adding YOLOv8 formats and tests

cvat-ai · Jul 30, 2024 · f5fefd2 · f5fefd2
1 parent d209741
commit f5fefd2
Show file tree

Hide file tree

Showing 6 changed files with 119 additions and 7 deletions.
diff --git a/cvat/apps/dataset_manager/formats/yolo.py b/cvat/apps/dataset_manager/formats/yolo.py
@@ -26,6 +26,7 @@ def _export(dst_file, temp_dir, instance_data, save_images=False):
 
     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):
     Archive(src_file.name).extractall(temp_dir)
@@ -52,3 +53,83 @@ def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs
     if load_data_callback is not None:
         load_data_callback(dataset, instance_data)
     import_dm_annotations(dataset, instance_data)
+
+
+@exporter(name='YOLOv8 Detection', ext='ZIP', version='1.0')
+def _export(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.export(temp_dir, 'yolov8', save_images=save_images)
+    make_zip_archive(temp_dir, dst_file)
+
+
+@exporter(name='YOLOv8 Oriented Bounding Boxes', ext='ZIP', version='1.0')
+def _export(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.export(temp_dir, 'yolov8_oriented_boxes', save_images=save_images)
+    make_zip_archive(temp_dir, dst_file)
+
+
+@exporter(name='YOLOv8 Segmentation', ext='ZIP', version='1.0')
+def _export(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.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(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.export(temp_dir, 'yolov8_pose', save_images=save_images)
+    make_zip_archive(temp_dir, dst_file)
+
+
+@importer(name='YOLOv8 Detection', ext="ZIP", version="1.0")
+def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs):
+    Archive(src_file.name).extractall(temp_dir)
+
+    detect_dataset(temp_dir, format_name='yolov8', importer=dm_env.importers.get('yolov8'))
+    dataset = Dataset.import_from(temp_dir, 'yolov8', env=dm_env)
+
+    if load_data_callback is not None:
+        load_data_callback(dataset, instance_data)
+    import_dm_annotations(dataset, instance_data)
+
+
+@importer(name='YOLOv8 Segmentation', ext="ZIP", version="1.0")
+def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs):
+    Archive(src_file.name).extractall(temp_dir)
+
+    detect_dataset(temp_dir, format_name='yolov8_segmentation', importer=dm_env.importers.get('yolov8_segmentation'))
+    dataset = Dataset.import_from(temp_dir, 'yolov8_segmentation', env=dm_env)
+
+    if load_data_callback is not None:
+        load_data_callback(dataset, instance_data)
+    import_dm_annotations(dataset, instance_data)
+
+
+@importer(name='YOLOv8 Oriented Bounding Boxes', ext="ZIP", version="1.0")
+def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs):
+    Archive(src_file.name).extractall(temp_dir)
+
+    detect_dataset(temp_dir, format_name='yolov8_oriented_boxes', importer=dm_env.importers.get('yolov8_oriented_boxes'))
+    dataset = Dataset.import_from(temp_dir, 'yolov8_oriented_boxes', env=dm_env)
+
+    if load_data_callback is not None:
+        load_data_callback(dataset, instance_data)
+    import_dm_annotations(dataset, instance_data)
+
+
+@importer(name='YOLOv8 Pose', ext="ZIP", version="1.0")
+def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs):
+    Archive(src_file.name).extractall(temp_dir)
+
+    detect_dataset(temp_dir, format_name='yolov8_pose', importer=dm_env.importers.get('yolov8_pose'))
+    dataset = Dataset.import_from(temp_dir, 'yolov8_pose', env=dm_env)
+
+    if load_data_callback is not None:
+        load_data_callback(dataset, instance_data)
+    import_dm_annotations(dataset, instance_data)
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'),
+            ('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:

diff --git a/cvat/apps/dataset_manager/tests/test_rest_api_formats.py b/cvat/apps/dataset_manager/tests/test_rest_api_formats.py
@@ -410,7 +410,7 @@ def test_api_v2_dump_and_upload_annotations_with_objects_type_is_shape(self):
                         "ImageNet 1.0", "MOTS PNG 1.0",
                         "PASCAL VOC 1.1", "Segmentation mask 1.1",
                         "VGGFace2 1.0",
-                        "WiderFace 1.0", "YOLO 1.1"
+                        "WiderFace 1.0", "YOLO 1.1", "YOLOv8 Detection 1.0",
                     ]:
                         self._create_annotations(task, dump_format_name, "default")
                     else:
@@ -517,7 +517,7 @@ def test_api_v2_dump_annotations_with_objects_type_is_track(self):
                             "Cityscapes 1.0", "ImageNet 1.0",
                             "MOTS PNG 1.0", "PASCAL VOC 1.1",
                             "Segmentation mask 1.1",
-                            "VGGFace2 1.0", "WiderFace 1.0", "YOLO 1.1"
+                            "VGGFace2 1.0", "WiderFace 1.0", "YOLO 1.1", "YOLOv8 Detection 1.0",
                     ]:
                         self._create_annotations(task, dump_format_name, "default")
                     else:
@@ -963,7 +963,7 @@ def test_api_v2_rewriting_annotations(self):
                         "MOT 1.1", "PASCAL VOC 1.1", "Segmentation mask 1.1",
                         "YOLO 1.1", "ImageNet 1.0",
                         "WiderFace 1.0", "VGGFace2 1.0",
-                        "Datumaro 1.0", "Open Images V6 1.0", "KITTI 1.0"
+                        "Datumaro 1.0", "Open Images V6 1.0", "KITTI 1.0", "YOLOv8 Detection 1.0",
                     ]:
                         self._create_annotations(task, dump_format_name, "default")
                     else:
@@ -1077,7 +1077,8 @@ def test_api_v2_tasks_annotations_dump_and_upload_with_datumaro(self):
                         "PASCAL VOC 1.1", "Segmentation mask 1.1",
                         "YOLO 1.1", "ImageNet 1.0",
                         "WiderFace 1.0", "VGGFace2 1.0", "LFW 1.0",
-                        "Open Images V6 1.0", "Datumaro 1.0", "KITTI 1.0"
+                        "Open Images V6 1.0", "Datumaro 1.0", "KITTI 1.0",
+                        "YOLOv8 Detection 1.0",
                     ]:
                         self._create_annotations(task, dump_format_name, "default")
                     else:
@@ -2051,7 +2052,7 @@ def test_api_v2_export_import_dataset(self):
                     "Cityscapes 1.0", "Datumaro 1.0", "ImageNet 1.0",
                     "MOT 1.1", "MOTS PNG 1.0", "PASCAL VOC 1.1",
                     "Segmentation mask 1.1", "VGGFace2 1.0",
-                    "WiderFace 1.0", "YOLO 1.1"
+                    "WiderFace 1.0", "YOLO 1.1", "YOLOv8 Detection 1.0",
                 ]:
                     self._create_annotations(task, dump_format_name, "default")
                 else:

diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py
@@ -6114,6 +6114,15 @@ def _get_initial_annotation(annotation_format):
             elif annotation_format == "YOLO 1.1":
                 annotations["shapes"] = rectangle_shapes_wo_attrs
 
+            elif annotation_format == "YOLOv8 Detection 1.0":
+                annotations["shapes"] = rectangle_shapes_wo_attrs
+
+            elif annotation_format == "YOLOv8 Oriented Bounding Boxes 1.0":
+                annotations["shapes"] = rectangle_shapes_wo_attrs
+
+            elif annotation_format == "YOLOv8 Segmentation 1.0":
+                annotations["shapes"] = polygon_shapes_wo_attrs
+
             elif annotation_format == "COCO 1.0":
                 annotations["shapes"] = polygon_shapes_wo_attrs
 
@@ -6471,7 +6480,7 @@ def etree_to_dict(t):
                     self.assertEqual(meta["task"]["name"], task["name"])
         elif format_name == "PASCAL VOC 1.1":
             self.assertTrue(zipfile.is_zipfile(content))
-        elif format_name == "YOLO 1.1":
+        elif format_name == ["YOLO 1.1", "YOLOv8 Detection 1.0", "YOLOv8 Segmentation 1.0", "YOLOv8 Oriented Bounding Boxes 1.0", "YOLOv8 Pose 1.0"]:
             self.assertTrue(zipfile.is_zipfile(content))
         elif format_name in ['Kitti Raw Format 1.0','Sly Point Cloud Format 1.0']:
             self.assertTrue(zipfile.is_zipfile(content))

diff --git a/tests/python/rest_api/test_projects.py b/tests/python/rest_api/test_projects.py
@@ -834,6 +834,10 @@ def test_can_export_and_import_dataset_after_deleting_related_storage(
             ("LFW 1.0", "{subset}/images/"),
             ("Cityscapes 1.0", "imgsFine/leftImg8bit/{subset}/"),
             ("Open Images V6 1.0", "images/{subset}/"),
+            ("YOLOv8 Detection 1.0", "images/{subset}/"),
+            ("YOLOv8 Oriented Bounding Boxes 1.0", "images/{subset}/"),
+            ("YOLOv8 Segmentation 1.0", "images/{subset}/"),
+            ("YOLOv8 Pose 1.0", "images/{subset}/"),
         ],
     )
     def test_creates_subfolders_for_subsets_on_export(

diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py
@@ -814,6 +814,7 @@ def test_export_dataset_after_deleting_related_cloud_storage(self, admin_user, t
         [
             ("Datumaro 1.0", "", "images/{subset}"),
             ("YOLO 1.1", "train", "obj_{subset}_data"),
+            ("YOLOv8 Detection 1.0", "train", "images/{subset}"),
         ],
     )
     def test_uses_subset_name(
@@ -2926,6 +2927,10 @@ def test_import_annotations_after_deleting_related_cloud_storage(
             "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_check_import_error_on_wrong_file_structure(self, tasks_with_shapes, format_name):