🛠️ fix ready for merging
This commit is contained in:
SkalskiP
parent
0f889935fd
commit
7de50e6b0b
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Reference in New Issue