Merge pull request #3 from ROBOTIS-move/feature-masking-error

Modified label masking function

labelme/cli/draw_segment_label.py

@@ -126,9 +126,12 @@ def multi_convert_json_to_mask(self, json_file):
                 copy.deepcopy(self.class_rgb),
                 self.segmentation_class)
 
-            utils.lblsave(
+            utils.lblsave_old(
                 os.path.join(self.masked_image_dir, file_name),
-                label)
+                label,
+                names,
+                copy.deepcopy(self.class_rgb),
+                self.segmentation_class)
 
             if self.merge_save:
                 PIL.Image.fromarray(image).save(os.path.join(

labelme/utils/__init__.py

@@ -1,6 +1,7 @@
 # flake8: noqa
 
 from ._io import lblsave
+from ._io import lblsave_old
 
 from .image import apply_exif_orientation
 from .image import img_arr_to_b64

labelme/utils/_io.py

@@ -21,3 +21,27 @@ def lblsave(filename, lbl):
             "[%s] Cannot save the pixel-wise class label as PNG. "
             "Please consider using the .npy format." % filename
         )
+
+def lblsave_old(filename, lbl, names, class_rgb, segmentation_class):
+    here = osp.dirname(osp.abspath(__file__))
+
+    if osp.splitext(filename)[1] != '.png':
+        filename += '.png'
+    if lbl.min() >= -1 and lbl.max() < 255:
+        lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
+        colormap = np.ones((255, 3), dtype=float)
+        colormap[0] = [0, 0, 0]
+        for num_classes, classes in enumerate(names):
+            color_val = class_rgb[segmentation_class.index(classes)]
+
+            for i in range(3):
+                color_val[i] = 256 - color_val[i]
+                colormap[num_classes+1] = color_val
+        lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
+        lbl_pil.save(filename)
+
+    else:
+        raise ValueError(
+            '[%s] Cannot save the pixel-wise class label as PNG. '
+            'Please consider using the .npy format.' % filename
+        )