From 453d4df0685fe80734706155018cf5c19f490a1f Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 14 Jun 2023 16:38:08 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=AA=20refactor=20+=20some=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 47 ++++---------------------- supervision/detection/utils.py | 46 +++++++++++++++++++++++++ test/detection/test_utils.py | 61 +++++++++++++++++++++++++++++++++- 3 files changed, 113 insertions(+), 41 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index e65082b5..97150e45 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -9,7 +9,7 @@ import numpy as np from supervision.detection.utils import ( extract_yolov8_masks, non_max_suppression, - polygon_to_mask, + process_roboflow_result, xywh_to_xyxy, ) from supervision.geometry.core import Position @@ -314,46 +314,13 @@ class Detections: >>> detections = sv.Detections.from_roboflow(roboflow_result, class_list) ``` """ - xyxy = [] - confidence = [] - class_id = [] - masks = [] - - img_width = int(roboflow_result["image"]["width"]) - img_height = int(roboflow_result["image"]["height"]) - - for prediction in roboflow_result["predictions"]: - x = prediction["x"] - y = prediction["y"] - width = prediction["width"] - height = prediction["height"] - x_min = x - width / 2 - y_min = y - height / 2 - x_max = x_min + width - y_max = y_min + height - xyxy.append([x_min, y_min, x_max, y_max]) - class_id.append(class_list.index(prediction["class"])) - confidence.append(prediction["confidence"]) - - if "points" not in prediction: - continue - - points = prediction["points"] - - polygon = np.array( - [(p["x"], p["y"]) for p in points], dtype=np.int32 - ).reshape((-1, 1, 2)) - - mask = polygon_to_mask(polygon, resolution_wh=(img_width, img_height)) - - masks.append(mask) - - masks = np.array(masks) if len(masks) > 0 else None - + xyxy, confidence, class_id, masks = process_roboflow_result( + roboflow_result=roboflow_result, class_list=class_list + ) return Detections( - xyxy=np.array(xyxy), - confidence=np.array(confidence), - class_id=np.array(class_id).astype(int), + xyxy=xyxy, + confidence=confidence, + class_id=class_id, mask=masks, ) diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 00d36901..63206ddc 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -292,3 +292,49 @@ def extract_yolov8_masks(yolov8_results) -> Optional[np.ndarray]: mask_maps.append(mask) return np.asarray(mask_maps, dtype=bool) + + +def process_roboflow_result( + roboflow_result: dict, class_list: List[str] +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]: + if not roboflow_result["predictions"]: + return np.empty((0, 4)), np.empty(0), np.empty(0), None + + xyxy = [] + confidence = [] + class_id = [] + masks = [] + + image_width = int(roboflow_result["image"]["width"]) + image_height = int(roboflow_result["image"]["height"]) + + for prediction in roboflow_result["predictions"]: + x = prediction["x"] + y = prediction["y"] + width = prediction["width"] + height = prediction["height"] + x_min = x - width / 2 + y_min = y - height / 2 + x_max = x_min + width + y_max = y_min + height + + xyxy.append([x_min, y_min, x_max, y_max]) + class_id.append(class_list.index(prediction["class"])) + confidence.append(prediction["confidence"]) + + if "points" not in prediction: + continue + + polygon = np.array( + [[point["x"], point["y"]] for point in prediction["points"]], dtype=int + ) + + mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height)) + masks.append(mask) + + xyxy = np.array(xyxy) + confidence = np.array(confidence) + class_id = np.array(class_id).astype(int) + masks = np.array(masks, dtype=bool) if len(masks) > 0 else None + + return xyxy, confidence, class_id, masks diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index cf7abfd3..80553854 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -5,7 +5,8 @@ import pytest import numpy as np -from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area +from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area, \ + process_roboflow_result @pytest.mark.parametrize( @@ -276,3 +277,61 @@ def test_filter_polygons_by_area( assert len(result) == len(expected_result) for result_polygon, expected_result_polygon in zip(result, expected_result): assert np.array_equal(result_polygon, expected_result_polygon) + + +@pytest.mark.parametrize( + "roboflow_result, class_list, expected_result, exception", + [ + ( + { + "predictions": [], + "image": {"width": 1000, "height": 1000} + }, + ["person", "car", "truck"], + ( + np.empty((0, 4)), + np.empty(0), + np.empty(0), + None + ), + DoesNotRaise() + ), # empty result + ( + { + "predictions": [ + { + "x": 200.0, + "y": 300.0, + "width": 50.0, + "height": 50.0, + "confidence": 0.9, + "class": "person" + } + ], + "image": {"width": 1000, "height": 1000} + }, + ["person", "car", "truck"], + ( + np.array([ + [175.0, 275.0, 225.0, 325.0] + ]), + np.array([0.9]), + np.array([0]), + None + ), + DoesNotRaise() + ), # single bounding box + ] +) +def test_process_roboflow_result( + roboflow_result: dict, + class_list: List[str], + expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]], + exception: Exception +) -> None: + with exception: + result = process_roboflow_result(roboflow_result=roboflow_result, class_list=class_list) + assert np.array_equal(result[0], expected_result[0]) + assert np.array_equal(result[1], expected_result[1]) + assert np.array_equal(result[2], expected_result[2]) + assert (result[3] is None and expected_result[3] is None) or (np.array_equal(result[3], expected_result[3]))