Make shape_to_mask support all shape types

examples/primitives/primitives.json

@@ -0,0 +1,140 @@
+{
+  "version": "3.5.0",
+  "flags": {},
+  "shapes": [
+    {
+      "label": "rectangle",
+      "line_color": null,
+      "fill_color": null,
+      "points": [
+        [
+          32,
+          35
+        ],
+        [
+          132,
+          135
+        ]
+      ],
+      "shape_type": "rectangle"
+    },
+    {
+      "label": "circle",
+      "line_color": null,
+      "fill_color": null,
+      "points": [
+        [
+          195,
+          84
+        ],
+        [
+          225,
+          125
+        ]
+      ],
+      "shape_type": "circle"
+    },
+    {
+      "label": "rectangle",
+      "line_color": null,
+      "fill_color": null,
+      "points": [
+        [
+          391,
+          33
+        ],
+        [
+          542,
+          135
+        ]
+      ],
+      "shape_type": "rectangle"
+    },
+    {
+      "label": "polygon",
+      "line_color": null,
+      "fill_color": null,
+      "points": [
+        [
+          69,
+          318
+        ],
+        [
+          45,
+          403
+        ],
+        [
+          173,
+          406
+        ],
+        [
+          198,
+          321
+        ]
+      ],
+      "shape_type": "polygon"
+    },
+    {
+      "label": "line",
+      "line_color": null,
+      "fill_color": null,
+      "points": [
+        [
+          188,
+          178
+        ],
+        [
+          160,
+          224
+        ]
+      ],
+      "shape_type": "line"
+    },
+    {
+      "label": "point",
+      "line_color": null,
+      "fill_color": null,
+      "points": [
+        [
+          345,
+          174
+        ]
+      ],
+      "shape_type": "point"
+    },
+    {
+      "label": "line_strip",
+      "line_color": null,
+      "fill_color": null,
+      "points": [
+        [
+          441,
+          181
+        ],
+        [
+          403,
+          274
+        ],
+        [
+          545,
+          275
+        ]
+      ],
+      "shape_type": "linestrip"
+    }
+  ],
+  "lineColor": [
+    0,
+    255,
+    0,
+    128
+  ],
+  "fillColor": [
+    255,
+    0,
+    0,
+    128
+  ],
+  "imagePath": "primitives.jpg",
+  "imageData": ""
+}

labelme/utils/shape.py

@@ -15,22 +15,32 @@ def polygons_to_mask(img_shape, polygons, shape_type=None):
     return shape_to_mask(img_shape, points=polygons, shape_type=shape_type)
 
 
-def shape_to_mask(img_shape, points, shape_type=None):
+def shape_to_mask(img_shape, points, shape_type=None,
+                  line_width=10, point_size=5):
     mask = np.zeros(img_shape[:2], dtype=np.uint8)
     mask = PIL.Image.fromarray(mask)
     draw = PIL.ImageDraw.Draw(mask)
-    if shape_type == 'circle' and len(points) == 2:
-        (cx, cy), (px, py) = points
+    xy = [tuple(point) for point in points]
+    if shape_type == 'circle':
+        assert len(xy) == 2, 'Shape of shape_type=circle must have 2 points'
+        (cx, cy), (px, py) = xy
         d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
         draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
-    elif shape_type == 'rectangle' and len(points) == 2:
-        xy = [tuple(point) for point in points]
+    elif shape_type == 'rectangle':
+        assert len(xy) == 2, 'Shape of shape_type=rectangle must have 2 points'
         draw.rectangle(xy, outline=1, fill=1)
+    elif shape_type == 'line':
+        assert len(xy) == 2, 'Shape of shape_type=line must have 2 points'
+        draw.line(xy=xy, fill=1, width=line_width)
     elif shape_type == 'linestrip':
-        xy = [tuple(point) for point in points]
-        draw.line(xy=xy, fill=1, width=10)
+        draw.line(xy=xy, fill=1, width=line_width)
+    elif shape_type == 'point':
+        assert len(xy) == 1, 'Shape of shape_type=point must have 1 points'
+        cx, cy = xy[0]
+        r = point_size
+        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
     else:
-        xy = [tuple(point) for point in points]
+        assert len(xy) > 2, 'Polygon must have points more than 2'
         draw.polygon(xy=xy, outline=1, fill=1)
     mask = np.array(mask, dtype=bool)
     return mask