[V7-1647] Make sure to include subs when converting bounding boxes to…

… coco (#206)

* Make sure to include subs when converting bounding boxes to coco

* changes from feedback

darwin/datatypes.py

@@ -8,6 +8,9 @@
 BoundingBox = Dict[str, float]
 Polygon = List[Point]
 ComplexPolygon = List[Polygon]
+Node = Dict[str, Any]
+EllipseData = Dict[str, Any]
+CuboidData = Dict[str, Any]
 
 
 @dataclass(frozen=True, eq=True)
@@ -33,6 +36,7 @@ def get_sub(self, annotation_type: str) -> Optional[SubAnnotation]:
         for sub in self.subs:
             if sub.annotation_type == annotation_type:
                 return sub
+        return None
 
 
 @dataclass(frozen=True, eq=True)
@@ -43,9 +47,6 @@ class VideoAnnotation:
     segments: List[List[int]]
     interpolated: bool
 
-    def get_frame(self, frame_index: int):
-        return frames[frame_index]
-
     def get_data(self, only_keyframes=True, post_processing=None):
         if not post_processing:
             post_processing = lambda annotation, data: data
@@ -86,43 +87,46 @@ def full_path(self) -> str:
         return construct_full_path(self.remote_path, self.filename)
 
 
-def make_bounding_box(class_name, x, y, w, h):
+def make_bounding_box(
+    class_name: str, x: float, y: float, w: float, h: float, subs: Optional[List[SubAnnotation]] = None
+):
     return Annotation(
         AnnotationClass(class_name, "bounding_box"),
         {"x": round(x, 3), "y": round(y, 3), "w": round(w, 3), "h": round(h, 3)},
+        subs or [],
     )
 
 
-def make_tag(class_name):
-    return Annotation(AnnotationClass(class_name, "tag"), {})
+def make_tag(class_name: str, subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "tag"), {}, subs or [])
 
 
-def make_polygon(class_name, point_path):
-    return Annotation(AnnotationClass(class_name, "polygon"), {"path": point_path})
+def make_polygon(class_name: str, point_path: List[Point], subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "polygon"), {"path": point_path}, subs or [])
 
 
-def make_complex_polygon(class_name, point_paths):
-    return Annotation(AnnotationClass(class_name, "complex_polygon", "polygon"), {"paths": point_paths})
+def make_complex_polygon(class_name: str, point_paths: List[List[Point]], subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "complex_polygon", "polygon"), {"paths": point_paths}, subs or [])
 
 
-def make_keypoint(class_name, x, y):
-    return Annotation(AnnotationClass(class_name, "keypoint"), {"x": x, "y": y})
+def make_keypoint(class_name: str, x: float, y: float, subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "keypoint"), {"x": x, "y": y}, subs or [])
 
 
-def make_line(class_name, path):
-    return Annotation(AnnotationClass(class_name, "line"), {"path": path})
+def make_line(class_name: str, path: List[Point], subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "line"), {"path": path}, subs or [])
 
 
-def make_skeleton(class_name, nodes):
-    return Annotation(AnnotationClass(class_name, "skeleton"), {"nodes": nodes})
+def make_skeleton(class_name: str, nodes: List[Node], subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "skeleton"), {"nodes": nodes}, subs or [])
 
 
-def make_ellipse(class_name, parameters):
-    return Annotation(AnnotationClass(class_name, "ellipse"), parameters)
+def make_ellipse(class_name: str, parameters: EllipseData, subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "ellipse"), parameters, subs or [])
 
 
-def make_cuboid(class_name, cuboid):
-    return Annotation(AnnotationClass(class_name, "cuboid"), cuboid)
+def make_cuboid(class_name: str, cuboid: CuboidData, subs: Optional[List[SubAnnotation]] = None):
+    return Annotation(AnnotationClass(class_name, "cuboid"), cuboid, subs or [])
 
 
 def make_instance_id(value):

darwin/exporter/formats/coco.py

@@ -175,12 +175,14 @@ def build_annotation(annotation_file, annotation_id, annotation: dt.Annotation,
         y = annotation.data["y"]
         w = annotation.data["w"]
         h = annotation.data["h"]
+
         return build_annotation(
             annotation_file,
             annotation_id,
             dt.make_polygon(
                 annotation.annotation_class.name,
                 [{"x": x, "y": y}, {"x": x + w, "y": y}, {"x": x + w, "y": y + h}, {"x": x, "y": y + h}],
+                annotation.subs,
             ),
             categories,
         )

tests/darwin/exporter/formats/coco_test.py

@@ -0,0 +1,33 @@
+from pathlib import Path
+
+import darwin.datatypes as dt
+import pytest
+from darwin.exporter.formats import coco
+
+
+def describe_build_annotation():
+    @pytest.fixture
+    def annotation_file() -> dt.AnnotationFile:
+        return dt.AnnotationFile(path=Path("test.json"), filename="test.json", annotation_classes=set(), annotations=[])
+
+    def polygon_include_extras(annotation_file: dt.AnnotationFile):
+        polygon = dt.Annotation(
+            dt.AnnotationClass("polygon_class", "polygon"),
+            {"path": [{"x": 1, "y": 1}, {"x": 2, "y": 2}, {"x": 1, "y": 2}]},
+            [dt.make_instance_id(1)],
+        )
+
+        categories = {"polygon_class": 1}
+
+        assert coco.build_annotation(annotation_file, "test-id", polygon, categories)["extra"] == {"instance_id": 1}
+
+    def bounding_boxes_include_extras(annotation_file: dt.AnnotationFile):
+        bbox = dt.Annotation(
+            dt.AnnotationClass("bbox_class", "bounding_box"),
+            {"x": 1, "y": 1, "w": 5, "h": 5},
+            [dt.make_instance_id(1)],
+        )
+
+        categories = {"bbox_class": 1}
+
+        assert coco.build_annotation(annotation_file, "test-id", bbox, categories)["extra"] == {"instance_id": 1}