Add MOT and CVAT video annotation import/export support

dagshub/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.9"
+__version__ = "0.6.10"
 from .logger import DAGsHubLogger, dagshub_logger
 from .common.init import init
 from .upload.wrapper import upload_files

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,22 +1,30 @@
 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.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 IRVideoBBoxFrameAnnotation
 
 from dagshub.common.api import UserAPI
 from dagshub.common.api.repo import PathNotFoundError
 from dagshub.common.helpers import log_message
+from dagshub.data_engine.annotation.video import build_video_sequence_from_annotations
 
 if TYPE_CHECKING:
     from dagshub.data_engine.model.datasource import Datasource
 
-AnnotationType = Literal["yolo", "cvat"]
+AnnotationType = Literal["yolo", "cvat", "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,145 @@ 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(result):
+                        annotation_dict = self._flatten_video_annotations(result)
+                    else:
+                        annotation_dict = result
+            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(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(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(result) -> bool:
+        """Check whether a loader result contains video annotations."""
+        from dagshub_annotation_converter.ir.video import IRVideoSequence
+
+        if isinstance(result, IRVideoSequence):
+            return True
+        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,
+        video_data,
+    ) -> Dict[str, Sequence[IRAnnotationBase]]:
+        """Flatten video annotations into a single entry keyed by video name."""
+        from dagshub_annotation_converter.ir.video import IRVideoSequence
+
+        video_name = self.additional_args.get("video_name", self.annotations_file.stem)
+        if isinstance(video_data, IRVideoSequence):
+            return {video_name: video_data.to_annotations()}
+
+        all_anns: List[IRAnnotationBase] = []
+        for frame_anns in video_data.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]]:
+        from dagshub_annotation_converter.ir.video import IRVideoSequence
+
+        flattened: Dict[str, List[IRAnnotationBase]] = {}
+        for rel_path, result in fs_annotations.items():
+            if isinstance(result, IRVideoSequence):
+                video_key = Path(rel_path).stem
+                flattened.setdefault(video_key, [])
+                flattened[video_key].extend(result.to_annotations())
+            elif isinstance(result, dict):
+                if self._is_video_annotation(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]]:
+        from dagshub_annotation_converter.ir.video import IRVideoSequence
+
+        flattened: Dict[str, List[IRAnnotationBase]] = {}
+        for rel_path, result in fs_annotations.items():
+            if not isinstance(result, tuple) or len(result) != 2:
+                continue
+            sequence_or_dict = result[0]
+            sequence_name = Path(rel_path).stem if rel_path not in (".", "") else self.annotations_file.stem
+            flattened.setdefault(sequence_name, [])
+            if isinstance(sequence_or_dict, IRVideoSequence):
+                flattened[sequence_name].extend(sequence_or_dict.to_annotations())
+            elif isinstance(sequence_or_dict, dict):
+                for frame_anns in sequence_or_dict.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 +246,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 == "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:
@@ -136,6 +280,9 @@ def remap_annotations(
             remap_func: Function that maps from an annotation path to a datapoint path. \
                 If None, we try to guess it by getting a datapoint and remapping that path
         """
+        if not annotations:
+            return {}
+
         if remap_func is None:
             first_ann = list(annotations.keys())[0]
             first_ann_filename = Path(first_ann).name
@@ -153,8 +300,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 +439,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 +449,21 @@ 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, IRVideoBBoxFrameAnnotation)]
+            if not video_anns:
+                continue
+            sequence = build_video_sequence_from_annotations(video_anns, filename=filename)
+            video_path = self.ds.source.raw_path(filename)
+            ls_tasks = video_ir_to_ls_video_tasks(sequence, 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):

dagshub/data_engine/annotation/video.py

@@ -0,0 +1,24 @@
+from typing import Optional, Sequence
+
+from dagshub_annotation_converter.ir.video import IRVideoBBoxFrameAnnotation, IRVideoSequence
+
+
+def build_video_sequence_from_annotations(
+    annotations: Sequence[IRVideoBBoxFrameAnnotation],
+    filename: Optional[str] = None,
+) -> IRVideoSequence:
+    sequence = IRVideoSequence.from_annotations(annotations, filename=filename)
+
+    resolved_width = sequence.resolved_video_width()
+    if sequence.video_width is None and resolved_width is not None:
+        sequence.video_width = resolved_width
+
+    resolved_height = sequence.resolved_video_height()
+    if sequence.video_height is None and resolved_height is not None:
+        sequence.video_height = resolved_height
+
+    resolved_length = sequence.resolved_sequence_length()
+    if sequence.sequence_length is None and resolved_length is not None:
+        sequence.sequence_length = resolved_length
+
+    return sequence