掩码转多边形坐标点
代码参考: X-AnyLabeling中的.\X-AnyLabeling\ tools\polygon_mask_conversion.py
。这里主要拆分出掩码转成多边形坐标点存储在json
文件的部分,同时增加了参数和支持带有中文名的路径。
代码示例
import os
import cv2
import json
import numpy as np
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(
img_file, mask_file, json_file,
version,
epsilon_factor=0.001,
shape_dict=None,
approx_len_threshold=5
):
custom_data = reset(version)
# binary_mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)
binary_mask = cv2.imdecode(
np.fromfile(mask_file, dtype=np.uint8),
0
)
contours, _ = cv2.findContours(
binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
)
image_name = os.path.basename(img_file)
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])
custom_data["shapes"].append(shape_dict)
image_width, image_height = get_image_size(img_file)
custom_data["imagePath"] = image_name
custom_data["imageHeight"] = image_height
custom_data["imageWidth"] = image_width
with open(json_file, "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 = "./jsons/test.json"
shape_dict = {
"label": "test",
"text": "test text",
"points": [],
"group_id": None,
"shape_type": "polygon",
"flags": {},
}
mask_to_polygon(
image_path,
mask_path,
json_path,
epsilon_factor=0.001,
version="2.4.0",
shape_dict=shape_dict,
approx_len_threshold=5
)