""" COCO格式数据转换为SegFormer训练格式 功能说明： 1. 读取CVAT导出的COCO格式标注文件 2. 将polygon标注转换为像素级mask（PNG图像） 3. 生成训练集和验证集的划分 4. 输出符合SegFormer训练要求的目录结构 输出目录结构： data/segformer_format/ ├── images/ │ ├── train/ │ │ ├── img1.jpg │ │ └── img2.jpg │ └── val/ │ └── img3.jpg └── annotations/ ├── train/ │ ├── img1.png # 像素值：0=背景, 1=食物区域 │ └── img2.png └── val/ └── img3.png """ import os import json import numpy as np from PIL import Image, ImageDraw from pathlib import Path from typing import List, Dict, Tuple import shutil class COCOToSegFormerConverter: """COCO格式到SegFormer格式的转换器""" def __init__(self, coco_json_path: str, coco_images_dir: str, output_dir: str): """ 初始化转换器 Args: coco_json_path: COCO标注JSON文件路径（例如：instances_default.json） coco_images_dir: COCO图像所在目录 output_dir: 输出目录（将创建segformer_format文件夹） """ self.coco_json_path = coco_json_path self.coco_images_dir = coco_images_dir self.output_dir = Path(output_dir) # 加载COCO标注 print(f"加载COCO标注文件: {coco_json_path}") with open(coco_json_path, 'r', encoding='utf-8') as f: self.coco_data = json.load(f) print(f"✓ 图像数量: {len(self.coco_data['images'])}") print(f"✓ 标注数量: {len(self.coco_data['annotations'])}") print(f"✓ 类别数量: {len(self.coco_data['categories'])}") # 创建输出目录结构 self._create_output_dirs() def _create_output_dirs(self): """创建输出目录结构""" dirs = [ self.output_dir / "images" / "train", self.output_dir / "images" / "val", self.output_dir / "annotations" / "train", self.output_dir / "annotations" / "val", ] for d in dirs: d.mkdir(parents=True, exist_ok=True) print(f"✓ 输出目录创建完成: {self.output_dir}") def _polygon_to_mask(self, segmentation: List, image_size: Tuple[int, int]) -> np.ndarray: """ 将COCO的polygon格式转换为像素级mask Args: segmentation: COCO的segmentation字段（polygon列表） image_size: 图像尺寸 (width, height) Returns: mask: 二值mask数组 (H, W)，1表示目标区域，0表示背景 """ width, height = image_size mask = Image.new('L', (width, height), 0) # 黑色背景 # COCO的segmentation可能包含多个polygon（例如：一个物体被遮挡分成多个部分） for polygon in segmentation: # polygon格式: [x1, y1, x2, y2, x3, y3, ...] # 转换为坐标点列表: [(x1,y1), (x2,y2), ...] points = [(polygon[i], polygon[i+1]) for i in range(0, len(polygon), 2)] # 在mask上绘制填充的多边形 ImageDraw.Draw(mask).polygon(points, outline=1, fill=1) return np.array(mask) def _rle_to_mask(self, rle: Dict, image_size: Tuple[int, int]) -> np.ndarray: """ 将COCO的RLE (Run-Length Encoding) 格式转换为像素级mask RLE格式说明： - counts: [v1, v2, v3, ...] 交替表示背景和前景的像素数量 - 从背景(0)开始，然后是前景(1)，再是背景(0)，依此类推 - 像素顺序：从左到右，从上到下（行优先） Args: rle: COCO的RLE字典，格式为 {"counts": [n1, n2, ...], "size": [height, width]} image_size: 图像尺寸 (width, height) Returns: mask: 二值mask数组 (H, W)，1表示目标区域，0表示背景 """ width, height = image_size # 获取RLE的counts if 'counts' in rle: counts = rle['counts'] else: print(f" ⚠️ 警告: RLE格式不包含counts字段") return np.zeros((height, width), dtype=np.uint8) # 创建一维数组 mask = np.zeros(height * width, dtype=np.uint8) current_pos = 0 current_value = 0 # 从背景开始 for count in counts: # 如果当前是前景(1)，填充mask if current_value == 1: mask[current_pos:current_pos + count] = 1 # 移动位置 current_pos += count # 切换值（0 -> 1 -> 0 -> 1 ...） current_value = 1 - current_value # 重塑为2D数组 (H, W) mask = mask.reshape((height, width)) return mask def _merge_annotations(self, image_id: int, image_size: Tuple[int, int]) -> np.ndarray: """ 合并一张图像的所有标注为单一mask 由于用户标注时所有食材区域都是同一类别，我们需要将同一图像的多个标注合并 支持两种格式： 1. Polygon格式: segmentation 是列表的列表 2. RLE格式: segmentation 是包含 'counts' 字段的字典 Args: image_id: COCO图像ID image_size: 图像尺寸 (width, height) Returns: merged_mask: 合并后的mask (H, W) 像素值: 0=背景（未标注区域）, 1=食物区域 """ width, height = image_size merged_mask = np.zeros((height, width), dtype=np.uint8) # 找到该图像的所有标注 annotations = [ann for ann in self.coco_data['annotations'] if ann['image_id'] == image_id] # 将所有标注合并到同一个mask for ann in annotations: if 'segmentation' not in ann: continue segmentation = ann['segmentation'] # 判断是RLE格式还是Polygon格式 if isinstance(segmentation, dict) and 'counts' in segmentation: # RLE格式 obj_mask = self._rle_to_mask(segmentation, image_size) elif isinstance(segmentation, list) and len(segmentation) > 0: # Polygon格式 obj_mask = self._polygon_to_mask(segmentation, image_size) else: print(f" ⚠️ 警告: 未知的segmentation格式: {type(segmentation)}") continue # 合并到总mask（取并集） merged_mask = np.maximum(merged_mask, obj_mask) return merged_mask def convert(self, train_ratio: float = 0.8, random_seed: int = 42): """ 执行转换流程 Args: train_ratio: 训练集占比（0.8表示80%训练，20%验证） random_seed: 随机种子，确保每次划分一致 """ print("\n" + "="*60) print("开始转换数据集") print("="*60) # 创建图像ID到文件名的映射 id_to_image = {img['id']: img for img in self.coco_data['images']} # 随机划分训练集和验证集 np.random.seed(random_seed) image_ids = list(id_to_image.keys()) np.random.shuffle(image_ids) split_idx = int(len(image_ids) * train_ratio) train_ids = image_ids[:split_idx] val_ids = image_ids[split_idx:] print(f"\n数据集划分:") print(f" 训练集: {len(train_ids)} 张图像") print(f" 验证集: {len(val_ids)} 张图像") # 处理训练集 print(f"\n处理训练集...") self._process_split(train_ids, id_to_image, split='train') # 处理验证集 print(f"\n处理验证集...") self._process_split(val_ids, id_to_image, split='val') # 保存数据集统计信息 self._save_dataset_info(train_ids, val_ids) print("\n" + "="*60) print("✓ 数据集转换完成!") print("="*60) print(f"\n输出目录: {self.output_dir}") print("\n下一步: 运行 2_visualize_data.py 查看转换结果") def _process_split(self, image_ids: List[int], id_to_image: Dict, split: str): """ 处理训练集或验证集 Args: image_ids: 图像ID列表 id_to_image: ID到图像信息的映射 split: 'train' 或 'val' """ for idx, img_id in enumerate(image_ids, 1): img_info = id_to_image[img_id] file_name = img_info['file_name'] width = img_info['width'] height = img_info['height'] print(f" [{idx}/{len(image_ids)}] {file_name}") # 1. 复制图像到目标目录 src_image_path = Path(self.coco_images_dir) / file_name dst_image_path = self.output_dir / "images" / split / file_name if not src_image_path.exists(): print(f" ⚠️ 警告: 图像文件不存在 {src_image_path}") continue shutil.copy2(src_image_path, dst_image_path) # 2. 生成mask并保存为PNG mask = self._merge_annotations(img_id, (width, height)) # 保存mask（像素值即为类别ID：0=背景, 1=食物） mask_filename = Path(file_name).stem + '.png' # 改为.png扩展名 mask_path = self.output_dir / "annotations" / split / mask_filename # 使用PIL保存，确保像素值不被压缩 Image.fromarray(mask, mode='L').save(mask_path) # 统计信息 food_pixels = np.sum(mask == 1) total_pixels = mask.size food_ratio = food_pixels / total_pixels * 100 print(f" ✓ 食物区域占比: {food_ratio:.1f}%") def _save_dataset_info(self, train_ids: List[int], val_ids: List[int]): """保存数据集统计信息""" info = { "dataset_name": "Food Segmentation Dataset", "num_classes": 2, # 背景 + 食物 "class_names": ["background", "food"], "train_size": len(train_ids), "val_size": len(val_ids), "total_size": len(train_ids) + len(val_ids), "image_format": "jpg/png", "annotation_format": "png (pixel value = class id)", } info_path = self.output_dir / "dataset_info.json" with open(info_path, 'w', encoding='utf-8') as f: json.dump(info, f, indent=2, ensure_ascii=False) print(f"\n✓ 数据集信息已保存: {info_path}") def main(): """ 主函数：配置路径并执行转换 使用前请修改以下路径： 1. COCO_JSON_PATH: CVAT导出的JSON文件路径 2. COCO_IMAGES_DIR: CVAT导出的图像目录 3. OUTPUT_DIR: 输出目录（将创建segformer_format文件夹） """ # ==================== 配置区 ==================== # 请根据您的实际路径修改以下三个变量 # CVAT导出的COCO标注文件（通常名为instances_default.json） COCO_JSON_PATH = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/raw_coco/annotations/instances_default.json" # CVAT导出的图像目录 COCO_IMAGES_DIR = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/raw_coco/images" # 输出目录（将在此目录下创建segformer_format文件夹） OUTPUT_DIR = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format" # 训练集/验证集划分比例（0.8表示80%训练，20%验证） TRAIN_RATIO = 0.8 # 随机种子（保证每次运行划分结果一致） RANDOM_SEED = 42 # =============================================== print("COCO数据集转换工具") print("目标格式: SegFormer训练格式\n") # 检查输入文件是否存在 if not os.path.exists(COCO_JSON_PATH): print(f"❌ 错误: COCO标注文件不存在") print(f" 路径: {COCO_JSON_PATH}") print(f"\n请检查:") print(f" 1. 是否已从CVAT导出COCO格式数据") print(f" 2. 标注文件路径是否正确") return if not os.path.exists(COCO_IMAGES_DIR): print(f"❌ 错误: 图像目录不存在") print(f" 路径: {COCO_IMAGES_DIR}") return # 创建转换器并执行转换 converter = COCOToSegFormerConverter( coco_json_path=COCO_JSON_PATH, coco_images_dir=COCO_IMAGES_DIR, output_dir=OUTPUT_DIR ) converter.convert(train_ratio=TRAIN_RATIO, random_seed=RANDOM_SEED) print("\n" + "="*60) print("转换完成! 接下来的步骤:") print("="*60) print("1. 运行 2_visualize_data.py 检查转换结果") print("2. 运行 3_train_minimal.py 开始训练") if __name__ == "__main__": main()