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

initial update of `process_roboflow_result` function

This commit is contained in:
SkalskiP 2024-01-02 21:00:17 +01:00
parent 8528a8a738
commit 706ff0fd46
2 changed files with 52 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
supervision/detection/utils.py
View File

@ -331,13 +331,21 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
def process_roboflow_result(
roboflow_result: dict,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray]:
) -> Tuple[
np.ndarray,
np.ndarray,
np.ndarray,
Optional[np.ndarray],
np.ndarray,
Dict[str, List[np.ndarray]]
]:
if not roboflow_result["predictions"]:
return np.empty((0, 4)), np.empty(0), np.empty(0), None, None
return np.empty((0, 4)), np.empty(0), np.empty(0), None, None, {"class_name": np.empty(0)}
xyxy = []
confidence = []
class_id = []
class_name = []
masks = []
tracker_ids = []
@ -357,6 +365,7 @@ def process_roboflow_result(
if "points" not in prediction:
xyxy.append([x_min, y_min, x_max, y_max])
class_id.append(prediction["class_id"])
class_name.append(prediction["class"])
confidence.append(prediction["confidence"])
if "tracker_id" in prediction:
tracker_ids.append(prediction["tracker_id"])
@ -367,6 +376,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(prediction["class_id"])
class_name.append(prediction["class"])
confidence.append(prediction["confidence"])
masks.append(mask)
if "tracker_id" in prediction:
@ -375,10 +385,12 @@ def process_roboflow_result(
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)
class_name = np.array(class_name) if len(class_name) > 0 else np.empty(0)
masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
tracker_id = np.array(tracker_ids).astype(int) if len(tracker_ids) > 0 else None
data = {"class_name": class_name}
return xyxy, confidence, class_id, masks, tracker_id
return xyxy, confidence, class_id, masks, tracker_id, data
def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:

40
test/detection/test_utils.py
View File

@ -268,7 +268,14 @@ def test_filter_polygons_by_area(
[
(
{"predictions": [], "image": {"width": 1000, "height": 1000}},
(np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
(
np.empty((0, 4)),
np.empty(0),
np.empty(0),
None,
None,
{"class_name": np.empty(0)}
),
DoesNotRaise(),
), # empty result
(
@ -292,6 +299,7 @@ def test_filter_polygons_by_area(
np.array([0]),
None,
None,
{"class_name": np.array(["person"])}
),
DoesNotRaise(),
), # single correct object detection result
@ -327,6 +335,7 @@ def test_filter_polygons_by_area(
np.array([0, 7]),
None,
np.array([1, 2]),
{"class_name": np.array(["person", "truck"])}
),
DoesNotRaise(),
), # two correct object detection result
@ -347,7 +356,14 @@ def test_filter_polygons_by_area(
],
"image": {"width": 1000, "height": 1000},
},
(np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
(
np.empty((0, 4)),
np.empty(0),
np.empty(0),
None,
None,
{"class_name": np.empty(0)}
),
DoesNotRaise(),
), # single incorrect instance segmentation result with no points
(
@ -366,7 +382,14 @@ def test_filter_polygons_by_area(
],
"image": {"width": 1000, "height": 1000},
},
(np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
(
np.empty((0, 4)),
np.empty(0),
np.empty(0),
None,
None,
{"class_name": np.empty(0)}
),
DoesNotRaise(),
), # single incorrect instance segmentation result with no enough points
(
@ -396,6 +419,7 @@ def test_filter_polygons_by_area(
np.array([0]),
TEST_MASK,
None,
{"class_name": np.array(["person"])}
),
DoesNotRaise(),
), # single incorrect instance segmentation result with no enough points
@ -436,6 +460,7 @@ def test_filter_polygons_by_area(
np.array([0]),
TEST_MASK,
None,
{"class_name": np.array(["person"])}
),
DoesNotRaise(),
), # two instance segmentation results - one correct, one incorrect
@ -459,6 +484,15 @@ def test_process_roboflow_result(
assert (result[4] is None and expected_result[4] is None) or (
np.array_equal(result[4], expected_result[4])
)
for key in result[5]:
if isinstance(result[5][key], np.ndarray):
assert np.array_equal(
result[5][key], expected_result[5][key]
), f"Mismatch in arrays for key {key}"
else:
assert (
result[5][key] == expected_result[5][key]
), f"Mismatch in non-array data for key {key}"
@pytest.mark.parametrize(