YOLOv8 support (#8240)

## Summary by CodeRabbit - **New Features** - Introduced support for YOLOv8 formats, enhancing object detection capabilities. - Added new export and import functions for YOLOv8 formats within the dataset manager. - Expanded documentation to cover YOLOv8 format specifications and export processes. - **Bug Fixes** - Improved handling of various YOLOv8 annotation formats to ensure accurate processing. - **Tests** - Enhanced test coverage for YOLOv8 formats in both dataset export/import and REST API functionalities. - **Documentation** - Updated existing links in the YOLO format documentation for clarity. - Added new documentation detailing YOLOv8 formats and their export processes. Co-authored-by: Roman Donchenko <roman@cvat.ai>
cvat-ai · Aug 23, 2024 · ac05a6d · ac05a6d
1 parent d9a525b
commit ac05a6d
Show file tree

Hide file tree

Showing 14 changed files with 495 additions and 45 deletions.
diff --git a/changelog.d/20240730_160939_dmitrii.lavrukhin_yolo8_support.md b/changelog.d/20240730_160939_dmitrii.lavrukhin_yolo8_support.md
@@ -0,0 +1,4 @@
+### Added
+
+- Added support for YOLOv8 formats
+  (<https://github.com/cvat-ai/cvat/pull/8240>)
diff --git a/cvat/apps/dataset_manager/formats/yolo.py b/cvat/apps/dataset_manager/formats/yolo.py
@@ -2,36 +2,54 @@
 # Copyright (C) 2023-2024 CVAT.ai Corporation
 #
 # SPDX-License-Identifier: MIT
-
 import os.path as osp
 from glob import glob
 
 from pyunpack import Archive
 
-from cvat.apps.dataset_manager.bindings import (GetCVATDataExtractor, detect_dataset,
-    import_dm_annotations, match_dm_item, find_dataset_root)
+from cvat.apps.dataset_manager.bindings import (
+    GetCVATDataExtractor,
+    detect_dataset,
+    import_dm_annotations,
+    match_dm_item,
+    find_dataset_root,
+)
 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 datumaro.plugins.yolo_format.extractor import YoloExtractor
 
 from .registry import dm_env, exporter, importer
 
 
-@exporter(name='YOLO', ext='ZIP', version='1.1')
-def _export(dst_file, temp_dir, instance_data, save_images=False):
+def _export_common(dst_file, temp_dir, instance_data, format_name, *, save_images=False):
     with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor:
         dataset = Dataset.from_extractors(extractor, env=dm_env)
-        dataset.export(temp_dir, 'yolo', save_images=save_images)
+        dataset.export(temp_dir, format_name, save_images=save_images)
 
     make_zip_archive(temp_dir, dst_file)
 
-@importer(name='YOLO', ext='ZIP', version='1.1')
-def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs):
+
+@exporter(name='YOLO', ext='ZIP', version='1.1')
+def _export_yolo(*args, **kwargs):
+    _export_common(*args, format_name='yolo', **kwargs)
+
+
+def _import_common(
+    src_file,
+    temp_dir,
+    instance_data,
+    format_name,
+    *,
+    load_data_callback=None,
+    import_kwargs=None,
+    **kwargs
+):
     Archive(src_file.name).extractall(temp_dir)
 
     image_info = {}
-    frames = [YoloExtractor.name_from_path(osp.relpath(p, temp_dir))
+    extractor = dm_env.extractors.get(format_name)
+    frames = [extractor.name_from_path(osp.relpath(p, temp_dir))
         for p in glob(osp.join(temp_dir, '**', '*.txt'), recursive=True)]
     root_hint = find_dataset_root(
         [DatasetItem(id=frame) for frame in frames], instance_data)
@@ -44,11 +62,75 @@ def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs
         except Exception: # nosec
             pass
         if frame_info is not None:
-            image_info[frame] = (frame_info['height'], frame_info['width'])
+            image_info[frame] = (frame_info["height"], frame_info["width"])
 
-    detect_dataset(temp_dir, format_name='yolo', importer=dm_env.importers.get('yolo'))
-    dataset = Dataset.import_from(temp_dir, 'yolo',
-        env=dm_env, image_info=image_info)
+    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 {}))
     if load_data_callback is not None:
         load_data_callback(dataset, instance_data)
     import_dm_annotations(dataset, instance_data)
+
+
+@importer(name='YOLO', ext='ZIP', version='1.1')
+def _import_yolo(*args, **kwargs):
+    _import_common(*args, format_name="yolo", **kwargs)
+
+
+@exporter(name='YOLOv8 Detection', ext='ZIP', version='1.0')
+def _export_yolov8_detection(*args, **kwargs):
+    _export_common(*args, format_name='yolov8_detection', **kwargs)
+
+
+@exporter(name='YOLOv8 Oriented Bounding Boxes', ext='ZIP', version='1.0')
+def _export_yolov8_oriented_boxes(*args, **kwargs):
+    _export_common(*args, format_name='yolov8_oriented_boxes', **kwargs)
+
+
+@exporter(name='YOLOv8 Segmentation', ext='ZIP', version='1.0')
+def _export_yolov8_segmentation(dst_file, temp_dir, instance_data, *, save_images=False):
+    with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor:
+        dataset = Dataset.from_extractors(extractor, env=dm_env)
+        dataset = dataset.transform('masks_to_polygons')
+        dataset.export(temp_dir, 'yolov8_segmentation', save_images=save_images)
+
+    make_zip_archive(temp_dir, dst_file)
+
+
+@exporter(name='YOLOv8 Pose', ext='ZIP', version='1.0')
+def _export_yolov8_pose(*args, **kwargs):
+    _export_common(*args, format_name='yolov8_pose', **kwargs)
+
+
+@importer(name='YOLOv8 Detection', ext="ZIP", version="1.0")
+def _import_yolov8_detection(*args, **kwargs):
+    _import_common(*args, format_name="yolov8_detection", **kwargs)
+
+
+@importer(name='YOLOv8 Segmentation', ext="ZIP", version="1.0")
+def _import_yolov8_segmentation(*args, **kwargs):
+    _import_common(*args, format_name="yolov8_segmentation", **kwargs)
+
+
+@importer(name='YOLOv8 Oriented Bounding Boxes', ext="ZIP", version="1.0")
+def _import_yolov8_oriented_boxes(*args, **kwargs):
+    _import_common(*args, format_name="yolov8_oriented_boxes", **kwargs)
+
+
+@importer(name='YOLOv8 Pose', ext="ZIP", version="1.0")
+def _import_yolov8_pose(src_file, temp_dir, instance_data, **kwargs):
+    with GetCVATDataExtractor(instance_data) as extractor:
+        point_categories = extractor.categories().get(AnnotationType.points)
+        label_categories = extractor.categories().get(AnnotationType.label)
+        true_skeleton_point_labels = {
+            label_categories[label_id].name: category.labels
+            for label_id, category in point_categories.items.items()
+        }
+    _import_common(
+        src_file,
+        temp_dir,
+        instance_data,
+        format_name="yolov8_pose",
+        import_kwargs=dict(skeleton_sub_labels=true_skeleton_point_labels),
+        **kwargs
+    )
diff --git a/cvat/apps/dataset_manager/project.py b/cvat/apps/dataset_manager/project.py
@@ -158,7 +158,7 @@ def import_dataset(self, dataset_file, importer, **options):
         os.makedirs(temp_dir_base, exist_ok=True)
         with TemporaryDirectory(dir=temp_dir_base) as temp_dir:
             try:
-                importer(dataset_file, temp_dir, project_data, self.load_dataset_data, **options)
+                importer(dataset_file, temp_dir, project_data, load_data_callback=self.load_dataset_data, **options)
             except (DatasetNotFoundError, CvatDatasetNotFoundError) as not_found:
                 if settings.CVAT_LOG_IMPORT_ERRORS:
                     dlogger.log_import_error(

diff --git a/cvat/apps/dataset_manager/tests/assets/annotations.json b/cvat/apps/dataset_manager/tests/assets/annotations.json
@@ -719,6 +719,131 @@
     ],
     "tracks": []
   },
+  "YOLOv8 Detection 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": []
+  },
+  "YOLOv8 Oriented Bounding Boxes 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",
+        "rotation": 30.0,
+        "attributes": []
+      }
+    ],
+    "tracks": []
+  },
+  "YOLOv8 Segmentation 1.0": {
+    "version": 0,
+    "tags": [],
+    "shapes": [
+      {
+        "type": "polygon",
+        "occluded": false,
+        "z_order": 0,
+        "points": [25.04, 13.7, 35.85, 20.2, 16.65, 19.8],
+        "frame": 0,
+        "label_id": null,
+        "group": 0,
+        "source": "manual",
+        "attributes": []
+      }
+    ],
+    "tracks": []
+  },
+  "YOLOv8 Pose 1.0": {
+    "version": 0,
+    "tags": [],
+    "shapes": [
+      {
+        "type": "skeleton",
+        "occluded": false,
+        "outside": false,
+        "z_order": 0,
+        "rotation": 0,
+        "points": [],
+        "frame": 0,
+        "label_id": 0,
+        "group": 0,
+        "source": "manual",
+        "attributes": [],
+        "elements": [
+          {
+            "type": "points",
+            "occluded": false,
+            "outside": true,
+            "z_order": 0,
+            "rotation": 0,
+            "points": [
+              223.02,
+              72.83
+            ],
+            "frame": 0,
+            "label_id": 1,
+            "group": 0,
+            "source": "manual",
+            "attributes": []
+          },
+          {
+            "type": "points",
+            "occluded": false,
+            "outside": false,
+            "z_order": 0,
+            "rotation": 0,
+            "points": [
+              232.98,
+              124.6
+            ],
+            "frame": 0,
+            "label_id": 2,
+            "group": 0,
+            "source": "manual",
+            "attributes": []
+          },
+          {
+            "type": "points",
+            "occluded": false,
+            "outside": false,
+            "z_order": 0,
+            "rotation": 0,
+            "points": [
+              281.22,
+              36.63
+            ],
+            "frame": 0,
+            "label_id": 3,
+            "group": 0,
+            "source": "manual",
+            "attributes": []
+          }
+        ]
+      }
+    ],
+    "tracks": []
+  },
   "VGGFace2 1.0": {
     "version": 0,
     "tags": [],

diff --git a/cvat/apps/dataset_manager/tests/assets/tasks.json b/cvat/apps/dataset_manager/tests/assets/tasks.json
@@ -592,5 +592,46 @@
         "svg": "<line x1=\"38.92810821533203\" y1=\"53.31378173828125\" x2=\"80.23341369628906\" y2=\"18.36313819885254\" stroke=\"black\" data-type=\"edge\" data-node-from=\"2\" stroke-width=\"0.5\" data-node-to=\"3\"></line><line x1=\"30.399484634399414\" y1=\"32.74474334716797\" x2=\"38.92810821533203\" y2=\"53.31378173828125\" stroke=\"black\" data-type=\"edge\" data-node-from=\"1\" stroke-width=\"0.5\" data-node-to=\"2\"></line><circle r=\"1.5\" stroke=\"black\" fill=\"#b3b3b3\" cx=\"30.399484634399414\" cy=\"32.74474334716797\" stroke-width=\"0.1\" data-type=\"element node\" data-element-id=\"1\" data-node-id=\"1\" data-label-name=\"1\"></circle><circle r=\"1.5\" stroke=\"black\" fill=\"#b3b3b3\" cx=\"38.92810821533203\" cy=\"53.31378173828125\" stroke-width=\"0.1\" data-type=\"element node\" data-element-id=\"2\" data-node-id=\"2\" data-label-name=\"2\"></circle><circle r=\"1.5\" stroke=\"black\" fill=\"#b3b3b3\" cx=\"80.23341369628906\" cy=\"18.36313819885254\" stroke-width=\"0.1\" data-type=\"element node\" data-element-id=\"3\" data-node-id=\"3\" data-label-name=\"3\"></circle>"
       }
     ]
+  },
+  "YOLOv8 Pose 1.0": {
+    "name": "YOLOv8 pose task",
+    "overlap": 0,
+    "segment_size": 100,
+    "labels": [
+      {
+        "name": "skeleton",
+        "color": "#2080c0",
+        "type": "skeleton",
+        "attributes": [
+          {
+            "name": "attr",
+            "mutable": false,
+            "input_type": "select",
+            "values": ["0", "1", "2"]
+          }
+        ],
+        "sublabels": [
+          {
+            "name": "1",
+            "color": "#d12345",
+            "attributes": [],
+            "type": "points"
+          },
+          {
+            "name": "2",
+            "color": "#350dea",
+            "attributes": [],
+            "type": "points"
+          },
+          {
+            "name": "3",
+            "color": "#479ffe",
+            "attributes": [],
+            "type": "points"
+          }
+        ],
+        "svg": "<line x1=\"38.92810821533203\" y1=\"53.31378173828125\" x2=\"80.23341369628906\" y2=\"18.36313819885254\" stroke=\"black\" data-type=\"edge\" data-node-from=\"2\" stroke-width=\"0.5\" data-node-to=\"3\"></line><line x1=\"30.399484634399414\" y1=\"32.74474334716797\" x2=\"38.92810821533203\" y2=\"53.31378173828125\" stroke=\"black\" data-type=\"edge\" data-node-from=\"1\" stroke-width=\"0.5\" data-node-to=\"2\"></line><circle r=\"1.5\" stroke=\"black\" fill=\"#b3b3b3\" cx=\"38.92810821533203\" cy=\"53.31378173828125\" stroke-width=\"0.1\" data-type=\"element node\" data-element-id=\"2\" data-node-id=\"2\" data-label-name=\"2\"></circle><circle r=\"1.5\" stroke=\"black\" fill=\"#b3b3b3\" cx=\"80.23341369628906\" cy=\"18.36313819885254\" stroke-width=\"0.1\" data-type=\"element node\" data-element-id=\"3\" data-node-id=\"3\" data-label-name=\"3\"></circle><circle r=\"1.5\" stroke=\"black\" fill=\"#b3b3b3\" cx=\"30.399484634399414\" cy=\"32.74474334716797\" stroke-width=\"0.1\" data-type=\"element node\" data-element-id=\"1\" data-node-id=\"1\" data-label-name=\"1\"></circle>"
+      }
+    ]
   }
 }
diff --git a/cvat/apps/dataset_manager/tests/test_formats.py b/cvat/apps/dataset_manager/tests/test_formats.py
@@ -309,7 +309,11 @@ def test_export_formats_query(self):
             'KITTI 1.0',
             'LFW 1.0',
             'Cityscapes 1.0',
-            'Open Images V6 1.0'
+            'Open Images V6 1.0',
+            'YOLOv8 Oriented Bounding Boxes 1.0',
+            'YOLOv8 Detection 1.0',
+            'YOLOv8 Pose 1.0',
+            'YOLOv8 Segmentation 1.0',
         })
 
     def test_import_formats_query(self):
@@ -342,6 +346,10 @@ def test_import_formats_query(self):
             'Open Images V6 1.0',
             'Datumaro 1.0',
             'Datumaro 3D 1.0',
+            'YOLOv8 Oriented Bounding Boxes 1.0',
+            'YOLOv8 Detection 1.0',
+            'YOLOv8 Pose 1.0',
+            'YOLOv8 Segmentation 1.0',
         })
 
     def test_exports(self):
@@ -391,6 +399,10 @@ def test_empty_images_are_exported(self):
             # ('KITTI 1.0', 'kitti') format does not support empty annotations
             ('LFW 1.0', 'lfw'),
             # ('Cityscapes 1.0', 'cityscapes'), does not support, empty annotations
+            ('YOLOv8 Oriented Bounding Boxes 1.0', 'yolov8_oriented_boxes'),
+            ('YOLOv8 Detection 1.0', 'yolov8_detection'),
+            ('YOLOv8 Pose 1.0', 'yolov8_pose'),
+            ('YOLOv8 Segmentation 1.0', 'yolov8_segmentation'),
         ]:
             with self.subTest(format=format_name):
                 if not dm.formats.registry.EXPORT_FORMATS[format_name].ENABLED: