X-AnyLabeling系列：多类别单通道掩码转多边形坐标点

多类别单通道掩码转多边形坐标点

整体与掩码转多边形坐标点类似，主要是单类别单通道掩码改为多类别单通道掩码。

示例代码

import os
import cv2
import json
import numpy as np
from copy import deepcopy

def get_image_size(image_file):
    image = cv2.imdecode(
      np.fromfile(image_file, dtype=np.uint8), 
      0
    )
    height, width = image.shape[:2]
    return width, height


def reset(version):
    custom_data = dict(
        version=version,
        flags={},
        shapes=[],
        imagePath="",
        imageData=None,
        imageHeight=-1,
        imageWidth=-1,
    )

    return custom_data


def mask_to_polygon(
    mask, shape_dict, 
    image_name,
    epsilon_factor,
    approx_len_threshold,
    point_interval_num=1
):
    contours, _ = cv2.findContours(
        mask, cv2.RETR_EXTERNAL, 
        cv2.CHAIN_APPROX_SIMPLE
    )
    for contour in contours:
        epsilon = epsilon_factor * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)
        if len(approx) < approx_len_threshold:
            print(
              f"{image_name}: \
                contour too small, len={len(approx)}"
            )
            continue
        for point in approx:
            x, y = point[0].tolist()
            shape_dict["points"].append([x, y])

    return shape_dict

def multiclass_mask_to_polygon(
  image_path, mask_path, json_path,
  mask_label_list, mask_color_list,
  shape_dict, 
  epsilon_factor, approx_len_threshold,
  point_interval_num=1,
):
  # mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
  mask = cv2.imdecode(
    np.fromfile(mask_path, dtype=np.uint8), 0
  )
  image_name = os.path.basename(image_path)
  image_width, image_height = get_image_size(image_path)
  custom_data = reset("2.4.0")
  custom_data["imagePath"] = image_name
  custom_data["imageHeight"] = image_height
  custom_data["imageWidth"] = image_width

  for label, color in zip(mask_label_list, mask_color_list):
    # 选择特定颜色的mask
    mask_label = cv2.inRange(mask, color, color)
    mask_label_shape_dict = deepcopy(shape_dict)
    mask_label_shape_dict["label"] = label
    mask_label_shape_dict = mask_to_polygon(
        mask_label, 
        mask_label_shape_dict, 
        image_name,
        epsilon_factor,
        approx_len_threshold,
        point_interval_num,

    )
    custom_data["shapes"].append(mask_label_shape_dict)

  with open(json_path, "w", encoding="utf-8") as f:
      json.dump(custom_data, f, indent=2, ensure_ascii=False)

if __name__ == '__main__':
  image_path = "./images/test.jpg"
  mask_path = "./masks/test.png"
  json_path = "./json/test.json"
  mask_label_list = ["label_1", "label_2"]
  mask_color_list = [255, 128]
  shape_dict = {
      "label": "",
      "description": "",
      "difficult": False,
      "points": [],
      "shape_type": "polygon",
  }
  multiclass_mask_to_polygon(
    image_path, mask_path, json_path,
    mask_label_list, mask_color_list,
    shape_dict, 
    epsilon_factor=0.001, approx_len_threshold=5,
    point_interval_num=2,
  )