ultralytics yolo support for tracks

cvat/apps/dataset_manager/bindings.py

@@ -2175,7 +2175,11 @@ def import_dm_annotations(dm_dataset: dm.Dataset, instance_data: Union[ProjectDa
         'coco',
         'coco_instances',
         'coco_person_keypoints',
-        'voc'
+        'voc',
+        'yolo_ultralytics_detection',
+        'yolo_ultralytics_segmentation',
+        'yolo_ultralytics_oriented_boxes',
+        'yolo_ultralytics_pose',
     ]
 
     label_cat = dm_dataset.categories()[dm.AnnotationType.label]

cvat/apps/dataset_manager/formats/transformations.py

@@ -37,6 +37,7 @@ def transform_item(self, item):
 
         return item.wrap(annotations=annotations)
 
+
 class MaskConverter:
     @staticmethod
     def cvat_rle_to_dm_rle(shape, img_h: int, img_w: int) -> dm.RleMask:
@@ -100,6 +101,7 @@ def rle(cls, arr: np.ndarray) -> list[int]:
 
         return cvat_rle
 
+
 class EllipsesToMasks:
     @staticmethod
     def convert_ellipse(ellipse, img_h, img_w):
@@ -115,6 +117,7 @@ def convert_ellipse(ellipse, img_h, img_w):
         return dm.RleMask(rle=rle, label=ellipse.label, z_order=ellipse.z_order,
             attributes=ellipse.attributes, group=ellipse.group)
 
+
 class MaskToPolygonTransformation:
     """
     Manages common logic for mask to polygons conversion in dataset import.
@@ -130,3 +133,13 @@ def convert_dataset(cls, dataset, **kwargs):
         if kwargs.get('conv_mask_to_poly', True):
             dataset.transform('masks_to_polygons')
         return dataset
+
+
+class SetKeyframeForEveryTrackShape(dm.ItemTransform):
+    def transform_item(self, item):
+        annotations = []
+        for ann in item.annotations:
+            if "track_id" in ann.attributes:
+                ann = ann.wrap(attributes=dict(ann.attributes, keyframe=True))
+            annotations.append(ann)
+        return item.wrap(annotations=annotations)

cvat/apps/dataset_manager/formats/yolo.py

@@ -4,28 +4,40 @@
 # SPDX-License-Identifier: MIT
 import os.path as osp
 from glob import glob
+from typing import Callable, Optional
 
 from pyunpack import Archive
 
 from cvat.apps.dataset_manager.bindings import (
-    GetCVATDataExtractor,
+    CommonData,
     detect_dataset,
+    find_dataset_root,
+    GetCVATDataExtractor,
     import_dm_annotations,
     match_dm_item,
-    find_dataset_root,
+    ProjectData,
 )
 from cvat.apps.dataset_manager.util import make_zip_archive
 from datumaro.components.annotation import AnnotationType
 from datumaro.components.extractor import DatasetItem
 from datumaro.components.project import Dataset
 
 from .registry import dm_env, exporter, importer
+from .transformations import SetKeyframeForEveryTrackShape
 
 
-def _export_common(dst_file, temp_dir, instance_data, format_name, *, save_images=False):
+def _export_common(
+    dst_file: str,
+    temp_dir: str,
+    instance_data: ProjectData | CommonData,
+    format_name: str,
+    *,
+    save_images: bool = False,
+    write_track_id: bool = False,
+):
     with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor:
         dataset = Dataset.from_extractors(extractor, env=dm_env)
-        dataset.export(temp_dir, format_name, save_images=save_images)
+        dataset.export(temp_dir, format_name, save_images=save_images, write_track_id=write_track_id)
 
     make_zip_archive(temp_dir, dst_file)
 
@@ -37,12 +49,12 @@ def _export_yolo(*args, **kwargs):
 
 def _import_common(
     src_file,
-    temp_dir,
-    instance_data,
-    format_name,
+    temp_dir: str,
+    instance_data: ProjectData | CommonData,
+    format_name: str,
     *,
-    load_data_callback=None,
-    import_kwargs=None,
+    load_data_callback: Optional[Callable] = None,
+    import_kwargs: dict | None = None,
     **kwargs
 ):
     Archive(src_file.name).extractall(temp_dir)
@@ -67,6 +79,7 @@ def _import_common(
     detect_dataset(temp_dir, format_name=format_name, importer=dm_env.importers.get(format_name))
     dataset = Dataset.import_from(temp_dir, format_name,
         env=dm_env, image_info=image_info, **(import_kwargs or {}))
+    dataset = dataset.transform(SetKeyframeForEveryTrackShape)
     if load_data_callback is not None:
         load_data_callback(dataset, instance_data)
     import_dm_annotations(dataset, instance_data)
@@ -82,6 +95,11 @@ def _export_yolo_ultralytics_detection(*args, **kwargs):
     _export_common(*args, format_name='yolo_ultralytics_detection', **kwargs)
 
 
+@exporter(name='Ultralytics YOLO Detection Track', ext='ZIP', version='1.0')
+def _export_yolo_ultralytics_detection(*args, **kwargs):
+    _export_common(*args, format_name='yolo_ultralytics_detection', write_track_id=True, **kwargs)
+
+
 @exporter(name='Ultralytics YOLO Oriented Bounding Boxes', ext='ZIP', version='1.0')
 def _export_yolo_ultralytics_oriented_boxes(*args, **kwargs):
     _export_common(*args, format_name='yolo_ultralytics_oriented_boxes', **kwargs)

cvat/apps/dataset_manager/tests/assets/annotations.json

@@ -1008,6 +1008,53 @@
     ],
     "tracks": []
   },
+  "Ultralytics YOLO Detection Track 1.0": {
+    "version": 0,
+    "tags": [],
+    "shapes": [
+      {
+        "type": "rectangle",
+        "occluded": false,
+        "z_order": 0,
+        "points": [8.3, 9.1, 19.2, 14.8],
+        "frame": 0,
+        "label_id": null,
+        "group": 0,
+        "source": "manual",
+        "attributes": []
+      }
+    ],
+    "tracks": [
+      {
+        "frame": 0,
+        "label_id": null,
+        "group": 0,
+        "source": "manual",
+        "shapes": [
+          {
+            "type": "rectangle",
+            "occluded": false,
+            "z_order": 0,
+            "points": [8.3, 9.1, 19.2, 14.8],
+            "frame": 0,
+            "outside": true,
+            "attributes": []
+          },
+          {
+            "type": "rectangle",
+            "occluded": false,
+            "z_order": 0,
+            "points": [8.3, 9.1, 19.2, 14.8],
+            "frame": 1,
+            "outside": false,
+            "attributes": [],
+            "keyframe": true
+          }
+        ],
+        "attributes": []
+      }
+    ]
+  },
   "Ultralytics YOLO Oriented Bounding Boxes 1.0": {
     "version": 0,
     "tags": [],

cvat/apps/dataset_manager/tests/test_formats.py

@@ -295,6 +295,7 @@ def test_export_formats_query(self):
             'Ultralytics YOLO Classification 1.0',
             'Ultralytics YOLO Oriented Bounding Boxes 1.0',
             'Ultralytics YOLO Detection 1.0',
+            'Ultralytics YOLO Detection Track 1.0',
             'Ultralytics YOLO Pose 1.0',
             'Ultralytics YOLO Segmentation 1.0',
         })

cvat/apps/dataset_manager/tests/test_rest_api_formats.py

@@ -58,6 +58,7 @@
     "Ultralytics YOLO Classification 1.0",
     "YOLO 1.1",
     "Ultralytics YOLO Detection 1.0",
+    "Ultralytics YOLO Detection Track 1.0",
     "Ultralytics YOLO Segmentation 1.0",
     "Ultralytics YOLO Oriented Bounding Boxes 1.0",
     "Ultralytics YOLO Pose 1.0",

cvat/requirements/base.in

@@ -12,7 +12,7 @@ azure-storage-blob==12.13.0
 boto3==1.17.61
 clickhouse-connect==0.6.8
 coreapi==2.3.3
-datumaro @ git+https://github.com/cvat-ai/datumaro.git@232c175ef1f3b7e55bd5162353df9c86a8116fde
+datumaro @ git+https://github.com/cvat-ai/datumaro.git@6247b61cf54438b3494634a0b5e01aa9a8e6b834
 dj-pagination==2.5.0
 # Despite direct indication allauth in requirements we should keep 'with_social' for dj-rest-auth
 # to avoid possible further versions conflicts (we use registration functionality)

cvat/requirements/base.txt

@@ -56,7 +56,7 @@ cryptography==44.0.0
     #   pyjwt
 cycler==0.12.1
     # via matplotlib
-datumaro @ git+https://github.com/cvat-ai/datumaro.git@232c175ef1f3b7e55bd5162353df9c86a8116fde
+datumaro @ git+https://github.com/cvat-ai/datumaro.git@6247b61cf54438b3494634a0b5e01aa9a8e6b834
     # via -r cvat/requirements/base.in
 defusedxml==0.7.1
     # via