feat(perception): add PointPillars transfer gate
This commit is contained in:
DCCONSTRUCTIONS
@@ -0,0 +1,321 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user