如何将分割的mask转为为分割标签

将分割的mask转换为分割标签通常涉及将每个像素的类别标识（在mask中以不同的灰度值或颜色表示）转换为整数标签。这些标签通常用于机器学习或深度学习模型的训练、验证和测试阶段。

使用方式，控制台或者命令行使用以下命令：

python polygon_mask_conversion.py --img_path xxx_folder --mask_path xxx_folder --mode mask2poly    

转换代码，来自X_anyLabeling的tool文件夹下的转换文件。

import argparse
import json
import os
import time
import cv2from PIL import Image
from tqdm import tqdm
from datetime import dateimport numpy as np
import matplotlib as pltimport syssys.path.append("./")
from anylabeling.app_info import __version__# ======================================================================= Usage ========================================================================#
#                                                                                                                                                      #
# -------------------------------------------------------------------- mask2poly  ----------------------------------------------------------------------#
# python tools/polygon_mask_conversion.py --img_path xxx_folder --mask_path xxx_folder --mode mask2poly                                                #
#                                                                                                                                                      #
# -------------------------------------------------------------------- poly2mask  ----------------------------------------------------------------------#
# [option1] python tools/polygon_mask_conversion.py --img_path xxx_folder --mask_path xxx_folder --mode poly2mask                                      #
# [option2] python tools/polygon_mask_conversion.py --img_path xxx_folder --mask_path xxx_folder --json_path xxx_folder --mode poly2mask               #
#                                                                                                                                                      #
# ======================================================================= Usage ========================================================================#VERSION = __version__
IMG_FORMATS = [".bmp",".dng",".jpeg",".jpg",".mpo",".png",".tif",".tiff",".webp",".pfm",
]class PolygonMaskConversion:def __init__(self, epsilon_factor=0.001):self.epsilon_factor = epsilon_factordef reset(self):self.custom_data = dict(version=VERSION,flags={},shapes=[],imagePath="",imageData=None,imageHeight=-1,imageWidth=-1,)def get_image_size(self, image_file):with Image.open(image_file) as img:width, height = img.sizereturn width, heightdef mask_to_polygon(self, img_file, mask_file, json_file):self.reset()binary_mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)for contour in contours:epsilon = self.epsilon_factor * cv2.arcLength(contour, True)approx = cv2.approxPolyDP(contour, epsilon, True)if len(approx) < 5:continueshape = {"label": "object","text": "","points": [],"group_id": None,"shape_type": "polygon","flags": {},}for point in approx:x, y = point[0].tolist()shape["points"].append([x, y])self.custom_data["shapes"].append(shape)image_width, image_height = self.get_image_size(img_file)self.custom_data["imagePath"] = os.path.basename(img_file)self.custom_data["imageHeight"] = image_heightself.custom_data["imageWidth"] = image_widthwith open(json_file, "w", encoding="utf-8") as f:json.dump(self.custom_data, f, indent=2, ensure_ascii=False)def polygon_to_mask(self, img_file, mask_file, json_file):with open(json_file, "r") as f:data = json.load(f)polygons = []for shape in data["shapes"]:points = shape["points"]polygon = []for point in points:x, y = pointpolygon.append((x, y))polygons.append(polygon)image_width, image_height = self.get_image_size(img_file)image_shape = (image_height, image_width)binary_mask = np.zeros(image_shape, dtype=np.uint8)for polygon_points in polygons:np_polygon = np.array(polygon_points, np.int32)np_polygon = np_polygon.reshape((-1, 1, 2))cv2.fillPoly(binary_mask, [np_polygon], color=255)cv2.imwrite(mask_file, binary_mask)def main():parser = argparse.ArgumentParser(description="Polygon Mask Conversion")parser.add_argument("--img_path", help="Path to image directory")parser.add_argument("--mask_path", help="Path to mask directory")parser.add_argument("--json_path", default="", help="Path to json directory")parser.add_argument("--epsilon_factor",default=0.001,type=float,help="Control the level of simplification when converting a polygon contour to a simplified version",)parser.add_argument("--mode",choices=["mask2poly", "poly2mask"],required=True,help="Choose the conversion mode what you need",)args = parser.parse_args()print(f"Starting conversion to {args.mode}...")start_time = time.time()converter = PolygonMaskConversion(args.epsilon_factor)if args.mode == "mask2poly":file_list = os.listdir(args.mask_path)for file_name in tqdm(file_list, desc="Converting files", unit="file", colour="blue"):img_file = os.path.join(args.img_path, file_name)mask_file = os.path.join(args.mask_path, file_name)json_file = os.path.join(args.img_path, os.path.splitext(file_name)[0] + ".json")converter.mask_to_polygon(img_file, mask_file, json_file)elif args.mode == "poly2mask":# Only binary mask transformations are supported.os.makedirs(args.mask_path, exist_ok=True)file_list = os.listdir(args.img_path)for file_name in tqdm(file_list, desc="Converting files", unit="file", colour="blue"):base_name, suffix = os.path.splitext(file_name)if suffix.lower() not in IMG_FORMATS:continueimg_file = os.path.join(args.img_path, file_name)if not args.json_path:json_file = os.path.join(args.img_path, base_name + ".json")else:json_file = os.path.join(args.json_path, base_name + ".json")mask_file = os.path.join(args.mask_path, base_name + ".png")converter.polygon_to_mask(img_file, mask_file, json_file)end_time = time.time()print(f"Conversion completed successfully!")print(f"Conversion time: {end_time - start_time:.2f} seconds")if __name__ == "__main__":main()