Merge annotations that share a common label (#180)

.github/workflows/tests-and-coverage.yml

@@ -2,11 +2,6 @@ name: Unit, functional, integration tests and code coverage
 
 on:
   push:
-    paths:
-      - "api/**"
-      - "client/**"
-      - "tests/**"
-      - "docs/**"
     branches: "**"
 
 jobs:

api/tests/unit-tests/test_schemas.py

@@ -771,7 +771,10 @@ def test_semantic_segmentation_validation():
             ],
         )
 
-    assert "appears more than" in str(e)
+    assert (
+        "semantic segmentation tasks can only have one annotation per label"
+        in str(e.value)
+    )
 
     with pytest.raises(ValueError) as e:
         schemas.GroundTruth(
@@ -794,7 +797,10 @@ def test_semantic_segmentation_validation():
             ],
         )
 
-    assert "appears more than" in str(e)
+    assert (
+        "semantic segmentation tasks can only have one annotation per label"
+        in str(e.value)
+    )
 
     # this is valid
     schemas.Prediction(
@@ -840,10 +846,10 @@ def test_semantic_segmentation_validation():
             ],
         )
 
-    assert "appears more than" in str(e)
-
-
-# velour_api.schemas.metadata
+    assert (
+        "semantic segmentation tasks can only have one annotation per label"
+        in str(e.value)
+    )
 
 
 # @TODO

api/velour_api/schemas/core.py

@@ -132,15 +132,17 @@ def _check_semantic_segmentations_single_label(
 ) -> None:
     # check that a label on appears once in the annotations for semenatic segmentations
     labels = []
-    for annotation in annotations:
+    indices = dict()
+    for index, annotation in enumerate(annotations):
         if annotation.task_type == enums.TaskType.SEMANTIC_SEGMENTATION:
             for label in annotation.labels:
                 if label in labels:
                     raise ValueError(
-                        f"Label {label} appears more than once but semantic segmentation "
-                        "tasks can only have at most one annotation per label."
+                        f"Label {label} appears in both annotation {index} and {indices[label]}, but semantic segmentation "
+                        "tasks can only have one annotation per label."
                     )
                 labels.append(label)
+                indices[label] = index
 
 
 class GroundTruth(BaseModel):

client/unit-tests/test_coco.py

@@ -1,6 +1,93 @@
 import PIL.Image
 
-from velour.integrations.coco import coco_rle_to_mask
+from velour.integrations.coco import coco_rle_to_mask, _merge_annotations
+from velour.enums import TaskType
+from velour.schemas import Label
+
+
+def test__merge_annotations():
+    """Check that we get the correct annotation set after merging semantic segmentions"""
+
+    initial_annotations = [
+        dict(
+            task_type=TaskType.SEMANTIC_SEGMENTATION,
+            labels=set([Label(key="k1", value="v1"), Label(key="k2", value="v2")]),
+            mask=[[True, False, False, False], [True, False, False, False]],
+        ),
+        dict(
+            task_type=TaskType.SEMANTIC_SEGMENTATION,
+            labels=set([Label(key="k1", value="v1"), Label(key="k3", value="v3")]),
+            mask=[[False, False, True, False], [False, False, True, False]],
+        ),
+        dict(
+            task_type=TaskType.SEMANTIC_SEGMENTATION,
+            labels=set(
+                [
+                    Label(key="k1", value="v1"),
+                    Label(key="k2", value="v2"),
+                    Label(key="k4", value="v4"),
+                ]
+            ),
+            mask=[[False, False, False, True], [False, False, False, True]],
+        ),
+        dict(
+            task_type=TaskType.INSTANCE_SEGMENTATION,
+            labels=set([Label(key="k1", value="v1"), Label(key="k3", value="v3")]),
+            mask=[[False, True, False, False], [False, True, False, False]],
+        ),
+    ]
+
+    expected = [
+        dict(
+            task_type=TaskType.SEMANTIC_SEGMENTATION,
+            labels=set([Label(key="k3", value="v3")]),
+            mask=[[False, False, True, False], [False, False, True, False]],
+        ),
+        dict(
+            task_type=TaskType.SEMANTIC_SEGMENTATION,
+            labels=set([Label(key="k4", value="v4")]),
+            mask=[[False, False, False, True], [False, False, False, True]],
+        ),
+        dict(
+            task_type=TaskType.INSTANCE_SEGMENTATION,
+            labels=set(
+                [
+                    Label(key="k1", value="v1"),
+                    Label(key="k3", value="v3"),
+                ]
+            ),
+            mask=[[False, True, False, False], [False, True, False, False]],
+        ),
+        dict(
+            task_type=TaskType.SEMANTIC_SEGMENTATION,
+            labels=set([Label(key="k1", value="v1")]),
+            mask=[[True, False, True, True], [True, False, True, True]],
+        ),
+        dict(
+            task_type=TaskType.SEMANTIC_SEGMENTATION,
+            labels=set([Label(key="k2", value="v2")]),
+            mask=[[True, False, False, True], [True, False, False, True]],
+        ),
+    ]
+
+    label_map = {
+        Label(key="k1", value="v1"): [0, 1, 2],
+        Label(key="k2", value="v2"): [0, 2],
+        Label(key="k3", value="v3"): [1],
+        Label(key="k4", value="v4"): [2],
+    }
+
+    merged_annotations = _merge_annotations(
+        annotation_list=initial_annotations, label_map=label_map
+    )
+
+    for i, v in enumerate(merged_annotations):
+        assert (
+            merged_annotations[i]["labels"] == expected[i]["labels"]
+        ), "Labels didn't merge as expected"
+        assert set(map(tuple, merged_annotations[i]["mask"])) == set(
+            map(tuple, expected[i]["mask"])
+        ), "Masks didn't merge as expected"
 
 
 def test_coco_rle_to_mask():

client/velour/integrations/coco.py

@@ -1,6 +1,8 @@
 import json
+from collections import defaultdict
 from pathlib import Path, PosixPath
 from typing import Any, Dict, List, Union
+from copy import deepcopy
 
 import numpy as np
 import PIL.Image
@@ -54,6 +56,128 @@ def coco_rle_to_mask(coco_rle_seg_dict: Dict[str, Any]) -> np.ndarray:
     return res
 
 
+def _get_task_type(isthing: bool) -> enums.TaskType:
+    """Get the correct TaskType for a given label"""
+    return (
+        enums.TaskType.INSTANCE_SEGMENTATION
+        if isthing
+        else enums.TaskType.SEMANTIC_SEGMENTATION
+    )
+
+
+def _is_semantic_task_type(task_type: enums.TaskType) -> bool:
+    """Check if a label is a semantic segmentation"""
+    return True if task_type == enums.TaskType.SEMANTIC_SEGMENTATION else False
+
+
+def _merge_annotations(annotation_list: list, label_map: dict):
+    """Aggregate masks of annotations that share a common label"""
+
+    # deepcopy since we use .remove()
+    annotation_list = deepcopy(annotation_list)
+
+    for label, indices in label_map.items():
+        if len(indices) > 1:
+            joined_mask = annotation_list[indices[0]]["mask"]
+            task_type = annotation_list[indices[0]]["task_type"]
+
+            # remove the label from the parent node
+            annotation_list[indices[0]]["labels"].remove(label)
+
+            for child_index in indices[1:]:
+                if indices[0] != child_index:
+                    child = annotation_list[child_index]
+                    joined_mask = np.logical_or(joined_mask, child["mask"])
+
+                    # remove the label from the child node
+                    annotation_list[child_index]["labels"].remove(label)
+
+            annotation_list.append(
+                dict(task_type=task_type, labels=set([label]), mask=joined_mask)
+            )
+
+    # delete any annotations without labels remaining (i.e., their mask is now incorporated into grouped annotations)
+    annotation_list = [
+        annotation for annotation in annotation_list if len(annotation["labels"]) > 0
+    ]
+
+    return annotation_list
+
+
+def _get_segs_groundtruth_for_single_image(
+    ann_dict: dict,
+    masks_path: str,
+    image_id_to_height: dict,
+    image_id_to_width: dict,
+    category_id_to_category: dict,
+) -> List[GroundTruth]:
+    mask = np.array(PIL.Image.open(masks_path / ann_dict["file_name"])).astype(int)
+    # convert the colors in the mask to ids
+    mask_ids = mask[:, :, 0] + 256 * mask[:, :, 1] + (256**2) * mask[:, :, 2]
+
+    # create datum
+    image_id = ann_dict["image_id"]
+    img = ImageMetadata(
+        uid=str(image_id),
+        height=image_id_to_height[image_id],
+        width=image_id_to_width[image_id],
+    ).to_datum()
+
+    # create initial list of annotations
+    annotation_list = []
+
+    semantic_labels = defaultdict(list)
+
+    for index, segment in enumerate(ann_dict["segments_info"]):
+        mask = mask_ids == segment["id"]
+        task_type = _get_task_type(
+            category_id_to_category[segment["category_id"]]["isthing"]
+        )
+        is_semantic = _is_semantic_task_type(task_type=task_type)
+
+        labels = set()
+
+        for k in ["supercategory", "name"]:
+            category_desc = str(category_id_to_category[segment["category_id"]][k])
+
+            label = Label(
+                key=k,
+                value=category_desc,
+            )
+
+            # identify the location of all semantic segmentation labels
+            if is_semantic:
+                semantic_labels[label].append(index)
+
+            labels.add(label)
+
+        annotation_list.append(dict(task_type=task_type, labels=labels, mask=mask))
+        is_crowd_label = Label(key="iscrowd", value=str(segment["iscrowd"]))
+
+    if is_semantic:
+        semantic_labels[is_crowd_label].append(len(ann_dict["segments_info"]) - 1)
+    else:
+        annotation_list.append(is_crowd_label)
+
+    # combine semantic segmentation masks by label
+    final_annotation_list = _merge_annotations(
+        annotation_list=annotation_list, label_map=semantic_labels
+    )
+
+    # create groundtruth
+    return GroundTruth(
+        datum=img,
+        annotations=[
+            Annotation(
+                task_type=annotation["task_type"],
+                labels=list(annotation["labels"]),
+                raster=Raster.from_numpy(annotation["mask"]),
+            )
+            for annotation in final_annotation_list
+        ],
+    )
+
+
 def upload_coco_panoptic(
     dataset: Dataset,
     annotations: Union[str, PosixPath, dict],
@@ -64,65 +188,20 @@ def upload_coco_panoptic(
         with open(annotations) as f:
             annotations = json.load(f)
 
-    category_id_to_category = {
-        cat["id"]: cat for cat in annotations["categories"]
-    }
+    category_id_to_category = {cat["id"]: cat for cat in annotations["categories"]}
 
     image_id_to_height, image_id_to_width, image_id_to_coco_url = {}, {}, {}
     for image in annotations["images"]:
         image_id_to_height[image["id"]] = image["height"]
         image_id_to_width[image["id"]] = image["width"]
         image_id_to_coco_url[image["id"]] = image["coco_url"]
 
-    def _get_segs_groundtruth_for_single_image(
-        ann_dict: dict,
-    ) -> List[GroundTruth]:
-        mask = np.array(
-            PIL.Image.open(masks_path / ann_dict["file_name"])
-        ).astype(int)
-        # convert the colors in the mask to ids
-        mask_ids = (
-            mask[:, :, 0] + 256 * mask[:, :, 1] + (256**2) * mask[:, :, 2]
-        )
-
-        # create datum
-        image_id = ann_dict["image_id"]
-        img = ImageMetadata(
-            uid=str(image_id),
-            height=image_id_to_height[image_id],
-            width=image_id_to_width[image_id],
-        ).to_datum()
-
-        # create groundtruth
-        return GroundTruth(
-            datum=img,
-            annotations=[
-                Annotation(
-                    task_type=(
-                        enums.TaskType.INSTANCE_SEGMENTATION
-                        if category_id_to_category[segment["category_id"]][
-                            "isthing"
-                        ]
-                        else enums.TaskType.SEMANTIC_SEGMENTATION
-                    ),
-                    labels=[
-                        Label(
-                            key=k,
-                            value=str(
-                                category_id_to_category[
-                                    segment["category_id"]
-                                ][k]
-                            ),
-                        )
-                        for k in ["supercategory", "name"]
-                    ]
-                    + [Label(key="iscrowd", value=str(segment["iscrowd"]))],
-                    raster=Raster.from_numpy(mask_ids == segment["id"]),
-                )
-                for segment in ann_dict["segments_info"]
-            ],
-        )
-
     for ann in tqdm(annotations["annotations"]):
-        gt = _get_segs_groundtruth_for_single_image(ann)
+        gt = _get_segs_groundtruth_for_single_image(
+            ann_dict=ann,
+            masks_path=masks_path,
+            image_id_to_height=image_id_to_height,
+            image_id_to_width=image_id_to_width,
+            category_id_to_category=category_id_to_category,
+        )
         dataset.add_groundtruth(gt)

integration_tests/test_client.py

@@ -1130,7 +1130,10 @@ def test_create_gt_segs_as_polys_or_masks(
 
         dataset.add_groundtruth(gts)
 
-    assert "appears more than once" in str(exc_info.value)
+    assert (
+        "semantic segmentation tasks can only have one annotation per label"
+        in str(exc_info.value)
+    )
 
     # fine with instance segmentation though
     gts = GroundTruth(