diff --git a/darwin/datatypes.py b/darwin/datatypes.py index b79b82f10..84cb7054e 100644 --- a/darwin/datatypes.py +++ b/darwin/datatypes.py @@ -101,12 +101,27 @@ def make_tag(class_name: str, subs: Optional[List[SubAnnotation]] = None): return Annotation(AnnotationClass(class_name, "tag"), {}, subs or []) -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_polygon( + class_name: str, point_path: List[Point], bounding_box: Optional[Dict], subs: Optional[List[SubAnnotation]] = None +): + return Annotation( + AnnotationClass(class_name, "polygon"), + _maybe_add_bounding_box_data({"path": point_path}, bounding_box), + subs or [], + ) -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_complex_polygon( + class_name: str, + point_paths: List[List[Point]], + bounding_box: Optional[Dict], + subs: Optional[List[SubAnnotation]] = None, +): + return Annotation( + AnnotationClass(class_name, "complex_polygon", "polygon"), + _maybe_add_bounding_box_data({"paths": point_paths}, bounding_box), + subs or [], + ) def make_keypoint(class_name: str, x: float, y: float, subs: Optional[List[SubAnnotation]] = None): @@ -169,3 +184,14 @@ def make_video_annotation(frames, keyframes, segments, interpolated): raise ValueError("invalid argument to make_video_annotation") return VideoAnnotation(first_annotation.annotation_class, frames, keyframes, segments, interpolated) + + +def _maybe_add_bounding_box_data(data: Dict[str, Any], bounding_box: Optional[Dict]) -> Dict[str, Any]: + if bounding_box: + data["bounding_box"] = { + "x": bounding_box["x"], + "y": bounding_box["y"], + "w": bounding_box["w"], + "h": bounding_box["h"], + } + return data diff --git a/darwin/exporter/formats/coco.py b/darwin/exporter/formats/coco.py index c2da2f376..cea9fb6b0 100644 --- a/darwin/exporter/formats/coco.py +++ b/darwin/exporter/formats/coco.py @@ -182,6 +182,7 @@ def build_annotation(annotation_file, annotation_id, annotation: dt.Annotation, 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}], + None, annotation.subs, ), categories, diff --git a/darwin/exporter/formats/darwin.py b/darwin/exporter/formats/darwin.py index 0a8827386..5c11ea688 100644 --- a/darwin/exporter/formats/darwin.py +++ b/darwin/exporter/formats/darwin.py @@ -8,6 +8,13 @@ def build_image_annotation(annotation_file: dt.AnnotationFile): annotation.annotation_class.annotation_type: build_annotation_data(annotation), "name": annotation.annotation_class.name, } + + if ( + annotation.annotation_class.annotation_type == "complex_polygon" + or annotation.annotation_class.annotation_type == "polygon" + ) and "bounding_box" in annotation.data: + payload["bounding_box"] = annotation.data["bounding_box"] + annotations.append(payload) return { @@ -25,4 +32,7 @@ def build_annotation_data(annotation: dt.Annotation): if annotation.annotation_class.annotation_type == "complex_polygon": return {"path": annotation.data["paths"]} + if annotation.annotation_class.annotation_type == "polygon": + return dict(filter(lambda item: item[0] != "bounding_box", annotation.data.items())) + return dict(annotation.data) diff --git a/darwin/utils.py b/darwin/utils.py index 3772f3035..12fb9d271 100644 --- a/darwin/utils.py +++ b/darwin/utils.py @@ -286,17 +286,19 @@ def parse_darwin_annotation(annotation: Dict[str, Any]): name = annotation["name"] main_annotation = None if "polygon" in annotation: + bounding_box = annotation.get("bounding_box") if "additional_paths" in annotation["polygon"]: paths = [annotation["polygon"]["path"]] + annotation["polygon"]["additional_paths"] - main_annotation = dt.make_complex_polygon(name, paths) + main_annotation = dt.make_complex_polygon(name, paths, bounding_box) else: - main_annotation = dt.make_polygon(name, annotation["polygon"]["path"]) + main_annotation = dt.make_polygon(name, annotation["polygon"]["path"], bounding_box) elif "complex_polygon" in annotation: + bounding_box = annotation.get("bounding_box") if "additional_paths" in annotation["complex_polygon"]: paths = annotation["complex_polygon"]["path"] + annotation["complex_polygon"]["additional_paths"] - main_annotation = dt.make_complex_polygon(name, paths) + main_annotation = dt.make_complex_polygon(name, paths, bounding_box) else: - main_annotation = dt.make_complex_polygon(name, annotation["complex_polygon"]["path"]) + main_annotation = dt.make_complex_polygon(name, annotation["complex_polygon"]["path"], bounding_box) elif "bounding_box" in annotation: bounding_box = annotation["bounding_box"] main_annotation = dt.make_bounding_box( @@ -350,7 +352,7 @@ def split_video_annotation(annotation): for i, frame_url in enumerate(annotation.frame_urls): annotations = [a.frames[i] for a in annotation.annotations if i in a.frames] annotation_classes = set([annotation.annotation_class for annotation in annotations]) - filename = f"{Path(annotation.filename).stem}/{i:07d}.jpg" + filename = f"{Path(annotation.filename).stem}/{i:07d}.png" frame_annotations.append( dt.AnnotationFile( diff --git a/tests/darwin/dataset/remote_dataset_test.py b/tests/darwin/dataset/remote_dataset_test.py index 6cc0c38e9..985852dd1 100644 --- a/tests/darwin/dataset/remote_dataset_test.py +++ b/tests/darwin/dataset/remote_dataset_test.py @@ -333,13 +333,13 @@ def it_works_on_videos( "annotations": [ {"name": "test_class", "polygon": {"path": [{"x": 0, "y": 0}, {"x": 1, "y": 1}, {"x": 1, "y": 0}]}} ], - "image": {"filename": "test_video/0000000.jpg", "height": 1080, "url": "frame_1.jpg", "width": 1920}, + "image": {"filename": "test_video/0000000.png", "height": 1080, "url": "frame_1.jpg", "width": 1920}, } with (video_path / "0000001.json").open() as f: assert json.load(f) == { "annotations": [], - "image": {"filename": "test_video/0000001.jpg", "height": 1080, "url": "frame_2.jpg", "width": 1920}, + "image": {"filename": "test_video/0000001.png", "height": 1080, "url": "frame_2.jpg", "width": 1920}, } with (video_path / "0000002.json").open() as f: @@ -347,7 +347,7 @@ def it_works_on_videos( "annotations": [ {"name": "test_class", "polygon": {"path": [{"x": 5, "y": 5}, {"x": 6, "y": 6}, {"x": 6, "y": 5}]}} ], - "image": {"filename": "test_video/0000002.jpg", "height": 1080, "url": "frame_3.jpg", "width": 1920}, + "image": {"filename": "test_video/0000002.png", "height": 1080, "url": "frame_3.jpg", "width": 1920}, } diff --git a/tests/darwin/utils_test.py b/tests/darwin/utils_test.py index 43497ea2b..ff13d542f 100644 --- a/tests/darwin/utils_test.py +++ b/tests/darwin/utils_test.py @@ -1,6 +1,6 @@ from unittest.mock import MagicMock, patch -from darwin.datatypes import AnnotationFile +import darwin.datatypes as dt from darwin.utils import ( is_extension_allowed, is_image_extension_allowed, @@ -114,7 +114,7 @@ def it_parses_darwin_images_correctly(tmp_path): import_file = directory / "darwin-file.json" import_file.write_text(content) - annotation_file: AnnotationFile = parse_darwin_json(import_file, None) + annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None) assert annotation_file.path == import_file assert annotation_file.filename == "P49-RediPad-ProPlayLEFTY_442.jpg" @@ -131,7 +131,7 @@ def it_parses_darwin_images_correctly(tmp_path): def it_parses_darwin_videos_correctly(tmp_path): content = """ - { + { "dataset": "my-dataset", "image": { "width": 3840, @@ -149,48 +149,48 @@ def it_parses_darwin_videos_correctly(tmp_path): }, "annotations": [ { - "frames": { - "3": { - "bounding_box": { - "h": 547.0, - "w": 400.0, - "x": 363.0, - "y": 701.0 - }, - "instance_id": { - "value": 119 - }, - "keyframe": true, - "polygon": { - "path": [ - { - "x": 748.0, - "y": 732.0 - }, - { - "x": 751.0, - "y": 735.0 - }, - { - "x": 748.0, - "y": 733.0 + "frames": { + "3": { + "bounding_box": { + "h": 547.0, + "w": 400.0, + "x": 363.0, + "y": 701.0 + }, + "instance_id": { + "value": 119 + }, + "keyframe": true, + "polygon": { + "path": [ + { + "x": 748.0, + "y": 732.0 + }, + { + "x": 751.0, + "y": 735.0 + }, + { + "x": 748.0, + "y": 733.0 + } + ] + } } + }, + "interpolate_algorithm": "linear-1.1", + "interpolated": true, + "name": "Hand", + "segments": [ + [ + 3, + 46 ] - } - } - }, - "interpolate_algorithm": "linear-1.1", - "interpolated": true, - "name": "Hand", - "segments": [ - [ - 3, - 46 ] - ] } ] - } + } """ directory = tmp_path / "imports" @@ -198,7 +198,7 @@ def it_parses_darwin_videos_correctly(tmp_path): import_file = directory / "darwin-file.json" import_file.write_text(content) - annotation_file: AnnotationFile = parse_darwin_json(import_file, None) + annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None) assert annotation_file.path == import_file assert annotation_file.filename == "above tractor.mp4" @@ -213,6 +213,29 @@ def it_parses_darwin_videos_correctly(tmp_path): assert annotation_file.frame_urls == ["https://my-website.com/api/videos/209/frames/0"] assert annotation_file.remote_path == "/" + assert annotation_file.annotations == [ + dt.VideoAnnotation( + annotation_class=dt.AnnotationClass( + name="Hand", annotation_type="polygon", annotation_internal_type=None + ), + frames={ + 3: dt.Annotation( + annotation_class=dt.AnnotationClass( + name="Hand", annotation_type="polygon", annotation_internal_type=None + ), + data={ + "path": [{"x": 748.0, "y": 732.0}, {"x": 751.0, "y": 735.0}, {"x": 748.0, "y": 733.0}], + "bounding_box": {"x": 363.0, "y": 701.0, "w": 400.0, "h": 547.0}, + }, + subs=[dt.SubAnnotation(annotation_type="instance_id", data=119)], + ) + }, + keyframes={3: True}, + segments=[[3, 46]], + interpolated=True, + ) + ] + def it_returns_None_if_no_annotations_exist(tmp_path): content = """ { @@ -232,7 +255,7 @@ def it_returns_None_if_no_annotations_exist(tmp_path): import_file = directory / "darwin-file.json" import_file.write_text(content) - annotation_file: AnnotationFile = parse_darwin_json(import_file, None) + annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None) assert not annotation_file @@ -260,11 +283,11 @@ def it_uses_a_default_path_if_one_is_missing(tmp_path): import_file = directory / "darwin-file.json" import_file.write_text(content) - annotation_file: AnnotationFile = parse_darwin_json(import_file, None) + annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None) assert annotation_file.remote_path == "/" - def it_imports_a_skeleteton(tmp_path): + def it_imports_a_skeleton(tmp_path): content = """ { "dataset": "cars", @@ -325,7 +348,7 @@ def it_imports_a_skeleteton(tmp_path): import_file = directory / "darwin-file.json" import_file.write_text(content) - annotation_file: AnnotationFile = parse_darwin_json(import_file, None) + annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None) assert annotation_file.annotations[0].annotation_class.annotation_type == "polygon" assert annotation_file.annotations[1].annotation_class.annotation_type == "skeleton" @@ -422,7 +445,7 @@ def it_imports_multiple_skeletetons(tmp_path): import_file = directory / "darwin-file.json" import_file.write_text(content) - annotation_file: AnnotationFile = parse_darwin_json(import_file, None) + annotation_file: dt.AnnotationFile = parse_darwin_json(import_file, None) assert annotation_file.annotations[0].annotation_class.annotation_type == "polygon" assert annotation_file.annotations[1].annotation_class.annotation_type == "skeleton"