Merge pull request #1227 from Bhavay-2001/OBB

Adding support for Yolo-OBB (Oriented Bounding Boxes) format

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,

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)
 

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