From 3fff16bfe75021a44d9c414733a87e894e1740d3 Mon Sep 17 00:00:00 2001 From: Andrea Azzini Date: Wed, 22 Sep 2021 09:38:02 +0200 Subject: [PATCH] [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 | 44 ++++++++++++---------- darwin/exporter/formats/coco.py | 2 + tests/darwin/exporter/formats/coco_test.py | 33 ++++++++++++++++ 3 files changed, 59 insertions(+), 20 deletions(-) create mode 100644 tests/darwin/exporter/formats/coco_test.py diff --git a/darwin/datatypes.py b/darwin/datatypes.py index 26a4935aa..b79b82f10 100644 --- a/darwin/datatypes.py +++ b/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): diff --git a/darwin/exporter/formats/coco.py b/darwin/exporter/formats/coco.py index 498ffa349..09258ddf0 100644 --- a/darwin/exporter/formats/coco.py +++ b/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, ) diff --git a/tests/darwin/exporter/formats/coco_test.py b/tests/darwin/exporter/formats/coco_test.py new file mode 100644 index 000000000..8535e4250 --- /dev/null +++ b/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}