diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ba0b70a60..2f2df75422 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 () - Skeleton support in datumaro format () +- Support for YOLOv8 formats + () ### Changed - `env.detect_dataset()` now returns a list of detected formats at all recursion levels diff --git a/datumaro/plugins/yolo_format/converter.py b/datumaro/plugins/yolo_format/converter.py index 71b759c3ff..a1c9be5294 100644 --- a/datumaro/plugins/yolo_format/converter.py +++ b/datumaro/plugins/yolo_format/converter.py @@ -1,21 +1,36 @@ # Copyright (C) 2019-2022 Intel Corporation +# Copyright (C) 2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT import logging as log +import math import os import os.path as osp from collections import OrderedDict - -from datumaro.components.annotation import AnnotationType, Bbox +from functools import cached_property +from itertools import cycle +from typing import Dict, List, Optional + +import yaml + +from datumaro.components.annotation import ( + Annotation, + AnnotationType, + Bbox, + Points, + PointsCategories, + Polygon, + Skeleton, +) from datumaro.components.converter import Converter -from datumaro.components.dataset import ItemStatus +from datumaro.components.dataset import DatasetPatch, ItemStatus from datumaro.components.errors import DatasetExportError, MediaTypeError from datumaro.components.extractor import DEFAULT_SUBSET_NAME, DatasetItem, IExtractor from datumaro.components.media import Image from datumaro.util import str_to_bool -from .format import YoloPath +from .format import YoloPath, YOLOv8Path def _make_yolo_bbox(img_size, box): @@ -29,9 +44,37 @@ def _make_yolo_bbox(img_size, box): return x, y, w, h +def _bbox_annotation_as_polygon(bbox: Bbox) -> List[float]: + points = bbox.as_polygon() + + def rotate_point(x: float, y: float): + new_x = ( + center_x + + math.cos(rotation_radians) * (x - center_x) + - math.sin(rotation_radians) * (y - center_y) + ) + new_y = ( + center_y + + math.sin(rotation_radians) * (x - center_x) + + math.cos(rotation_radians) * (y - center_y) + ) + return new_x, new_y + + if rotation_radians := math.radians(bbox.attributes.get("rotation", 0)): + center_x = bbox.x + bbox.w / 2 + center_y = bbox.y + bbox.h / 2 + points = [ + coordinate + for x, y in zip(points[::2], points[1::2]) + for coordinate in rotate_point(x, y) + ] + return points + + class YoloConverter(Converter): # https://github.com/AlexeyAB/darknet#how-to-train-to-detect-your-custom-objects DEFAULT_IMAGE_EXT = ".jpg" + RESERVED_CONFIG_KEYS = YoloPath.RESERVED_CONFIG_KEYS @classmethod def build_cmdline_parser(cls, **kwargs): @@ -52,7 +95,6 @@ def __init__( self._prefix = "data" if add_path_prefix else "" def apply(self): - extractor = self._extractor save_dir = self._save_dir if self._extractor.media_type() and not issubclass(self._extractor.media_type(), Image): @@ -63,11 +105,6 @@ def apply(self): if self._save_dataset_meta: self._save_meta_file(self._save_dir) - label_categories = extractor.categories()[AnnotationType.label] - label_ids = {label.name: idx for idx, label in enumerate(label_categories.items)} - with open(osp.join(save_dir, "obj.names"), "w", encoding="utf-8") as f: - f.writelines("%s\n" % l[0] for l in sorted(label_ids.items(), key=lambda x: x[1])) - subset_lists = OrderedDict() subsets = self._extractor.subsets() @@ -75,50 +112,42 @@ def apply(self): for (subset_name, subset), pbar in zip(subsets.items(), pbars): if not subset_name or subset_name == DEFAULT_SUBSET_NAME: subset_name = YoloPath.DEFAULT_SUBSET_NAME - elif subset_name in YoloPath.RESERVED_CONFIG_KEYS: + elif subset_name in self.RESERVED_CONFIG_KEYS: raise DatasetExportError( f"Can't export '{subset_name}' subset in YOLO format, this word is reserved." ) - subset_dir = osp.join(save_dir, "obj_%s_data" % subset_name) - os.makedirs(subset_dir, exist_ok=True) + subset_image_dir = self._make_image_subset_folder(save_dir, subset_name) + subset_anno_dir = self._make_annotation_subset_folder(save_dir, subset_name) + os.makedirs(subset_image_dir, exist_ok=True) + os.makedirs(subset_anno_dir, exist_ok=True) image_paths = OrderedDict() for item in pbar.iter(subset, desc=f"Exporting '{subset_name}'"): try: - if not item.media or not (item.media.has_data or item.media.has_size): - raise Exception( - "Failed to export item '%s': " "item has no image info" % item.id - ) - - image_name = self._make_image_filename(item) - if self._save_media: - if item.media: - self._save_image(item, osp.join(subset_dir, image_name)) - else: - log.warning("Item '%s' has no image" % item.id) + image_fpath = self._export_media(item, subset_image_dir) + image_name = osp.relpath(image_fpath, subset_image_dir) image_paths[item.id] = osp.join( - self._prefix, osp.basename(subset_dir), image_name + self._prefix, osp.relpath(subset_image_dir, save_dir), image_name ) - yolo_annotation = self._export_item_annotation(item) - annotation_path = osp.join(subset_dir, "%s.txt" % item.id) - os.makedirs(osp.dirname(annotation_path), exist_ok=True) - with open(annotation_path, "w", encoding="utf-8") as f: - f.write(yolo_annotation) + self._export_item_annotation(item, subset_anno_dir) + except Exception as e: self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset)) - subset_list_name = f"{subset_name}.txt" - subset_list_path = osp.join(save_dir, subset_list_name) - if self._patch and subset_name in self._patch.updated_subsets and not image_paths: - if osp.isfile(subset_list_path): - os.remove(subset_list_path) - continue + if subset_list_name := self._make_subset_list_file(subset_name, image_paths): + subset_lists[subset_name] = subset_list_name - subset_lists[subset_name] = subset_list_name - with open(subset_list_path, "w", encoding="utf-8") as f: - f.writelines("%s\n" % s.replace("\\", "/") for s in image_paths.values()) + self._save_config_files(subset_lists) + + def _save_config_files(self, subset_lists: Dict[str, str]): + extractor = self._extractor + save_dir = self._save_dir + label_categories = extractor.categories()[AnnotationType.label] + label_ids = {label.name: idx for idx, label in enumerate(label_categories.items)} + with open(osp.join(save_dir, "obj.names"), "w", encoding="utf-8") as f: + f.writelines("%s\n" % l[0] for l in sorted(label_ids.items(), key=lambda x: x[1])) with open(osp.join(save_dir, "obj.data"), "w", encoding="utf-8") as f: f.write(f"classes = {len(label_ids)}\n") @@ -132,23 +161,79 @@ def apply(self): f.write("names = %s\n" % osp.join(self._prefix, "obj.names")) f.write("backup = backup/\n") - def _export_item_annotation(self, item): - height, width = item.media.size + def _make_subset_list_file(self, subset_name, image_paths): + subset_list_name = f"{subset_name}{YoloPath.SUBSET_LIST_EXT}" + subset_list_path = osp.join(self._save_dir, subset_list_name) + if self._patch and subset_name in self._patch.updated_subsets and not image_paths: + if osp.isfile(subset_list_path): + os.remove(subset_list_path) + return + + with open(subset_list_path, "w", encoding="utf-8") as f: + f.writelines("%s\n" % s.replace("\\", "/") for s in image_paths.values()) + return subset_list_name + + def _export_media(self, item: DatasetItem, subset_img_dir: str) -> str: + try: + if not item.media or not (item.media.has_data or item.media.has_size): + raise DatasetExportError( + "Failed to export item '%s': " "item has no image info" % item.id + ) - yolo_annotation = "" + image_name = self._make_image_filename(item) + image_fpath = osp.join(subset_img_dir, image_name) - for bbox in item.annotations: - if not isinstance(bbox, Bbox) or bbox.label is None: - continue + if self._save_media: + if item.media: + self._save_image(item, image_fpath) + else: + log.warning("Item '%s' has no image" % item.id) + + return image_fpath + + except Exception as e: + self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset)) + + def _export_item_annotation(self, item: DatasetItem, subset_dir: str) -> None: + try: + height, width = item.media.size + + yolo_annotation = "" - yolo_bb = _make_yolo_bbox((width, height), bbox.points) - yolo_bb = " ".join("%.6f" % p for p in yolo_bb) - yolo_annotation += "%s %s\n" % (bbox.label, yolo_bb) + for bbox in item.annotations: + annotation_line = self._make_annotation_line(width, height, bbox) + if annotation_line: + yolo_annotation += annotation_line - return yolo_annotation + annotation_path = osp.join(subset_dir, f"{item.id}{YoloPath.LABELS_EXT}") + os.makedirs(osp.dirname(annotation_path), exist_ok=True) + + with open(annotation_path, "w", encoding="utf-8") as f: + f.write(yolo_annotation) + + except Exception as e: + self._ctx.error_policy.report_item_error(e, item_id=(item.id, item.subset)) + + def _make_annotation_line(self, width: int, height: int, anno: Annotation) -> Optional[str]: + if not isinstance(anno, Bbox) or anno.label is None: + return + if anno.attributes.get("rotation", 0) % 360.0 > 0.00001: + return + + values = _make_yolo_bbox((width, height), anno.points) + string_values = " ".join("%.6f" % p for p in values) + return "%s %s\n" % (anno.label, string_values) + + @staticmethod + def _make_image_subset_folder(save_dir: str, subset: str) -> str: + return osp.join(save_dir, f"obj_{subset}_data") + + @staticmethod + def _make_annotation_subset_folder(save_dir: str, subset: str) -> str: + return osp.join(save_dir, f"obj_{subset}_data") @classmethod - def patch(cls, dataset, patch, save_dir, **kwargs): + def patch(cls, dataset: IExtractor, patch: DatasetPatch, save_dir: str, **kwargs): conv = cls(dataset, save_dir=save_dir, **kwargs) conv._patch = patch conv.apply() @@ -164,12 +249,148 @@ def patch(cls, dataset, patch, save_dir, **kwargs): if subset == DEFAULT_SUBSET_NAME: subset = YoloPath.DEFAULT_SUBSET_NAME - subset_dir = osp.join(save_dir, "obj_%s_data" % subset) + subset_image_dir = cls._make_image_subset_folder(save_dir, subset) + subset_anno_dir = cls._make_annotation_subset_folder(save_dir, subset) - image_path = osp.join(subset_dir, conv._make_image_filename(item)) + image_path = osp.join(subset_image_dir, conv._make_image_filename(item)) if osp.isfile(image_path): os.remove(image_path) - ann_path = osp.join(subset_dir, "%s.txt" % item.id) + ann_path = osp.join(subset_anno_dir, f"{item.id}{YoloPath.LABELS_EXT}") if osp.isfile(ann_path): os.remove(ann_path) + + +class YOLOv8Converter(YoloConverter): + RESERVED_CONFIG_KEYS = YOLOv8Path.RESERVED_CONFIG_KEYS + + def __init__( + self, + extractor: IExtractor, + save_dir: str, + *, + add_path_prefix: bool = True, + config_file=None, + **kwargs, + ) -> None: + super().__init__(extractor, save_dir, add_path_prefix=add_path_prefix, **kwargs) + self._config_filename = config_file or YOLOv8Path.DEFAULT_CONFIG_FILE + + @classmethod + def build_cmdline_parser(cls, **kwargs): + parser = super().build_cmdline_parser(**kwargs) + parser.add_argument( + "--config-file", + default=YOLOv8Path.DEFAULT_CONFIG_FILE, + type=str, + help="config file name (default: %(default)s)", + ) + return parser + + def _save_config_files(self, subset_lists: Dict[str, str]): + extractor = self._extractor + save_dir = self._save_dir + with open(osp.join(save_dir, self._config_filename), "w", encoding="utf-8") as f: + label_categories = extractor.categories()[AnnotationType.label] + data = dict( + path=".", + names={idx: label.name for idx, label in enumerate(label_categories.items)}, + **subset_lists, + ) + yaml.dump(data, f) + + @staticmethod + def _make_image_subset_folder(save_dir: str, subset: str) -> str: + return osp.join(save_dir, YOLOv8Path.IMAGES_FOLDER_NAME, subset) + + @staticmethod + def _make_annotation_subset_folder(save_dir: str, subset: str) -> str: + return osp.join(save_dir, YOLOv8Path.LABELS_FOLDER_NAME, subset) + + +class YOLOv8SegmentationConverter(YOLOv8Converter): + def _make_annotation_line(self, width: int, height: int, anno: Annotation) -> Optional[str]: + if anno.label is None or not isinstance(anno, Polygon): + return + values = [value / size for value, size in zip(anno.points, cycle((width, height)))] + string_values = " ".join("%.6f" % p for p in values) + return "%s %s\n" % (anno.label, string_values) + + +class YOLOv8OrientedBoxesConverter(YOLOv8Converter): + def _make_annotation_line(self, width: int, height: int, anno: Annotation) -> Optional[str]: + if anno.label is None or not isinstance(anno, Bbox): + return + points = _bbox_annotation_as_polygon(anno) + values = [value / size for value, size in zip(points, cycle((width, height)))] + string_values = " ".join("%.6f" % p for p in values) + return "%s %s\n" % (anno.label, string_values) + + +class YOLOv8PoseConverter(YOLOv8Converter): + @cached_property + def _map_labels_for_save(self): + point_categories = self._extractor.categories().get( + AnnotationType.points, PointsCategories.from_iterable([]) + ) + return {label_id: index for index, label_id in enumerate(sorted(point_categories.items))} + + def _save_config_files(self, subset_lists: Dict[str, str]): + extractor = self._extractor + save_dir = self._save_dir + + point_categories = extractor.categories().get( + AnnotationType.points, PointsCategories.from_iterable([]) + ) + if len(set(len(cat.labels) for cat in point_categories.items.values())) > 1: + raise DatasetExportError( + "Can't export: skeletons should have the same number of points" + ) + n_of_points = ( + len(next(iter(point_categories.items.values())).labels) + if len(point_categories) > 0 + else 0 + ) + + with open(osp.join(save_dir, self._config_filename), "w", encoding="utf-8") as f: + label_categories = extractor.categories()[AnnotationType.label] + parent_categories = { + self._map_labels_for_save[label_id]: label_categories.items[label_id].name + for label_id in point_categories.items + } + assert set(parent_categories.keys()) == set(range(len(parent_categories))) + data = dict( + path=".", + names=parent_categories, + kpt_shape=[n_of_points, 3], + **subset_lists, + ) + yaml.dump(data, f) + + def _make_annotation_line(self, width: int, height: int, skeleton: Annotation) -> Optional[str]: + if skeleton.label is None or not isinstance(skeleton, Skeleton): + return + + x, y, w, h = skeleton.get_bbox() + bbox_values = _make_yolo_bbox((width, height), [x, y, x + w, y + h]) + bbox_string_values = " ".join("%.6f" % p for p in bbox_values) + + point_label_ids = [ + self._extractor.categories()[AnnotationType.label].find( + name=child_label, + parent=self._extractor.categories()[AnnotationType.label][skeleton.label].name, + )[0] + for child_label in self._extractor.categories()[AnnotationType.points] + .items[skeleton.label] + .labels + ] + + points_values = [f"0.0, 0.0, {Points.Visibility.absent.value}"] * len(point_label_ids) + for element in skeleton.elements: + assert len(element.points) == 2 and len(element.visibility) == 1 + position = point_label_ids.index(element.label) + x = element.points[0] / width + y = element.points[1] / height + points_values[position] = f"{x:.6f} {y:.6f} {element.visibility[0].value}" + + return f"{self._map_labels_for_save[skeleton.label]} {bbox_string_values} {' '.join(points_values)}\n" diff --git a/datumaro/plugins/yolo_format/extractor.py b/datumaro/plugins/yolo_format/extractor.py index dd07324b79..817260c842 100644 --- a/datumaro/plugins/yolo_format/extractor.py +++ b/datumaro/plugins/yolo_format/extractor.py @@ -1,23 +1,37 @@ # Copyright (C) 2019-2022 Intel Corporation -# Copyright (C) 2023 CVAT.ai Corporation +# Copyright (C) 2023-2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT from __future__ import annotations +import math +import os import os.path as osp import re from collections import OrderedDict -from typing import Dict, List, Optional, Tuple, Type, TypeVar, Union - -from datumaro.components.annotation import Annotation, AnnotationType, Bbox, LabelCategories +from functools import cached_property +from itertools import cycle +from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union + +import yaml + +from datumaro.components.annotation import ( + Annotation, + AnnotationType, + Bbox, + LabelCategories, + Points, + PointsCategories, + Polygon, + Skeleton, +) from datumaro.components.errors import ( DatasetImportError, InvalidAnnotationError, UndeclaredLabelError, ) -from datumaro.components.extractor import DatasetItem, Extractor, Importer, SourceExtractor -from datumaro.components.format_detection import FormatDetectionContext +from datumaro.components.extractor import DatasetItem, Extractor, SourceExtractor from datumaro.components.media import Image from datumaro.util.image import ( DEFAULT_IMAGE_META_FILE_NAME, @@ -25,15 +39,18 @@ load_image, load_image_meta_file, ) -from datumaro.util.meta_file_util import has_meta_file, parse_meta_file +from datumaro.util.meta_file_util import get_meta_file, has_meta_file, parse_meta_file from datumaro.util.os_util import split_path -from .format import YoloPath +from ...util import parse_json_file, take_by +from .format import YoloPath, YOLOv8Path, YOLOv8PoseFormat T = TypeVar("T") class YoloExtractor(SourceExtractor): + RESERVED_CONFIG_KEYS = YoloPath.RESERVED_CONFIG_KEYS + class Subset(Extractor): def __init__(self, name: str, parent: YoloExtractor): super().__init__() @@ -65,6 +82,7 @@ def __init__( super().__init__(**kwargs) rootpath = osp.dirname(config_path) + self._config_path = config_path self._path = rootpath assert image_info is None or isinstance(image_info, (str, dict)) @@ -77,11 +95,7 @@ def __init__( self._image_info = image_info - config = self._parse_config(config_path) - - names_path = config.get("names") - if not names_path: - raise InvalidAnnotationError(f"Failed to parse names file path from config") + self._categories = self._load_categories() # The original format is like this: # @@ -93,31 +107,29 @@ def __init__( # # To support more subset names, we disallow subsets # called 'classes', 'names' and 'backup'. - subsets = {k: v for k, v in config.items() if k not in YoloPath.RESERVED_CONFIG_KEYS} + subsets = {k: v for k, v in self._config.items() if k not in self.RESERVED_CONFIG_KEYS} for subset_name, list_path in subsets.items(): - list_path = osp.join(self._path, self.localize_path(list_path)) - if not osp.isfile(list_path): - raise InvalidAnnotationError(f"Can't find '{subset_name}' subset list file") - subset = YoloExtractor.Subset(subset_name, self) - with open(list_path, "r", encoding="utf-8") as f: - subset.items = OrderedDict( - (self.name_from_path(p), self.localize_path(p)) for p in f if p.strip() - ) + subset.items = OrderedDict( + (self.name_from_path(p), self.localize_path(p)) + for p in self._iterate_over_image_paths(subset_name, list_path) + ) subsets[subset_name] = subset self._subsets: Dict[str, YoloExtractor.Subset] = subsets - self._categories = { - AnnotationType.label: self._load_categories( - osp.join(self._path, self.localize_path(names_path)) - ) - } + def _iterate_over_image_paths(self, subset_name: str, list_path: str): + list_path = osp.join(self._path, self.localize_path(list_path)) + if not osp.isfile(list_path): + raise InvalidAnnotationError(f"Can't find '{subset_name}' subset list file") - @staticmethod - def _parse_config(path: str) -> Dict[str, str]: - with open(path, "r", encoding="utf-8") as f: + with open(list_path, "r", encoding="utf-8") as f: + yield from (p for p in f if p.strip()) + + @cached_property + def _config(self) -> Dict[str, str]: + with open(self._config_path, "r", encoding="utf-8") as f: config_lines = f.readlines() config = {} @@ -166,6 +178,9 @@ def name_from_path(cls, path: str) -> str: def _image_loader(cls, *args, **kwargs): return load_image(*args, **kwargs, keep_exif=True) + def _get_labels_path_from_image_path(self, image_path: str) -> str: + return osp.splitext(image_path)[0] + YoloPath.LABELS_EXT + def _get(self, item_id: str, subset_name: str) -> Optional[DatasetItem]: subset = self._subsets[subset_name] item = subset.items[item_id] @@ -180,7 +195,7 @@ def _get(self, item_id: str, subset_name: str) -> Optional[DatasetItem]: else: image = Image(path=image_path, data=self._image_loader) - anno_path = osp.splitext(image.path)[0] + ".txt" + anno_path = self._get_labels_path_from_image_path(image.path) annotations = self._parse_annotations( anno_path, image, item_id=(item_id, subset_name) ) @@ -227,41 +242,54 @@ def _parse_annotations( for line in lines: try: - parts = line.split() - if len(parts) != 5: - raise InvalidAnnotationError( - f"Unexpected field count {len(parts)} in the bbox description. " - "Expected 5 fields (label, xc, yc, w, h)." - ) - label_id, xc, yc, w, h = parts - - label_id = self._parse_field(label_id, int, "bbox label id") - if label_id not in self._categories[AnnotationType.label]: - raise UndeclaredLabelError(str(label_id)) - - w = self._parse_field(w, float, "bbox width") - h = self._parse_field(h, float, "bbox height") - x = self._parse_field(xc, float, "bbox center x") - w * 0.5 - y = self._parse_field(yc, float, "bbox center y") - h * 0.5 - annotations.append( - Bbox( - x * image_width, - y * image_height, - w * image_width, - h * image_height, - label=label_id, - ) + self._load_one_annotation(line.split(), image_height, image_width) ) except Exception as e: self._ctx.error_policy.report_annotation_error(e, item_id=item_id) return annotations - @staticmethod - def _load_categories(names_path): + def _load_one_annotation( + self, parts: List[str], image_height: int, image_width: int + ) -> Annotation: + if len(parts) != 5: + raise InvalidAnnotationError( + f"Unexpected field count {len(parts)} in the bbox description. " + "Expected 5 fields (label, xc, yc, w, h)." + ) + label_id, xc, yc, w, h = parts + + label_id = self._parse_field(label_id, int, "bbox label id") + if label_id not in self._categories[AnnotationType.label]: + raise UndeclaredLabelError(str(label_id)) + + w = self._parse_field(w, float, "bbox width") + h = self._parse_field(h, float, "bbox height") + x = self._parse_field(xc, float, "bbox center x") - w * 0.5 + y = self._parse_field(yc, float, "bbox center y") - h * 0.5 + + return Bbox( + x * image_width, + y * image_height, + w * image_width, + h * image_height, + label=label_id, + ) + + def _load_categories(self) -> Dict[AnnotationType, LabelCategories]: + names_path = self._config.get("names") + if not names_path: + raise InvalidAnnotationError(f"Failed to parse names file path from config") + + names_path = osp.join(self._path, self.localize_path(names_path)) + if has_meta_file(osp.dirname(names_path)): - return LabelCategories.from_iterable(parse_meta_file(osp.dirname(names_path)).keys()) + return { + AnnotationType.label: LabelCategories.from_iterable( + parse_meta_file(osp.dirname(names_path)).keys() + ) + } label_categories = LabelCategories() @@ -271,7 +299,7 @@ def _load_categories(names_path): if label: label_categories.add(label) - return label_categories + return {AnnotationType.label: label_categories} def __iter__(self): subsets = self._subsets @@ -287,11 +315,321 @@ def get_subset(self, name): return self._subsets[name] -class YoloImporter(Importer): - @classmethod - def detect(cls, context: FormatDetectionContext) -> None: - context.require_file("obj.data") +class YOLOv8Extractor(YoloExtractor): + RESERVED_CONFIG_KEYS = YOLOv8Path.RESERVED_CONFIG_KEYS + + def __init__( + self, + *args, + config_file=None, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + + @cached_property + def _config(self) -> Dict[str, Any]: + with open(self._config_path) as stream: + try: + return yaml.safe_load(stream) + except yaml.YAMLError: + raise InvalidAnnotationError("Failed to parse config file") + + def _load_categories(self) -> Dict[AnnotationType, LabelCategories]: + if has_meta_file(self._path): + return { + AnnotationType.label: LabelCategories.from_iterable( + parse_meta_file(self._path).keys() + ) + } + + if (names := self._config.get("names")) is not None: + if isinstance(names, dict): + return { + AnnotationType.label: LabelCategories.from_iterable( + [names[i] for i in range(len(names))] + ) + } + if isinstance(names, list): + return {AnnotationType.label: LabelCategories.from_iterable(names)} + + raise InvalidAnnotationError(f"Failed to parse names from config") + + def _get_labels_path_from_image_path(self, image_path: str) -> str: + relative_image_path = osp.relpath( + image_path, osp.join(self._path, YOLOv8Path.IMAGES_FOLDER_NAME) + ) + relative_labels_path = osp.splitext(relative_image_path)[0] + YOLOv8Path.LABELS_EXT + return osp.join(self._path, YOLOv8Path.LABELS_FOLDER_NAME, relative_labels_path) @classmethod - def find_sources(cls, path): - return cls._find_sources_recursive(path, ".data", "yolo") + def name_from_path(cls, path: str) -> str: + """ + Obtains from the path like [data/]images//.ext + + can be , so it is + more involved than just calling "basename()". + """ + path = cls.localize_path(path) + + parts = split_path(path) + if 2 < len(parts) and not osp.isabs(path): + path = osp.join(*parts[2:]) # pylint: disable=no-value-for-parameter + return osp.splitext(path)[0] + + def _iterate_over_image_paths( + self, subset_name: str, subset_images_source: Union[str, List[str]] + ): + if isinstance(subset_images_source, str): + if subset_images_source.endswith(YoloPath.SUBSET_LIST_EXT): + yield from super()._iterate_over_image_paths(subset_name, subset_images_source) + else: + path = osp.join(self._path, self.localize_path(subset_images_source)) + if not osp.isdir(path): + raise InvalidAnnotationError(f"Can't find '{subset_name}' subset image folder") + yield from ( + osp.relpath(osp.join(root, file), self._path) + for root, dirs, files in os.walk(path) + for file in files + if osp.isfile(osp.join(root, file)) + ) + else: + yield from subset_images_source + + +class YOLOv8SegmentationExtractor(YOLOv8Extractor): + def _load_segmentation_annotation( + self, parts: List[str], image_height: int, image_width: int + ) -> Polygon: + label_id = self._parse_field(parts[0], int, "Polygon label id") + if label_id not in self._categories[AnnotationType.label]: + raise UndeclaredLabelError(str(label_id)) + + points = [ + self._parse_field( + value, float, f"polygon point {idx // 2} {'x' if idx % 2 == 0 else 'y'}" + ) + for idx, value in enumerate(parts[1:]) + ] + scaled_points = [ + value * size for value, size in zip(points, cycle((image_width, image_height))) + ] + return Polygon(scaled_points, label=label_id) + + def _load_one_annotation( + self, parts: List[str], image_height: int, image_width: int + ) -> Annotation: + if len(parts) > 5 and len(parts) % 2 == 1: + return self._load_segmentation_annotation(parts, image_height, image_width) + raise InvalidAnnotationError( + f"Unexpected field count {len(parts)} in the polygon description. " + "Expected odd number > 5 of fields for segment annotation (label, x1, y1, x2, y2, x3, y3, ...)" + ) + + +class YOLOv8OrientedBoxesExtractor(YOLOv8Extractor): + @staticmethod + def _check_is_rectangle(p1, p2, p3, p4): + p12_angle = math.atan2(p2[0] - p1[0], p2[1] - p1[1]) + p23_angle = math.atan2(p3[0] - p2[0], p3[1] - p2[1]) + p43_angle = math.atan2(p3[0] - p4[0], p3[1] - p4[1]) + p14_angle = math.atan2(p4[0] - p1[0], p4[1] - p1[1]) + + if abs(p12_angle - p43_angle) > 0.001 or abs(p23_angle - p14_angle) > 0.001: + raise InvalidAnnotationError( + "Given points do not form a rectangle: opposite sides have different slope angles." + ) + + def _load_one_annotation( + self, parts: List[str], image_height: int, image_width: int + ) -> Annotation: + if len(parts) != 9: + raise InvalidAnnotationError( + f"Unexpected field count {len(parts)} in the bbox description. " + "Expected 9 fields (label, x1, y1, x2, y2, x3, y3, x4, y4)." + ) + label_id = self._parse_field(parts[0], int, "bbox label id") + if label_id not in self._categories[AnnotationType.label]: + raise UndeclaredLabelError(str(label_id)) + points = [ + ( + self._parse_field(x, float, f"bbox point {idx} x") * image_width, + self._parse_field(y, float, f"bbox point {idx} y") * image_height, + ) + for idx, (x, y) in enumerate(take_by(parts[1:], 2)) + ] + self._check_is_rectangle(*points) + + (x1, y1), (x2, y2), (x3, y3), (x4, y4) = points + + width = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2) + height = math.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2) + rotation = math.atan2(y2 - y1, x2 - x1) + if rotation < 0: + rotation += math.pi * 2 + + center_x = (x1 + x2 + x3 + x4) / 4 + center_y = (y1 + y2 + y3 + y4) / 4 + + return Bbox( + x=center_x - width / 2, + y=center_y - height / 2, + w=width, + h=height, + label=label_id, + attributes=(dict(rotation=math.degrees(rotation)) if abs(rotation) > 0.00001 else {}), + ) + + +class YOLOv8PoseExtractor(YOLOv8Extractor): + @cached_property + def _kpt_shape(self): + if YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME not in self._config: + raise InvalidAnnotationError( + f"Failed to parse {YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME} from config" + ) + kpt_shape = self._config[YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME] + if not isinstance(kpt_shape, list) or len(kpt_shape) != 2: + raise InvalidAnnotationError( + f"Failed to parse {YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME} from config" + ) + if kpt_shape[1] not in [2, 3]: + raise InvalidAnnotationError( + f"Unexpected values per point {kpt_shape[1]} in field" + f"{YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME}. Expected 2 or 3." + ) + if not isinstance(kpt_shape[0], int) or kpt_shape[0] < 0: + raise InvalidAnnotationError( + f"Unexpected number of points {kpt_shape[0]} in field " + f"{YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME}. Expected non-negative integer." + ) + + return kpt_shape + + @cached_property + def _skeleton_id_to_label_id(self): + point_categories = self._categories.get( + AnnotationType.points, PointsCategories.from_iterable([]) + ) + return {index: label_id for index, label_id in enumerate(sorted(point_categories.items))} + + def _load_categories(self) -> Dict[AnnotationType, LabelCategories]: + if "names" not in self._config: + raise InvalidAnnotationError(f"Failed to parse names from config") + + if has_meta_file(self._path): + dataset_meta = parse_json_file(get_meta_file(self._path)) + point_categories = PointsCategories.from_iterable( + dataset_meta.get("point_categories", []) + ) + categories = { + AnnotationType.label: LabelCategories.from_iterable( + dataset_meta["label_categories"] + ) + } + if len(point_categories) > 0: + categories[AnnotationType.points] = point_categories + return categories + + number_of_points, _ = self._kpt_shape + names = self._config["names"] + if isinstance(names, dict): + if set(names.keys()) != set(range(len(names))): + raise InvalidAnnotationError( + f"Failed to parse names from config: non-sequential label ids" + ) + skeleton_labels = [names[i] for i in range(len(names))] + elif isinstance(names, list): + skeleton_labels = names + else: + raise InvalidAnnotationError(f"Failed to parse names from config") + + def make_children_names(skeleton_label): + return [ + f"{skeleton_label}_point_{point_index}" for point_index in range(number_of_points) + ] + + point_labels = [ + (child_name, skeleton_label) + for skeleton_label in skeleton_labels + for child_name in make_children_names(skeleton_label) + ] + + point_categories = PointsCategories.from_iterable( + [ + ( + index, + make_children_names(skeleton_label), + set(), + ) + for index, skeleton_label in enumerate(skeleton_labels) + ] + ) + categories = { + AnnotationType.label: LabelCategories.from_iterable(skeleton_labels + point_labels) + } + if len(point_categories) > 0: + categories[AnnotationType.points] = point_categories + + return categories + + def _load_one_annotation( + self, parts: List[str], image_height: int, image_width: int + ) -> Annotation: + number_of_points, values_per_point = self._kpt_shape + if len(parts) != 5 + number_of_points * values_per_point: + raise InvalidAnnotationError( + f"Unexpected field count {len(parts)} in the skeleton description. " + "Expected 5 fields (label, xc, yc, w, h)" + f"and then {values_per_point} for each of {number_of_points} points" + ) + + skeleton_id = self._parse_field(parts[0], int, "skeleton label id") + label_id = self._skeleton_id_to_label_id.get(skeleton_id, -1) + if label_id not in self._categories[AnnotationType.label]: + raise UndeclaredLabelError(str(skeleton_id)) + if self._categories[AnnotationType.label][label_id].parent != "": + raise InvalidAnnotationError("WTF") + + point_labels = self._categories[AnnotationType.points][label_id].labels + point_label_ids = [ + self._categories[AnnotationType.label].find( + name=point_label, + parent=self._categories[AnnotationType.label][label_id].name, + )[0] + for point_label in point_labels + ] + + points = [ + Points( + [ + self._parse_field(parts[x_index], float, f"skeleton point {point_index} x") + * image_width, + self._parse_field(parts[y_index], float, f"skeleton point {point_index} y") + * image_height, + ], + ( + [ + self._parse_field( + parts[visibility_index], + int, + f"skeleton point {point_index} visibility", + ), + ] + if values_per_point == 3 + else [Points.Visibility.visible.value] + ), + label=label_id, + ) + for point_index, label_id in enumerate(point_label_ids) + for x_index, y_index, visibility_index in [ + ( + 5 + point_index * values_per_point, + 5 + point_index * values_per_point + 1, + 5 + point_index * values_per_point + 2, + ), + ] + ] + return Skeleton( + points, + label=label_id, + ) diff --git a/datumaro/plugins/yolo_format/format.py b/datumaro/plugins/yolo_format/format.py index ed80025b21..a1f7b1563b 100644 --- a/datumaro/plugins/yolo_format/format.py +++ b/datumaro/plugins/yolo_format/format.py @@ -1,4 +1,5 @@ # Copyright (C) 2019-2021 Intel Corporation +# Copyright (C) 2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT @@ -7,3 +8,21 @@ class YoloPath: DEFAULT_SUBSET_NAME = "train" SUBSET_NAMES = ["train", "valid"] RESERVED_CONFIG_KEYS = ["backup", "classes", "names"] + LABELS_EXT = ".txt" + SUBSET_LIST_EXT = ".txt" + + +class YOLOv8Path(YoloPath): + CONFIG_FILE_EXT = ".yaml" + DEFAULT_CONFIG_FILE = "data.yaml" + RESERVED_CONFIG_KEYS = YoloPath.RESERVED_CONFIG_KEYS + [ + "path", + "kpt_shape", + "flip_idx", + ] + IMAGES_FOLDER_NAME = "images" + LABELS_FOLDER_NAME = "labels" + + +class YOLOv8PoseFormat: + KPT_SHAPE_FIELD_NAME = "kpt_shape" diff --git a/datumaro/plugins/yolo_format/importer.py b/datumaro/plugins/yolo_format/importer.py new file mode 100644 index 0000000000..4a6e50e23e --- /dev/null +++ b/datumaro/plugins/yolo_format/importer.py @@ -0,0 +1,109 @@ +# Copyright (C) 2023 Intel Corporation +# Copyright (C) 2024 CVAT.ai Corporation +# +# SPDX-License-Identifier: MIT + +from __future__ import annotations + +from os import path as osp +from typing import Any, Dict, List + +import yaml + +from datumaro import Importer +from datumaro.components.format_detection import FormatDetectionContext +from datumaro.plugins.yolo_format.extractor import ( + YOLOv8Extractor, + YOLOv8OrientedBoxesExtractor, + YOLOv8PoseExtractor, + YOLOv8SegmentationExtractor, +) +from datumaro.plugins.yolo_format.format import YOLOv8Path, YOLOv8PoseFormat + + +class YoloImporter(Importer): + @classmethod + def detect(cls, context: FormatDetectionContext) -> None: + context.require_file("obj.data") + + @classmethod + def find_sources(cls, path) -> List[Dict[str, Any]]: + return cls._find_sources_recursive(path, ".data", "yolo") + + +class YOLOv8Importer(Importer): + EXTRACTOR = YOLOv8Extractor + + @classmethod + def build_cmdline_parser(cls, **kwargs): + parser = super().build_cmdline_parser(**kwargs) + parser.add_argument( + "--config-file", + help="The name of the file to read dataset config from", + ) + return parser + + @classmethod + def _check_config_file(cls, context, config_file): + with context.probe_text_file( + config_file, + f"must not have '{YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME}' field", + ) as f: + try: + config = yaml.safe_load(f) + if YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME in config: + raise Exception + except yaml.YAMLError: + raise Exception + + @classmethod + def detect(cls, context: FormatDetectionContext) -> None: + context.require_file(f"*{YOLOv8Path.CONFIG_FILE_EXT}") + sources = cls.find_sources_with_params(context.root_path) + if not sources or len(sources) > 1: + context.fail("Cannot choose config file") + + cls._check_config_file(context, osp.relpath(sources[0]["url"], context.root_path)) + + @classmethod + def find_sources_with_params( + cls, path, config_file=None, **extra_params + ) -> List[Dict[str, Any]]: + sources = cls._find_sources_recursive( + path, YOLOv8Path.CONFIG_FILE_EXT, cls.EXTRACTOR.NAME, max_depth=1 + ) + + if config_file: + return [source for source in sources if source["url"] == osp.join(path, config_file)] + if len(sources) <= 1: + return sources + return [ + source + for source in sources + if source["url"] == osp.join(path, YOLOv8Path.DEFAULT_CONFIG_FILE) + ] + + +class YOLOv8SegmentationImporter(YOLOv8Importer): + EXTRACTOR = YOLOv8SegmentationExtractor + + +class YOLOv8OrientedBoxesImporter(YOLOv8Importer): + EXTRACTOR = YOLOv8OrientedBoxesExtractor + + +class YOLOv8PoseImporter(YOLOv8Importer): + EXTRACTOR = YOLOv8PoseExtractor + + @classmethod + def _check_config_file(cls, context, config_file): + with context.probe_text_file( + config_file, + f"must have '{YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME}' field", + ) as f: + try: + config = yaml.safe_load(f) + if YOLOv8PoseFormat.KPT_SHAPE_FIELD_NAME not in config: + raise Exception + except yaml.YAMLError: + raise Exception diff --git a/datumaro/util/meta_file_util.py b/datumaro/util/meta_file_util.py index b10effdd2e..8e0aba99f6 100644 --- a/datumaro/util/meta_file_util.py +++ b/datumaro/util/meta_file_util.py @@ -1,4 +1,5 @@ # Copyright (C) 2022 Intel Corporation +# Copyright (C) 2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT @@ -50,6 +51,14 @@ def save_meta_file(path, categories): labels = [label.name for label in categories[AnnotationType.label]] dataset_meta["labels"] = labels + dataset_meta["label_categories"] = [ + (label.name, label.parent) for label in categories[AnnotationType.label] + ] + if categories.get(AnnotationType.points): + dataset_meta["point_categories"] = [ + (label_id, cat.labels, list(cat.joints)) + for label_id, cat in categories[AnnotationType.points].items.items() + ] if categories.get(AnnotationType.mask): label_map = {} diff --git a/site/content/en/docs/formats/yolo_v8.md b/site/content/en/docs/formats/yolo_v8.md new file mode 100644 index 0000000000..addb960b17 --- /dev/null +++ b/site/content/en/docs/formats/yolo_v8.md @@ -0,0 +1,345 @@ +--- +title: 'YOLOv8' +linkTitle: 'YOLOv8' +description: '' +--- + +## Format specification + +The YOLOv8 format family allows you to define the dataset root directory, the relative paths +to training/validation/testing image directories or *.txt files containing image paths, +and a dictionary of class names. + +Family consists of four formats: +- [Detection](https://docs.ultralytics.com/datasets/detect/) +- [Oriented bounding Box](https://docs.ultralytics.com/datasets/obb/) +- [Segmentation](https://docs.ultralytics.com/datasets/segment/) +- [Pose](https://docs.ultralytics.com/datasets/pose/) + +Supported annotation types and formats: +- `Bbox` + - Detection (only not rotated) + - Oriented Bounding Box, +- `Polygon` + - Segmentation +- `Skeleton` + - Pose + +The format supports arbitrary subset names, except `classes`, `names`, `backup`, `path`, `kpt_shape`, `flip_idx`. + +> Note, that by default, the YOLO framework does not expect any subset names, + except `train` and `val`, Datumaro supports this as an extension. + If there is no subset separation in a project, the data + will be saved in the `train` subset. + +## Import YOLOv8 dataset +To create a Datumaro project with a YOLOv8 source, use the following commands: + +```bash +datum create +datum import --format yolov8 # for Detection dataset +datum import --format yolov8_oriented_boxes # for Oriented Bounding Box dataset +datum import --format yolov8_segmentation # for Segmentation dataset +datum import --format yolov8_pose # for Pose dataset +``` + +The YOLOv8 dataset directory should have the following structure: + +```bash +└─ yolo_dataset/ + │ # a list of non-format labels (optional) # file with list of classes + ├── data.yaml # file with dataset information + ├── train.txt # list of image paths in train subset [Optional] + ├── val.txt # list of image paths in valid subset [Optional] + │ + ├── images/ + │ ├── train/ # directory with images for train subset + │ │ ├── image1.jpg + │ │ ├── image2.jpg + │ │ ├── image3.jpg + │ │ └── ... + │ ├── valid/ # directory with images for validation subset + │ │ ├── image11.jpg + │ │ ├── image12.jpg + │ │ ├── image13.jpg + │ │ └── ... + ├── labels/ + │ ├── train/ # directory with annotations for train subset + │ │ ├── image1.txt + │ │ ├── image2.txt + │ │ ├── image3.txt + │ │ └── ... + │ ├── valid/ # directory with annotations for validation subset + │ │ ├── image11.txt + │ │ ├── image12.txt + │ │ ├── image13.txt + │ │ └── ... +``` + +`data.yaml` should have the following content: + +```yaml +path: ./ # dataset root dir +train: train.txt # train images (relative to 'path') 4 images +val: val.txt # val images (relative to 'path') 4 images + +# YOLOv8 Pose specific field +# First number is a number of points in skeleton +# Second number defines a format of point info in an annotation txt files +kpt_shape: [17, 3] + +# Classes +names: + 0: person + 1: bicycle + 2: car + # ... + 77: teddy bear + 78: hair drier + 79: toothbrush +``` +> Note, that though by default YOLOv8 framework expects `data.yaml`, + Datumaro allows this file to have arbitrary name. + +`data.yaml` can specify what images a subset contains in 3 ways: +1. subset can point to a folder, in which case all images in the folder will belong to the subset: + ```yaml + val: images/valid + ``` +2. subset can be described as a list of images: + ```yaml + val: + - images/valid/image1.jpg + - images/valid/image2.jpg + ``` +3. subset can point at a `.txt` file which contains a list of images: + ```yaml + val: val.txt + ``` + `val.txt` should have the following structure: + ```txt + + + ... + ``` + +Files in directories `labels/train/` and `labels/valid/` should +contain information about labels for images in `images/train` and `images/valid` respectively. +If there are no objects in an image, no `.txt` file is required. + +Content of the `.txt` file depends on format. + +For **Detection** it contains bounding boxes: +```txt +# labels/train/image1.txt: +# +0 0.250000 0.400000 0.300000 0.400000 +3 0.600000 0.400000 0.400000 0.266667 +... +``` + +For **Oriented Bounding Box** it contains coordinates of four corners of oriented bounding box: +```txt +# labels/train/image1.txt: +# +0 0.146731 0.151795 0.319936 0.301795 0.186603 0.648205 0.013397 0.498205 +3 0.557735 0.090192 0.357735 0.609808 0.242265 0.509808 0.442265 -0.009808 +... +``` + +For **Segmentation** it contains coordinates of all points of polygon. +A polygon can have three or more points: +```txt +# labels/train/image1.txt: +# ... +0 0.146731 0.151795 0.319936 0.301795 0.186603 0.648205 +3 0.557735 0.090192 0.357735 0.609808 0.242265 0.509808 0.442265 -0.009808 0.400000 0.266667 +... +``` + +For **Pose** it contains bounding boxes and then description of points in one of two forms. +- If the second number in kpt_shape field is 2, + then the line contains two values for every point - `x`, `y`. +- If the second number in kpt_shape field is 3, + then the line contains three values for every point - `x`, `y`, `visibility`, + where `visibility` can have one of three values: + - 0: The keypoint is not visible. + - 1: The keypoint is partially visible. + - 2: The keypoint is fully visible. + +```txt +# labels/image1.txt: +# ... +0 0.250000 0.400000 0.300000 0.400000 0.250000 0.400000 2 0.350000 0.500000 0 ... +3 0.600000 0.400000 0.400000 0.266667 0.250000 0.400000 1 0.440000 0.550000 2 ... +... +``` + +All coordinates must be normalized and be in range \[0, 1\]. +It can be achieved by dividing x coordinates and widths by image width, +and y coordinates and heights by image height. + + +## Export to other formats + +Datumaro can convert a YOLOv8 dataset into any other format Datumaro supports. +To get the expected result, convert the dataset to formats +that support the same annotations as YOLOv8 format you have. + +```bash +datum create +datum add -f yolov8 +datum export -f coco_instances -o +``` +or +```bash +datum convert -if yolov8 -i -f coco_instances -o +``` + +Extra options for importing YOLOv8 format: +- `--config-file` allows to specify config file name to use instead of default `data.yaml` + +Alternatively, using the Python API: + +```python +from datumaro.components.dataset import Dataset + +data_path = 'path/to/dataset' +data_format = 'yolov8' + +dataset = Dataset.import_from(data_path, data_format) +dataset.export('save_dir', 'coco_instances') +``` + +## Export to YOLOv8 format +Datumaro can convert an existing dataset to YOLOv8 format +if it supports annotations from source format. + +Example: + +```bash +datum create +datum import -f coco_instances +datum export -f yolov8 -o +``` + +Extra options for exporting to YOLOv8 format: +- `--save-media` allow to export dataset with saving media files + (default: `False`) +- `--image-ext ` allow to specify image extension + for exporting dataset (default: use original or `.jpg`, if none) +- `--add-path-prefix` allows to specify, whether to include the + `data/` path prefix in the annotation files or not (default: `True`) +- `--config-file` allows to specify config file name to use instead of default `data.yaml` + +## Examples + +### Example 1. Create a custom dataset in YOLOv8 Detection format + +```python +import numpy as np +import datumaro as dm + +dataset = dm.Dataset.from_iterable( + [ + dm.DatasetItem( + id=1, + subset="train", + media=dm.Image(data=np.ones((8, 8, 3))), + annotations=[ + dm.Bbox(0, 2, 4, 2, label=2), + dm.Bbox(0, 1, 2, 3, label=4), + ], + ), + ], + categories=["label_" + str(i) for i in range(10)], +) +dataset.export('../yolov8_dataset', format='yolov8') +``` + +### Example 2. Create a custom dataset in YOLOv8 Oriented Bounding Box format + +Orientation of bounding boxes is controlled through `rotation` attribute of `Bbox` annotation. +Its value is a counter-clockwise angle in degrees. + +```python +import numpy as np +import datumaro as dm + +dataset = dm.Dataset.from_iterable( + [ + dm.DatasetItem( + id=1, + subset="train", + media=dm.Image(data=np.ones((8, 8, 3))), + annotations=[ + dm.Bbox(0, 2, 4, 2, label=2), + dm.Bbox(0, 1, 2, 3, label=4, attributes={"rotation": 30.0}), + ], + ), + ], + categories=["label_" + str(i) for i in range(10)], +) +dataset.export('../yolov8_dataset', format='yolov8_oriented_boxes') +``` + +### Example 3. Create a custom dataset in YOLOv8 Segmentation format + +```python +import numpy as np +import datumaro as dm + +dataset = dm.Dataset.from_iterable( + [ + dm.DatasetItem( + id=1, + subset="train", + media=dm.Image(data=np.ones((8, 8, 3))), + annotations=[ + dm.Polygon([3.0, 1.5, 6.0, 1.5, 6.0, 7.5, 4.5, 7.5, 3.75, 3.0], label=4), + ], + ), + ], + categories=["label_" + str(i) for i in range(10)], +) +dataset.export('../yolov8_dataset', format='yolov8_segmentation') +``` + +### Example 4. Create a custom dataset in YOLOv8 Pose format + +```python +import numpy as np +import datumaro as dm + +dataset = dm.Dataset.from_iterable( + [ + dm.DatasetItem( + id="1", + subset="train", + media=dm.Image(data=np.ones((5, 10, 3))), + annotations=[ + dm.Skeleton( + [ + dm.Points([1.5, 2.0], [2], label=1), + dm.Points([4.5, 4.0], [2], label=2), + dm.Points([7.5, 6.0], [1], label=3), + ], + label=0, + ), + ], + ), + ], + categories={ + dm.AnnotationType.label: dm.LabelCategories.from_iterable([ + "skeleton_label", + ("point_0", "skeleton_label"), + ("point_1", "skeleton_label"), + ("point_2", "skeleton_label"), + ]), + dm.AnnotationType.points: dm.PointsCategories.from_iterable([ + (0, ["point_0", "point_1", "point_2"], set()) + ]), + }, +) +dataset.export('../yolov8_dataset', format='yolov8_pose') +``` diff --git a/site/content/en/docs/user-manual/supported_formats.md b/site/content/en/docs/user-manual/supported_formats.md index cea5d41f96..794c4a5e2f 100644 --- a/site/content/en/docs/user-manual/supported_formats.md +++ b/site/content/en/docs/user-manual/supported_formats.md @@ -161,6 +161,23 @@ List of supported formats: - [Format specification](https://github.com/AlexeyAB/darknet#how-to-train-pascal-voc-data) - [Dataset example](https://github.com/cvat-ai/datumaro/tree/develop/tests/assets/yolo_dataset) - [Format documentation](/docs/formats/yolo) +- YOLOv8 (`detection`, `segmentation`, `pose`, `oriented_box`) + - Detection + - [Format specification](https://docs.ultralytics.com/datasets/detect/) + - [Dataset example](https://docs.ultralytics.com/datasets/detect/coco8/) + - [Format documentation](/docs/formats/yolo_v8) + - Segmentation + - [Format specification](https://docs.ultralytics.com/datasets/segment/) + - [Dataset example](https://docs.ultralytics.com/datasets/segment/coco8-seg/) + - [Format documentation](/docs/formats/yolo_v8) + - Pose + - [Format specification](https://docs.ultralytics.com/datasets/pose/) + - [Dataset example](https://docs.ultralytics.com/datasets/pose/coco8-pose/) + - [Format documentation](/docs/formats/yolo_v8) + - Oriented box + - [Format specification](https://docs.ultralytics.com/datasets/obb/) + - [Dataset example](https://docs.ultralytics.com/datasets/obb/dota8/) + - [Format documentation](/docs/formats/yolo_v8) ### Supported annotation types diff --git a/tests/assets/yolo_dataset/obj.data b/tests/assets/yolo_dataset/yolo/obj.data similarity index 100% rename from tests/assets/yolo_dataset/obj.data rename to tests/assets/yolo_dataset/yolo/obj.data diff --git a/tests/assets/yolo_dataset/obj.names b/tests/assets/yolo_dataset/yolo/obj.names similarity index 100% rename from tests/assets/yolo_dataset/obj.names rename to tests/assets/yolo_dataset/yolo/obj.names diff --git a/tests/assets/yolo_dataset/obj_train_data/1.jpg b/tests/assets/yolo_dataset/yolo/obj_train_data/1.jpg similarity index 100% rename from tests/assets/yolo_dataset/obj_train_data/1.jpg rename to tests/assets/yolo_dataset/yolo/obj_train_data/1.jpg diff --git a/tests/assets/yolo_dataset/obj_train_data/1.txt b/tests/assets/yolo_dataset/yolo/obj_train_data/1.txt similarity index 100% rename from tests/assets/yolo_dataset/obj_train_data/1.txt rename to tests/assets/yolo_dataset/yolo/obj_train_data/1.txt diff --git a/tests/assets/yolo_dataset/train.txt b/tests/assets/yolo_dataset/yolo/train.txt similarity index 100% rename from tests/assets/yolo_dataset/train.txt rename to tests/assets/yolo_dataset/yolo/train.txt diff --git a/tests/assets/yolo_dataset/yolov8/data.yaml b/tests/assets/yolo_dataset/yolov8/data.yaml new file mode 100644 index 0000000000..7e9546ac9f --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8/data.yaml @@ -0,0 +1,14 @@ +path: . +train: images/train + +names: + 0: label_0 + 1: label_1 + 2: label_2 + 3: label_3 + 4: label_4 + 5: label_5 + 6: label_6 + 7: label_7 + 8: label_8 + 9: label_9 diff --git a/tests/assets/yolo_dataset/yolov8/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8/labels/train/1.txt new file mode 100644 index 0000000000..1f507909e2 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8/labels/train/1.txt @@ -0,0 +1,2 @@ +2 0.133333 0.300000 0.266667 0.200000 +4 0.266667 0.450000 0.133333 0.300000 diff --git a/tests/assets/yolo_dataset/yolov8_oriented_boxes/data.yaml b/tests/assets/yolo_dataset/yolov8_oriented_boxes/data.yaml new file mode 100644 index 0000000000..7e9546ac9f --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_oriented_boxes/data.yaml @@ -0,0 +1,14 @@ +path: . +train: images/train + +names: + 0: label_0 + 1: label_1 + 2: label_2 + 3: label_3 + 4: label_4 + 5: label_5 + 6: label_6 + 7: label_7 + 8: label_8 + 9: label_9 diff --git a/tests/assets/yolo_dataset/yolov8_oriented_boxes/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8_oriented_boxes/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8_oriented_boxes/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8_oriented_boxes/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8_oriented_boxes/labels/train/1.txt new file mode 100644 index 0000000000..f5103ac08e --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_oriented_boxes/labels/train/1.txt @@ -0,0 +1,2 @@ +2 0.146731 0.151795 0.319936 0.301795 0.186603 0.648205 0.013397 0.498205 +4 0.557735 0.090192 0.357735 0.609808 0.242265 0.509808 0.442265 -0.009808 diff --git a/tests/assets/yolo_dataset/yolov8_pose/data.yaml b/tests/assets/yolo_dataset/yolov8_pose/data.yaml new file mode 100644 index 0000000000..1ace833a40 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_pose/data.yaml @@ -0,0 +1,7 @@ +path: . +train: images/train + +kpt_shape: [3, 3] + +names: + 0: skeleton_label diff --git a/tests/assets/yolo_dataset/yolov8_pose/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8_pose/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8_pose/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8_pose/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8_pose/labels/train/1.txt new file mode 100644 index 0000000000..fd13aac108 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_pose/labels/train/1.txt @@ -0,0 +1 @@ +0 0.133333 0.300000 0.266667 0.200000 0.1 0.2 2 0.3 0.4 2 0.5 0.6 2 diff --git a/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/data.yaml b/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/data.yaml new file mode 100644 index 0000000000..9f5d940a88 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/data.yaml @@ -0,0 +1,7 @@ +path: . +train: images/train + +kpt_shape: [3, 2] + +names: + 0: skeleton_label diff --git a/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/labels/train/1.txt new file mode 100644 index 0000000000..a43f2cad89 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_pose_two_values_per_point/labels/train/1.txt @@ -0,0 +1 @@ +0 0.133333 0.300000 0.266667 0.200000 0.1 0.2 0.3 0.4 0.5 0.6 diff --git a/tests/assets/yolo_dataset/yolov8_segmentation/data.yaml b/tests/assets/yolo_dataset/yolov8_segmentation/data.yaml new file mode 100644 index 0000000000..7e9546ac9f --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_segmentation/data.yaml @@ -0,0 +1,14 @@ +path: . +train: images/train + +names: + 0: label_0 + 1: label_1 + 2: label_2 + 3: label_3 + 4: label_4 + 5: label_5 + 6: label_6 + 7: label_7 + 8: label_8 + 9: label_9 diff --git a/tests/assets/yolo_dataset/yolov8_segmentation/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8_segmentation/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8_segmentation/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8_segmentation/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8_segmentation/labels/train/1.txt new file mode 100644 index 0000000000..4b8be40618 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_segmentation/labels/train/1.txt @@ -0,0 +1,2 @@ +2 0.100000 0.100000 0.400000 0.100000 0.400000 0.500000 +4 0.200000 0.150000 0.400000 0.150000 0.400000 0.750000 0.300000 0.750000 0.250000 0.300000 diff --git a/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/data.yaml b/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/data.yaml new file mode 100644 index 0000000000..552e53814a --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/data.yaml @@ -0,0 +1,15 @@ +path: . +train: + - images/train/1.jpg + +names: + 0: label_0 + 1: label_1 + 2: label_2 + 3: label_3 + 4: label_4 + 5: label_5 + 6: label_6 + 7: label_7 + 8: label_8 + 9: label_9 diff --git a/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/labels/train/1.txt new file mode 100644 index 0000000000..1f507909e2 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_with_list_of_imgs/labels/train/1.txt @@ -0,0 +1,2 @@ +2 0.133333 0.300000 0.266667 0.200000 +4 0.266667 0.450000 0.133333 0.300000 diff --git a/tests/assets/yolo_dataset/yolov8_with_list_of_names/data.yaml b/tests/assets/yolo_dataset/yolov8_with_list_of_names/data.yaml new file mode 100644 index 0000000000..0dee2d196c --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_with_list_of_names/data.yaml @@ -0,0 +1,14 @@ +path: . +train: images/train + +names: + - label_0 + - label_1 + - label_2 + - label_3 + - label_4 + - label_5 + - label_6 + - label_7 + - label_8 + - label_9 diff --git a/tests/assets/yolo_dataset/yolov8_with_list_of_names/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8_with_list_of_names/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8_with_list_of_names/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8_with_list_of_names/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8_with_list_of_names/labels/train/1.txt new file mode 100644 index 0000000000..1f507909e2 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_with_list_of_names/labels/train/1.txt @@ -0,0 +1,2 @@ +2 0.133333 0.300000 0.266667 0.200000 +4 0.266667 0.450000 0.133333 0.300000 diff --git a/tests/assets/yolo_dataset/yolov8_with_subset_txt/data.yaml b/tests/assets/yolo_dataset/yolov8_with_subset_txt/data.yaml new file mode 100644 index 0000000000..ce3f710c76 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_with_subset_txt/data.yaml @@ -0,0 +1,14 @@ +path: . +train: train.txt + +names: + 0: label_0 + 1: label_1 + 2: label_2 + 3: label_3 + 4: label_4 + 5: label_5 + 6: label_6 + 7: label_7 + 8: label_8 + 9: label_9 diff --git a/tests/assets/yolo_dataset/yolov8_with_subset_txt/images/train/1.jpg b/tests/assets/yolo_dataset/yolov8_with_subset_txt/images/train/1.jpg new file mode 100644 index 0000000000..8689b95631 Binary files /dev/null and b/tests/assets/yolo_dataset/yolov8_with_subset_txt/images/train/1.jpg differ diff --git a/tests/assets/yolo_dataset/yolov8_with_subset_txt/labels/train/1.txt b/tests/assets/yolo_dataset/yolov8_with_subset_txt/labels/train/1.txt new file mode 100644 index 0000000000..1f507909e2 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_with_subset_txt/labels/train/1.txt @@ -0,0 +1,2 @@ +2 0.133333 0.300000 0.266667 0.200000 +4 0.266667 0.450000 0.133333 0.300000 diff --git a/tests/assets/yolo_dataset/yolov8_with_subset_txt/train.txt b/tests/assets/yolo_dataset/yolov8_with_subset_txt/train.txt new file mode 100644 index 0000000000..8b44a247e3 --- /dev/null +++ b/tests/assets/yolo_dataset/yolov8_with_subset_txt/train.txt @@ -0,0 +1 @@ +data/images/train/1.jpg diff --git a/tests/conftest.py b/tests/conftest.py index e52bdafeff..971ecb61e4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,11 @@ # Copyright (C) 2021 Intel Corporation -# Copyright (C) 2022 CVAT.ai Corporation +# Copyright (C) 2022-2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT +from datumaro.util.test_utils import TestDir from .fixtures import * +from .utils.test_utils import TestCaseHelper def pytest_configure(config): @@ -19,3 +21,14 @@ def pytest_configure(config): config.addinivalue_line("markers", "components(ids): link a test with a component") config.addinivalue_line("markers", "reqids(ids): link a test with a requirement") config.addinivalue_line("markers", "bugs(ids): link a test with a bug") + + +@pytest.fixture(scope="function") +def test_dir(): + with TestDir() as test_dir: + yield test_dir + + +@pytest.fixture(scope="class") +def helper_tc(): + return TestCaseHelper() diff --git a/tests/test_yolo_format.py b/tests/test_yolo_format.py deleted file mode 100644 index be102c100e..0000000000 --- a/tests/test_yolo_format.py +++ /dev/null @@ -1,515 +0,0 @@ -import os -import os.path as osp -import pickle # nosec - disable B403:import_pickle check -import shutil -from unittest import TestCase - -import numpy as np -from PIL import Image as PILImage - -from datumaro.components.annotation import Bbox -from datumaro.components.dataset import Dataset -from datumaro.components.environment import Environment -from datumaro.components.errors import ( - AnnotationImportError, - DatasetExportError, - DatasetImportError, - InvalidAnnotationError, - ItemImportError, - UndeclaredLabelError, -) -from datumaro.components.extractor import DatasetItem -from datumaro.components.media import Image -from datumaro.plugins.yolo_format.converter import YoloConverter -from datumaro.plugins.yolo_format.extractor import YoloExtractor, YoloImporter -from datumaro.util.image import save_image -from datumaro.util.test_utils import TestDir, compare_datasets, compare_datasets_strict - -from .requirements import Requirements, mark_requirement - - -class YoloConvertertTest(TestCase): - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_save_and_load(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=1, - subset="train", - media=Image(data=np.ones((8, 8, 3))), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(0, 1, 2, 3, label=4), - ], - ), - DatasetItem( - id=2, - subset="train", - media=Image(data=np.ones((10, 10, 3))), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(3, 3, 2, 3, label=4), - Bbox(2, 1, 2, 3, label=4), - ], - ), - DatasetItem( - id=3, - subset="valid", - media=Image(data=np.ones((8, 8, 3))), - annotations=[ - Bbox(0, 1, 5, 2, label=2), - Bbox(0, 2, 3, 2, label=5), - Bbox(0, 2, 4, 2, label=6), - Bbox(0, 7, 3, 2, label=7), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir, save_media=True) - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - compare_datasets(self, source_dataset, parsed_dataset) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_save_dataset_with_image_info(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=1, - subset="train", - media=Image(path="1.jpg", size=(10, 15)), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(3, 3, 2, 3, label=4), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir) - - save_image( - osp.join(test_dir, "obj_train_data", "1.jpg"), np.ones((10, 15, 3)) - ) # put the image for dataset - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - compare_datasets(self, source_dataset, parsed_dataset) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_load_dataset_with_exact_image_info(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=1, - subset="train", - media=Image(path="1.jpg", size=(10, 15)), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(3, 3, 2, 3, label=4), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir) - - parsed_dataset = Dataset.import_from(test_dir, "yolo", image_info={"1": (10, 15)}) - - compare_datasets(self, source_dataset, parsed_dataset) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_save_dataset_with_cyrillic_and_spaces_in_filename(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem( - id="кириллица с пробелом", - subset="train", - media=Image(data=np.ones((8, 8, 3))), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(0, 1, 2, 3, label=4), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir, save_media=True) - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - compare_datasets(self, source_dataset, parsed_dataset, require_media=True) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_relative_paths(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem(id="1", subset="train", media=Image(data=np.ones((4, 2, 3)))), - DatasetItem(id="subdir1/1", subset="train", media=Image(data=np.ones((2, 6, 3)))), - DatasetItem(id="subdir2/1", subset="train", media=Image(data=np.ones((5, 4, 3)))), - ], - categories=[], - ) - - for save_media in {True, False}: - with self.subTest(save_media=save_media): - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir, save_media=save_media) - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - compare_datasets(self, source_dataset, parsed_dataset) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_save_and_load_image_with_arbitrary_extension(self): - dataset = Dataset.from_iterable( - [ - DatasetItem( - "q/1", subset="train", media=Image(path="q/1.JPEG", data=np.zeros((4, 3, 3))) - ), - DatasetItem( - "a/b/c/2", - subset="valid", - media=Image(path="a/b/c/2.bmp", data=np.zeros((3, 4, 3))), - ), - ], - categories=[], - ) - - with TestDir() as test_dir: - YoloConverter.convert(dataset, test_dir, save_media=True) - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - compare_datasets(self, dataset, parsed_dataset, require_media=True) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_inplace_save_writes_only_updated_data(self): - expected = Dataset.from_iterable( - [ - DatasetItem(1, subset="train", media=Image(data=np.ones((2, 4, 3)))), - DatasetItem(2, subset="train", media=Image(data=np.ones((3, 2, 3)))), - ], - categories=[], - ) - - with TestDir() as path: - dataset = Dataset.from_iterable( - [ - DatasetItem(1, subset="train", media=Image(data=np.ones((2, 4, 3)))), - DatasetItem(2, subset="train", media=Image(path="2.jpg", size=(3, 2))), - DatasetItem(3, subset="valid", media=Image(data=np.ones((2, 2, 3)))), - ], - categories=[], - ) - dataset.export(path, "yolo", save_media=True) - - dataset.put(DatasetItem(2, subset="train", media=Image(data=np.ones((3, 2, 3))))) - dataset.remove(3, "valid") - dataset.save(save_media=True) - - self.assertEqual( - {"1.txt", "2.txt", "1.jpg", "2.jpg"}, - set(os.listdir(osp.join(path, "obj_train_data"))), - ) - self.assertEqual(set(), set(os.listdir(osp.join(path, "obj_valid_data")))) - compare_datasets(self, expected, Dataset.import_from(path, "yolo"), require_media=True) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_save_and_load_with_meta_file(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=1, - subset="train", - media=Image(data=np.ones((8, 8, 3))), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(0, 1, 2, 3, label=4), - ], - ), - DatasetItem( - id=2, - subset="train", - media=Image(data=np.ones((10, 10, 3))), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(3, 3, 2, 3, label=4), - Bbox(2, 1, 2, 3, label=4), - ], - ), - DatasetItem( - id=3, - subset="valid", - media=Image(data=np.ones((8, 8, 3))), - annotations=[ - Bbox(0, 1, 5, 2, label=2), - Bbox(0, 2, 3, 2, label=5), - Bbox(0, 2, 4, 2, label=6), - Bbox(0, 7, 3, 2, label=7), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir, save_media=True, save_dataset_meta=True) - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - self.assertTrue(osp.isfile(osp.join(test_dir, "dataset_meta.json"))) - compare_datasets(self, source_dataset, parsed_dataset) - - @mark_requirement(Requirements.DATUM_565) - def test_can_save_and_load_with_custom_subset_name(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=3, - subset="anything", - media=Image(data=np.ones((8, 8, 3))), - annotations=[ - Bbox(0, 1, 5, 2, label=2), - Bbox(0, 2, 3, 2, label=5), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir, save_media=True) - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - compare_datasets(self, source_dataset, parsed_dataset) - - @mark_requirement(Requirements.DATUM_565) - def test_cant_save_with_reserved_subset_name(self): - for subset in ["backup", "classes"]: - dataset = Dataset.from_iterable( - [ - DatasetItem( - id=3, - subset=subset, - media=Image(data=np.ones((8, 8, 3))), - ), - ], - categories=["a"], - ) - - with TestDir() as test_dir: - with self.assertRaisesRegex(DatasetExportError, f"Can't export '{subset}' subset"): - YoloConverter.convert(dataset, test_dir) - - @mark_requirement(Requirements.DATUM_609) - def test_can_save_and_load_without_path_prefix(self): - source_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=3, - subset="valid", - media=Image(data=np.ones((8, 8, 3))), - annotations=[ - Bbox(0, 1, 5, 2, label=1), - ], - ), - ], - categories=["a", "b"], - ) - - with TestDir() as test_dir: - YoloConverter.convert(source_dataset, test_dir, save_media=True, add_path_prefix=False) - parsed_dataset = Dataset.import_from(test_dir, "yolo") - - with open(osp.join(test_dir, "obj.data"), "r") as f: - lines = f.readlines() - self.assertIn("valid = valid.txt\n", lines) - - with open(osp.join(test_dir, "valid.txt"), "r") as f: - lines = f.readlines() - self.assertIn("obj_valid_data/3.jpg\n", lines) - - compare_datasets(self, source_dataset, parsed_dataset) - - -DUMMY_DATASET_DIR = osp.join(osp.dirname(__file__), "assets", "yolo_dataset") - - -class YoloImporterTest(TestCase): - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_detect(self): - detected_formats = Environment().detect_dataset(DUMMY_DATASET_DIR) - self.assertEqual([YoloImporter.NAME], detected_formats) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_import(self): - expected_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=1, - subset="train", - media=Image(data=np.ones((10, 15, 3))), - annotations=[ - Bbox(0, 2, 4, 2, label=2), - Bbox(3, 3, 2, 3, label=4), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - dataset = Dataset.import_from(DUMMY_DATASET_DIR, "yolo") - - compare_datasets(self, expected_dataset, dataset) - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_import_with_exif_rotated_images(self): - expected_dataset = Dataset.from_iterable( - [ - DatasetItem( - id=1, - subset="train", - media=Image(data=np.ones((15, 10, 3))), - annotations=[ - Bbox(0, 3, 2.67, 3.0, label=2), - Bbox(2, 4.5, 1.33, 4.5, label=4), - ], - ), - ], - categories=["label_" + str(i) for i in range(10)], - ) - - with TestDir() as test_dir: - dataset_path = osp.join(test_dir, "dataset") - shutil.copytree(DUMMY_DATASET_DIR, dataset_path) - - # Add exif rotation for image - image_path = osp.join(dataset_path, "obj_train_data", "1.jpg") - img = PILImage.open(image_path) - exif = img.getexif() - exif.update([(296, 3), (282, 28.0), (531, 1), (274, 6), (283, 28.0)]) - img.save(image_path, exif=exif) - - dataset = Dataset.import_from(dataset_path, "yolo") - - compare_datasets(self, expected_dataset, dataset, require_media=True) - - @mark_requirement(Requirements.DATUM_673) - def test_can_pickle(self): - source = Dataset.import_from(DUMMY_DATASET_DIR, format="yolo") - - parsed = pickle.loads(pickle.dumps(source)) # nosec - - compare_datasets_strict(self, source, parsed) - - -class YoloExtractorTest(TestCase): - def _prepare_dataset(self, path: str) -> Dataset: - dataset = Dataset.from_iterable( - [ - DatasetItem( - "a", - subset="train", - media=Image(np.ones((5, 10, 3))), - annotations=[Bbox(1, 1, 2, 4, label=0)], - ) - ], - categories=["test"], - ) - dataset.export(path, "yolo", save_images=True) - - return dataset - - @mark_requirement(Requirements.DATUM_GENERAL_REQ) - def test_can_parse(self): - with TestDir() as test_dir: - expected = self._prepare_dataset(test_dir) - - actual = Dataset.import_from(test_dir, "yolo") - compare_datasets(self, expected, actual) - - @mark_requirement(Requirements.DATUM_ERROR_REPORTING) - def test_can_report_invalid_data_file(self): - with TestDir() as test_dir: - with self.assertRaisesRegex(DatasetImportError, "Can't read dataset descriptor file"): - YoloExtractor(test_dir) - - @mark_requirement(Requirements.DATUM_ERROR_REPORTING) - def test_can_report_invalid_ann_line_format(self): - with TestDir() as test_dir: - self._prepare_dataset(test_dir) - with open(osp.join(test_dir, "obj_train_data", "a.txt"), "w") as f: - f.write("1 2 3\n") - - with self.assertRaises(AnnotationImportError) as capture: - Dataset.import_from(test_dir, "yolo").init_cache() - self.assertIsInstance(capture.exception.__cause__, InvalidAnnotationError) - self.assertIn("Unexpected field count", str(capture.exception.__cause__)) - - @mark_requirement(Requirements.DATUM_ERROR_REPORTING) - def test_can_report_invalid_label(self): - with TestDir() as test_dir: - self._prepare_dataset(test_dir) - with open(osp.join(test_dir, "obj_train_data", "a.txt"), "w") as f: - f.write("10 0.5 0.5 0.5 0.5\n") - - with self.assertRaises(AnnotationImportError) as capture: - Dataset.import_from(test_dir, "yolo").init_cache() - self.assertIsInstance(capture.exception.__cause__, UndeclaredLabelError) - self.assertEqual(capture.exception.__cause__.id, "10") - - @mark_requirement(Requirements.DATUM_ERROR_REPORTING) - def test_can_report_invalid_field_type(self): - for field, field_name in [ - (1, "bbox center x"), - (2, "bbox center y"), - (3, "bbox width"), - (4, "bbox height"), - ]: - with self.subTest(field_name=field_name): - with TestDir() as test_dir: - self._prepare_dataset(test_dir) - with open(osp.join(test_dir, "obj_train_data", "a.txt"), "w") as f: - values = [0, 0.5, 0.5, 0.5, 0.5] - values[field] = "a" - f.write(" ".join(str(v) for v in values)) - - with self.assertRaises(AnnotationImportError) as capture: - Dataset.import_from(test_dir, "yolo").init_cache() - self.assertIsInstance(capture.exception.__cause__, InvalidAnnotationError) - self.assertIn(field_name, str(capture.exception.__cause__)) - - @mark_requirement(Requirements.DATUM_ERROR_REPORTING) - def test_can_report_missing_ann_file(self): - with TestDir() as test_dir: - self._prepare_dataset(test_dir) - os.remove(osp.join(test_dir, "obj_train_data", "a.txt")) - - with self.assertRaises(ItemImportError) as capture: - Dataset.import_from(test_dir, "yolo").init_cache() - self.assertIsInstance(capture.exception.__cause__, FileNotFoundError) - - @mark_requirement(Requirements.DATUM_ERROR_REPORTING) - def test_can_report_missing_image_info(self): - with TestDir() as test_dir: - self._prepare_dataset(test_dir) - os.remove(osp.join(test_dir, "obj_train_data", "a.jpg")) - - with self.assertRaises(ItemImportError) as capture: - Dataset.import_from(test_dir, "yolo").init_cache() - self.assertIsInstance(capture.exception.__cause__, DatasetImportError) - self.assertIn("Can't find image info", str(capture.exception.__cause__)) - - @mark_requirement(Requirements.DATUM_ERROR_REPORTING) - def test_can_report_missing_subset_info(self): - with TestDir() as test_dir: - self._prepare_dataset(test_dir) - os.remove(osp.join(test_dir, "train.txt")) - - with self.assertRaisesRegex(InvalidAnnotationError, "subset list file"): - Dataset.import_from(test_dir, "yolo").init_cache() diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/data_formats/__init__.py b/tests/unit/data_formats/__init__.py new file mode 100644 index 0000000000..ff847f0120 --- /dev/null +++ b/tests/unit/data_formats/__init__.py @@ -0,0 +1,3 @@ +# Copyright (C) 2023 Intel Corporation +# +# SPDX-License-Identifier: MIT diff --git a/tests/unit/data_formats/test_yolo_format.py b/tests/unit/data_formats/test_yolo_format.py new file mode 100644 index 0000000000..94f699733c --- /dev/null +++ b/tests/unit/data_formats/test_yolo_format.py @@ -0,0 +1,1223 @@ +# Copyright (C) 2024 CVAT.ai Corporation +# +# SPDX-License-Identifier: MIT + +import copy +import os +import os.path as osp +import pickle # nosec - disable B403:import_pickle check +import random +import shutil + +import numpy as np +import pytest +import yaml +from PIL import ExifTags +from PIL import Image as PILImage + +from datumaro.components.annotation import ( + AnnotationType, + Bbox, + LabelCategories, + Points, + PointsCategories, + Polygon, + Skeleton, +) +from datumaro.components.dataset import Dataset +from datumaro.components.environment import Environment +from datumaro.components.errors import ( + AnnotationImportError, + DatasetExportError, + DatasetImportError, + DatasetNotFoundError, + InvalidAnnotationError, + ItemImportError, + UndeclaredLabelError, +) +from datumaro.components.extractor import DatasetItem +from datumaro.components.format_detection import FormatDetectionContext, FormatRequirementsUnmet +from datumaro.components.media import Image +from datumaro.plugins.yolo_format.converter import ( + YoloConverter, + YOLOv8Converter, + YOLOv8OrientedBoxesConverter, + YOLOv8PoseConverter, + YOLOv8SegmentationConverter, +) +from datumaro.plugins.yolo_format.extractor import ( + YoloExtractor, + YOLOv8Extractor, + YOLOv8OrientedBoxesExtractor, + YOLOv8PoseExtractor, + YOLOv8SegmentationExtractor, +) +from datumaro.plugins.yolo_format.importer import ( + YoloImporter, + YOLOv8Importer, + YOLOv8OrientedBoxesImporter, + YOLOv8PoseImporter, + YOLOv8SegmentationImporter, +) +from datumaro.util.image import save_image +from datumaro.util.test_utils import compare_datasets, compare_datasets_strict + +from ...requirements import Requirements, mark_requirement +from ...utils.assets import get_test_asset_path + + +@pytest.fixture(autouse=True) +def seed_random(): + random.seed(1234) + + +def randint(a, b): + return random.randint(a, b) # nosec B311 NOSONAR + + +class CompareDatasetMixin: + @pytest.fixture(autouse=True) + def setup(self, helper_tc): + self.helper_tc = helper_tc + + def compare_datasets(self, expected, actual, **kwargs): + compare_datasets(self.helper_tc, expected, actual, **kwargs) + + +class CompareDatasetsRotationMixin(CompareDatasetMixin): + def compare_datasets(self, expected, actual, **kwargs): + actual_copy = copy.deepcopy(actual) + compare_datasets(self.helper_tc, expected, actual, ignored_attrs=["rotation"], **kwargs) + for item_a, item_b in zip(expected, actual_copy): + for ann_a, ann_b in zip(item_a.annotations, item_b.annotations): + assert ("rotation" in ann_a.attributes) == ("rotation" in ann_b.attributes) + assert ( + abs(ann_a.attributes.get("rotation", 0) - ann_b.attributes.get("rotation", 0)) + < 0.01 + ) + + +class YoloConverterTest(CompareDatasetMixin): + CONVERTER = YoloConverter + IMPORTER = YoloImporter + + def _generate_random_bbox(self, n_of_labels=10, **kwargs): + return Bbox( + x=randint(0, 4), + y=randint(0, 4), + w=randint(1, 4), + h=randint(1, 4), + label=randint(0, n_of_labels - 1), + attributes=kwargs, + ) + + def _generate_random_annotation(self, n_of_labels=10): + return self._generate_random_bbox(n_of_labels=n_of_labels) + + @staticmethod + def _make_image_path(test_dir: str, subset_name: str, image_id: str): + return osp.join(test_dir, f"obj_{subset_name}_data", image_id) + + def _generate_random_dataset(self, recipes, n_of_labels=10): + items = [ + DatasetItem( + id=recipe.get("id", index + 1), + subset=recipe.get("subset", "train"), + media=recipe.get( + "media", + Image(data=np.ones((randint(8, 10), randint(8, 10), 3))), + ), + annotations=[ + self._generate_random_annotation(n_of_labels=n_of_labels) + for _ in range(recipe.get("annotations", 1)) + ], + ) + for index, recipe in enumerate(recipes) + ] + return Dataset.from_iterable( + items, + categories=["label_" + str(i) for i in range(n_of_labels)], + ) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_save_and_load(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + {"annotations": 2}, + {"annotations": 3}, + {"annotations": 4}, + ] + ) + + self.CONVERTER.convert(source_dataset, test_dir, save_media=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_save_dataset_with_image_info(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + { + "annotations": 2, + "media": Image(path="1.jpg", size=(10, 15)), + }, + ] + ) + + self.CONVERTER.convert(source_dataset, test_dir) + + save_image( + self._make_image_path(test_dir, "train", "1.jpg"), np.ones((10, 15, 3)) + ) # put the image for dataset + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_load_dataset_with_exact_image_info(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + { + "annotations": 2, + "media": Image(path="1.jpg", size=(10, 15)), + }, + ] + ) + + self.CONVERTER.convert(source_dataset, test_dir) + parsed_dataset = Dataset.import_from( + test_dir, self.IMPORTER.NAME, image_info={"1": (10, 15)} + ) + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_save_dataset_with_cyrillic_and_spaces_in_filename(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + { + "id": "кириллица с пробелом", + "annotations": 2, + }, + ] + ) + + self.CONVERTER.convert(source_dataset, test_dir, save_media=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + self.compare_datasets(source_dataset, parsed_dataset, require_media=True) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + @pytest.mark.parametrize("save_media", [True, False]) + def test_relative_paths(self, save_media, test_dir): + source_dataset = Dataset.from_iterable( + [ + DatasetItem(id="1", subset="train", media=Image(data=np.ones((4, 2, 3)))), + DatasetItem(id="subdir1/1", subset="train", media=Image(data=np.ones((2, 6, 3)))), + DatasetItem(id="subdir2/1", subset="train", media=Image(data=np.ones((5, 4, 3)))), + ], + categories=[], + ) + + self.CONVERTER.convert(source_dataset, test_dir, save_media=save_media) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_save_and_load_image_with_arbitrary_extension(self, test_dir): + dataset = Dataset.from_iterable( + [ + DatasetItem( + "q/1", subset="train", media=Image(path="q/1.JPEG", data=np.zeros((4, 3, 3))) + ), + DatasetItem( + "a/b/c/2", + subset="valid", + media=Image(path="a/b/c/2.bmp", data=np.zeros((3, 4, 3))), + ), + ], + categories=[], + ) + + self.CONVERTER.convert(dataset, test_dir, save_media=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + self.compare_datasets(dataset, parsed_dataset, require_media=True) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_inplace_save_writes_only_updated_data(self, test_dir): + expected = Dataset.from_iterable( + [ + DatasetItem(1, subset="train", media=Image(data=np.ones((2, 4, 3)))), + DatasetItem(2, subset="train", media=Image(data=np.ones((3, 2, 3)))), + ], + categories=[], + ) + + dataset = Dataset.from_iterable( + [ + DatasetItem(1, subset="train", media=Image(data=np.ones((2, 4, 3)))), + DatasetItem(2, subset="train", media=Image(path="2.jpg", size=(3, 2))), + DatasetItem(3, subset="valid", media=Image(data=np.ones((2, 2, 3)))), + ], + categories=[], + ) + dataset.export(test_dir, self.CONVERTER.NAME, save_media=True) + + dataset.put(DatasetItem(2, subset="train", media=Image(data=np.ones((3, 2, 3))))) + dataset.remove(3, "valid") + dataset.save(save_media=True) + + self._check_inplace_save_writes_only_updated_data(test_dir, expected) + + def _check_inplace_save_writes_only_updated_data(self, test_dir, expected): + assert set(os.listdir(osp.join(test_dir, "obj_train_data"))) == { + "1.txt", + "2.txt", + "1.jpg", + "2.jpg", + } + assert set(os.listdir(osp.join(test_dir, "obj_valid_data"))) == set() + self.compare_datasets( + expected, + Dataset.import_from(test_dir, "yolo"), + require_media=True, + ) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_save_and_load_with_meta_file(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + {"annotations": 2}, + {"annotations": 3}, + {"annotations": 4}, + ] + ) + + self.CONVERTER.convert(source_dataset, test_dir, save_media=True, save_dataset_meta=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + + assert osp.isfile(osp.join(test_dir, "dataset_meta.json")) + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_565) + def test_can_save_and_load_with_custom_subset_name(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + {"annotations": 2, "subset": "anything", "id": 3}, + ] + ) + + self.CONVERTER.convert(source_dataset, test_dir, save_media=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_565) + @pytest.mark.parametrize("subset", ["backup", "classes"]) + def test_cant_save_with_reserved_subset_name(self, test_dir, subset): + self._check_cant_save_with_reserved_subset_name(test_dir, subset) + + def _check_cant_save_with_reserved_subset_name(self, test_dir, subset): + dataset = Dataset.from_iterable( + [ + DatasetItem( + id=3, + subset=subset, + media=Image(data=np.ones((8, 8, 3))), + ), + ], + categories=["a"], + ) + + with pytest.raises(DatasetExportError, match=f"Can't export '{subset}' subset"): + self.CONVERTER.convert(dataset, test_dir) + + @mark_requirement(Requirements.DATUM_609) + def test_can_save_and_load_without_path_prefix(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + {"subset": "valid", "id": 3}, + ], + n_of_labels=2, + ) + + self.CONVERTER.convert(source_dataset, test_dir, save_media=True, add_path_prefix=False) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + + with open(osp.join(test_dir, "obj.data"), "r") as f: + lines = f.readlines() + assert "valid = valid.txt\n" in lines + + with open(osp.join(test_dir, "valid.txt"), "r") as f: + lines = f.readlines() + assert "obj_valid_data/3.jpg\n" in lines + + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_export_rotated_bbox(self, test_dir): + expected_dataset = Dataset.from_iterable( + [ + DatasetItem( + id=3, + subset="valid", + media=Image(data=np.ones((8, 8, 3))), + annotations=[ + self._generate_random_bbox(n_of_labels=2), + self._generate_random_bbox(n_of_labels=2), + ], + ), + ], + categories=["a", "b"], + ) + source_dataset = Dataset.from_iterable( + [ + DatasetItem( + id=3, + subset="valid", + media=Image(data=np.ones((8, 8, 3))), + annotations=list(expected_dataset)[0].annotations + + [ + self._generate_random_bbox(n_of_labels=2, rotation=30.0), + ], + ), + ], + categories=["a", "b"], + ) + source_dataset.export(test_dir, self.CONVERTER.NAME, save_media=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + self.compare_datasets(expected_dataset, parsed_dataset) + + +class YOLOv8ConverterTest(YoloConverterTest): + CONVERTER = YOLOv8Converter + IMPORTER = YOLOv8Importer + + @staticmethod + def _make_image_path(test_dir: str, subset_name: str, image_id: str): + return osp.join(test_dir, "images", subset_name, image_id) + + @mark_requirement(Requirements.DATUM_565) + @pytest.mark.parametrize("subset", ["backup", "classes", "path", "names"]) + def test_cant_save_with_reserved_subset_name(self, test_dir, subset): + self._check_cant_save_with_reserved_subset_name(test_dir, subset) + + @mark_requirement(Requirements.DATUM_609) + def test_can_save_and_load_with_custom_config_file(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + {"subset": "valid", "id": 3}, + ], + n_of_labels=2, + ) + filename = "custom_config_name.yaml" + self.CONVERTER.convert( + source_dataset, test_dir, save_media=True, add_path_prefix=False, config_file=filename + ) + assert not osp.exists(osp.join(test_dir, "data.yaml")) + assert osp.isfile(osp.join(test_dir, filename)) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME, config_file=filename) + self.compare_datasets(source_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_609) + def test_can_save_and_load_without_path_prefix(self, test_dir): + source_dataset = self._generate_random_dataset( + [ + {"subset": "valid", "id": 3}, + ], + n_of_labels=2, + ) + + self.CONVERTER.convert(source_dataset, test_dir, save_media=True, add_path_prefix=False) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + + with open(osp.join(test_dir, "data.yaml"), "r") as f: + config = yaml.safe_load(f) + assert config.get("valid") == "valid.txt" + + with open(osp.join(test_dir, "valid.txt"), "r") as f: + lines = f.readlines() + assert "images/valid/3.jpg\n" in lines + + self.compare_datasets(source_dataset, parsed_dataset) + + def _check_inplace_save_writes_only_updated_data(self, test_dir, expected): + assert set(os.listdir(osp.join(test_dir, "images", "train"))) == { + "1.jpg", + "2.jpg", + } + assert set(os.listdir(osp.join(test_dir, "labels", "train"))) == { + "1.txt", + "2.txt", + } + assert set(os.listdir(osp.join(test_dir, "images", "valid"))) == set() + assert set(os.listdir(osp.join(test_dir, "labels", "valid"))) == set() + self.compare_datasets( + expected, + Dataset.import_from(test_dir, self.IMPORTER.NAME), + require_media=True, + ) + + +class YOLOv8SegmentationConverterTest(YOLOv8ConverterTest): + CONVERTER = YOLOv8SegmentationConverter + IMPORTER = YOLOv8SegmentationImporter + + def _generate_random_annotation(self, n_of_labels=10): + return Polygon( + points=[randint(0, 6) for _ in range(randint(3, 7) * 2)], + label=randint(0, n_of_labels - 1), + ) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_export_rotated_bbox(self, test_dir): + pass + + +class YOLOv8OrientedBoxesConverterTest(CompareDatasetsRotationMixin, YOLOv8ConverterTest): + CONVERTER = YOLOv8OrientedBoxesConverter + IMPORTER = YOLOv8OrientedBoxesImporter + + def _generate_random_annotation(self, n_of_labels=10): + return self._generate_random_bbox(n_of_labels=n_of_labels, rotation=randint(10, 350)) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_export_rotated_bbox(self, test_dir): + source_dataset = Dataset.from_iterable( + [ + DatasetItem( + id=3, + subset="valid", + media=Image(data=np.ones((8, 8, 3))), + annotations=[ + self._generate_random_bbox(n_of_labels=2, rotation=30.0), + ], + ), + ], + categories=["a", "b"], + ) + source_dataset.export(test_dir, self.CONVERTER.NAME, save_media=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + assert abs(list(parsed_dataset)[0].annotations[0].attributes["rotation"] - 30) < 0.001 + self.compare_datasets(source_dataset, parsed_dataset) + + +class YOLOv8PoseConverterTest(YOLOv8ConverterTest): + CONVERTER = YOLOv8PoseConverter + IMPORTER = YOLOv8PoseImporter + + def _generate_random_skeleton_annotation(self, skeleton_label_to_point_labels, n_of_labels=10): + label_id = random.choice(list(skeleton_label_to_point_labels.keys())) # nosec B311 NOSONAR + return Skeleton( + [ + Points( + [randint(1, 7), randint(1, 7)], + [randint(0, 2)], + label=label, + ) + for label in skeleton_label_to_point_labels[label_id] + ], + label=label_id, + ) + + def _generate_random_dataset(self, recipes, n_of_labels=10): + n_of_points_in_skeleton = randint(3, 8) + labels = [f"skeleton_label_{index}" for index in range(n_of_labels)] + [ + (f"skeleton_label_{parent_index}_point_{point_index}", f"skeleton_label_{parent_index}") + for parent_index in range(n_of_labels) + for point_index in range(n_of_points_in_skeleton) + ] + skeleton_label_to_point_labels = { + skeleton_label_id: [ + label_id + for label_id, label in enumerate(labels) + if isinstance(label, tuple) and label[1] == f"skeleton_label_{skeleton_label_id}" + ] + for skeleton_label_id, skeleton_label in enumerate(labels) + if isinstance(skeleton_label, str) + } + items = [ + DatasetItem( + id=recipe.get("id", index + 1), + subset=recipe.get("subset", "train"), + media=recipe.get( + "media", + Image(data=np.ones((randint(8, 10), randint(8, 10), 3))), + ), + annotations=[ + self._generate_random_skeleton_annotation( + skeleton_label_to_point_labels, + n_of_labels=n_of_labels, + ) + for _ in range(recipe.get("annotations", 1)) + ], + ) + for index, recipe in enumerate(recipes) + ] + + point_categories = PointsCategories.from_iterable( + [ + ( + index, + [ + f"skeleton_label_{index}_point_{point_index}" + for point_index in range(n_of_points_in_skeleton) + ], + set(), + ) + for index in range(n_of_labels) + ] + ) + + return Dataset.from_iterable( + items, + categories={ + AnnotationType.label: LabelCategories.from_iterable(labels), + AnnotationType.points: point_categories, + }, + ) + + def test_export_rotated_bbox(self, test_dir): + pass + + @staticmethod + def _make_dataset_with_edges_and_point_labels(): + items = [ + DatasetItem( + id="1", + subset="train", + media=Image(data=np.ones((5, 10, 3))), + annotations=[ + Skeleton( + [ + Points([1.5, 2.0], [2], label=4), + Points([4.5, 4.0], [2], label=5), + ], + label=3, + ), + ], + ), + ] + return Dataset.from_iterable( + items, + categories={ + AnnotationType.label: LabelCategories.from_iterable( + [ + "skeleton_label_1", + ("point_label_1", "skeleton_label_1"), + ("point_label_2", "skeleton_label_1"), + "skeleton_label_2", + ("point_label_3", "skeleton_label_2"), + ("point_label_4", "skeleton_label_2"), + ] + ), + AnnotationType.points: PointsCategories.from_iterable( + [ + (0, ["point_label_1", "point_label_2"], {(0, 1)}), + (3, ["point_label_3", "point_label_4"], {}), + ], + ), + }, + ) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_loses_some_info_on_save_load_without_meta_file(self, test_dir): + # loses point labels + # loses edges + # loses label ids - groups skeleton labels to the start + source_dataset = self._make_dataset_with_edges_and_point_labels() + expected_dataset = Dataset.from_iterable( + [ + DatasetItem( + id="1", + subset="train", + media=Image(data=np.ones((5, 10, 3))), + annotations=[ + Skeleton( + [ + Points([1.5, 2.0], [2], label=4), + Points([4.5, 4.0], [2], label=5), + ], + label=1, + ), + ], + ), + ], + categories={ + AnnotationType.label: LabelCategories.from_iterable( + [ + "skeleton_label_1", + "skeleton_label_2", + ("skeleton_label_1_point_0", "skeleton_label_1"), + ("skeleton_label_1_point_1", "skeleton_label_1"), + ("skeleton_label_2_point_0", "skeleton_label_2"), + ("skeleton_label_2_point_1", "skeleton_label_2"), + ] + ), + AnnotationType.points: PointsCategories.from_iterable( + [ + (0, ["skeleton_label_1_point_0", "skeleton_label_1_point_1"], set()), + (1, ["skeleton_label_2_point_0", "skeleton_label_2_point_1"], set()), + ], + ), + }, + ) + self.CONVERTER.convert(source_dataset, test_dir, save_media=True) + + # check that annotation with label 3 was saved as 1 + with open(osp.join(test_dir, "labels", "train", "1.txt"), "r") as f: + assert f.readlines()[0].startswith("1 ") + + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + self.compare_datasets(expected_dataset, parsed_dataset) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_save_and_load_with_meta_file(self, test_dir): + source_dataset = self._make_dataset_with_edges_and_point_labels() + self.CONVERTER.convert(source_dataset, test_dir, save_media=True, save_dataset_meta=True) + parsed_dataset = Dataset.import_from(test_dir, self.IMPORTER.NAME) + assert osp.isfile(osp.join(test_dir, "dataset_meta.json")) + self.compare_datasets(source_dataset, parsed_dataset) + + +class YoloImporterTest(CompareDatasetMixin): + IMPORTER = YoloImporter + ASSETS = ["yolo"] + + def test_can_detect(self): + dataset_dir = get_test_asset_path("yolo_dataset", "yolo") + detected_formats = Environment().detect_dataset(dataset_dir) + assert detected_formats == [self.IMPORTER.NAME] + + @staticmethod + def _asset_dataset(): + return Dataset.from_iterable( + [ + DatasetItem( + id=1, + subset="train", + media=Image(data=np.ones((10, 15, 3))), + annotations=[ + Bbox(0, 2, 4, 2, label=2), + Bbox(3, 3, 2, 3, label=4), + ], + ), + ], + categories=["label_" + str(i) for i in range(10)], + ) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_import(self): + expected_dataset = self._asset_dataset() + for asset in self.ASSETS: + dataset_dir = get_test_asset_path("yolo_dataset", asset) + dataset = Dataset.import_from(dataset_dir, self.IMPORTER.NAME) + self.compare_datasets(expected_dataset, dataset) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_import_with_exif_rotated_images(self, test_dir): + expected_dataset = Dataset.from_iterable( + [ + DatasetItem( + id=1, + subset="train", + media=Image(data=np.ones((15, 10, 3))), + annotations=[ + Bbox(0, 3, 2.67, 3.0, label=2), + Bbox(2, 4.5, 1.33, 4.5, label=4), + ], + ), + ], + categories=["label_" + str(i) for i in range(10)], + ) + + dataset_path = osp.join(test_dir, "dataset") + shutil.copytree(get_test_asset_path("yolo_dataset", "yolo"), dataset_path) + + # Add exif rotation for image + image_path = osp.join(dataset_path, "obj_train_data", "1.jpg") + img = PILImage.open(image_path) + exif = img.getexif() + exif.update( + [ + (ExifTags.Base.ResolutionUnit, 3), + (ExifTags.Base.XResolution, 28.0), + (ExifTags.Base.YCbCrPositioning, 1), + (ExifTags.Base.Orientation, 6), + (ExifTags.Base.YResolution, 28.0), + ] + ) + img.save(image_path, exif=exif) + + dataset = Dataset.import_from(dataset_path, "yolo") + + self.compare_datasets(expected_dataset, dataset, require_media=True) + + @mark_requirement(Requirements.DATUM_673) + def test_can_pickle(self, helper_tc): + for asset in self.ASSETS: + dataset_dir = get_test_asset_path("yolo_dataset", asset) + source = Dataset.import_from(dataset_dir, format=self.IMPORTER.NAME) + parsed = pickle.loads(pickle.dumps(source)) # nosec + compare_datasets_strict(helper_tc, source, parsed) + + +class YOLOv8ImporterTest(YoloImporterTest): + IMPORTER = YOLOv8Importer + ASSETS = [ + "yolov8", + "yolov8_with_list_of_imgs", + "yolov8_with_subset_txt", + "yolov8_with_list_of_names", + ] + + def test_can_detect(self): + for asset in self.ASSETS: + dataset_dir = get_test_asset_path("yolo_dataset", asset) + detected_formats = Environment().detect_dataset(dataset_dir) + assert set(detected_formats) == { + YOLOv8Importer.NAME, + YOLOv8SegmentationImporter.NAME, + YOLOv8OrientedBoxesImporter.NAME, + } + + def test_can_detect_and_import_with_any_yaml_as_config(self, test_dir): + expected_dataset = self._asset_dataset() + dataset_path = osp.join(test_dir, "dataset") + shutil.copytree(get_test_asset_path("yolo_dataset", self.ASSETS[0]), dataset_path) + os.rename( + osp.join(dataset_path, "data.yaml"), osp.join(dataset_path, "custom_file_name.yaml") + ) + + self.IMPORTER.detect(FormatDetectionContext(dataset_path)) + dataset = Dataset.import_from(dataset_path, self.IMPORTER.NAME) + self.compare_datasets(expected_dataset, dataset) + + def test_can_detect_and_import_if_multiple_yamls_with_default_among_them(self, test_dir): + expected_dataset = self._asset_dataset() + dataset_path = osp.join(test_dir, "dataset") + shutil.copytree(get_test_asset_path("yolo_dataset", self.ASSETS[0]), dataset_path) + shutil.copyfile( + osp.join(dataset_path, "data.yaml"), osp.join(dataset_path, "custom_file_name.yaml") + ) + + self.IMPORTER.detect(FormatDetectionContext(dataset_path)) + dataset = Dataset.import_from(dataset_path, self.IMPORTER.NAME) + self.compare_datasets(expected_dataset, dataset) + + def test_can_not_detect_or_import_if_multiple_yamls_but_no_default_among_them(self, test_dir): + dataset_path = osp.join(test_dir, "dataset") + shutil.copytree(get_test_asset_path("yolo_dataset", self.ASSETS[0]), dataset_path) + shutil.copyfile( + osp.join(dataset_path, "data.yaml"), + osp.join(dataset_path, "custom_file_name1.yaml"), + ) + os.rename( + osp.join(dataset_path, "data.yaml"), + osp.join(dataset_path, "custom_file_name2.yaml"), + ) + + with pytest.raises(FormatRequirementsUnmet): + self.IMPORTER.detect(FormatDetectionContext(dataset_path)) + with pytest.raises(DatasetNotFoundError): + Dataset.import_from(dataset_path, self.IMPORTER.NAME) + + def test_can_import_despite_multiple_yamls_if_config_file_provided_as_argument(self, test_dir): + expected_dataset = self._asset_dataset() + dataset_path = osp.join(test_dir, "dataset") + shutil.copytree(get_test_asset_path("yolo_dataset", self.ASSETS[0]), dataset_path) + shutil.copyfile( + osp.join(dataset_path, "data.yaml"), + osp.join(dataset_path, "custom_file_name1.yaml"), + ) + os.rename( + osp.join(dataset_path, "data.yaml"), + osp.join(dataset_path, "custom_file_name2.yaml"), + ) + + dataset = Dataset.import_from( + dataset_path, self.IMPORTER.NAME, config_file="custom_file_name1.yaml" + ) + self.compare_datasets(expected_dataset, dataset) + + +class YOLOv8SegmentationImporterTest(YOLOv8ImporterTest): + IMPORTER = YOLOv8SegmentationImporter + ASSETS = [ + "yolov8_segmentation", + ] + + @staticmethod + def _asset_dataset(): + return Dataset.from_iterable( + [ + DatasetItem( + id=1, + subset="train", + media=Image(data=np.ones((10, 15, 3))), + annotations=[ + Polygon([1.5, 1.0, 6.0, 1.0, 6.0, 5.0], label=2), + Polygon([3.0, 1.5, 6.0, 1.5, 6.0, 7.5, 4.5, 7.5, 3.75, 3.0], label=4), + ], + ), + ], + categories=["label_" + str(i) for i in range(10)], + ) + + +class YOLOv8OrientedBoxesImporterTest(CompareDatasetsRotationMixin, YOLOv8ImporterTest): + IMPORTER = YOLOv8OrientedBoxesImporter + ASSETS = ["yolov8_oriented_boxes"] + + @staticmethod + def _asset_dataset(): + return Dataset.from_iterable( + [ + DatasetItem( + id=1, + subset="train", + media=Image(data=np.ones((10, 15, 3))), + annotations=[ + Bbox(1, 2, 3, 4, label=2, attributes=dict(rotation=30)), + Bbox(3, 2, 6, 2, label=4, attributes=dict(rotation=120)), + ], + ), + ], + categories=["label_" + str(i) for i in range(10)], + ) + + +class YOLOv8PoseImporterTest(YOLOv8ImporterTest): + IMPORTER = YOLOv8PoseImporter + ASSETS = [ + "yolov8_pose", + "yolov8_pose_two_values_per_point", + ] + + def test_can_detect(self): + for asset in self.ASSETS: + dataset_dir = get_test_asset_path("yolo_dataset", asset) + detected_formats = Environment().detect_dataset(dataset_dir) + assert detected_formats == [self.IMPORTER.NAME] + + @staticmethod + def _asset_dataset(): + return Dataset.from_iterable( + [ + DatasetItem( + id="1", + subset="train", + media=Image(data=np.ones((5, 10, 3))), + annotations=[ + Skeleton( + [ + Points([1.5, 2.0], [2], label=1), + Points([4.5, 4.0], [2], label=2), + Points([7.5, 6.0], [2], label=3), + ], + label=0, + ), + ], + ), + ], + categories={ + AnnotationType.label: LabelCategories.from_iterable( + [ + "skeleton_label", + ("skeleton_label_point_0", "skeleton_label"), + ("skeleton_label_point_1", "skeleton_label"), + ("skeleton_label_point_2", "skeleton_label"), + ] + ), + AnnotationType.points: PointsCategories.from_iterable( + [ + ( + 0, + [ + "skeleton_label_point_0", + "skeleton_label_point_1", + "skeleton_label_point_2", + ], + set(), + ) + ], + ), + }, + ) + + +class YoloExtractorTest: + IMPORTER = YoloImporter + EXTRACTOR = YoloExtractor + + def _prepare_dataset(self, path: str, anno=None) -> Dataset: + if anno is None: + anno = Bbox(1, 1, 2, 4, label=0) + dataset = Dataset.from_iterable( + [ + DatasetItem( + "a", + subset="train", + media=Image(np.ones((5, 10, 3))), + annotations=[anno], + ) + ], + categories=["test"], + ) + dataset.export(path, self.EXTRACTOR.NAME, save_media=True) + return dataset + + @staticmethod + def _get_annotation_dir(subset="train"): + return f"obj_{subset}_data" + + @staticmethod + def _get_image_dir(subset="train"): + return f"obj_{subset}_data" + + @staticmethod + def _make_some_annotation_values(): + return [0.5, 0.5, 0.5, 0.5] + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_parse(self, helper_tc, test_dir): + expected = self._prepare_dataset(test_dir) + actual = Dataset.import_from(test_dir, self.IMPORTER.NAME) + compare_datasets(helper_tc, expected, actual) + + def test_can_report_invalid_data_file(self, test_dir): + with pytest.raises(DatasetImportError, match="Can't read dataset descriptor file"): + self.EXTRACTOR(test_dir) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_can_report_invalid_ann_line_format(self, test_dir): + self._prepare_dataset(test_dir) + with open(osp.join(test_dir, self._get_annotation_dir(), "a.txt"), "w") as f: + f.write("1 2 3\n") + + with pytest.raises(AnnotationImportError) as capture: + Dataset.import_from(test_dir, self.IMPORTER.NAME).init_cache() + assert isinstance(capture.value.__cause__, InvalidAnnotationError) + assert "Unexpected field count" in str(capture.value.__cause__) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_can_report_invalid_label(self, test_dir): + self._prepare_dataset(test_dir) + with open(osp.join(test_dir, self._get_annotation_dir(), "a.txt"), "w") as f: + f.write(" ".join(str(v) for v in [10] + self._make_some_annotation_values())) + + with pytest.raises(AnnotationImportError) as capture: + Dataset.import_from(test_dir, self.IMPORTER.NAME).init_cache() + assert isinstance(capture.value.__cause__, UndeclaredLabelError) + assert capture.value.__cause__.id == "10" + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + @pytest.mark.parametrize( + "field, field_name", + [ + (1, "bbox center x"), + (2, "bbox center y"), + (3, "bbox width"), + (4, "bbox height"), + ], + ) + def test_can_report_invalid_field_type(self, field, field_name, test_dir): + self._check_can_report_invalid_field_type(field, field_name, test_dir) + + def _check_can_report_invalid_field_type(self, field, field_name, test_dir): + self._prepare_dataset(test_dir) + with open(osp.join(test_dir, self._get_annotation_dir(), "a.txt"), "w") as f: + values = [0] + self._make_some_annotation_values() + values[field] = "a" + f.write(" ".join(str(v) for v in values)) + + with pytest.raises(AnnotationImportError) as capture: + Dataset.import_from(test_dir, self.IMPORTER.NAME).init_cache() + assert isinstance(capture.value.__cause__, InvalidAnnotationError) + assert field_name in str(capture.value.__cause__) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_can_report_missing_ann_file(self, test_dir): + self._prepare_dataset(test_dir) + os.remove(osp.join(test_dir, self._get_annotation_dir(), "a.txt")) + + with pytest.raises(ItemImportError) as capture: + Dataset.import_from(test_dir, self.IMPORTER.NAME).init_cache() + assert isinstance(capture.value.__cause__, FileNotFoundError) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_can_report_missing_image_info(self, test_dir): + self._prepare_dataset(test_dir) + os.remove(osp.join(test_dir, self._get_image_dir(), "a.jpg")) + + with pytest.raises(ItemImportError) as capture: + Dataset.import_from(test_dir, self.IMPORTER.NAME).init_cache() + assert isinstance(capture.value.__cause__, DatasetImportError) + assert "Can't find image info" in str(capture.value.__cause__) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_can_report_missing_subset_info(self, test_dir): + self._prepare_dataset(test_dir) + os.remove(osp.join(test_dir, "train.txt")) + + with pytest.raises(InvalidAnnotationError, match="subset list file"): + Dataset.import_from(test_dir, self.IMPORTER.NAME).init_cache() + + +class YOLOv8ExtractorTest(YoloExtractorTest): + IMPORTER = YOLOv8Importer + EXTRACTOR = YOLOv8Extractor + + @staticmethod + def _get_annotation_dir(subset="train"): + return osp.join("labels", subset) + + @staticmethod + def _get_image_dir(subset="train"): + return osp.join("images", subset) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_can_report_missing_subset_folder(self, test_dir): + dataset_path = osp.join(test_dir, "dataset") + shutil.copytree(get_test_asset_path("yolo_dataset", self.IMPORTER.NAME), dataset_path) + shutil.rmtree(osp.join(dataset_path, "images", "train")) + + with pytest.raises(InvalidAnnotationError, match="subset image folder"): + Dataset.import_from(dataset_path, self.IMPORTER.NAME).init_cache() + + +class YOLOv8SegmentationExtractorTest(YOLOv8ExtractorTest): + IMPORTER = YOLOv8SegmentationImporter + EXTRACTOR = YOLOv8SegmentationExtractor + + def _prepare_dataset(self, path: str, anno=None) -> Dataset: + return super()._prepare_dataset( + path, anno=Polygon(points=[1, 1, 2, 4, 4, 2, 8, 8], label=0) + ) + + @staticmethod + def _make_some_annotation_values(): + return [0.5, 0.5] * 3 + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + @pytest.mark.parametrize( + "field, field_name", + [ + (1, "polygon point 0 x"), + (2, "polygon point 0 y"), + (3, "polygon point 1 x"), + (4, "polygon point 1 y"), + (5, "polygon point 2 x"), + (6, "polygon point 2 y"), + ], + ) + def test_can_report_invalid_field_type(self, field, field_name, test_dir): + self._check_can_report_invalid_field_type(field, field_name, test_dir) + + +class YOLOv8OrientedBoxesExtractorTest(YOLOv8ExtractorTest): + IMPORTER = YOLOv8OrientedBoxesImporter + EXTRACTOR = YOLOv8OrientedBoxesExtractor + + def _prepare_dataset(self, path: str, anno=None) -> Dataset: + return super()._prepare_dataset( + path, anno=Bbox(1, 1, 2, 4, label=0, attributes=dict(rotation=30)) + ) + + @staticmethod + def _make_some_annotation_values(): + return [0.5, 0.5] * 4 + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_parse(self, helper_tc, test_dir): + expected = self._prepare_dataset(test_dir) + actual = Dataset.import_from(test_dir, self.IMPORTER.NAME) + assert abs(list(actual)[0].annotations[0].attributes["rotation"] - 30) < 0.001 + compare_datasets(helper_tc, expected, actual, ignored_attrs=["rotation"]) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + @pytest.mark.parametrize( + "field, field_name", + [ + (1, "bbox point 0 x"), + (2, "bbox point 0 y"), + (3, "bbox point 1 x"), + (4, "bbox point 1 y"), + (5, "bbox point 2 x"), + (6, "bbox point 2 y"), + ], + ) + def test_can_report_invalid_field_type(self, field, field_name, test_dir): + self._check_can_report_invalid_field_type(field, field_name, test_dir) + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + def test_can_report_invalid_shape(self, test_dir): + self._prepare_dataset(test_dir) + with open(osp.join(test_dir, self._get_annotation_dir(), "a.txt"), "w") as f: + f.write("0 0.1 0.1 0.5 0.1 0.5 0.5 0.5 0.2") + + with pytest.raises(AnnotationImportError) as capture: + Dataset.import_from(test_dir, self.IMPORTER.NAME).init_cache() + assert isinstance(capture.value.__cause__, InvalidAnnotationError) + assert "Given points do not form a rectangle" in str(capture.value.__cause__) + + +class YOLOv8PoseExtractorTest(YOLOv8ExtractorTest): + IMPORTER = YOLOv8PoseImporter + EXTRACTOR = YOLOv8PoseExtractor + + def _prepare_dataset(self, path: str, anno=None) -> Dataset: + dataset = Dataset.from_iterable( + [ + DatasetItem( + "a", + subset="train", + media=Image(np.ones((5, 10, 3))), + annotations=[ + Skeleton( + [ + Points([1, 2], [Points.Visibility.visible.value], label=1), + Points([3, 6], [Points.Visibility.visible.value], label=2), + Points([4, 5], [Points.Visibility.visible.value], label=3), + Points([8, 7], [Points.Visibility.visible.value], label=4), + ], + label=0, + ) + ], + ) + ], + categories={ + AnnotationType.label: LabelCategories.from_iterable( + [ + "test", + ("test_point_0", "test"), + ("test_point_1", "test"), + ("test_point_2", "test"), + ("test_point_3", "test"), + ] + ), + AnnotationType.points: PointsCategories.from_iterable( + [(0, ["test_point_0", "test_point_1", "test_point_2", "test_point_3"], set())] + ), + }, + ) + dataset.export(path, self.EXTRACTOR.NAME, save_media=True) + return dataset + + @staticmethod + def _make_some_annotation_values(): + return [0.5, 0.5, 0.5, 0.5] + [0.5, 0.5, 2] * 4 + + @mark_requirement(Requirements.DATUM_ERROR_REPORTING) + @pytest.mark.parametrize( + "field, field_name", + [ + (5, "skeleton point 0 x"), + (6, "skeleton point 0 y"), + (7, "skeleton point 0 visibility"), + (8, "skeleton point 1 x"), + (9, "skeleton point 1 y"), + (10, "skeleton point 1 visibility"), + ], + ) + def test_can_report_invalid_field_type(self, field, field_name, test_dir): + self._check_can_report_invalid_field_type(field, field_name, test_dir) diff --git a/tests/utils/assets.py b/tests/utils/assets.py new file mode 100644 index 0000000000..d8aaa820ea --- /dev/null +++ b/tests/utils/assets.py @@ -0,0 +1,10 @@ +# Copyright (C) 2023 Intel Corporation +# +# SPDX-License-Identifier: MIT + +import os.path as osp + + +def get_test_asset_path(*args): + cur_dir = osp.dirname(__file__) + return osp.abspath(osp.join(cur_dir, "..", "assets", *args)) diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py new file mode 100644 index 0000000000..5eee14dcb4 --- /dev/null +++ b/tests/utils/test_utils.py @@ -0,0 +1,30 @@ +# Copyright (C) 2019-2024 Intel Corporation +# +# SPDX-License-Identifier: MIT +from typing import Any, List + +import pytest + + +class TestCaseHelper: + """This class will exist until we complete the migration from unittest to pytest. + It is designed to mimic unittest.TestCase behaviors to minimize the migration work labor cost. + """ + + def assertTrue(self, boolean: bool, err_msg: str = ""): + assert boolean, err_msg + + def assertFalse(self, boolean: bool, err_msg: str = ""): + assert not boolean, err_msg + + def assertEqual(self, item1: Any, item2: Any, err_msg: str = ""): + assert item1 == item2, err_msg + + def assertListEqual(self, list1: List[Any], list2: List[Any], err_msg: str = ""): + assert isinstance(list1, list) and isinstance(list2, list), err_msg + assert len(list1) == len(list2), err_msg + for item1, item2 in zip(list1, list2): + self.assertEqual(item1, item2, err_msg) + + def fail(self, msg): + pytest.fail(reason=msg)