Browse Source

Create _io.py

LaraStuStu 5 years ago
parent
commit
beff5c9082
1 changed files with 23 additions and 0 deletions
  1. 23 0
      DataAnnotation/labelme/labelme/utils/_io.py

+ 23 - 0
DataAnnotation/labelme/labelme/utils/_io.py

@@ -0,0 +1,23 @@
+import os.path as osp
+
+import numpy as np
+import PIL.Image
+
+from labelme.utils.draw import label_colormap
+
+
+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:
+        raise ValueError(
+            '[%s] Cannot save the pixel-wise class label as PNG. '
+            'Please consider using the .npy format.' % filename
+        )