diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 1c0ad8af..b13f0533 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -172,7 +172,7 @@ class OrientedBoxAnnotator(BaseAnnotator): return scene for detection_idx in range(len(detections)): - bbox = np.int0(detections.data.get(ORIENTED_BOX_COORDINATES)[detection_idx]) + bbox = np.intp(detections.data.get(ORIENTED_BOX_COORDINATES)[detection_idx]) color = resolve_color( color=self.color, detections=detections, diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index c8863df3..e6d69f76 100644 --- a/supervision/dataset/core.py +++ b/supervision/dataset/core.py @@ -266,6 +266,7 @@ class DetectionDataset(BaseDataset): annotations_directory_path: str, data_yaml_path: str, force_masks: bool = False, + is_obb: bool = False, ) -> DetectionDataset: """ Creates a Dataset instance from YOLO formatted data. @@ -280,6 +281,9 @@ class DetectionDataset(BaseDataset): force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. + is_obb (bool, optional): If True, loads the annotations in OBB format. + OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`, + where pairs of [x, y] are box corners. Returns: DetectionDataset: A DetectionDataset instance @@ -312,6 +316,7 @@ class DetectionDataset(BaseDataset): annotations_directory_path=annotations_directory_path, data_yaml_path=data_yaml_path, force_masks=force_masks, + is_obb=is_obb, ) return DetectionDataset(classes=classes, images=images, annotations=annotations) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index c8136205..f634d4fa 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -5,6 +5,7 @@ from typing import Dict, List, Optional, Tuple import cv2 import numpy as np +from supervision.config import ORIENTED_BOX_COORDINATES from supervision.dataset.utils import approximate_mask_with_polygons from supervision.detection.core import Detections from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy @@ -70,12 +71,15 @@ def _image_name_to_annotation_name(image_name: str) -> str: def yolo_annotations_to_detections( - lines: List[str], resolution_wh: Tuple[int, int], with_masks: bool + lines: List[str], + resolution_wh: Tuple[int, int], + with_masks: bool, + is_obb: bool = False, ) -> Detections: if len(lines) == 0: return Detections.empty() - class_id, relative_xyxy, relative_polygon = [], [], [] + class_id, relative_xyxy, relative_polygon, relative_xyxyxyxy = [], [], [], [] w, h = resolution_wh for line in lines: values = line.split() @@ -88,21 +92,30 @@ def yolo_annotations_to_detections( elif len(values) > 5: polygon = _parse_polygon(values=values[1:]) relative_xyxy.append(polygon_to_xyxy(polygon=polygon)) + if is_obb: + relative_xyxyxyxy.append(np.array(values[1:])) if with_masks: relative_polygon.append(polygon) class_id = np.array(class_id, dtype=int) relative_xyxy = np.array(relative_xyxy, dtype=np.float32) xyxy = relative_xyxy * np.array([w, h, w, h], dtype=np.float32) + data = {} + + if is_obb: + relative_xyxyxyxy = np.array(relative_xyxyxyxy, dtype=np.float32) + xyxyxyxy = relative_xyxyxyxy.reshape(-1, 4, 2) + xyxyxyxy *= np.array([w, h], dtype=np.float32) + data[ORIENTED_BOX_COORDINATES] = xyxyxyxy if not with_masks: - return Detections(class_id=class_id, xyxy=xyxy) + return Detections(class_id=class_id, xyxy=xyxy, data=data) polygons = [ (polygon * np.array(resolution_wh)).astype(int) for polygon in relative_polygon ] mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh) - return Detections(class_id=class_id, xyxy=xyxy, mask=mask) + return Detections(class_id=class_id, xyxy=xyxy, data=data, mask=mask) def load_yolo_annotations( @@ -110,6 +123,7 @@ def load_yolo_annotations( annotations_directory_path: str, data_yaml_path: str, force_masks: bool = False, + is_obb: bool = False, ) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: """ Loads YOLO annotations and returns class names, images, @@ -123,6 +137,9 @@ def load_yolo_annotations( YAML file containing class information. force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present. + is_obb (bool, optional): If True, loads the annotations in OBB format. + OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`, + where pairs of [x, y] are box corners. Returns: Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: @@ -156,9 +173,11 @@ def load_yolo_annotations( with_masks = _with_mask(lines=lines) with_masks = force_masks if force_masks else with_masks annotation = yolo_annotations_to_detections( - lines=lines, resolution_wh=resolution_wh, with_masks=with_masks + lines=lines, + resolution_wh=resolution_wh, + with_masks=with_masks, + is_obb=is_obb, ) - images[image_path] = image annotations[image_path] = annotation return classes, images, annotations