anomlib_custom/split_data.py at main · lijin662/anomlib_custom · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
import os
import shutil
import random
from pathlib import Path
from sklearn.model_selection import train_test_split

def split_anomaly_detection_dataset(
    base_dir="./",
    output_dir="./split_data",
    train_ratio=0.7,
    val_ratio=0.15,
    test_ratio=0.15,
    random_seed=42
):
    """
    划分工业异常检测/分割数据集。
    严格保证 normal、anomaly 图片的独立分層划分，并同步联动对应 masks 目录下的真值分割标签。

    默认目录结构：
    base_dir/
      ├── normal/     (正常图片)
      ├── anomaly/    (异常/缺陷图片)
      └── masks/      (异常图片对应的 Segmentation Mask 标签)
    """
    # 比例校验
    assert abs((train_ratio + val_ratio + test_ratio) - 1.0) < 1e-5, "划分比例之和必须为 1.0！"

    random.seed(random_seed)

    # 路径实例化
    normal_dir = Path(base_dir) / "normal"
    anomaly_dir = Path(base_dir) / "anomaly"
    masks_dir = Path(base_dir) / "masks"

    # 支持的常见图像后缀
    valid_suffixes = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.npy'}

    # 读取过滤后的文件路径列表
    normal_img_paths = sorted([p for p in normal_dir.glob("*") if p.suffix.lower() in valid_suffixes])
    anomaly_img_paths = sorted([p for p in anomaly_dir.glob("*") if p.suffix.lower() in valid_suffixes])

    print(f"[*] 检测到正常样本 (Normal): {len(normal_img_paths)} 张")
    print(f"[*] 检测到异常样本 (Anomaly): {len(anomaly_img_paths)} 张")

    # 1. 划分正常样本 (Normal)
    if len(normal_img_paths) > 0:
        norm_train, norm_temp = train_test_split(normal_img_paths, test_size=(val_ratio + test_ratio), random_state=random_seed)
        norm_val, norm_test = train_test_split(norm_temp, test_size=(test_ratio / (val_ratio + test_ratio)), random_state=random_seed)
    else:
        norm_train, norm_val, norm_test = [], [], []

    # 2. 划分异常样本 (Anomaly) -> 保持与 Masks 同步
    if len(anomaly_img_paths) > 0:
        anom_train, anom_temp = train_test_split(anomaly_img_paths, test_size=(val_ratio + test_ratio), random_state=random_seed)
        anom_val, anom_test = train_test_split(anom_temp, test_size=(test_ratio / (val_ratio + test_ratio)), random_state=random_seed)
    else:
        anom_train, anom_val, anom_test = [], [], []

    # 构建结构化字典，方便循环分发
    splits = {
        "train": {"normal": norm_train, "anomaly": anom_train},
        "val": {"normal": norm_val, "anomaly": anom_val},
        "test": {"normal": norm_test, "anomaly": anom_test}
    }

    print("\n[*] 开始物理复制与对齐文件...")

    # 3. 动态创建标准结构并执行分发
    for split_name, categories in splits.items():
        for cat_name, paths in categories.items():
            # 目标图像子目录 (例如: ./split_data/train/normal)
            dest_img_dir = Path(output_dir) / split_name / cat_name
            dest_img_dir.mkdir(parents=True, exist_ok=True)

            # 目标掩码子目录 (例如: ./split_data/train/masks)
            dest_mask_dir = Path(output_dir) / split_name / "masks"
            if cat_name == "anomaly" and len(paths) > 0:
                dest_mask_dir.mkdir(parents=True, exist_ok=True)

            for img_path in paths:
                # 复制原图
                shutil.copy2(img_path, dest_img_dir / img_path.name)

                # 联动复制对应的 Mask 标签
                if cat_name == "anomaly":
                    # 方案 A: 严格同名同后缀匹配
                    mask_path = masks_dir / img_path.name
                    if mask_path.exists():
                        shutil.copy2(mask_path, dest_mask_dir / mask_path.name)
                    else:
                        # 方案 B: 容错处理（文件名相同但后缀不同，如原图 .jpg 标签 .png）
                        possible_masks = list(masks_dir.glob(f"{img_path.stem}.*"))
                        if possible_masks:
                            shutil.copy2(possible_masks[0], dest_mask_dir / possible_masks[0].name)
                        else:
                            print(f" [Warning] 未找到异常图片 {img_path.name} 对应的 Mask 标签文件！")

    # 打印最终划分统计
    print("\n" + "="*50)
    print(" 数据集划分与对齐任务圆满结束！")
    print("="*50)
    for split in ["train", "val", "test"]:
        n_count = len(splits[split]["normal"])
        a_count = len(splits[split]["anomaly"])
        print(f" -> [{split.upper():5s}] 正常集: {n_count:5d} | 异常集: {a_count:5d} | 总图片数: {n_count + a_count:5d}")
    print("="*50)

if __name__ == '__main__':
    # 执行划分，默认比例为 7:1.5:1.5
    split_anomaly_detection_dataset(
        base_dir="./",
        output_dir="./my_data",
        train_ratio=0.7,
        val_ratio=0.15,
        test_ratio=0.15,
        random_seed=42
    )