mmdetection-实例分割

import os
import json
import random
import shutil
import numpy as np
from PIL import Image, ImageDraw
from tqdm import tqdm

# ==== 配置路径 ====
json_path = r"Annotations\coco_info.json"       # 你的 COCO 标注文件
images_dir = r"\Images"              # 所有原始图片路径
output_dir = r"\dataset"               # 输出的根目录
train_ratio = 0.7                      # 训练集比例

# ==== 创建输出目录 ====
images_train_dir = os.path.join(output_dir, "images/train2017")
images_val_dir = os.path.join(output_dir, "images/val2017")
stuff_train_dir = os.path.join(output_dir, "stuffthingmaps/train2017")
stuff_val_dir = os.path.join(output_dir, "stuffthingmaps/val2017")
ann_dir = os.path.join(output_dir, "annotations")

for d in [images_train_dir, images_val_dir, stuff_train_dir, stuff_val_dir, ann_dir]:
    os.makedirs(d, exist_ok=True)

# ==== 读取 JSON ====
with open(json_path, "r", encoding="utf-8") as f:
    coco_data = json.load(f)

images_info = coco_data["images"]
annotations = coco_data["annotations"]
categories = coco_data["categories"]

# ==== 按图片分组标注 ====
image_to_anns = {}
for ann in annotations:
    image_to_anns.setdefault(ann["image_id"], []).append(ann)

# ==== 随机划分 ====
random.seed(42)
image_ids = [img["id"] for img in images_info]
random.shuffle(image_ids)

split_idx = int(len(image_ids) * train_ratio)
train_ids = set(image_ids[:split_idx])
val_ids = set(image_ids[split_idx:])

train_images, val_images = [], []
train_anns, val_anns = [], []

# ==== 生成 stuffthingmaps 并划分 ====
for img_info in tqdm(images_info, desc="Processing images"):
    file_name = img_info["file_name"]
    img_id = img_info["id"]
    width, height = img_info["width"], img_info["height"]

    # 创建单通道mask（全背景0）
    mask = np.zeros((height, width), dtype=np.uint8)
    anns = image_to_anns.get(img_id, [])

    for ann in anns:
        if ann["iscrowd"] == 1:
            continue
        if ann["segmentation"]:
            img_mask = Image.new("L", (width, height), 0)
            draw = ImageDraw.Draw(img_mask)
            seg = ann["segmentation"]
            if isinstance(seg, list):  # polygon
                for poly in seg:
                    if len(poly) >= 6:
                        xy = [(poly[i], poly[i+1]) for i in range(0, len(poly), 2)]
                        draw.polygon(xy, fill=1)  # 单类：前景设为1
            elif isinstance(seg, dict) and "counts" in seg:  # RLE
                import pycocotools.mask as maskUtils
                m = maskUtils.decode(seg)
                mask[m > 0] = 1
            mask = np.maximum(mask, np.array(img_mask))  # 取最大值，保持前景1

    # 保存mask
    if img_id in train_ids:
        mask_path = os.path.join(stuff_train_dir, os.path.splitext(file_name)[0] + ".png")
        img_out_path = os.path.join(images_train_dir, file_name)
        train_images.append(img_info)
        train_anns.extend(anns)
    else:
        mask_path = os.path.join(stuff_val_dir, os.path.splitext(file_name)[0] + ".png")
        img_out_path = os.path.join(images_val_dir, file_name)
        val_images.append(img_info)
        val_anns.extend(anns)

    Image.fromarray(mask).save(mask_path)
    shutil.copy(os.path.join(images_dir, file_name), img_out_path)

# ==== 保存新的JSON ====
train_json = {
    "images": train_images,
    "annotations": train_anns,
    "categories": categories
}
val_json = {
    "images": val_images,
    "annotations": val_anns,
    "categories": categories
}

with open(os.path.join(ann_dir, "instances_train2017.json"), "w", encoding="utf-8") as f:
    json.dump(train_json, f, ensure_ascii=False)
with open(os.path.join(ann_dir, "instances_val2017.json"), "w", encoding="utf-8") as f:
    json.dump(val_json, f, ensure_ascii=False)

print("数据集处理完成！")

METAINFO = {
        'classes':
        ( 'rice'),
        # palette is a list of color tuples, which is used for visualization.
        'palette':
        [ (255, 0, 0)]
    }