import os import json # ======================== 你的实际路径配置（无需修改） ======================== # 脚本所在目录（train目录） SCRIPT_ROOT = os.path.dirname(os.path.abspath(__file__)) # 数据集根路径（train/data） BASE_PATH = os.path.join(SCRIPT_ROOT, "data") # 元数据保存路径（框架固定目录） JSONL_SAVE_DIR = os.path.join(SCRIPT_ROOT, "dataset", "metadata", "opencvxuetang") # 原图目录：data/test/[类别名]/ IMG_ROOT = os.path.join(BASE_PATH, "test") # Mask根目录：data/ground_truth/[类别名]/ MASK_ROOT = os.path.join(BASE_PATH, "ground_truth") # 类别映射： # key=原图文件夹名 → (label值, 是否有mask, mask所在子文件夹名) # mask子文件夹名和原图类别名一致（比如contamination的mask在ground_truth/contamination/） CLASS_MAP = { "good": (0, False, ""), # 正常样本，无mask "contamination": (1, True, "contamination"), # 异常，mask在ground_truth/contamination/ "broken_small": (1, True, "broken_small"), # 异常，mask在ground_truth/broken_small/ "broken_large": (1, True, "broken_large"), # 异常，mask在ground_truth/broken_large/ } # Mask文件命名规则：原图名 → mask名（比如 xxx.jpg → xxx_mask.png） MASK_SUFFIX = "_mask.png" # ============================================================================= # 创建元数据保存目录 os.makedirs(JSONL_SAVE_DIR, exist_ok=True) JSONL_FILE_PATH = os.path.join(JSONL_SAVE_DIR, "full-shot.jsonl") meta_list = [] print("===== 开始生成元数据 =====\n") # 遍历所有类别 for class_name, (label, has_mask, mask_subdir) in CLASS_MAP.items(): # 1. 拼接原图类别文件夹路径 img_class_dir = os.path.join(IMG_ROOT, class_name) if not os.path.exists(img_class_dir): print(f"⚠️ 原图文件夹不存在：{img_class_dir} → 跳过该类别") continue print(f"🔍 处理类别：{class_name}") print(f" 原图路径：{img_class_dir}") # 2. 遍历该类别下的所有图片 img_files = [f for f in os.listdir(img_class_dir) if f.lower().endswith((".jpg", ".png", ".jpeg"))] if len(img_files) == 0: print(f" ❌ 该类别下无有效图片（jpg/png/jpeg）→ 跳过\n") continue print(f" 找到有效图片数：{len(img_files)}") for img_filename in img_files: # 构建原图的相对路径（框架要求：test/[类别名]/[图片名]） img_rel_path = os.path.join("test", class_name, img_filename).replace("\\", "/") sample = { "image_path": img_rel_path, "label": label, "class_name": class_name, "mask_path": "" } # 3. 处理异常样本的mask（核心修正：拼接mask子文件夹路径） if has_mask: # 拼接mask类别文件夹路径 mask_class_dir = os.path.join(MASK_ROOT, mask_subdir) if not os.path.exists(mask_class_dir): print(f" ❌ mask文件夹不存在：{mask_class_dir} → 跳过图片{img_filename}") continue # 提取原图名（去掉后缀）+ 拼接mask文件名 img_base_name = os.path.splitext(img_filename)[0] mask_filename = img_base_name + MASK_SUFFIX # 拼接mask的绝对路径（检查是否存在） mask_abs_path = os.path.join(mask_class_dir, mask_filename) if os.path.exists(mask_abs_path): # 构建mask的相对路径（框架要求：ground_truth/[类别名]/[mask名]） mask_rel_path = os.path.join("ground_truth", mask_subdir, mask_filename).replace("\\", "/") sample["mask_path"] = mask_rel_path meta_list.append(sample) print(f" ✅ 图片{img_filename} → mask{mask_filename} → 添加成功") else: print(f" ❌ 图片{img_filename}的mask不存在：{mask_abs_path} → 跳过") else: # good类直接添加 meta_list.append(sample) print(f" ✅ 正常图片{img_filename} → 添加成功") print("\n") # 写入jsonl文件 with open(JSONL_FILE_PATH, "w", encoding="utf-8") as f: for sample in meta_list: f.write(json.dumps(sample, ensure_ascii=False) + "\n") # 最终验证 print("===== 生成完成 ======") print(f"📁 元数据文件路径：{JSONL_FILE_PATH}") print(f"📊 总有效样本数：{len(meta_list)}") # 打印前5个样本示例 if len(meta_list) > 0: print("\n📝 前5个样本示例：") for i in range(min(5, len(meta_list))): print(f" 样本{i + 1}：{meta_list[i]}")
