diff --git a/supervision/detection/core.py b/supervision/detection/core.py index f7f1bb5f..0ba28491 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -432,7 +432,7 @@ class Detections: ) @classmethod - def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections: + def from_roboflow(cls, roboflow_result: dict) -> Detections: """ Create a Detections object from the [Roboflow](https://roboflow.com/) API inference result. @@ -440,8 +440,6 @@ class Detections: Args: roboflow_result (dict): The result from the Roboflow API containing predictions. - class_list (List[str]): A list of class names - corresponding to the class IDs in the API result. Returns: (Detections): A Detections object containing the bounding boxes, class IDs, @@ -458,19 +456,19 @@ class Detections: ... "y": 0.5, ... "width": 0.2, ... "height": 0.3, + ... "class_id": 0, ... "class": "person", ... "confidence": 0.9 ... }, ... # ... more predictions ... ... ] ... } - >>> class_list = ["person", "car", "dog"] - >>> detections = sv.Detections.from_roboflow(roboflow_result, class_list) + >>> detections = sv.Detections.from_roboflow(roboflow_result) ``` """ xyxy, confidence, class_id, masks = process_roboflow_result( - roboflow_result=roboflow_result, class_list=class_list + roboflow_result=roboflow_result ) if np.asarray(xyxy).shape[0] == 0: diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index e882173b..2beff657 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -332,7 +332,7 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]: def process_roboflow_result( - roboflow_result: dict, class_list: List[str] + roboflow_result: dict, ) -> 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 @@ -357,7 +357,7 @@ def process_roboflow_result( if "points" not in prediction: xyxy.append([x_min, y_min, x_max, y_max]) - class_id.append(class_list.index(prediction["class"])) + class_id.append(prediction["class_id"]) confidence.append(prediction["confidence"]) elif len(prediction["points"]) >= 3: polygon = np.array( @@ -365,7 +365,7 @@ def process_roboflow_result( ) mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height)) xyxy.append([x_min, y_min, x_max, y_max]) - class_id.append(class_list.index(prediction["class"])) + class_id.append(prediction["class_id"]) confidence.append(prediction["confidence"]) masks.append(mask) diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py index 63a2c33c..529914ae 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -260,11 +260,10 @@ def test_filter_polygons_by_area( @pytest.mark.parametrize( - "roboflow_result, class_list, expected_result, exception", + "roboflow_result, 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 @@ -277,12 +276,12 @@ def test_filter_polygons_by_area( "width": 50.0, "height": 50.0, "confidence": 0.9, + "class_id": 0, "class": "person", } ], "image": {"width": 1000, "height": 1000}, }, - ["person", "car", "truck"], ( np.array([[175.0, 275.0, 225.0, 325.0]]), np.array([0.9]), @@ -300,6 +299,7 @@ def test_filter_polygons_by_area( "width": 50.0, "height": 50.0, "confidence": 0.9, + "class_id": 0, "class": "person", }, { @@ -308,16 +308,16 @@ def test_filter_polygons_by_area( "width": 100.0, "height": 100.0, "confidence": 0.8, + "class_id": 7, "class": "truck", }, ], "image": {"width": 1000, "height": 1000}, }, - ["person", "car", "truck"], ( np.array([[175.0, 275.0, 225.0, 325.0], [450.0, 450.0, 550.0, 550.0]]), np.array([0.9, 0.8]), - np.array([0, 2]), + np.array([0, 7]), None, ), DoesNotRaise(), @@ -331,13 +331,13 @@ def test_filter_polygons_by_area( "width": 50.0, "height": 50.0, "confidence": 0.9, + "class_id": 0, "class": "person", "points": [], } ], "image": {"width": 1000, "height": 1000}, }, - ["person", "car", "truck"], (np.empty((0, 4)), np.empty(0), np.empty(0), None), DoesNotRaise(), ), # single incorrect instance segmentation result with no points @@ -350,13 +350,13 @@ def test_filter_polygons_by_area( "width": 50.0, "height": 50.0, "confidence": 0.9, + "class_id": 0, "class": "person", "points": [{"x": 200.0, "y": 300.0}, {"x": 250.0, "y": 300.0}], } ], "image": {"width": 1000, "height": 1000}, }, - ["person", "car", "truck"], (np.empty((0, 4)), np.empty(0), np.empty(0), None), DoesNotRaise(), ), # single incorrect instance segmentation result with no enough points @@ -369,6 +369,7 @@ def test_filter_polygons_by_area( "width": 50.0, "height": 50.0, "confidence": 0.9, + "class_id": 0, "class": "person", "points": [ {"x": 200.0, "y": 300.0}, @@ -380,7 +381,6 @@ def test_filter_polygons_by_area( ], "image": {"width": 1000, "height": 1000}, }, - ["person", "car", "truck"], ( np.array([[175.0, 275.0, 225.0, 325.0]]), np.array([0.9]), @@ -398,6 +398,7 @@ def test_filter_polygons_by_area( "width": 50.0, "height": 50.0, "confidence": 0.9, + "class_id": 0, "class": "person", "points": [ {"x": 200.0, "y": 300.0}, @@ -412,13 +413,13 @@ def test_filter_polygons_by_area( "width": 100.0, "height": 100.0, "confidence": 0.8, + "class_id": 7, "class": "truck", "points": [], }, ], "image": {"width": 1000, "height": 1000}, }, - ["person", "car", "truck"], ( np.array([[175.0, 275.0, 225.0, 325.0]]), np.array([0.9]), @@ -431,14 +432,11 @@ def test_filter_polygons_by_area( ) 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 - ) + result = process_roboflow_result(roboflow_result=roboflow_result) 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])