From afb11bb64da58fa4cdc5464acbaea51e3bb8af44 Mon Sep 17 00:00:00 2001
From: eungiCho <ceg@robotis.com>
Date: Thu, 12 May 2022 16:52:00 +0900
Subject: [PATCH] Modified label masking function

Signed-off-by: eungiCho <ceg@robotis.com>
---
 labelme/cli/draw_segment_label.py |  7 +++++--
 labelme/utils/__init__.py         |  1 +
 labelme/utils/_io.py              | 24 ++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/labelme/cli/draw_segment_label.py b/labelme/cli/draw_segment_label.py
index e67eb5dfe..2679f6713 100644
--- a/labelme/cli/draw_segment_label.py
+++ b/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(
diff --git a/labelme/utils/__init__.py b/labelme/utils/__init__.py
index a0ec16af4..b6bd866f6 100644
--- a/labelme/utils/__init__.py
+++ b/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
diff --git a/labelme/utils/_io.py b/labelme/utils/_io.py
index cf869e1e7..3cf7ca721 100644
--- a/labelme/utils/_io.py
+++ b/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
+        )