Add labelme.utils.lblsave

examples/instance_segmentation/labelme2voc.py

@@ -90,43 +90,18 @@ def main():
             ins[cls == -1] = 0  # ignore it.
 
             # class label
-
-            # Assume class label ranses [-1, 254] for int32,
-            # and [0, 255] for uint8 as VOC.
-            if cls.min() >= -1 and cls.max() < 255:
-                cls_pil = PIL.Image.fromarray(cls.astype(np.uint8), mode='P')
-                cls_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
-                cls_pil.save(out_clsp_file)
-            else:
-                labelme.logger.warn(
-                    '[%s] Cannot save the pixel-wise class label as PNG, '
-                    'so please use the npy file.' % label_file
-                )
-
+            labelme.utils.lblsave(out_clsp_file, cls)
             np.save(out_cls_file, cls)
             clsv = labelme.utils.draw_label(
                 cls, img, class_names, colormap=colormap)
             PIL.Image.fromarray(clsv).save(out_clsv_file)
 
             # instance label
-
-            # Assume instance label ranses [-1, 254] for int32,
-            # and [0, 255] for uint8 as VOC.
-            if ins.min() >= -1 and ins.max() < 255:
-                ins_pil = PIL.Image.fromarray(ins.astype(np.uint8), mode='P')
-                ins_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
-                ins_pil.save(out_insp_file)
-            else:
-                labelme.logger.warn(
-                    '[%s] Cannot save the pixel-wise instance label as PNG, '
-                    'so please use the npy file.' % label_file
-                )
-
+            labelme.utils.lblsave(out_insp_file, ins)
             np.save(out_ins_file, ins)
             instance_ids = np.unique(ins)
             instance_names = [str(i) for i in range(max(instance_ids) + 1)]
-            insv = labelme.utils.draw_label(
-                ins, img, instance_names)
+            insv = labelme.utils.draw_label(ins, img, instance_names)
             PIL.Image.fromarray(insv).save(out_insv_file)
 
 

examples/semantic_segmentation/labelme2voc.py

@@ -77,18 +77,7 @@ def main():
                 shapes=data['shapes'],
                 label_name_to_value=class_name_to_id,
             )
-
-            # Assume label ranses [-1, 254] for int32,
-            # and [0, 255] for uint8 as VOC.
-            if lbl.min() >= -1 and lbl.max() < 255:
-                lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
-                lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
-                lbl_pil.save(out_png_file)
-            else:
-                labelme.logger.warn(
-                    '[%s] Cannot save the pixel-wise class label as PNG, '
-                    'so please use the npy file.' % label_file
-                )
+            labelme.utils.lblsave(out_png_file, lbl)
 
             np.save(out_lbl_file, lbl)
 

labelme/cli/json_to_dataset.py

@@ -66,7 +66,7 @@ def main():
     lbl_viz = utils.draw_label(lbl, img, captions)
 
     PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
-    PIL.Image.fromarray(lbl).save(osp.join(out_dir, 'label.png'))
+    utils.lblsave(osp.join(out_dir, 'label.png'), lbl)
     PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))
 
     with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:

labelme/utils.py

@@ -1,11 +1,14 @@
 import base64
 import io
+import os.path as osp
 import warnings
 
 import numpy as np
 import PIL.Image
 import PIL.ImageDraw
 
+from labelme import logger
+
 
 def label_colormap(N=256):
 
@@ -177,3 +180,20 @@ def labelme_shapes_to_label(img_shape, shapes):
 
     lbl = shapes_to_label(img_shape, shapes, label_name_to_value)
     return lbl, label_name_to_value
+
+
+def lblsave(filename, lbl):
+    if osp.splitext(filename)[1] != '.png':
+        filename += '.png'
+    # Assume label ranses [-1, 254] for int32,
+    # and [0, 255] for uint8 as VOC.
+    if lbl.min() >= -1 and lbl.max() < 255:
+        lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
+        colormap = label_colormap(255)
+        lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
+        lbl_pil.save(filename)
+    else:
+        logger.warn(
+            '[%s] Cannot save the pixel-wise class label as PNG, '
+            'so please use the npy file.' % filename
+        )