From 3d0aec03e03559a6cd71686c74b40cb4153ce09e Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 8 May 2023 13:51:49 +0200 Subject: [PATCH] yolo_annotations_to_detections function full implementation --- supervision/dataset/formats/yolo.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 23e3062e..ffbb68b0 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -2,7 +2,7 @@ import numpy as np from typing import List, Tuple -from supervision import polygon_to_mask +from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy from supervision.detection.core import Detections @@ -26,7 +26,7 @@ def box_to_polygon(box: np.ndarray) -> np.ndarray: def parse_polygon(values: List[str]) -> np.ndarray: - pass + return np.array(values, dtype=np.float32).reshape(-1, 2) def polygons_to_masks(polygons: List[np.ndarray], resolution_wh: Tuple[int, int]) -> np.ndarray: @@ -44,7 +44,7 @@ def yolo_annotations_to_detections( ) -> Detections: if len(lines) == 0: return Detections.empty() - class_id, relative_xyxy, relative_polygons = [], [], [] + class_id, relative_xyxy, relative_polygon = [], [], [] w, h = resolution_wh for line in lines: values = line.split() @@ -55,11 +55,11 @@ def yolo_annotations_to_detections( elif len(values) == 5 and force_segmentations: 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 + relative_polygon.append(box_to_polygon(box=box)) + else: + polygon = parse_polygon(values=values[1:]) + relative_polygon.append(polygon) + relative_xyxy.append(polygon_to_xyxy(polygon=polygon)) class_id = np.array(class_id, dtype=int) relative_xyxy = np.array(relative_xyxy, dtype=np.float32) @@ -69,15 +69,14 @@ def yolo_annotations_to_detections( return Detections(class_id=class_id, xyxy=xyxy) polygons = [ - (relative_polygon * np.array(resolution_wh)).astype(int) - for relative_polygon - in relative_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) - def detections_to_yolo_annotations(detections: Detections, resolution_wh: Tuple[int, int]) -> List[str]: if len(detections): return []