Merge pull request #300 from ShubhamKanitkar32/null-check-fillPoly
Fix issue #299
This commit is contained in:
Piotr Skalski
GitHub
commit
2ec33dbe9a
|
|
@ -20,6 +20,7 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n
|
|||
"""
|
||||
width, height = resolution_wh
|
||||
mask = np.zeros((height, width))
|
||||
|
||||
cv2.fillPoly(mask, [polygon], color=1)
|
||||
return mask
|
||||
|
||||
|
|
@ -354,23 +355,23 @@ 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
|
||||
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)
|
||||
|
||||
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)
|
||||
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
|
||||
|
|
|
|||
|
|
@ -11,6 +11,9 @@ from supervision.detection.utils import (
|
|||
process_roboflow_result,
|
||||
)
|
||||
|
||||
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",
|
||||
|
|
@ -286,7 +289,143 @@ 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