如何将RGB掩码转换为适配Mask R-CNN的多边形标注JSON文件

前置依赖 #

需要安装3个第三方库，执行以下命令安装即可：
pip install pillow opencv-python numpy

完整实现代码 #

代码保留了你已编写的create_sub_masks子掩码提取逻辑，补充了轮廓提取、坐标偏移修正、格式组装、文件输出的全流程：

from PIL import Image
import cv2
import numpy as np
import json
import warnings


def create_sub_masks(mask_image):
    width, height = mask_image.size
    # 初始化按RGB值索引的子掩码字典
    sub_masks = {}
    for x in range(width):
        for y in range(height):
            # 读取像素RGB值
            pixel = mask_image.getpixel((x,y))[:3]
            # 跳过黑色背景
            if pixel != (0, 0, 0):
                pixel_str = str(pixel)
                sub_mask = sub_masks.get(pixel_str)
                if sub_mask is None:
                    # 子掩码添加1像素边缘padding，适配contours模块边界处理
                    sub_masks[pixel_str] = Image.new('1', (width+2, height+2))
                # 对应padding偏移写入像素值
                sub_masks[pixel_str].putpixel((x+1, y+1), 1)
    return sub_masks


def mask_to_labelme_json(mask_path, output_json_path, color_label_map):
    # 读取输入掩码图像
    mask_img = Image.open(mask_path).convert("RGB")
    width, height = mask_img.size
    # 拆分各颜色对应的二值子掩码
    sub_masks = create_sub_masks(mask_img)
    
    # 初始化符合要求的标注结构
    annotation = {
        "version": "5.0.1",
        "flags": {},
        "shapes": [],
        "imagePath": mask_path.split("\\")[-1].split("/")[-1],
        "imageData": None,
        "imageHeight": height,
        "imageWidth": width
    }

    for pixel_color_str, sub_mask in sub_masks.items():
        # 转换字符串格式的颜色值为元组
        color = eval(pixel_color_str)
        # 匹配对应类别
        if color not in color_label_map:
            warnings.warn(f"检测到未配置映射的颜色值{color}，已自动跳过对应区域")
            continue
        label = color_label_map[color]

        # PIL格式子掩码转OpenCV可处理的数组格式
        sub_mask_np = np.array(sub_mask, dtype=np.uint8) * 255
        # 提取外轮廓，忽略内部孔洞
        contours, _ = cv2.findContours(
            sub_mask_np,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )

        # 遍历同类别下的所有独立连通域轮廓
        for contour in contours:
            # 过滤面积极小的噪点轮廓，阈值可根据实际需求调整
            if cv2.contourArea(contour) < 2:
                continue
            # 转换轮廓点格式，同时抵消1像素padding带来的坐标偏移
            points = []
            contour_points = contour.squeeze()
            # 处理单点轮廓的异常情况
            if len(contour_points.shape) < 2:
                continue
            for point in contour_points:
                x = float(point[0]) - 1
                y = float(point[1]) - 1
                points.append([x, y])
            # 点数小于3无法构成有效多边形，跳过
            if len(points) < 3:
                continue
            # 组装单条标注
            shape_item = {
                "label": label,
                "points": points,
                "group_id": None,
                "shape_type": "polygon",
                "flags": {}
            }
            annotation["shapes"].append(shape_item)
    
    # 写入最终JSON文件
    with open(output_json_path, "w", encoding="utf-8") as f:
        json.dump(annotation, f, ensure_ascii=False, indent=2)
    print(f"标注文件生成完成，保存路径：{output_json_path}")


if __name__ == "__main__":
    # --------------- 按需修改以下配置 ---------------
    # 颜色(RGB)到类别名的映射表，必须和掩码中的颜色值完全对应
    COLOR_TO_LABEL = {
        (255, 0, 0): "window",
        (0, 255, 0): "door",
        (0, 0, 255): "wall"
    }
    # 输入掩码文件路径、输出JSON文件路径
    INPUT_MASK_PATH = "sample.png"
    OUTPUT_JSON_PATH = "sample.json"
    # -----------------------------------------------
    mask_to_labelme_json(
        mask_path=INPUT_MASK_PATH,
        output_json_path=OUTPUT_JSON_PATH,
        color_label_map=COLOR_TO_LABEL
    )

使用说明 #

• 使用前必须根据自身掩码的实际RGB颜色值修改COLOR_TO_LABEL映射配置，否则无法正确匹配类别标签
• 代码默认过滤了面积小于2像素的噪点轮廓，如果存在极小尺寸目标，可自行调整cv2.contourArea判断的阈值
• 生成的JSON为标准Labelme格式，既可以用Labelme打开人工校验，也可以直接转换为COCO格式供Mask R-CNN训练使用
• 轮廓提取默认使用CHAIN_APPROX_SIMPLE算法压缩冗余顶点，如果需要更高精度的轮廓，可以将该参数替换为cv2.CHAIN_APPROX_NONE保留所有轮廓点
• 代码自动兼容Windows、Linux路径格式，自动提取文件名填入标注结构

内容的提问来源于stack exchange，提问作者Bingo Man