okooooooo/supervision
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #399 from akashAD98/update_detection_sv

Browse Source 
Remove class list dependency for roboflow to detection conversion
This commit is contained in:
Piotr Skalski authored and GitHub committed 2023-10-04 12:08:04 +02:00
2 parents 492f033882 e911516180
commit 67fa66c7e7
3 files changed
+18 -22
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

supervision/detection/core.py
+4 -6
View File
@@ -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:
supervision/detection/utils.py
+3 -3
View File
@@ -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)
test/detection/test_utils.py
+11 -13
View File
@@ -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])