多类别三通道掩码转多边形坐标点
整体与掩码转多边形坐标点类似,主要是单类别单通道掩码改为多类别三通道掩码。
代码示例
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[::point_interval_num]:
x, y = point[0].tolist()
shape_dict["points"].append([x, y])
# 可选,确保添加最后一个元素
if len(approx) % point_interval_num!= 0:
shape_dict["points"].append(approx[-1][0].tolist())
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),
1
)
image_name = os.path.basename(image_path)
image_width, image_height = get_image_size(image_path)
custom_data = reset(__version__)
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_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"]
# bgr
mask_color_list = [(5, 178, 124), (69, 150, 255)]
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,
)