Fix importing CVAT image XML files with annotation_type="cvat"

dagshub/data_engine/annotation/importer.py

@@ -61,6 +61,8 @@ def __init__(
         self.annotations_file = Path(annotations_file)
         self.load_from = load_from if load_from is not None else self.determine_load_location(ds, annotations_file)
         self.additional_args = kwargs
+        # Set when a "cvat" import turns out to contain video annotations
+        self._loaded_cvat_video = False
 
         if self.annotations_type == "yolo":
             if "yolo_type" not in kwargs:
@@ -71,7 +73,7 @@ def __init__(
 
     @property
     def is_video_format(self) -> bool:
-        return self.annotations_type in ("mot", "cvat_video")
+        return self.annotations_type in ("mot", "cvat_video") or self._loaded_cvat_video
 
     def import_annotations(self) -> Mapping[str, Sequence[IRTaskAnnotation]]:
         # Double check that the annotation file exists
@@ -99,15 +101,6 @@ def import_annotations(self) -> Mapping[str, Sequence[IRTaskAnnotation]]:
                 annotation_dict, _ = load_yolo_from_fs(
                     annotation_type=self.additional_args["yolo_type"], meta_file=annotations_file
                 )
-            elif self.annotations_type == "cvat":
-                if annotations_file.is_dir():
-                    annotation_dict = self._key_cvat_fs_annotations_by_filename(load_cvat_from_fs(annotations_file))
-                else:
-                    result = load_cvat_from_zip(annotations_file)
-                    if self._determine_cvat_annotation(result) == "video":
-                        annotation_dict = self._key_video_annotations_by_filename(result)
-                    else:
-                        annotation_dict = result
             elif self.annotations_type == "coco":
                 annotation_dict, _ = load_coco_from_file(annotations_file)
             elif self.annotations_type == "mot":
@@ -140,24 +133,28 @@ def import_annotations(self) -> Mapping[str, Sequence[IRTaskAnnotation]]:
                 else:
                     video_anns, _ = load_mot_from_dir(annotations_file, **mot_kwargs)
                     annotation_dict = self._key_video_annotations_by_filename(video_anns)
-            elif self.annotations_type == "cvat_video":
+            elif self.annotations_type in ("cvat", "cvat_video"):
                 cvat_kwargs = {}
                 if "image_width" in self.additional_args:
                     cvat_kwargs["image_width"] = self.additional_args["image_width"]
                 if "image_height" in self.additional_args:
                     cvat_kwargs["image_height"] = self.additional_args["image_height"]
                 if annotations_file.is_dir():
                     raw = load_cvat_from_fs(annotations_file, **cvat_kwargs)
+                    if any(self._determine_cvat_annotation(r) == "video" for r in raw.values()):
+                        self._loaded_cvat_video = True
                     annotation_dict = self._key_cvat_fs_annotations_by_filename(raw)
                 elif annotations_file.suffix == ".zip":
                     result = load_cvat_from_zip(annotations_file, **cvat_kwargs)
                     if self._determine_cvat_annotation(result) == "video":
+                        self._loaded_cvat_video = True
                         annotation_dict = self._key_video_annotations_by_filename(result)
                     else:
                         annotation_dict = result
                 else:
                     result = load_cvat_from_xml_file(annotations_file, **cvat_kwargs)
                     if self._determine_cvat_annotation(result) == "video":
+                        self._loaded_cvat_video = True
                         annotation_dict = self._key_video_annotations_by_filename(result)
                     else:
                         annotation_dict = result

tests/data_engine/annotation_import/test_cvat.py

@@ -0,0 +1,34 @@
+import zipfile
+
+import pytest
+from dagshub_annotation_converter.ir.image import IRBBoxImageAnnotation
+
+from dagshub.data_engine.annotation.importer import AnnotationImporter
+
+_CVAT_IMAGE_XML = b"""<?xml version="1.0" encoding="utf-8"?>
+<annotations>
+  <version>1.1</version>
+  <image id="0" name="img.jpg" width="100" height="50">
+    <box label="cat" source="manual" occluded="0" xtl="10" ytl="5" xbr="30" ybr="25" z_order="0"></box>
+  </image>
+</annotations>
+"""
+
+
+@pytest.mark.parametrize("as_zip", [False, True])
+def test_import_cvat_image_annotations(ds, tmp_path, as_zip):
+    if as_zip:
+        annotations_file = tmp_path / "annotations.zip"
+        with zipfile.ZipFile(annotations_file, "w") as z:
+            z.writestr("annotations.xml", _CVAT_IMAGE_XML)
+    else:
+        annotations_file = tmp_path / "annotations.xml"
+        annotations_file.write_bytes(_CVAT_IMAGE_XML)
+
+    importer = AnnotationImporter(ds, "cvat", annotations_file, load_from="disk")
+    result = importer.import_annotations()
+
+    assert list(result.keys()) == ["img.jpg"]
+    assert len(result["img.jpg"]) == 1
+    assert isinstance(result["img.jpg"][0], IRBBoxImageAnnotation)
+    assert not importer.is_video_format

tests/data_engine/annotation_import/test_cvat_video.py

@@ -26,13 +26,17 @@ def mock_source_prefix(ds):
 # --- import ---
 
 
-def test_import_cvat_video(ds, tmp_path):
+@pytest.mark.parametrize("annotation_type", ["cvat_video", "cvat"])
+def test_import_cvat_video(ds, tmp_path, annotation_type):
     xml_file = tmp_path / "annotations.xml"
     xml_file.write_bytes(_make_cvat_video_xml())
 
-    importer = AnnotationImporter(ds, "cvat_video", xml_file, load_from="disk")
+    importer = AnnotationImporter(ds, annotation_type, xml_file, load_from="disk")
     result = importer.import_annotations()
 
+    # "cvat" detects the video data, so Label Studio conversion makes video tasks
+    assert importer.is_video_format
+
     assert len(result) == 1
     anns = list(result.values())[0]
     assert len(anns) == 2