增加图像分割相应的代码
This commit is contained in:
@@ -0,0 +1,303 @@
|
||||
"""
|
||||
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 _merge_annotations(self, image_id: int, image_size: Tuple[int, int]) -> np.ndarray:
|
||||
"""
|
||||
合并一张图像的所有标注为单一mask
|
||||
|
||||
由于用户标注时所有食材区域都是同一类别，我们需要将同一图像的多个标注合并
|
||||
|
||||
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' in ann and isinstance(ann['segmentation'], list):
|
||||
# 转换polygon为mask
|
||||
obj_mask = self._polygon_to_mask(ann['segmentation'], image_size)
|
||||
# 合并到总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文件夹）
|
||||
"""
|
||||
|
||||
# ==================== 配置区 ====================
|
||||
# TODO: 请根据您的实际路径修改以下三个变量
|
||||
|
||||
# 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()
|
||||
Reference in New Issue
Block a user