Add function to create multiple annotations from the same image. Closes #29

examples/multi-annotations/multi_annotation_infos.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+import json
+
+from PIL import Image
+import numpy as np
+
+from pycococreatortools import pycococreatortools
+
+SAMPLE_IMAGE_PATH = Path(__file__).parent / "sample.png"
+assert SAMPLE_IMAGE_PATH.exists()
+
+raw_image_bitmask = Image.open(SAMPLE_IMAGE_PATH)
+image_bitmask = np.asarray(raw_image_bitmask.convert("1")).astype(np.uint8)
+
+annotation_infos = pycococreatortools.create_annotation_infos(
+    1, 1, {"id": 1}, image_bitmask, raw_image_bitmask.size, tolerance=2, connectivity=1
+)
+
+print(f"Found {len(annotation_infos)} annotations in the image:")
+print(json.dumps(annotation_infos, indent=4))

pycococreatortools/pycococreatortools.py

@@ -53,13 +53,13 @@ def binary_mask_to_polygon(binary_mask, tolerance=0):
             continue
         contour = np.flip(contour, axis=1)
         segmentation = contour.ravel().tolist()
-        # after padding and subtracting 1 we may get -0.5 points in our segmentation 
+        # after padding and subtracting 1 we may get -0.5 points in our segmentation
         segmentation = [0 if i < 0 else i for i in segmentation]
         polygons.append(segmentation)
 
     return polygons
 
-def create_image_info(image_id, file_name, image_size, 
+def create_image_info(image_id, file_name, image_size,
                       date_captured=datetime.datetime.utcnow().isoformat(' '),
                       license_id=1, coco_url="", flickr_url=""):
 
@@ -76,7 +76,7 @@ def create_image_info(image_id, file_name, image_size,
 
     return image_info
 
-def create_annotation_info(annotation_id, image_id, category_info, binary_mask, 
+def create_annotation_info(annotation_id, image_id, category_info, binary_mask,
                            image_size=None, tolerance=2, bounding_box=None):
 
     if image_size is not None:
@@ -91,10 +91,10 @@ def create_annotation_info(annotation_id, image_id, category_info, binary_mask,
     if bounding_box is None:
         bounding_box = mask.toBbox(binary_mask_encoded)
 
-    if category_info["is_crowd"]:
+    if "is_crowd" in category_info and category_info["is_crowd"]:
         is_crowd = 1
         segmentation = binary_mask_to_rle(binary_mask)
-    else :
+    else:
         is_crowd = 0
         segmentation = binary_mask_to_polygon(binary_mask, tolerance)
         if not segmentation:
@@ -110,6 +110,67 @@ def create_annotation_info(annotation_id, image_id, category_info, binary_mask,
         "segmentation": segmentation,
         "width": binary_mask.shape[1],
         "height": binary_mask.shape[0],
-    } 
+    }
 
     return annotation_info
+
+
+def create_annotation_infos(
+    start_annotation_id, image_id, category_info, binary_mask,
+    image_size=None, tolerance=2, create_labels=True, connectivity=None
+):
+    """Create multiple annotation infos for each connected component in the given binary mask
+
+    The same category is used for each annotation.
+    The annotation ids start from `start_annotation_id` and
+    are incremented for each annotation.
+
+    The `create_labels` argument can be used to control whether or
+    not labels should be automatically created or not.
+    See:
+    https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label
+
+    The `connectivity` argument can be used to specify the connectivity
+    for the labels according to:
+    https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label
+
+    Limitations:
+        * Crowds are not supported
+    """
+    if "is_crowd" in category_info and category_info["is_crowd"]:
+        raise NotImplementedError("Creating multiple crowd annotations from a single binary mask is not supported")
+
+    if image_size is not None:
+        binary_mask = resize_binary_mask(binary_mask, image_size)
+
+    # label connected components in binary mask image
+    if create_labels:
+        label_image = measure.label(binary_mask, connectivity=connectivity)
+    else:
+        label_image = binary_mask
+
+    region_props = measure.regionprops(label_image)
+
+    # create a binary mask image per region property
+    binary_masks = []
+    for region_label, region_bbox in ((r.label, r.bbox) for r in region_props):
+        region_binary_mask = np.zeros_like(binary_mask, dtype=np.bool)
+
+        # copy the region into the region binary mask
+        bbox_slice = (
+            slice(region_bbox[0], region_bbox[2]),
+            slice(region_bbox[1], region_bbox[3]),
+        )
+        region_binary_mask[bbox_slice] = label_image[bbox_slice] == region_label
+        binary_masks.append(region_binary_mask)
+
+    # create annotations for each binary mask
+    annotation_infos = []
+    for annotation_id, region_binary_mask in enumerate(binary_masks, start=start_annotation_id):
+        annotation_info = create_annotation_info(
+            annotation_id, image_id, category_info, region_binary_mask,
+            image_size=image_size, tolerance=tolerance
+        )
+        annotation_infos.append(annotation_info)
+
+    return annotation_infos