367 lines
13 KiB
Python
367 lines
13 KiB
Python
"""
|
||
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()
|