diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 7e8a4ee9..95d1fc26 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -355,24 +355,25 @@ def process_roboflow_result( 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 - - if len(prediction["points"]) >= 3: - polygon = np.array( - [[point["x"], point["y"]] for point in prediction["points"]], dtype=int - ) - + xyxy.append([x_min, y_min, x_max, y_max]) + class_id.append(class_list.index(prediction["class"])) + confidence.append(prediction["confidence"]) + elif len(prediction["points"]) >= 3: + 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)) + xyxy.append([x_min, y_min, x_max, y_max]) + class_id.append(class_list.index(prediction["class"])) + confidence.append(prediction["confidence"]) masks.append(mask) - xyxy = np.array(xyxy) - confidence = np.array(confidence) - class_id = np.array(class_id).astype(int) + xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4)) + confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0) + class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0) 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 b096fd08..316f23e5 100644 --- a/test/detection/test_utils.py +++ b/test/detection/test_utils.py @@ -12,6 +12,10 @@ from supervision.detection.utils import ( ) +TEST_MASK = np.zeros((1, 1000, 1000), dtype=bool) +TEST_MASK[:, 300:351, 200:251] = True + + @pytest.mark.parametrize( "predictions, iou_threshold, expected_result, exception", [ @@ -259,9 +263,17 @@ def test_filter_polygons_by_area( "roboflow_result, class_list, expected_result, exception", [ ( - {"predictions": [], "image": {"width": 1000, "height": 1000}}, + { + "predictions": [], + "image": {"width": 1000, "height": 1000} + }, ["person", "car", "truck"], - (np.empty((0, 4)), np.empty(0), np.empty(0), None), + ( + np.empty((0, 4)), + np.empty(0), + np.empty(0), + None + ), DoesNotRaise(), ), # empty result ( @@ -286,7 +298,156 @@ def test_filter_polygons_by_area( None, ), DoesNotRaise(), - ), # single bounding box + ), # single correct object detection result +( + { + "predictions": [ + { + "x": 200.0, + "y": 300.0, + "width": 50.0, + "height": 50.0, + "confidence": 0.9, + "class": "person", + }, + { + "x": 500.0, + "y": 500.0, + "width": 100.0, + "height": 100.0, + "confidence": 0.8, + "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]), + None, + ), + DoesNotRaise(), + ), # two correct object detection result + ( + { + "predictions": [ + { + "x": 200.0, + "y": 300.0, + "width": 50.0, + "height": 50.0, + "confidence": 0.9, + "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 + ( + { + "predictions": [ + { + "x": 200.0, + "y": 300.0, + "width": 50.0, + "height": 50.0, + "confidence": 0.9, + "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 + ( + { + "predictions": [ + { + "x": 200.0, + "y": 300.0, + "width": 50.0, + "height": 50.0, + "confidence": 0.9, + "class": "person", + "points": [ + {"x": 200.0, "y": 300.0}, + {"x": 250.0, "y": 300.0}, + {"x": 250.0, "y": 350.0}, + {"x": 200.0, "y": 350.0}, + ] + } + ], + "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]), + TEST_MASK + ), + DoesNotRaise(), + ), # single incorrect instance segmentation result with no enough points + ( + { + "predictions": [ + { + "x": 200.0, + "y": 300.0, + "width": 50.0, + "height": 50.0, + "confidence": 0.9, + "class": "person", + "points": [ + {"x": 200.0, "y": 300.0}, + {"x": 250.0, "y": 300.0}, + {"x": 250.0, "y": 350.0}, + {"x": 200.0, "y": 350.0}, + ] + }, + { + "x": 500.0, + "y": 500.0, + "width": 100.0, + "height": 100.0, + "confidence": 0.8, + "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]), + np.array([0]), + TEST_MASK + ), + DoesNotRaise(), + ), # two instance segmentation results - one correct, one incorrect ], ) def test_process_roboflow_result(