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

🧪 refactor + some tests

This commit is contained in:
SkalskiP 2023-06-14 16:38:08 +02:00
parent dd4cae328b
commit 453d4df068
3 changed files with 113 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
supervision/detection/core.py
View File

@ -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,
)

46
supervision/detection/utils.py
View File

@ -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

61
test/detection/test_utils.py
View File

@ -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]))