yolo annotations box_to_polygon initial implementation

supervision/dataset/formats/yolo.py

@@ -2,6 +2,7 @@ import numpy as np
 
 from typing import List, Tuple
 
+from supervision import polygon_to_mask
 from supervision.detection.core import Detections
 
 
@@ -15,10 +16,27 @@ def parse_box(values: List[str]) -> np.ndarray:
     ], dtype=np.float32)
 
 
-def parse_segments(values: List[str]) -> np.ndarray:
+def box_to_polygon(box: np.ndarray) -> np.ndarray:
+    return np.array([
+        [box[0], box[1]],
+        [box[2], box[1]],
+        [box[2], box[3]],
+        [box[0], box[3]]
+    ])
+
+
+def parse_polygon(values: List[str]) -> np.ndarray:
     pass
 
 
+def polygons_to_masks(polygons: List[np.ndarray], resolution_wh: Tuple[int, int]) -> np.ndarray:
+    return np.array([
+        polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
+        for polygon
+        in polygons
+    ], dtype=bool)
+
+
 def yolo_annotations_to_detections(
     lines: List[str],
     resolution_wh: Tuple[int, int],
@@ -26,27 +44,37 @@ def yolo_annotations_to_detections(
 ) -> Detections:
     if len(lines) == 0:
         return Detections.empty()
-    class_id, xyxy, mask = [], [], []
+    class_id, relative_xyxy, relative_polygons = [], [], []
     w, h = resolution_wh
     for line in lines:
         values = line.split()
         class_id.append(int(values[0]))
         if len(values) == 5 and not force_segmentations:
-            xyxy.append(parse_box(values=values[1:]))
+            box = parse_box(values=values[1:])
+            relative_xyxy.append(box)
         elif len(values) == 5 and force_segmentations:
-            pass
+            box = parse_box(values=values[1:])
+            relative_xyxy.append(box)
+            relative_polygons.append(box_to_polygon(box=box))
         elif len(values) > 5 and not force_segmentations:
             pass
         elif len(values) > 5 and force_segmentations:
             pass
 
     class_id = np.array(class_id, dtype=int)
-    xyxy = np.array(xyxy, dtype=np.float32)
-    xyxy = xyxy * np.array([w, h, w, h], dtype=np.float32)
-    return Detections(
-        class_id=class_id,
-        xyxy=xyxy
-    )
+    relative_xyxy = np.array(relative_xyxy, dtype=np.float32)
+    xyxy = relative_xyxy * np.array([w, h, w, h], dtype=np.float32)
+
+    if not force_segmentations:
+        return Detections(class_id=class_id, xyxy=xyxy)
+
+    polygons = [
+        (relative_polygon * np.array(resolution_wh)).astype(int)
+        for relative_polygon
+        in relative_polygons
+    ]
+    mask = polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
+    return Detections(class_id=class_id, xyxy=xyxy, mask=mask)