Browse Source

Update labelme2voc.py and labelme2coco.py accordingly

Kentaro Wada 5 years ago
parent
commit
612b40df6f

+ 16 - 8
examples/instance_segmentation/labelme2coco.py

@@ -8,6 +8,7 @@ import json
 import os
 import os.path as osp
 import sys
+import uuid
 
 import numpy as np
 import PIL.Image
@@ -111,21 +112,28 @@ def main():
         for shape in label_data['shapes']:
             points = shape['points']
             label = shape['label']
-            shape_type = shape.get('shape_type', None)
+            group_id = shape.get('group_id')
+            shape_type = shape.get('shape_type')
             mask = labelme.utils.shape_to_mask(
                 img.shape[:2], points, shape_type
             )
 
-            if label in masks:
-                masks[label] = masks[label] | mask
+            if group_id is None:
+                group_id = uuid.uuid1()
+
+            instance = (label, group_id)
+
+            if instance in masks:
+                masks[instance] = masks[instance] | mask
             else:
-                masks[label] = mask
+                masks[instance] = mask
 
             points = np.asarray(points).flatten().tolist()
-            segmentations[label].append(points)
+            segmentations[instance].append(points)
+        segmentations = dict(segmentations)
 
-        for label, mask in masks.items():
-            cls_name = label.split('-')[0]
+        for instance, mask in masks.items():
+            cls_name, group_id = instance
             if cls_name not in class_name_to_id:
                 continue
             cls_id = class_name_to_id[cls_name]
@@ -139,7 +147,7 @@ def main():
                 id=len(data['annotations']),
                 image_id=image_id,
                 category_id=cls_id,
-                segmentation=segmentations[label],
+                segmentation=segmentations[instance],
                 area=area,
                 bbox=bbox,
                 iscrowd=0,

+ 0 - 1
examples/instance_segmentation/labelme2voc.py

@@ -103,7 +103,6 @@ def main():
                 img_shape=img.shape,
                 shapes=data['shapes'],
                 label_name_to_value=class_name_to_id,
-                type='instance',
             )
             ins[cls == -1] = 0  # ignore it.
 

+ 1 - 1
examples/semantic_segmentation/labelme2voc.py

@@ -81,7 +81,7 @@ def main():
             f.write(label_file.imageData)
         img = labelme.utils.img_data_to_arr(label_file.imageData)
 
-        lbl = labelme.utils.shapes_to_label(
+        lbl, _ = labelme.utils.shapes_to_label(
             img_shape=img.shape,
             shapes=label_file.shapes,
             label_name_to_value=class_name_to_id,

+ 3 - 1
labelme/cli/draw_json.py

@@ -41,7 +41,9 @@ def main():
         else:
             label_value = len(label_name_to_value)
             label_name_to_value[label_name] = label_value
-    lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
+    lbl, _ = utils.shapes_to_label(
+        img.shape, data['shapes'], label_name_to_value
+    )
 
     label_names = [None] * (max(label_name_to_value.values()) + 1)
     for name, value in label_name_to_value.items():

+ 3 - 1
labelme/cli/json_to_dataset.py

@@ -50,7 +50,9 @@ def main():
         else:
             label_value = len(label_name_to_value)
             label_name_to_value[label_name] = label_value
-    lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
+    lbl, _ = utils.shapes_to_label(
+        img.shape, data['shapes'], label_name_to_value
+    )
 
     label_names = [None] * (max(label_name_to_value.values()) + 1)
     for name, value in label_name_to_value.items():

+ 18 - 19
labelme/utils/shape.py

@@ -1,4 +1,5 @@
 import math
+import uuid
 
 import numpy as np
 import PIL.Image
@@ -46,33 +47,31 @@ def shape_to_mask(img_shape, points, shape_type=None,
     return mask
 
 
-def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):
-    assert type in ['class', 'instance']
-
+def shapes_to_label(img_shape, shapes, label_name_to_value):
     cls = np.zeros(img_shape[:2], dtype=np.int32)
-    if type == 'instance':
-        ins = np.zeros(img_shape[:2], dtype=np.int32)
-        instance_names = ['_background_']
+    ins = np.zeros_like(cls)
+    instances = []
     for shape in shapes:
         points = shape['points']
         label = shape['label']
+        group_id = shape.get('group_id')
+        if group_id is None:
+            group_id = uuid.uuid1()
         shape_type = shape.get('shape_type', None)
-        if type == 'class':
-            cls_name = label
-        elif type == 'instance':
-            cls_name = label.split('-')[0]
-            if label not in instance_names:
-                instance_names.append(label)
-            ins_id = instance_names.index(label)
+
+        cls_name = label
+        instance = (cls_name, group_id)
+
+        if instance not in instances:
+            instances.append(instance)
+        ins_id = instances.index(instance) + 1
         cls_id = label_name_to_value[cls_name]
+
         mask = shape_to_mask(img_shape[:2], points, shape_type)
         cls[mask] = cls_id
-        if type == 'instance':
-            ins[mask] = ins_id
+        ins[mask] = ins_id
 
-    if type == 'instance':
-        return cls, ins
-    return cls
+    return cls, ins
 
 
 def labelme_shapes_to_label(img_shape, shapes):
@@ -88,7 +87,7 @@ def labelme_shapes_to_label(img_shape, shapes):
             label_value = len(label_name_to_value)
             label_name_to_value[label_name] = label_value
 
-    lbl = shapes_to_label(img_shape, shapes, label_name_to_value)
+    lbl, _ = shapes_to_label(img_shape, shapes, label_name_to_value)
     return lbl, label_name_to_value
 
 

+ 1 - 1
tests/labelme_tests/utils_tests/test_shape.py

@@ -10,7 +10,7 @@ def test_shapes_to_label():
         label_name = shape['label']
         label_value = len(label_name_to_value)
         label_name_to_value[label_name] = label_value
-    cls = shape_module.shapes_to_label(
+    cls, _ = shape_module.shapes_to_label(
         img.shape, data['shapes'], label_name_to_value)
     assert cls.shape == img.shape[:2]
 

+ 1 - 1
tests/labelme_tests/utils_tests/util.py

@@ -31,7 +31,7 @@ def get_img_and_lbl():
     for label_name, label_value in label_name_to_value.items():
         label_names[label_value] = label_name
 
-    lbl = shape_module.shapes_to_label(
+    lbl, _ = shape_module.shapes_to_label(
         img.shape, data['shapes'], label_name_to_value
     )
     return img, lbl, label_names