Add support for video (MOT + CVAT Video) and COCO image (BBOX and Segmentation) import and export

dagshub/auth/token_auth.py

@@ -37,7 +37,7 @@ def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
 
     def can_renegotiate(self):
         # Env var tokens cannot renegotiate, every other token type can
-        return not type(self._token) is EnvVarDagshubToken
+        return type(self._token) is not EnvVarDagshubToken
 
     def renegotiate_token(self):
         if not self._token_storage.is_valid_token(self._token, self._host):

dagshub/data_engine/annotation/importer.py

@@ -1,13 +1,21 @@
 from difflib import SequenceMatcher
 from pathlib import Path, PurePosixPath, PurePath
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Literal, Optional, Union, Sequence, Mapping, Callable, List
-
-from dagshub_annotation_converter.converters.cvat import load_cvat_from_zip
+from typing import TYPE_CHECKING, Dict, Literal, Optional, Union, Sequence, Mapping, Callable, List
+
+from dagshub_annotation_converter.converters.coco import load_coco_from_file
+from dagshub_annotation_converter.converters.cvat import (
+    load_cvat_from_fs,
+    load_cvat_from_zip,
+    load_cvat_from_xml_file,
+)
+from dagshub_annotation_converter.converters.mot import load_mot_from_dir, load_mot_from_fs, load_mot_from_zip
 from dagshub_annotation_converter.converters.yolo import load_yolo_from_fs
+from dagshub_annotation_converter.converters.label_studio_video import video_ir_to_ls_video_tasks
 from dagshub_annotation_converter.formats.label_studio.task import LabelStudioTask
 from dagshub_annotation_converter.formats.yolo import YoloContext
 from dagshub_annotation_converter.ir.image.annotations.base import IRAnnotationBase
+from dagshub_annotation_converter.ir.video import IRVideoBBoxAnnotation
 
 from dagshub.common.api import UserAPI
 from dagshub.common.api.repo import PathNotFoundError
@@ -16,7 +24,7 @@
 if TYPE_CHECKING:
     from dagshub.data_engine.model.datasource import Datasource
 
-AnnotationType = Literal["yolo", "cvat"]
+AnnotationType = Literal["yolo", "cvat", "coco", "mot", "cvat_video"]
 AnnotationLocation = Literal["repo", "disk"]
 
 
@@ -57,6 +65,10 @@ def __init__(
                     'Add `yolo_type="bbox"|"segmentation"|pose"` to the arguments.'
                 )
 
+    @property
+    def is_video_format(self) -> bool:
+        return self.annotations_type in ("mot", "cvat_video")
+
     def import_annotations(self) -> Mapping[str, Sequence[IRAnnotationBase]]:
         # Double check that the annotation file exists
         if self.load_from == "disk":
@@ -84,15 +96,130 @@ def import_annotations(self) -> Mapping[str, Sequence[IRAnnotationBase]]:
                     annotation_type=self.additional_args["yolo_type"], meta_file=annotations_file
                 )
             elif self.annotations_type == "cvat":
-                annotation_dict = load_cvat_from_zip(annotations_file)
+                if annotations_file.is_dir():
+                    annotation_dict = self._flatten_cvat_fs_annotations(load_cvat_from_fs(annotations_file))
+                else:
+                    result = load_cvat_from_zip(annotations_file)
+                    if self._is_video_annotation_dict(result):
+                        annotation_dict = self._flatten_video_annotations(result)
+                    else:
+                        annotation_dict = result
+            elif self.annotations_type == "coco":
+                annotation_dict, _ = load_coco_from_file(annotations_file)
+            elif self.annotations_type == "mot":
+                mot_kwargs = {}
+                if "image_width" in self.additional_args:
+                    mot_kwargs["image_width"] = self.additional_args["image_width"]
+                if "image_height" in self.additional_args:
+                    mot_kwargs["image_height"] = self.additional_args["image_height"]
+                if "video_name" in self.additional_args:
+                    mot_kwargs["video_file"] = self.additional_args["video_name"]
+                if annotations_file.is_dir():
+                    video_files = self.additional_args.get("video_files")
+                    raw_datasource_path = self.additional_args.get("datasource_path")
+                    if raw_datasource_path is None:
+                        raw_datasource_path = self.ds.source.source_prefix
+                    datasource_path = PurePosixPath(raw_datasource_path).as_posix().lstrip("/")
+                    if datasource_path == ".":
+                        datasource_path = ""
+                    mot_results = load_mot_from_fs(
+                        annotations_file,
+                        image_width=mot_kwargs.get("image_width"),
+                        image_height=mot_kwargs.get("image_height"),
+                        video_files=video_files,
+                        datasource_path=datasource_path,
+                    )
+                    annotation_dict = self._flatten_mot_fs_annotations(mot_results)
+                elif annotations_file.suffix == ".zip":
+                    video_anns, _ = load_mot_from_zip(annotations_file, **mot_kwargs)
+                    annotation_dict = self._flatten_video_annotations(video_anns)
+                else:
+                    video_anns, _ = load_mot_from_dir(annotations_file, **mot_kwargs)
+                    annotation_dict = self._flatten_video_annotations(video_anns)
+            elif self.annotations_type == "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)
+                    annotation_dict = self._flatten_cvat_fs_annotations(raw)
+                elif annotations_file.suffix == ".zip":
+                    result = load_cvat_from_zip(annotations_file, **cvat_kwargs)
+                    if self._is_video_annotation_dict(result):
+                        annotation_dict = self._flatten_video_annotations(result)
+                    else:
+                        annotation_dict = result
+                else:
+                    result = load_cvat_from_xml_file(annotations_file, **cvat_kwargs)
+                    if self._is_video_annotation_dict(result):
+                        annotation_dict = self._flatten_video_annotations(result)
+                    else:
+                        annotation_dict = result
+            else:
+                raise ValueError(f"Unsupported annotation type: {self.annotations_type}")
 
             return annotation_dict
 
+    @staticmethod
+    def _is_video_annotation_dict(result) -> bool:
+        """Check if the result from a CVAT loader is video annotations (int keys) vs image annotations (str keys)."""
+        if not isinstance(result, dict) or len(result) == 0:
+            return False
+        first_key = next(iter(result.keys()))
+        return isinstance(first_key, int)
+
+    def _flatten_video_annotations(
+        self,
+        frame_annotations: Dict[int, Sequence[IRAnnotationBase]],
+    ) -> Dict[str, Sequence[IRAnnotationBase]]:
+        """Flatten frame-indexed video annotations into a single entry keyed by video name."""
+        video_name = self.additional_args.get("video_name", self.annotations_file.stem)
+        all_anns: List[IRAnnotationBase] = []
+        for frame_anns in frame_annotations.values():
+            all_anns.extend(frame_anns)
+        return {video_name: all_anns}
+
+    def _flatten_cvat_fs_annotations(
+        self, fs_annotations: Mapping[str, object]
+    ) -> Dict[str, Sequence[IRAnnotationBase]]:
+        flattened: Dict[str, List[IRAnnotationBase]] = {}
+        for rel_path, result in fs_annotations.items():
+            if not isinstance(result, dict):
+                continue
+            if self._is_video_annotation_dict(result):
+                video_key = Path(rel_path).stem
+                flattened.setdefault(video_key, [])
+                for frame_anns in result.values():
+                    flattened[video_key].extend(frame_anns)
+            else:
+                for filename, anns in result.items():
+                    flattened.setdefault(filename, [])
+                    flattened[filename].extend(anns)
+        return flattened
+
+    def _flatten_mot_fs_annotations(
+        self,
+        fs_annotations: Mapping[str, object],
+    ) -> Dict[str, Sequence[IRAnnotationBase]]:
+        flattened: Dict[str, List[IRAnnotationBase]] = {}
+        for rel_path, result in fs_annotations.items():
+            if not isinstance(result, tuple) or len(result) != 2:
+                continue
+            frame_annotations = result[0]
+            if not isinstance(frame_annotations, dict):
+                continue
+            sequence_name = Path(rel_path).stem if rel_path not in (".", "") else self.annotations_file.stem
+            flattened.setdefault(sequence_name, [])
+            for frame_anns in frame_annotations.values():
+                flattened[sequence_name].extend(frame_anns)
+        return flattened
+
     def download_annotations(self, dest_dir: Path):
         log_message("Downloading annotations from repository")
         repoApi = self.ds.source.repoApi
-        if self.annotations_type == "cvat":
-            # Download just the annotation file
+        if self.annotations_type in ("cvat", "cvat_video"):
             repoApi.download(self.annotations_file.as_posix(), dest_dir, keep_source_prefix=True)
         elif self.annotations_type == "yolo":
             # Download the dataset .yaml file and the images + annotations
@@ -104,6 +231,8 @@ def download_annotations(self, dest_dir: Path):
             # Download the annotation data
             assert context.path is not None
             repoApi.download(self.annotations_file.parent / context.path, dest_dir, keep_source_prefix=True)
+        elif self.annotations_type in ("coco", "mot"):
+            repoApi.download(self.annotations_file.as_posix(), dest_dir, keep_source_prefix=True)
 
     @staticmethod
     def determine_load_location(ds: "Datasource", annotations_path: Union[str, Path]) -> AnnotationLocation:
@@ -153,8 +282,12 @@ def remap_annotations(
                 )
                 continue
             for ann in anns:
-                assert ann.filename is not None
-                ann.filename = remap_func(ann.filename)
+                if ann.filename is not None:
+                    ann.filename = remap_func(ann.filename)
+                else:
+                    if not self.is_video_format:
+                        raise ValueError(f"Non-video annotation has no filename: {ann}")
+                    ann.filename = new_filename
             remapped[new_filename] = anns
 
         return remapped
@@ -288,6 +421,8 @@ def convert_to_ls_tasks(self, annotations: Mapping[str, Sequence[IRAnnotationBas
         """
         Converts the annotations to Label Studio tasks.
         """
+        if self.is_video_format:
+            return self._convert_to_ls_video_tasks(annotations)
         current_user_id = UserAPI.get_current_user(self.ds.source.repoApi.host).user_id
         tasks = {}
         for filename, anns in annotations.items():
@@ -296,3 +431,20 @@ def convert_to_ls_tasks(self, annotations: Mapping[str, Sequence[IRAnnotationBas
             t.add_ir_annotations(anns)
             tasks[filename] = t.model_dump_json().encode("utf-8")
         return tasks
+
+    def _convert_to_ls_video_tasks(
+        self, annotations: Mapping[str, Sequence[IRAnnotationBase]]
+    ) -> Mapping[str, bytes]:
+        """
+        Converts video annotations to Label Studio video tasks.
+        """
+        tasks = {}
+        for filename, anns in annotations.items():
+            video_anns = [a for a in anns if isinstance(a, IRVideoBBoxAnnotation)]
+            if not video_anns:
+                continue
+            video_path = self.ds.source.raw_path(filename)
+            ls_tasks = video_ir_to_ls_video_tasks(video_anns, video_path=video_path)
+            if ls_tasks:
+                tasks[filename] = ls_tasks[0].model_dump_json().encode("utf-8")
+        return tasks

dagshub/data_engine/annotation/metadata.py

@@ -22,6 +22,11 @@
 
     from dagshub.data_engine.model.datapoint import Datapoint
 
+from dagshub_annotation_converter.formats.label_studio.videorectangle import VideoRectangleAnnotation
+from dagshub_annotation_converter.formats.label_studio.task import task_lookup as _task_lookup
+
+_task_lookup["videorectangle"] = VideoRectangleAnnotation
+
 
 class AnnotationMetaDict(dict):
     def __init__(self, annotation: "MetadataAnnotations", *args, **kwargs):
@@ -271,6 +276,28 @@ def add_image_pose(
         self.annotations.append(ann)
         self._update_datapoint()
 
+    def add_coco_annotation(
+        self,
+        coco_json: str,
+    ):
+        """
+        Add annotations from a COCO-format JSON string.
+
+        Args:
+            coco_json: A COCO-format JSON string with ``categories``, ``images``, and ``annotations`` keys.
+        """
+        from dagshub_annotation_converter.converters.coco import load_coco_from_json_string
+
+        grouped, _ = load_coco_from_json_string(coco_json)
+        new_anns: list[IRAnnotationBase] = []
+        for anns in grouped.values():
+            for ann in anns:
+                ann.filename = self.datapoint.path
+                new_anns.append(ann)
+        self.annotations.extend(new_anns)
+        log_message(f"Added {len(new_anns)} COCO annotation(s) to datapoint {self.datapoint.path}")
+        self._update_datapoint()
+
     def add_yolo_annotation(
         self,
         annotation_type: Literal["bbox", "segmentation", "pose"],