Skip to content

Commit 4fdc2df

Browse files
EldiesSpecLad
authored and
Bradley Schultz
committed
YOLOv8 support (cvat-ai#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 <[email protected]>
1 parent f3c6944 commit 4fdc2df

File tree

14 files changed

+495
-45
lines changed

14 files changed

+495
-45
lines changed

‎changelog.d/20240730_160939_dmitrii.lavrukhin_yolo8_support.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
### Added
2+
3+
- Added support for YOLOv8 formats
4+
(<https://github.com/cvat-ai/cvat/pull/8240>)

‎cvat/apps/dataset_manager/formats/yolo.py

Lines changed: 96 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,36 +2,54 @@
22
# Copyright (C) 2023-2024 CVAT.ai Corporation
33
#
44
# SPDX-License-Identifier: MIT
5-
65
import os.path as osp
76
from glob import glob
87

98
from pyunpack import Archive
109

11-
from cvat.apps.dataset_manager.bindings import (GetCVATDataExtractor, detect_dataset,
12-
import_dm_annotations, match_dm_item, find_dataset_root)
10+
from cvat.apps.dataset_manager.bindings import (
11+
GetCVATDataExtractor,
12+
detect_dataset,
13+
import_dm_annotations,
14+
match_dm_item,
15+
find_dataset_root,
16+
)
1317
from cvat.apps.dataset_manager.util import make_zip_archive
18+
from datumaro.components.annotation import AnnotationType
1419
from datumaro.components.extractor import DatasetItem
1520
from datumaro.components.project import Dataset
16-
from datumaro.plugins.yolo_format.extractor import YoloExtractor
1721

1822
from .registry import dm_env, exporter, importer
1923

2024

21-
@exporter(name='YOLO', ext='ZIP', version='1.1')
22-
def _export(dst_file, temp_dir, instance_data, save_images=False, all_images=True):
25+
def _export_common(dst_file, temp_dir, instance_data, format_name, *, save_images=False, all_images=True):
2326
with GetCVATDataExtractor(instance_data, include_images=save_images, all_images=all_images) as extractor:
2427
dataset = Dataset.from_extractors(extractor, env=dm_env)
25-
dataset.export(temp_dir, 'yolo', save_images=save_images)
28+
dataset.export(temp_dir, format_name, save_images=save_images)
2629

2730
make_zip_archive(temp_dir, dst_file)
2831

29-
@importer(name='YOLO', ext='ZIP', version='1.1')
30-
def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs):
32+
33+
@exporter(name='YOLO', ext='ZIP', version='1.1')
34+
def _export_yolo(*args, **kwargs):
35+
_export_common(*args, format_name='yolo', **kwargs)
36+
37+
38+
def _import_common(
39+
src_file,
40+
temp_dir,
41+
instance_data,
42+
format_name,
43+
*,
44+
load_data_callback=None,
45+
import_kwargs=None,
46+
**kwargs
47+
):
3148
Archive(src_file.name).extractall(temp_dir)
3249

3350
image_info = {}
34-
frames = [YoloExtractor.name_from_path(osp.relpath(p, temp_dir))
51+
extractor = dm_env.extractors.get(format_name)
52+
frames = [extractor.name_from_path(osp.relpath(p, temp_dir))
3553
for p in glob(osp.join(temp_dir, '**', '*.txt'), recursive=True)]
3654
root_hint = find_dataset_root(
3755
[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
4462
except Exception: # nosec
4563
pass
4664
if frame_info is not None:
47-
image_info[frame] = (frame_info['height'], frame_info['width'])
65+
image_info[frame] = (frame_info["height"], frame_info["width"])
4866

49-
detect_dataset(temp_dir, format_name='yolo', importer=dm_env.importers.get('yolo'))
50-
dataset = Dataset.import_from(temp_dir, 'yolo',
51-
env=dm_env, image_info=image_info)
67+
detect_dataset(temp_dir, format_name=format_name, importer=dm_env.importers.get(format_name))
68+
dataset = Dataset.import_from(temp_dir, format_name,
69+
env=dm_env, image_info=image_info, **(import_kwargs or {}))
5270
if load_data_callback is not None:
5371
load_data_callback(dataset, instance_data)
5472
import_dm_annotations(dataset, instance_data)
73+
74+
75+
@importer(name='YOLO', ext='ZIP', version='1.1')
76+
def _import_yolo(*args, **kwargs):
77+
_import_common(*args, format_name="yolo", **kwargs)
78+
79+
80+
@exporter(name='YOLOv8 Detection', ext='ZIP', version='1.0')
81+
def _export_yolov8_detection(*args, **kwargs):
82+
_export_common(*args, format_name='yolov8_detection', **kwargs)
83+
84+
85+
@exporter(name='YOLOv8 Oriented Bounding Boxes', ext='ZIP', version='1.0')
86+
def _export_yolov8_oriented_boxes(*args, **kwargs):
87+
_export_common(*args, format_name='yolov8_oriented_boxes', **kwargs)
88+
89+
90+
@exporter(name='YOLOv8 Segmentation', ext='ZIP', version='1.0')
91+
def _export_yolov8_segmentation(dst_file, temp_dir, instance_data, *, save_images=False):
92+
with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor:
93+
dataset = Dataset.from_extractors(extractor, env=dm_env)
94+
dataset = dataset.transform('masks_to_polygons')
95+
dataset.export(temp_dir, 'yolov8_segmentation', save_images=save_images)
96+
97+
make_zip_archive(temp_dir, dst_file)
98+
99+
100+
@exporter(name='YOLOv8 Pose', ext='ZIP', version='1.0')
101+
def _export_yolov8_pose(*args, **kwargs):
102+
_export_common(*args, format_name='yolov8_pose', **kwargs)
103+
104+
105+
@importer(name='YOLOv8 Detection', ext="ZIP", version="1.0")
106+
def _import_yolov8_detection(*args, **kwargs):
107+
_import_common(*args, format_name="yolov8_detection", **kwargs)
108+
109+
110+
@importer(name='YOLOv8 Segmentation', ext="ZIP", version="1.0")
111+
def _import_yolov8_segmentation(*args, **kwargs):
112+
_import_common(*args, format_name="yolov8_segmentation", **kwargs)
113+
114+
115+
@importer(name='YOLOv8 Oriented Bounding Boxes', ext="ZIP", version="1.0")
116+
def _import_yolov8_oriented_boxes(*args, **kwargs):
117+
_import_common(*args, format_name="yolov8_oriented_boxes", **kwargs)
118+
119+
120+
@importer(name='YOLOv8 Pose', ext="ZIP", version="1.0")
121+
def _import_yolov8_pose(src_file, temp_dir, instance_data, **kwargs):
122+
with GetCVATDataExtractor(instance_data) as extractor:
123+
point_categories = extractor.categories().get(AnnotationType.points)
124+
label_categories = extractor.categories().get(AnnotationType.label)
125+
true_skeleton_point_labels = {
126+
label_categories[label_id].name: category.labels
127+
for label_id, category in point_categories.items.items()
128+
}
129+
_import_common(
130+
src_file,
131+
temp_dir,
132+
instance_data,
133+
format_name="yolov8_pose",
134+
import_kwargs=dict(skeleton_sub_labels=true_skeleton_point_labels),
135+
**kwargs
136+
)

‎cvat/apps/dataset_manager/project.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ def import_dataset(self, dataset_file, importer, **options):
158158
os.makedirs(temp_dir_base, exist_ok=True)
159159
with TemporaryDirectory(dir=temp_dir_base) as temp_dir:
160160
try:
161-
importer(dataset_file, temp_dir, project_data, self.load_dataset_data, **options)
161+
importer(dataset_file, temp_dir, project_data, load_data_callback=self.load_dataset_data, **options)
162162
except (DatasetNotFoundError, CvatDatasetNotFoundError) as not_found:
163163
if settings.CVAT_LOG_IMPORT_ERRORS:
164164
dlogger.log_import_error(

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

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -719,6 +719,131 @@
719719
],
720720
"tracks": []
721721
},
722+
"YOLOv8 Detection 1.0": {
723+
"version": 0,
724+
"tags": [],
725+
"shapes": [
726+
{
727+
"type": "rectangle",
728+
"occluded": false,
729+
"z_order": 0,
730+
"points": [8.3, 9.1, 19.2, 14.8],
731+
"frame": 0,
732+
"label_id": null,
733+
"group": 0,
734+
"source": "manual",
735+
"attributes": []
736+
}
737+
],
738+
"tracks": []
739+
},
740+
"YOLOv8 Oriented Bounding Boxes 1.0": {
741+
"version": 0,
742+
"tags": [],
743+
"shapes": [
744+
{
745+
"type": "rectangle",
746+
"occluded": false,
747+
"z_order": 0,
748+
"points": [8.3, 9.1, 19.2, 14.8],
749+
"frame": 0,
750+
"label_id": null,
751+
"group": 0,
752+
"source": "manual",
753+
"rotation": 30.0,
754+
"attributes": []
755+
}
756+
],
757+
"tracks": []
758+
},
759+
"YOLOv8 Segmentation 1.0": {
760+
"version": 0,
761+
"tags": [],
762+
"shapes": [
763+
{
764+
"type": "polygon",
765+
"occluded": false,
766+
"z_order": 0,
767+
"points": [25.04, 13.7, 35.85, 20.2, 16.65, 19.8],
768+
"frame": 0,
769+
"label_id": null,
770+
"group": 0,
771+
"source": "manual",
772+
"attributes": []
773+
}
774+
],
775+
"tracks": []
776+
},
777+
"YOLOv8 Pose 1.0": {
778+
"version": 0,
779+
"tags": [],
780+
"shapes": [
781+
{
782+
"type": "skeleton",
783+
"occluded": false,
784+
"outside": false,
785+
"z_order": 0,
786+
"rotation": 0,
787+
"points": [],
788+
"frame": 0,
789+
"label_id": 0,
790+
"group": 0,
791+
"source": "manual",
792+
"attributes": [],
793+
"elements": [
794+
{
795+
"type": "points",
796+
"occluded": false,
797+
"outside": true,
798+
"z_order": 0,
799+
"rotation": 0,
800+
"points": [
801+
223.02,
802+
72.83
803+
],
804+
"frame": 0,
805+
"label_id": 1,
806+
"group": 0,
807+
"source": "manual",
808+
"attributes": []
809+
},
810+
{
811+
"type": "points",
812+
"occluded": false,
813+
"outside": false,
814+
"z_order": 0,
815+
"rotation": 0,
816+
"points": [
817+
232.98,
818+
124.6
819+
],
820+
"frame": 0,
821+
"label_id": 2,
822+
"group": 0,
823+
"source": "manual",
824+
"attributes": []
825+
},
826+
{
827+
"type": "points",
828+
"occluded": false,
829+
"outside": false,
830+
"z_order": 0,
831+
"rotation": 0,
832+
"points": [
833+
281.22,
834+
36.63
835+
],
836+
"frame": 0,
837+
"label_id": 3,
838+
"group": 0,
839+
"source": "manual",
840+
"attributes": []
841+
}
842+
]
843+
}
844+
],
845+
"tracks": []
846+
},
722847
"VGGFace2 1.0": {
723848
"version": 0,
724849
"tags": [],

‎cvat/apps/dataset_manager/tests/assets/tasks.json

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -592,5 +592,46 @@
592592
"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>"
593593
}
594594
]
595+
},
596+
"YOLOv8 Pose 1.0": {
597+
"name": "YOLOv8 pose task",
598+
"overlap": 0,
599+
"segment_size": 100,
600+
"labels": [
601+
{
602+
"name": "skeleton",
603+
"color": "#2080c0",
604+
"type": "skeleton",
605+
"attributes": [
606+
{
607+
"name": "attr",
608+
"mutable": false,
609+
"input_type": "select",
610+
"values": ["0", "1", "2"]
611+
}
612+
],
613+
"sublabels": [
614+
{
615+
"name": "1",
616+
"color": "#d12345",
617+
"attributes": [],
618+
"type": "points"
619+
},
620+
{
621+
"name": "2",
622+
"color": "#350dea",
623+
"attributes": [],
624+
"type": "points"
625+
},
626+
{
627+
"name": "3",
628+
"color": "#479ffe",
629+
"attributes": [],
630+
"type": "points"
631+
}
632+
],
633+
"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>"
634+
}
635+
]
595636
}
596637
}

‎cvat/apps/dataset_manager/tests/test_formats.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,11 @@ def test_export_formats_query(self):
309309
'KITTI 1.0',
310310
'LFW 1.0',
311311
'Cityscapes 1.0',
312-
'Open Images V6 1.0'
312+
'Open Images V6 1.0',
313+
'YOLOv8 Oriented Bounding Boxes 1.0',
314+
'YOLOv8 Detection 1.0',
315+
'YOLOv8 Pose 1.0',
316+
'YOLOv8 Segmentation 1.0',
313317
})
314318

315319
def test_import_formats_query(self):
@@ -342,6 +346,10 @@ def test_import_formats_query(self):
342346
'Open Images V6 1.0',
343347
'Datumaro 1.0',
344348
'Datumaro 3D 1.0',
349+
'YOLOv8 Oriented Bounding Boxes 1.0',
350+
'YOLOv8 Detection 1.0',
351+
'YOLOv8 Pose 1.0',
352+
'YOLOv8 Segmentation 1.0',
345353
})
346354

347355
def test_exports(self):
@@ -391,6 +399,10 @@ def test_empty_images_are_exported(self):
391399
# ('KITTI 1.0', 'kitti') format does not support empty annotations
392400
('LFW 1.0', 'lfw'),
393401
# ('Cityscapes 1.0', 'cityscapes'), does not support, empty annotations
402+
('YOLOv8 Oriented Bounding Boxes 1.0', 'yolov8_oriented_boxes'),
403+
('YOLOv8 Detection 1.0', 'yolov8_detection'),
404+
('YOLOv8 Pose 1.0', 'yolov8_pose'),
405+
('YOLOv8 Segmentation 1.0', 'yolov8_segmentation'),
394406
]:
395407
with self.subTest(format=format_name):
396408
if not dm.formats.registry.EXPORT_FORMATS[format_name].ENABLED:

0 commit comments

Comments
 (0)