Added support for exporting PNG masks in grayscale and RGB formats, a…

…llowing unique class representation through either grayscale values or distinct RGB colors (CVHub520#131)

anylabeling/resources/resources.qrc

@@ -29,6 +29,7 @@
         <file>images/format_dota.png</file>
         <file>images/format_yolo.png</file>
         <file>images/format_coco.png</file>
+        <file>images/format_mask.png</file>
         <file>images/help.png</file>
         <file>images/icon.icns</file>
         <file>images/icon.ico</file>

anylabeling/views/labeling/label_converter.py

@@ -1,4 +1,5 @@
 import os
+import cv2
 import csv
 import json
 import natsort
@@ -13,11 +14,15 @@
 
 
 class LabelConverter:
-    def __init__(self, classes_file):
+    def __init__(self, classes_file, mapping_file):
         self.classes = []
         if classes_file:
             with open(classes_file, "r", encoding="utf-8") as f:
                 self.classes = f.read().splitlines()
+        self.mapping_table = None
+        if mapping_file:
+            with open(mapping_file, "r", encoding="utf-8") as f:
+                self.mapping_table = json.load(f)
 
     @staticmethod
     def calculate_polygon_area(segmentation):
@@ -311,3 +316,50 @@ def custom_to_coco(self, root_path, output_file, formats):
         # Save the coco result
         with open(output_file, "w", encoding="utf-8") as f:
             json.dump(coco_meta_data, f, indent=4, ensure_ascii=False)
+
+    def custom_polygon_to_mask(self, data, output_file):
+        image_width = data["imageWidth"]
+        image_height = data["imageHeight"]
+        image_shape = (image_height, image_width)
+
+        polygons = {}
+        for shape in data["shapes"]:
+            points = shape["points"]
+            polygon = []
+            for point in points:
+                x, y = point
+                polygon.append((int(x), int(y)))  # Convert to integers
+            polygons[shape["label"]] = polygon
+
+        output_format = self.mapping_table["type"]
+        if output_format not in  ["grayscale", "rgb"]:
+            raise ValueError("Invalid output format specified")
+        mapping_color = self.mapping_table["colors"]
+
+        if output_format == "grayscale":
+            binary_mask = np.zeros(image_shape, dtype=np.uint8)
+            for label, polygon in polygons.items():
+                mask = np.zeros(image_shape, dtype=np.uint8)
+                cv2.fillPoly(mask, [np.array(polygon, dtype=np.int32)], 1)
+                if label in mapping_color:
+                    mask_mapped = mask * mapping_color[label]
+                else:
+                    mask_mapped = mask
+                binary_mask += mask_mapped
+            cv2.imwrite(output_file, binary_mask)
+        elif output_format == "rgb":
+            # Initialize rgb_mask
+            color_mask = np.zeros((image_height, image_width, 3), dtype=np.uint8)
+            for label, polygon in polygons.items():
+                # Create a mask for each polygon
+                mask = np.zeros(image_shape[:2], dtype=np.uint8)
+                cv2.fillPoly(mask, [np.array(polygon, dtype=np.int32)], 1)
+                # Initialize mask_mapped with a default value
+                mask_mapped = mask
+                # Map the mask values using the provided mapping table
+                if label in mapping_color:
+                    color = mapping_color[label]
+                    mask_mapped = np.zeros_like(color_mask)
+                    cv2.fillPoly(mask_mapped, [np.array(polygon, dtype=np.int32)], color)
+                    color_mask = cv2.addWeighted(color_mask, 1, mask_mapped, 1, 0)
+            cv2.imwrite(output_file, cv2.cvtColor(color_mask, cv2.COLOR_BGR2RGB))

anylabeling/views/labeling/label_file.py

@@ -166,6 +166,7 @@ def save(
         flags=None,
         output_format="defalut",
         classes_file=None,
+        mapping_file=None,
     ):
         if image_data is not None:
             image_data = base64.b64encode(image_data).decode("utf-8")
@@ -194,15 +195,17 @@ def save(
             with io_open(filename, "w") as f:
                 json.dump(data, f, ensure_ascii=False, indent=2)
             self.filename = filename
-            _ = self.save_other_mode(data, output_format, classes_file)
+            _ = self.save_other_mode(
+                data, output_format, classes_file, mapping_file
+            )
         except Exception as e:  # noqa
             raise LabelFileError(e) from e
 
     @staticmethod
     def is_label_file(filename):
         return osp.splitext(filename)[1].lower() == LabelFile.suffix
 
-    def save_other_mode(self, data, mode, classes_file=None):
+    def save_other_mode(self, data, mode, classes_file=None, mapping_file=None):
         target_formats = ["polygon", "rectangle", "rotation"]
         shape_type = self.get_shape_type(data, target_formats)
         if mode == "default" or not shape_type:
@@ -227,8 +230,15 @@ def save_other_mode(self, data, mode, classes_file=None):
             os.makedirs(save_path, exist_ok=True)
         elif mode == "mot":
             dst_file = root_path + "/" + base_name.rsplit("_", 1)[0] + ".csv"
+        elif mode == "mask":
+            save_path = root_path + "/masks"
+            dst_file = save_path + "/" + base_name + ".png"
+            os.makedirs(save_path, exist_ok=True)
 
-        converter = LabelConverter(classes_file=classes_file)
+        converter = LabelConverter(
+            classes_file=classes_file, 
+            mapping_file=mapping_file,
+        )
         if mode == "yolo" and shape_type == "rectangle":
             converter.custom_to_yolo_rectangle(data, dst_file)
             return True
@@ -247,6 +257,9 @@ def save_other_mode(self, data, mode, classes_file=None):
         elif mode == "mot" and shape_type == "rectangle":
             converter.custom_to_mot_rectangle(data, dst_file, base_name)
             return True
+        elif mode == "mask" and shape_type == "polygon":
+            converter.custom_polygon_to_mask(data, dst_file)
+            return True
         else:
             return False
 

anylabeling/views/labeling/label_widget.py

@@ -84,7 +84,8 @@ def __init__(
         self.image_data = None
         self.label_file = None
         self.other_data = {}
-        self.classes = None
+        self.classes_file = None
+        self.mapping_file = None
         self.attributes = {}
         self.current_category = None
 
@@ -762,6 +763,14 @@ def __init__(
             checked=self._config["save_mode"] == "mot",
             enabled=self._config["save_mode"] != "mot",
         )
+        select_mask_format = action(
+            "MASK",
+            functools.partial(self.set_output_format, "mask"),
+            icon="format_mask",
+            checkable=True,
+            checked=self._config["save_mode"] == "mask",
+            enabled=self._config["save_mode"] != "mask",
+        )
 
         # Group zoom controls into a list for easier toggling.
         zoom_actions = (
@@ -854,6 +863,7 @@ def __init__(
             select_voc_format=select_voc_format,
             select_dota_format=select_dota_format,
             select_mot_format=select_mot_format,
+            select_mask_format=select_mask_format,
             zoom=zoom,
             zoom_in=zoom_in,
             zoom_out=zoom_out,
@@ -992,6 +1002,7 @@ def __init__(
                 select_voc_format,
                 select_dota_format,
                 select_mot_format,
+                select_mask_format,
             ),
         )
         utils.add_actions(
@@ -1276,13 +1287,13 @@ def set_output_format(self, mode):
         confirm_flag = True
         if mode in ["yolo", "coco", "mot"]:
             filter = "Classes Files (*.txt);;All Files (*)"
-            self.classes, _ = QtWidgets.QFileDialog.getOpenFileName(
+            self.classes_file, _ = QtWidgets.QFileDialog.getOpenFileName(
                 self,
                 self.tr("Select a specific classes file"),
                 "",
                 filter,
             )
-            if not self.classes:
+            if not self.classes_file:
                 QtWidgets.QMessageBox.warning(
                     self,
                     self.tr("Warning"),
@@ -1293,7 +1304,7 @@ def set_output_format(self, mode):
                 self._config["save_mode"] = "default"
                 return
             else:
-                with open(self.classes, "r", encoding="utf-8") as f:
+                with open(self.classes_file, "r", encoding="utf-8") as f:
                     classes = f.read().splitlines()
                     for label in classes:
                         if not self.unique_label_list.find_items_by_label(
@@ -1318,7 +1329,7 @@ def set_output_format(self, mode):
                     ] + [f"*{LabelFile.suffix}"]
                     if "*.json" in formats:
                         formats.remove("*.json")
-                    converter = LabelConverter(classes_file=self.classes)
+                    converter = LabelConverter(classes_file=self.classes_file)
                     root_path = osp.split(self.filename)[0]
                     save_path = root_path + "/annotations"
                     os.makedirs(save_path, exist_ok=True)
@@ -1333,6 +1344,42 @@ def set_output_format(self, mode):
                     msg_box.exec_()
                     confirm_flag = False
                     self._config["save_mode"] = "default"
+        elif mode == "mask":
+            filter = "JSON Files (*.json);;All Files (*)"
+            self.mapping_file, _ = QtWidgets.QFileDialog.getOpenFileName(
+                self,
+                self.tr("Select a specific color_map file"),
+                "",
+                filter,
+            )
+            if not self.mapping_file:
+                QtWidgets.QMessageBox.warning(
+                    self,
+                    self.tr("Warning"),
+                    self.tr("Please select a specific color_map file!"),
+                    QtWidgets.QMessageBox.Ok,
+                )
+                confirm_flag = False
+                self._config["save_mode"] = "default"
+                return
+            else:
+                with open(self.mapping_file, "r", encoding="utf-8") as f:
+                    mapping_table = json.load(f)
+                    classes = list(mapping_table["colors"].keys())
+                    for label in classes:
+                        if not self.unique_label_list.find_items_by_label(
+                            label
+                        ):
+                            item = (
+                                self.unique_label_list.create_item_from_label(
+                                    label
+                                )
+                            )
+                            self.unique_label_list.addItem(item)
+                            rgb = self._get_rgb_by_label(label)
+                            self.unique_label_list.set_item_label(
+                                item, label, rgb
+                            )
 
         # Show dialog to restart application
         if confirm_flag and self._config["save_mode"] != "default":
@@ -1349,34 +1396,47 @@ def set_output_format(self, mode):
             self.actions.select_voc_format.setEnabled(True)
             self.actions.select_dota_format.setEnabled(True)
             self.actions.select_mot_format.setEnabled(True)
+            self.actions.select_mask_format.setEnabled(True)
         elif self._config["save_mode"] == "yolo":
             self.actions.select_default_format.setEnabled(True)
             self.actions.select_yolo_format.setEnabled(False)
             self.actions.select_coco_format.setEnabled(True)
             self.actions.select_voc_format.setEnabled(True)
             self.actions.select_dota_format.setEnabled(True)
             self.actions.select_mot_format.setEnabled(True)
+            self.actions.select_mask_format.setEnabled(True)
         elif self._config["save_mode"] == "voc":
             self.actions.select_default_format.setEnabled(True)
             self.actions.select_yolo_format.setEnabled(True)
             self.actions.select_coco_format.setEnabled(True)
             self.actions.select_voc_format.setEnabled(False)
             self.actions.select_dota_format.setEnabled(True)
             self.actions.select_mot_format.setEnabled(True)
+            self.actions.select_mask_format.setEnabled(True)
         elif self._config["save_mode"] == "dota":
             self.actions.select_default_format.setEnabled(True)
             self.actions.select_yolo_format.setEnabled(True)
             self.actions.select_coco_format.setEnabled(True)
             self.actions.select_voc_format.setEnabled(True)
             self.actions.select_dota_format.setEnabled(False)
             self.actions.select_mot_format.setEnabled(True)
+            self.actions.select_mask_format.setEnabled(True)
         elif self._config["save_mode"] == "mot":
             self.actions.select_default_format.setEnabled(True)
             self.actions.select_yolo_format.setEnabled(True)
             self.actions.select_coco_format.setEnabled(True)
             self.actions.select_voc_format.setEnabled(True)
             self.actions.select_dota_format.setEnabled(True)
             self.actions.select_mot_format.setEnabled(False)
+            self.actions.select_mask_format.setEnabled(True)
+        elif self._config["save_mode"] == "mask":
+            self.actions.select_default_format.setEnabled(True)
+            self.actions.select_yolo_format.setEnabled(True)
+            self.actions.select_coco_format.setEnabled(True)
+            self.actions.select_voc_format.setEnabled(True)
+            self.actions.select_dota_format.setEnabled(True)
+            self.actions.select_mot_format.setEnabled(True)
+            self.actions.select_mask_format.setEnabled(False)
 
     def get_labeling_instruction(self):
         text_mode = self.tr("Mode:")
@@ -1912,7 +1972,8 @@ def format_shape(s):
                 other_data=self.other_data,
                 flags=flags,
                 output_format=self._config["save_mode"],
-                classes_file=self.classes,
+                classes_file=self.classes_file,
+                mapping_file=self.mapping_file,
             )
             self.label_file = label_file
             items = self.file_list_widget.findItems(
@@ -2145,7 +2206,8 @@ def format_shape(s):
                 other_data=self.other_data,
                 flags=flags,
                 output_format=self._config["save_mode"],
-                classes_file=self.classes,
+                classes_file=self.classes_file,
+                mapping_file=self.mapping_file,
             )
             self.label_file = label_file
             items = self.file_list_widget.findItems(

assets/mask_color_map.json

@@ -0,0 +1,24 @@
+{
+    "type": "rgb",
+    "colors": {
+        "Road": [128, 64, 128],
+        "Sidewalk": [244, 35, 232],
+        "Building": [70, 70, 70],
+        "Wall": [102, 102, 156],
+        "Fence": [190, 153, 153],
+        "Pole": [153, 153, 153],
+        "Traffic Light": [250, 170, 30],
+        "Traffic Sign": [220, 220, 0],
+        "Vegetation": [107, 142, 35],
+        "Terrain": [152, 251, 152],
+        "Sky": [70, 130, 180],
+        "Person": [220, 20, 60],
+        "Rider": [255, 0, 0],
+        "Car": [0, 0, 142],
+        "Truck": [0, 0, 70],
+        "Bus": [0, 60, 100],
+        "Train": [0, 80, 100],
+        "Motorcycle": [0, 0, 230],
+        "Bicycle": [119, 11, 32]
+    }
+}

assets/mask_grayscale_map.json

@@ -0,0 +1,24 @@
+{
+    "type": "grayscale",
+    "colors": {
+        "Road": 128,
+        "Sidewalk": 200,
+        "Building": 50,
+        "Wall": 120,
+        "Fence": 180,
+        "Pole": 150,
+        "Traffic Light": 220,
+        "Traffic Sign": 190,
+        "Vegetation": 70,
+        "Terrain": 210,
+        "Sky": 100,
+        "Person": 10,
+        "Rider": 5,
+        "Car": 60,
+        "Truck": 30,
+        "Bus": 40,
+        "Train": 50,
+        "Motorcycle": 80,
+        "Bicycle": 90
+    }
+}

docs/Q&A.md

@@ -61,16 +61,16 @@ A: X-AnyLabeling 工具目前内置了多种主流数据格式的导出，包括
 2. 此处 `--classes` 参数指定的  `*.txt` 文件是用户预定义的类别文件，每一行代表一个类别，类别编号按从上到下的顺序编排，可参考此文件[classes.txt](../assets/classes.txt)。</br>
 
 Q: **语义分割任务如何将输出的标签文件转换为 \*.png 格式输出？**</br>
-A: 针对工具本身自定义(`custom`)的格式，我们可以使用工程目录下的 `tools/polygon_mask_conversion.py` 脚本轻松转换，以下是参考的转换指令：
+A: 当前 `X-AnyLabeling` 中同样支持了 png 掩码图的导出，仅需在标注之前准备好一份自定义颜色映射表文件（具体可参考[mask_color_map.json](../assets/mask_color_map.json)和[mask_grayscale_map.json](../assets/mask_grayscale_map.json)文件，分别用于将当前分割结果映射为相应地RGB格式活灰度图格式掩码图，可按需选取），并将当前导出格式设置为 `MASK` 格式并导入事先准备好的文件即可，掩码图默认保存在与图像文件同级目录下。</br>
+当然，针对工具本身自定义(`custom`)的格式，我们可以使用工程目录下的 `tools/polygon_mask_conversion.py` 脚本 （仅支持二分类转换）轻松转换，以下是参考的转换指令：
 
 ```bash
 python tools/polygon_mask_conversion.py --img_path xxx_folder --mask_path xxx_folder --mode poly2mask
-
 # [option] 如果标签和图像不在同一目录下，请使用以下命令：
 python tools/polygon_mask_conversion.py --img_path xxx_folder --mask_path xxx_folder --json_path xxx_folder --mode poly2mask
 ```
 
-同样地，也支持将掩码图一键转换为自定义格式导入 `X-AnyLabeling` 中进行修正，输出的 `*.json` 文件默认保存至 'img_path' 目录：
+此外，也支持将掩码图一键转换为自定义格式导入 `X-AnyLabeling` 中进行修正，输出的 `*.json` 文件默认保存至 'img_path' 目录：
 
 ```bash
 python tools/polygon_mask_conversion.py --img_path xxx_folder --mask_path xxx_folder --mode mask2poly