diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 24a66204..23e3062e 100644 --- a/supervision/dataset/formats/yolo.py +++ b/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)