diff --git a/darwin/importer/importer.py b/darwin/importer/importer.py index 0ce221f6e..9d00b8aa5 100644 --- a/darwin/importer/importer.py +++ b/darwin/importer/importer.py @@ -1586,6 +1586,7 @@ def _handle_subs( data: dt.DictFreeForm, annotation_class_id: str, attributes: Dict[str, dt.UnknownType], + include_empty_attributes: Optional[bool] = False, ) -> dt.DictFreeForm: for sub in annotation.subs: if sub.annotation_type == "text": @@ -1609,6 +1610,9 @@ def _handle_subs( else: data[sub.annotation_type] = sub.data + if not data.get("attributes") and include_empty_attributes: + data["attributes"] = {"attributes": []} + return data @@ -1707,6 +1711,7 @@ def _get_annotation_data( _format_polygon_for_import(annotation, data), annotation_class_id, attributes, + include_empty_attributes=True, ), ) else: diff --git a/e2e_tests/data/import/video_annotations_with_subtypes/mini_uct.json b/e2e_tests/data/import/video_annotations_with_subtypes/mini_uct.json index cf6a3c709..2eb80bfc5 100644 --- a/e2e_tests/data/import/video_annotations_with_subtypes/mini_uct.json +++ b/e2e_tests/data/import/video_annotations_with_subtypes/mini_uct.json @@ -105,7 +105,7 @@ "instance_id": { "value": 14 }, - "keyframe": true, + "keyframe": false, "skeleton": { "nodes": [ { @@ -190,7 +190,7 @@ "instance_id": { "value": 15 }, - "keyframe": true, + "keyframe": false, "line": { "path": [ { @@ -263,7 +263,7 @@ "instance_id": { "value": 16 }, - "keyframe": true, + "keyframe": false, "keypoint": { "x": 8.2353, "y": 19.951 @@ -362,7 +362,7 @@ "instance_id": { "value": 17 }, - "keyframe": true, + "keyframe": false, "text": { "text": "test_text" } @@ -476,7 +476,7 @@ "instance_id": { "value": 18 }, - "keyframe": true, + "keyframe": false, "polygon": { "paths": [ [ @@ -578,7 +578,7 @@ "instance_id": { "value": 19 }, - "keyframe": true, + "keyframe": false, "text": { "text": "test_text" } @@ -687,7 +687,7 @@ } }, "2": { - "keyframe": true, + "keyframe": false, "tag": {}, "text": { "text": "test_text" diff --git a/tests/darwin/importer/importer_test.py b/tests/darwin/importer/importer_test.py index 9b55d9cc5..2aaf76812 100644 --- a/tests/darwin/importer/importer_test.py +++ b/tests/darwin/importer/importer_test.py @@ -520,6 +520,37 @@ def test_handle_subs() -> None: assert result == expected_result +def test__handle_subs_empty_attributes() -> None: + from darwin.importer.importer import _handle_subs + + annotation = dt.Annotation( + annotation_class=dt.AnnotationClass( + name="bbox1", annotation_type="bounding_box" + ), + data={"x": 451.525, "y": 213.559, "w": 913.22, "h": 538.983}, + subs=[], + slot_names=[], + id="a25e4613-718c-4cc8-9170-1bf372853f22", + ) + + initial_data = { + "bounding_box": {"x": 451.525, "y": 213.559, "w": 913.22, "h": 538.983} + } + + result = _handle_subs( + annotation=annotation, + data=initial_data, + annotation_class_id="bbox1", + attributes={}, + include_empty_attributes=True, + ) + + assert result == { + "bounding_box": {"x": 451.525, "y": 213.559, "w": 913.22, "h": 538.983}, + "attributes": {"attributes": []}, + } + + def test__format_polygon_for_import() -> None: from darwin.importer.importer import _format_polygon_for_import @@ -643,6 +674,57 @@ def test__get_annotation_data() -> None: assert mock_hcp.call_args_list[0][0][1] == {"TEST_TYPE": "TEST DATA"} +def test__get_annotation_data_video_annotation_with_attributes_that_become_empty() -> ( + None +): + from darwin.importer.importer import _get_annotation_data + + video_annotation_class = dt.AnnotationClass("video_class", "bounding_box") + video_annotation = dt.VideoAnnotation(video_annotation_class, {}, {}, [], False) + video_annotation.keyframes = {1: True, 2: True, 3: True, 4: True} + video_annotation.frames = { + 1: dt.Annotation( + annotation_class=video_annotation_class, + data={"x": 1, "y": 2, "w": 3, "h": 4}, + subs=[], + slot_names=[], + ), + 2: dt.Annotation( + annotation_class=video_annotation_class, + data={"x": 1, "y": 2, "w": 3, "h": 4}, + subs=[ + dt.SubAnnotation( + annotation_type="attributes", data=["attribute_1", "attribute_2"] + ) + ], + slot_names=[], + ), + 3: dt.Annotation( + annotation_class=video_annotation_class, + data={"x": 1, "y": 2, "w": 3, "h": 4}, + subs=[ + dt.SubAnnotation( + annotation_type="attributes", + data=["attribute_1"], + ) + ], + slot_names=[], + ), + 4: dt.Annotation( + annotation_class=video_annotation_class, + data={"x": 1, "y": 2, "w": 3, "h": 4}, + subs=[], + slot_names=[], + ), + } + attributes = {"video_class_id": {"attribute_1": "id_1", "attribute_2": "id_2"}} + result = _get_annotation_data(video_annotation, "video_class_id", attributes) + assert result["frames"][1]["attributes"] == {"attributes": []} + assert result["frames"][2]["attributes"] == {"attributes": ["id_1", "id_2"]} + assert result["frames"][3]["attributes"] == {"attributes": ["id_1"]} + assert result["frames"][4]["attributes"] == {"attributes": []} + + def __expectation_factory(i: int, slot_names: List[str]) -> dt.Annotation: annotation = dt.Annotation( dt.AnnotationClass(f"class_{i}", f"TEST_TYPE_{i}"), {}, [], []