137 lines
3.7 KiB
Python
137 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from k1link.compute.pointpillars_postprocess import (
|
|
PointPillarsBox,
|
|
PointPillarsPostprocessError,
|
|
decode_pointpillars_output,
|
|
oriented_3d_iou,
|
|
oriented_bev_iou,
|
|
)
|
|
|
|
|
|
def _outputs(rows: list[list[float]]) -> tuple[np.ndarray, np.ndarray]:
|
|
boxes = np.zeros((1, 393_216, 9), dtype=np.float32)
|
|
boxes[0, : len(rows)] = np.asarray(rows, dtype=np.float32)
|
|
return boxes, np.asarray([len(rows)], dtype=np.int32)
|
|
|
|
|
|
def _row(
|
|
*,
|
|
x: float,
|
|
y: float = 0.0,
|
|
length: float = 4.0,
|
|
width: float = 2.0,
|
|
yaw: float = 0.0,
|
|
class_id: int = 0,
|
|
score: float = 0.9,
|
|
) -> list[float]:
|
|
return [x, y, 0.0, length, width, 1.5, yaw, float(class_id), score]
|
|
|
|
|
|
def _box(*, yaw: float = 0.0, x: float = 0.0) -> PointPillarsBox:
|
|
return PointPillarsBox(
|
|
x_m=x,
|
|
y_m=0.0,
|
|
z_m=0.0,
|
|
length_m=4.0,
|
|
width_m=2.0,
|
|
height_m=1.5,
|
|
yaw_rad=yaw,
|
|
class_id=0,
|
|
model_class="Vehicle",
|
|
score=0.9,
|
|
)
|
|
|
|
|
|
def test_decodes_native_label_order_and_sorts_by_score() -> None:
|
|
output_boxes, num_boxes = _outputs(
|
|
[
|
|
_row(x=20.0, class_id=2, score=0.6),
|
|
_row(x=0.0, class_id=0, score=0.9),
|
|
_row(x=10.0, class_id=1, score=0.8),
|
|
]
|
|
)
|
|
|
|
decoded = decode_pointpillars_output(output_boxes, num_boxes)
|
|
|
|
assert [box.model_class for box in decoded] == [
|
|
"Vehicle",
|
|
"Pedestrian",
|
|
"Cyclist",
|
|
]
|
|
assert [box.score for box in decoded] == pytest.approx([0.9, 0.8, 0.6])
|
|
|
|
|
|
def test_nms_reproduces_nvidia_sample_class_agnostic_suppression() -> None:
|
|
output_boxes, num_boxes = _outputs(
|
|
[
|
|
_row(x=0.0, class_id=0, score=0.9),
|
|
_row(x=0.1, class_id=1, score=0.8),
|
|
_row(x=20.0, class_id=1, score=0.7),
|
|
]
|
|
)
|
|
|
|
decoded = decode_pointpillars_output(output_boxes, num_boxes)
|
|
|
|
assert [(box.x_m, box.model_class) for box in decoded] == [
|
|
(0.0, "Vehicle"),
|
|
(20.0, "Pedestrian"),
|
|
]
|
|
|
|
|
|
def test_pre_nms_cap_is_applied_after_stable_score_ordering() -> None:
|
|
output_boxes, num_boxes = _outputs(
|
|
[
|
|
_row(x=0.0, score=0.7),
|
|
_row(x=10.0, score=0.9),
|
|
_row(x=20.0, score=0.8),
|
|
]
|
|
)
|
|
|
|
decoded = decode_pointpillars_output(
|
|
output_boxes,
|
|
num_boxes,
|
|
pre_nms_top_n=2,
|
|
)
|
|
|
|
assert [box.x_m for box in decoded] == [10.0, 20.0]
|
|
|
|
|
|
def test_oriented_bev_iou_handles_rotation_and_separation() -> None:
|
|
assert oriented_bev_iou(_box(), _box()) == pytest.approx(1.0)
|
|
assert oriented_bev_iou(_box(), _box(yaw=math.pi / 2.0)) == pytest.approx(
|
|
1.0 / 3.0
|
|
)
|
|
assert oriented_bev_iou(_box(), _box(x=20.0)) == 0.0
|
|
assert oriented_3d_iou(_box(), _box()) == pytest.approx(1.0)
|
|
assert oriented_3d_iou(_box(), _box(x=20.0)) == 0.0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("row", "message"),
|
|
[
|
|
(_row(x=0.0, class_id=3), "class id"),
|
|
(_row(x=0.0, length=0.0), "dimensions or score"),
|
|
(_row(x=0.0, score=0.09), "dimensions or score"),
|
|
(_row(x=0.0, score=1.1), "dimensions or score"),
|
|
],
|
|
)
|
|
def test_invalid_candidate_fails_closed(row: list[float], message: str) -> None:
|
|
output_boxes, num_boxes = _outputs([row])
|
|
|
|
with pytest.raises(PointPillarsPostprocessError, match=message):
|
|
decode_pointpillars_output(output_boxes, num_boxes)
|
|
|
|
|
|
def test_output_tensor_contract_is_exact() -> None:
|
|
with pytest.raises(PointPillarsPostprocessError, match="output_boxes contract"):
|
|
decode_pointpillars_output(
|
|
np.zeros((1, 1, 9), dtype=np.float32),
|
|
np.asarray([0], dtype=np.int32),
|
|
)
|