322 lines
10 KiB
Python
322 lines
10 KiB
Python
"""Validated NVIDIA PointPillars candidate decoding and sample-compatible NMS."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Final
|
|
|
|
import numpy as np
|
|
|
|
POINTPILLARS_MODEL_CLASSES: Final = ("Vehicle", "Pedestrian", "Cyclist")
|
|
POINTPILLARS_OUTPUT_FIELDS: Final = (
|
|
"x",
|
|
"y",
|
|
"z",
|
|
"length",
|
|
"width",
|
|
"height",
|
|
"yaw",
|
|
"class_id",
|
|
"score",
|
|
)
|
|
POINTPILLARS_NMS_IOU_THRESHOLD: Final = 0.01
|
|
POINTPILLARS_PRE_NMS_TOP_N: Final = 4_096
|
|
POINTPILLARS_EMBEDDED_SCORE_THRESHOLD: Final = 0.1
|
|
POINTPILLARS_MODEL_POINT_CLOUD_RANGE: Final = (
|
|
-51.20000076293945,
|
|
-51.20000076293945,
|
|
-1.399999976158142,
|
|
51.20000076293945,
|
|
51.20000076293945,
|
|
4.400000095367432,
|
|
)
|
|
NVIDIA_REFERENCE_COMMIT: Final = "a540badc47812a17a94e924b537d49ad3969b5a8"
|
|
_EPSILON: Final = 1e-8
|
|
|
|
|
|
class PointPillarsPostprocessError(RuntimeError):
|
|
"""PointPillars output does not satisfy the frozen NVIDIA contract."""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PointPillarsBox:
|
|
x_m: float
|
|
y_m: float
|
|
z_m: float
|
|
length_m: float
|
|
width_m: float
|
|
height_m: float
|
|
yaw_rad: float
|
|
class_id: int
|
|
model_class: str
|
|
score: float
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Geometry:
|
|
corners: tuple[tuple[float, float], ...]
|
|
area: float
|
|
minimum_x: float
|
|
maximum_x: float
|
|
minimum_y: float
|
|
maximum_y: float
|
|
|
|
|
|
def decode_pointpillars_output(
|
|
output_boxes: np.ndarray,
|
|
num_boxes: np.ndarray,
|
|
*,
|
|
pre_nms_top_n: int = POINTPILLARS_PRE_NMS_TOP_N,
|
|
nms_iou_threshold: float = POINTPILLARS_NMS_IOU_THRESHOLD,
|
|
) -> tuple[PointPillarsBox, ...]:
|
|
"""Decode validated rows and reproduce the NVIDIA sample's class-agnostic NMS."""
|
|
|
|
boxes = np.asarray(output_boxes)
|
|
counts = np.asarray(num_boxes)
|
|
if boxes.shape != (1, 393_216, 9) or boxes.dtype != np.float32:
|
|
raise PointPillarsPostprocessError("PointPillars output_boxes contract changed")
|
|
if counts.shape != (1,) or counts.dtype != np.int32:
|
|
raise PointPillarsPostprocessError("PointPillars num_boxes contract changed")
|
|
count = int(counts[0])
|
|
if count < 0 or count > boxes.shape[1]:
|
|
raise PointPillarsPostprocessError("PointPillars candidate count is invalid")
|
|
if (
|
|
isinstance(pre_nms_top_n, bool)
|
|
or not isinstance(pre_nms_top_n, int)
|
|
or pre_nms_top_n < 1
|
|
or not math.isfinite(nms_iou_threshold)
|
|
or not 0.0 <= nms_iou_threshold <= 1.0
|
|
):
|
|
raise PointPillarsPostprocessError("PointPillars NMS parameters are invalid")
|
|
if count == 0:
|
|
return ()
|
|
|
|
candidates = boxes[0, :count]
|
|
if not np.isfinite(candidates).all():
|
|
raise PointPillarsPostprocessError("PointPillars candidate is non-finite")
|
|
decoded = tuple(_decode_row(row) for row in candidates)
|
|
ordered = tuple(
|
|
sorted(
|
|
decoded,
|
|
key=lambda box: box.score,
|
|
reverse=True,
|
|
)[:pre_nms_top_n]
|
|
)
|
|
return _class_agnostic_nms(ordered, nms_iou_threshold)
|
|
|
|
|
|
def oriented_bev_iou(one: PointPillarsBox, another: PointPillarsBox) -> float:
|
|
"""Return BEV IoU for two validated oriented boxes."""
|
|
|
|
one_geometry = _geometry(one)
|
|
another_geometry = _geometry(another)
|
|
overlap = _intersection_area(one_geometry, another_geometry)
|
|
denominator = one_geometry.area + another_geometry.area - overlap
|
|
return overlap / max(denominator, _EPSILON)
|
|
|
|
|
|
def oriented_3d_iou(one: PointPillarsBox, another: PointPillarsBox) -> float:
|
|
"""Return oriented 3D IoU for two center-based LiDAR-frame boxes."""
|
|
|
|
one_geometry = _geometry(one)
|
|
another_geometry = _geometry(another)
|
|
bev_overlap = _intersection_area(one_geometry, another_geometry)
|
|
one_minimum_z = one.z_m - one.height_m / 2.0
|
|
one_maximum_z = one.z_m + one.height_m / 2.0
|
|
another_minimum_z = another.z_m - another.height_m / 2.0
|
|
another_maximum_z = another.z_m + another.height_m / 2.0
|
|
height_overlap = max(
|
|
0.0,
|
|
min(one_maximum_z, another_maximum_z)
|
|
- max(one_minimum_z, another_minimum_z),
|
|
)
|
|
intersection = bev_overlap * height_overlap
|
|
one_volume = one_geometry.area * one.height_m
|
|
another_volume = another_geometry.area * another.height_m
|
|
return intersection / max(one_volume + another_volume - intersection, _EPSILON)
|
|
|
|
|
|
def _decode_row(row: np.ndarray) -> PointPillarsBox:
|
|
class_value = float(row[7])
|
|
class_id = int(round(class_value))
|
|
if (
|
|
abs(class_value - class_id) > 1e-5
|
|
or class_id < 0
|
|
or class_id >= len(POINTPILLARS_MODEL_CLASSES)
|
|
):
|
|
raise PointPillarsPostprocessError("PointPillars class id is invalid")
|
|
length_m, width_m, height_m = (float(row[index]) for index in (3, 4, 5))
|
|
score = float(row[8])
|
|
if (
|
|
length_m <= 0.0
|
|
or width_m <= 0.0
|
|
or height_m <= 0.0
|
|
or score < POINTPILLARS_EMBEDDED_SCORE_THRESHOLD
|
|
or score > 1.0
|
|
):
|
|
raise PointPillarsPostprocessError(
|
|
"PointPillars dimensions or score are invalid"
|
|
)
|
|
x_m, y_m, z_m = (float(row[index]) for index in (0, 1, 2))
|
|
return PointPillarsBox(
|
|
x_m=x_m,
|
|
y_m=y_m,
|
|
z_m=z_m,
|
|
length_m=length_m,
|
|
width_m=width_m,
|
|
height_m=height_m,
|
|
yaw_rad=float(row[6]),
|
|
class_id=class_id,
|
|
model_class=POINTPILLARS_MODEL_CLASSES[class_id],
|
|
score=score,
|
|
)
|
|
|
|
|
|
def _class_agnostic_nms(
|
|
boxes: tuple[PointPillarsBox, ...],
|
|
threshold: float,
|
|
) -> tuple[PointPillarsBox, ...]:
|
|
geometries = tuple(_geometry(box) for box in boxes)
|
|
suppressed = [False] * len(boxes)
|
|
accepted: list[PointPillarsBox] = []
|
|
for index, box in enumerate(boxes):
|
|
if suppressed[index]:
|
|
continue
|
|
accepted.append(box)
|
|
geometry = geometries[index]
|
|
for candidate_index in range(index + 1, len(boxes)):
|
|
if suppressed[candidate_index]:
|
|
continue
|
|
another = geometries[candidate_index]
|
|
if not _aabbs_overlap(geometry, another):
|
|
continue
|
|
overlap = _intersection_area(geometry, another)
|
|
iou = overlap / max(geometry.area + another.area - overlap, _EPSILON)
|
|
if iou >= threshold:
|
|
suppressed[candidate_index] = True
|
|
return tuple(accepted)
|
|
|
|
|
|
def _geometry(box: PointPillarsBox) -> _Geometry:
|
|
half_length = box.length_m / 2.0
|
|
half_width = box.width_m / 2.0
|
|
cosine = math.cos(box.yaw_rad)
|
|
sine = math.sin(box.yaw_rad)
|
|
corners = tuple(
|
|
(
|
|
box.x_m + local_x * cosine - local_y * sine,
|
|
box.y_m + local_x * sine + local_y * cosine,
|
|
)
|
|
for local_x, local_y in (
|
|
(-half_length, -half_width),
|
|
(half_length, -half_width),
|
|
(half_length, half_width),
|
|
(-half_length, half_width),
|
|
)
|
|
)
|
|
x_values = [point[0] for point in corners]
|
|
y_values = [point[1] for point in corners]
|
|
return _Geometry(
|
|
corners=corners,
|
|
area=box.length_m * box.width_m,
|
|
minimum_x=min(x_values),
|
|
maximum_x=max(x_values),
|
|
minimum_y=min(y_values),
|
|
maximum_y=max(y_values),
|
|
)
|
|
|
|
|
|
def _aabbs_overlap(one: _Geometry, another: _Geometry) -> bool:
|
|
return not (
|
|
one.maximum_x < another.minimum_x
|
|
or another.maximum_x < one.minimum_x
|
|
or one.maximum_y < another.minimum_y
|
|
or another.maximum_y < one.minimum_y
|
|
)
|
|
|
|
|
|
def _intersection_area(one: _Geometry, another: _Geometry) -> float:
|
|
if not _aabbs_overlap(one, another):
|
|
return 0.0
|
|
polygon = list(one.corners)
|
|
clip = another.corners
|
|
for index, edge_start in enumerate(clip):
|
|
edge_end = clip[(index + 1) % len(clip)]
|
|
polygon = _clip_polygon(polygon, edge_start, edge_end)
|
|
if not polygon:
|
|
return 0.0
|
|
return abs(
|
|
sum(
|
|
one_point[0] * another_point[1]
|
|
- another_point[0] * one_point[1]
|
|
for one_point, another_point in zip(
|
|
polygon,
|
|
(*polygon[1:], polygon[0]),
|
|
strict=True,
|
|
)
|
|
)
|
|
) / 2.0
|
|
|
|
|
|
def _clip_polygon(
|
|
polygon: list[tuple[float, float]],
|
|
edge_start: tuple[float, float],
|
|
edge_end: tuple[float, float],
|
|
) -> list[tuple[float, float]]:
|
|
if not polygon:
|
|
return []
|
|
result: list[tuple[float, float]] = []
|
|
previous = polygon[-1]
|
|
previous_inside = _inside(previous, edge_start, edge_end)
|
|
for current in polygon:
|
|
current_inside = _inside(current, edge_start, edge_end)
|
|
if current_inside:
|
|
if not previous_inside:
|
|
result.append(_line_intersection(previous, current, edge_start, edge_end))
|
|
result.append(current)
|
|
elif previous_inside:
|
|
result.append(_line_intersection(previous, current, edge_start, edge_end))
|
|
previous = current
|
|
previous_inside = current_inside
|
|
return result
|
|
|
|
|
|
def _inside(
|
|
point: tuple[float, float],
|
|
edge_start: tuple[float, float],
|
|
edge_end: tuple[float, float],
|
|
) -> bool:
|
|
return _cross(
|
|
edge_end[0] - edge_start[0],
|
|
edge_end[1] - edge_start[1],
|
|
point[0] - edge_start[0],
|
|
point[1] - edge_start[1],
|
|
) >= -_EPSILON
|
|
|
|
|
|
def _line_intersection(
|
|
segment_start: tuple[float, float],
|
|
segment_end: tuple[float, float],
|
|
edge_start: tuple[float, float],
|
|
edge_end: tuple[float, float],
|
|
) -> tuple[float, float]:
|
|
segment_x = segment_end[0] - segment_start[0]
|
|
segment_y = segment_end[1] - segment_start[1]
|
|
edge_x = edge_end[0] - edge_start[0]
|
|
edge_y = edge_end[1] - edge_start[1]
|
|
denominator = _cross(segment_x, segment_y, edge_x, edge_y)
|
|
if abs(denominator) <= _EPSILON:
|
|
return segment_end
|
|
offset_x = edge_start[0] - segment_start[0]
|
|
offset_y = edge_start[1] - segment_start[1]
|
|
ratio = _cross(offset_x, offset_y, edge_x, edge_y) / denominator
|
|
return (
|
|
segment_start[0] + ratio * segment_x,
|
|
segment_start[1] + ratio * segment_y,
|
|
)
|
|
|
|
|
|
def _cross(one_x: float, one_y: float, another_x: float, another_y: float) -> float:
|
|
return one_x * another_y - one_y * another_x
|