[AI-1287] fixed the siutation when a user pass only a image (#696)

* fixed the siutation when a user pass only a image

* changed the check a bit

* black

* fixed exporter test

* black formatting

* added fix for when there are no annoations (empty dict)

* added tests and extra logic for when annotations is empty, one option is to throw an error and not let the user include images without annotations

* black

* refactored for readability, area is it's own check now

* black

Author: ChristofferEdlund

darwin/torch/transforms.py

@@ -143,7 +143,9 @@ class ColorJitter(transforms.ColorJitter):
     def __call__(
         self, image: PILImage.Image, target: Optional[TargetType] = None
     ) -> Union[PILImage.Image, Tuple[PILImage.Image, TargetType]]:
-        transform = self.get_params(self.brightness, self.contrast, self.saturation, self.hue)
+        transform = self.get_params(
+            self.brightness, self.contrast, self.saturation, self.hue
+        )
         image = transform(image)
         if target is None:
             return image
@@ -198,7 +200,9 @@ class ConvertPolygonsToInstanceMasks(object):
     Converts given polygon to an ``InstanceMask``.
     """
 
-    def __call__(self, image: PILImage.Image, target: TargetType) -> Tuple[PILImage.Image, TargetType]:
+    def __call__(
+        self, image: PILImage.Image, target: TargetType
+    ) -> Tuple[PILImage.Image, TargetType]:
         w, h = image.size
 
         image_id = target["image_id"]
@@ -255,7 +259,9 @@ class ConvertPolygonsToSemanticMask(object):
     Converts given polygon to an ``SemanticMask``.
     """
 
-    def __call__(self, image: PILImage.Image, target: TargetType) -> Tuple[PILImage.Image, TargetType]:
+    def __call__(
+        self, image: PILImage.Image, target: TargetType
+    ) -> Tuple[PILImage.Image, TargetType]:
         w, h = image.size
         image_id = target["image_id"]
         image_id = torch.tensor([image_id])
@@ -282,7 +288,9 @@ class ConvertPolygonToMask(object):
     Converts given polygon to a ``Mask``.
     """
 
-    def __call__(self, image: PILImage.Image, annotation: Dict[str, Any]) -> Tuple[PILImage.Image, PILImage.Image]:
+    def __call__(
+        self, image: PILImage.Image, annotation: Dict[str, Any]
+    ) -> Tuple[PILImage.Image, PILImage.Image]:
         w, h = image.size
         segmentations = [obj["segmentation"] for obj in annotation]
         cats = [obj["category_id"] for obj in annotation]
@@ -327,10 +335,18 @@ def from_dict(cls, alb_dict: dict) -> "AlbumentationsTransform":
     def __call__(self, image, annotation: dict = None) -> tuple:
         np_image = np.array(image)
         if annotation is None:
-            annotation = {}
-        albu_data = self._pre_process(np_image, annotation)
+            annotation_dict = {}
+        else:
+            annotation_dict = annotation
+
+        albu_data = self._pre_process(np_image, annotation_dict)
         transformed_data = self.transform(**albu_data)
-        image, transformed_annotation = self._post_process(transformed_data, annotation)
+        image, transformed_annotation = self._post_process(
+            transformed_data, annotation_dict
+        )
+
+        if annotation is None:
+            return image
 
         return image, transformed_annotation
 
@@ -349,7 +365,10 @@ def _pre_process(self, image: np.ndarray, annotation: dict) -> dict:
             albumentation_dict["labels"] = labels.tolist()
 
         masks = annotation.get("masks")
-        if masks is not None:
+        if (
+            masks is not None and masks.numel() > 0
+        ):  # using numel() to check if tensor is non-empty
+            print("WE GOT MASKS")
             albumentation_dict["masks"] = masks.numpy()
 
         return albumentation_dict
@@ -364,8 +383,6 @@ def _post_process(self, albumentation_output: dict, annotation: dict) -> tuple:
         bboxes = albumentation_output.get("bboxes")
         if bboxes is not None:
             output_annotation["boxes"] = torch.tensor(bboxes)
-            if "area" in annotation and "masks" not in albumentation_output:
-                output_annotation["area"] = output_annotation["boxes"][:, 2] * output_annotation["boxes"][:, 3]
 
         labels = albumentation_output.get("labels")
         if labels is not None:
@@ -377,8 +394,20 @@ def _post_process(self, albumentation_output: dict, annotation: dict) -> tuple:
                 output_annotation["masks"] = torch.tensor(np.array(masks))
             else:
                 output_annotation["masks"] = torch.stack(masks)
-            if "area" in annotation:
-                output_annotation["area"] = torch.sum(output_annotation["masks"], dim=[1, 2])
+        elif "masks" in annotation:
+            output_annotation["masks"] = torch.tensor([])
+
+        if "area" in annotation:
+            if "masks" in output_annotation and output_annotation["masks"].numel() > 0:
+                output_annotation["area"] = torch.sum(
+                    output_annotation["masks"], dim=[1, 2]
+                )
+            elif "boxes" in output_annotation and len(output_annotation["boxes"]) > 0:
+                output_annotation["area"] = (
+                    output_annotation["boxes"][:, 2] * output_annotation["boxes"][:, 3]
+                )
+            else:
+                output_annotation["area"] = torch.tensor([])
 
         # Copy other metadata from original annotation
         for key, value in annotation.items():

tests/darwin/exporter/formats/export_darwin_1_0_test.py

@@ -1,13 +1,17 @@
 from pathlib import Path
 
-
 import darwin.datatypes as dt
 from darwin.exporter.formats.darwin_1_0 import _build_json
 
 
 class TestBuildJson:
     def test_empty_annotation_file(self):
-        annotation_file = dt.AnnotationFile(path=Path("test.json"), filename="test.json", annotation_classes=[], annotations=[])
+        annotation_file = dt.AnnotationFile(
+            path=Path("test.json"),
+            filename="test.json",
+            annotation_classes=[],
+            annotations=[],
+        )
 
         assert _build_json(annotation_file) == {
             "image": {
@@ -58,7 +62,9 @@ def test_complete_annotation_file(self):
         ]
 
         annotation_class = dt.AnnotationClass(name="test", annotation_type="polygon")
-        annotation = dt.Annotation(annotation_class=annotation_class, data={"path": polygon_path}, subs=[])
+        annotation = dt.Annotation(
+            annotation_class=annotation_class, data={"path": polygon_path}, subs=[]
+        )
 
         annotation_file = dt.AnnotationFile(
             path=Path("test.json"),
@@ -82,14 +88,24 @@ def test_complete_annotation_file(self):
                 "path": None,
                 "workview_url": None,
             },
-            "annotations": [{"polygon": {"path": polygon_path}, "name": "test", "slot_names": []}],
+            "annotations": [
+                {"polygon": {"path": polygon_path}, "name": "test", "slot_names": []}
+            ],
             "dataset": "None",
         }
 
     def test_complex_polygon(self):
         polygon_path = [
-            [{"x": 230.06, "y": 174.04}, {"x": 226.39, "y": 170.36}, {"x": 224.61, "y": 166.81}],
-            [{"x": 238.98, "y": 171.69}, {"x": 236.97, "y": 174.04}, {"x": 238.67, "y": 174.04}],
+            [
+                {"x": 230.06, "y": 174.04},
+                {"x": 226.39, "y": 170.36},
+                {"x": 224.61, "y": 166.81},
+            ],
+            [
+                {"x": 238.98, "y": 171.69},
+                {"x": 236.97, "y": 174.04},
+                {"x": 238.67, "y": 174.04},
+            ],
             [
                 {"x": 251.75, "y": 169.77},
                 {"x": 251.75, "y": 154.34},
@@ -98,8 +114,12 @@ def test_complex_polygon(self):
             ],
         ]
 
-        annotation_class = dt.AnnotationClass(name="test", annotation_type="complex_polygon")
-        annotation = dt.Annotation(annotation_class=annotation_class, data={"paths": polygon_path}, subs=[])
+        annotation_class = dt.AnnotationClass(
+            name="test", annotation_type="complex_polygon"
+        )
+        annotation = dt.Annotation(
+            annotation_class=annotation_class, data={"paths": polygon_path}, subs=[]
+        )
 
         annotation_file = dt.AnnotationFile(
             path=Path("test.json"),
@@ -123,7 +143,13 @@ def test_complex_polygon(self):
                 "path": None,
                 "workview_url": None,
             },
-            "annotations": [{"complex_polygon": {"path": polygon_path}, "name": "test", "slot_names": []}],
+            "annotations": [
+                {
+                    "complex_polygon": {"path": polygon_path},
+                    "name": "test",
+                    "slot_names": [],
+                }
+            ],
             "dataset": "None",
         }
 
@@ -137,7 +163,11 @@ def test_polygon_annotation_file_with_bbox(self):
         bounding_box = {"x": 557.66, "y": 428.98, "w": 160.76, "h": 315.3}
 
         annotation_class = dt.AnnotationClass(name="test", annotation_type="polygon")
-        annotation = dt.Annotation(annotation_class=annotation_class, data={"path": polygon_path, "bounding_box": bounding_box}, subs=[])
+        annotation = dt.Annotation(
+            annotation_class=annotation_class,
+            data={"path": polygon_path, "bounding_box": bounding_box},
+            subs=[],
+        )
 
         annotation_file = dt.AnnotationFile(
             path=Path("test.json"),
@@ -161,14 +191,29 @@ def test_polygon_annotation_file_with_bbox(self):
                 "path": None,
                 "workview_url": None,
             },
-            "annotations": [{"polygon": {"path": polygon_path}, "name": "test", "slot_names": [], "bounding_box": bounding_box}],
+            "annotations": [
+                {
+                    "polygon": {"path": polygon_path},
+                    "name": "test",
+                    "slot_names": [],
+                    "bounding_box": bounding_box,
+                }
+            ],
             "dataset": "None",
         }
 
     def test_complex_polygon_with_bbox(self):
         polygon_path = [
-            [{"x": 230.06, "y": 174.04}, {"x": 226.39, "y": 170.36}, {"x": 224.61, "y": 166.81}],
-            [{"x": 238.98, "y": 171.69}, {"x": 236.97, "y": 174.04}, {"x": 238.67, "y": 174.04}],
+            [
+                {"x": 230.06, "y": 174.04},
+                {"x": 226.39, "y": 170.36},
+                {"x": 224.61, "y": 166.81},
+            ],
+            [
+                {"x": 238.98, "y": 171.69},
+                {"x": 236.97, "y": 174.04},
+                {"x": 238.67, "y": 174.04},
+            ],
             [
                 {"x": 251.75, "y": 169.77},
                 {"x": 251.75, "y": 154.34},
@@ -179,8 +224,14 @@ def test_complex_polygon_with_bbox(self):
 
         bounding_box = {"x": 557.66, "y": 428.98, "w": 160.76, "h": 315.3}
 
-        annotation_class = dt.AnnotationClass(name="test", annotation_type="complex_polygon")
-        annotation = dt.Annotation(annotation_class=annotation_class, data={"paths": polygon_path, "bounding_box": bounding_box}, subs=[])
+        annotation_class = dt.AnnotationClass(
+            name="test", annotation_type="complex_polygon"
+        )
+        annotation = dt.Annotation(
+            annotation_class=annotation_class,
+            data={"paths": polygon_path, "bounding_box": bounding_box},
+            subs=[],
+        )
 
         annotation_file = dt.AnnotationFile(
             path=Path("test.json"),
@@ -204,14 +255,25 @@ def test_complex_polygon_with_bbox(self):
                 "path": None,
                 "workview_url": None,
             },
-            "annotations": [{"complex_polygon": {"path": polygon_path}, "name": "test", "slot_names": [], "bounding_box": bounding_box}],
+            "annotations": [
+                {
+                    "complex_polygon": {"path": polygon_path},
+                    "name": "test",
+                    "slot_names": [],
+                    "bounding_box": bounding_box,
+                }
+            ],
             "dataset": "None",
         }
 
     def test_bounding_box(self):
         bounding_box_data = {"x": 100, "y": 150, "w": 50, "h": 30}
-        annotation_class = dt.AnnotationClass(name="bbox_test", annotation_type="bounding_box")
-        annotation = dt.Annotation(annotation_class=annotation_class, data=bounding_box_data, subs=[])
+        annotation_class = dt.AnnotationClass(
+            name="bbox_test", annotation_type="bounding_box"
+        )
+        annotation = dt.Annotation(
+            annotation_class=annotation_class, data=bounding_box_data, subs=[]
+        )
 
         annotation_file = dt.AnnotationFile(
             path=Path("test.json"),
@@ -235,14 +297,22 @@ def test_bounding_box(self):
                 "path": None,
                 "workview_url": None,
             },
-            "annotations": [{"bounding_box": bounding_box_data, "name": "bbox_test", "slot_names": []}],
+            "annotations": [
+                {
+                    "bounding_box": bounding_box_data,
+                    "name": "bbox_test",
+                    "slot_names": [],
+                }
+            ],
             "dataset": "None",
         }
 
     def test_tags(self):
         tag_data = "sample_tag"
         annotation_class = dt.AnnotationClass(name="tag_test", annotation_type="tag")
-        annotation = dt.Annotation(annotation_class=annotation_class, data=tag_data, subs=[])
+        annotation = dt.Annotation(
+            annotation_class=annotation_class, data=tag_data, subs=[]
+        )
 
         annotation_file = dt.AnnotationFile(
             path=Path("test.json"),