235 lines
8.3 KiB
Python
235 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
target_math.py — 纯 numpy 的"像素 -> 三维目标"定位核心（不依赖 ROS，方便离线测试）。
|
||
|
||
输入：一帧点云（雷达系）+ 查询像素
|
||
输出：目标在相机系/雷达系下的坐标、距离、置信度、是否有多个深度层等
|
||
"""
|
||
import numpy as np
|
||
|
||
import calib
|
||
|
||
|
||
def _fit_plane_ray(P, ray, max_residual=0.05):
|
||
"""PCA 拟合平面，再让射线与平面求交。返回 (点, 残差)。"""
|
||
if P.shape[0] < 6:
|
||
return None, None
|
||
c = P.mean(axis=0)
|
||
_, _, vt = np.linalg.svd(P - c)
|
||
n = vt[2]
|
||
denom = float(n @ ray)
|
||
if abs(denom) < 1e-6:
|
||
return None, None
|
||
lam = float(n @ c) / denom
|
||
if lam <= 0:
|
||
return None, None
|
||
p = lam * ray
|
||
residual = float(np.sqrt(np.mean(((P - c) @ n) ** 2)))
|
||
if residual > max_residual:
|
||
return None, residual
|
||
return p, residual
|
||
|
||
|
||
def locate(xyz, query_u, query_v, conf=None, radius=10.0, depth_band=0.6,
|
||
fit_plane=True, max_residual=0.05, min_conf=None, min_z=0.05):
|
||
"""在点云里找查询像素方向上最近的可见点。
|
||
|
||
参数
|
||
xyz : (N,3) 雷达系点云
|
||
query_u/v : 查询像素
|
||
conf : (N,) 置信度，可选
|
||
radius : 候选像素半径
|
||
返回 dict（未找到时 found=False）
|
||
"""
|
||
xyz = np.asarray(xyz, dtype=np.float64)
|
||
good = np.isfinite(xyz).all(axis=1) & (np.abs(xyz).sum(axis=1) > 1e-6)
|
||
if conf is not None and min_conf is not None:
|
||
good &= np.asarray(conf) >= min_conf
|
||
idx = np.where(good)[0]
|
||
pts = xyz[good]
|
||
|
||
p_cam = calib.lidar_to_cam(pts)
|
||
u, v, ok = calib.project_cam_to_pixel(p_cam)
|
||
|
||
dpx = np.hypot(u - query_u, v - query_v)
|
||
cand = ok & (dpx <= radius) & (p_cam[:, 2] > min_z)
|
||
|
||
rng = np.linalg.norm(p_cam, axis=1)
|
||
res = {
|
||
'found': False, 'n_valid': int(pts.shape[0]), 'n_cand': int(cand.sum()),
|
||
'radius': float(radius), 'query_u': float(query_u), 'query_v': float(query_v),
|
||
}
|
||
if not cand.any():
|
||
return res
|
||
|
||
cand_idx = np.where(cand)[0]
|
||
order = cand_idx[np.argsort(rng[cand_idx])]
|
||
sel = int(order[0])
|
||
|
||
ray = calib.pixel_to_ray(query_u, query_v)
|
||
p_near = p_cam[sel]
|
||
|
||
# 深度分层：候选距离跨度大 => 可能擦过边缘
|
||
r_sorted = np.sort(rng[cand_idx])
|
||
multi = bool(len(r_sorted) > 3 and (r_sorted[-1] - r_sorted[0]) > 2.0)
|
||
|
||
p_final = p_near
|
||
plane_used = False
|
||
residual = None
|
||
if fit_plane:
|
||
band = p_cam[cand_idx][rng[cand_idx] <= rng[sel] + depth_band]
|
||
p_fit, residual = _fit_plane_ray(band, ray, max_residual)
|
||
if p_fit is not None:
|
||
p_final = p_fit
|
||
plane_used = True
|
||
|
||
res.update({
|
||
'found': True,
|
||
'p_cam': p_final,
|
||
'range': float(np.linalg.norm(p_final)),
|
||
'z': float(p_final[2]),
|
||
'p_lidar': calib.cam_to_lidar(p_final)[0],
|
||
'hit_u': float(u[sel]),
|
||
'hit_v': float(v[sel]),
|
||
'px_err': float(dpx[sel]),
|
||
'lidar_index': int(idx[sel]),
|
||
'plane_used': plane_used,
|
||
'plane_residual': residual,
|
||
'multi_layer': multi,
|
||
'cand_ranges': rng[cand_idx],
|
||
'confidence': float(conf[idx[sel]]) if conf is not None else None,
|
||
})
|
||
return res
|
||
|
||
|
||
class TargetIndex:
|
||
"""一帧点云的「像素格子 -> 点索引」倒排表。
|
||
|
||
为什么需要：
|
||
locate() 每次点击都会把 N 个点全量投影一遍。N≈5 万时一次几毫秒到
|
||
几十毫秒，点一次无所谓；但同一帧反复点、或在板子上跑时，这部分是
|
||
重复劳动。这里在【每帧】只做一次全量投影 + 分桶，之后每次点击只
|
||
检查查询像素周围几个格子里的点，平面拟合也只在局部候选里做。
|
||
|
||
用法：
|
||
idx = TargetIndex(xyz, conf=conf) # 每帧建一次
|
||
res = idx.locate(u, v, radius=10) # 每次点击，返回结构和 locate 一样
|
||
|
||
注意：
|
||
- 桶按点的像素坐标分，格子越大，查询时扫到的点越多；默认 32 px。
|
||
- 只索引落在图像内的点；图像外的点本来也不会被点到。
|
||
"""
|
||
|
||
def __init__(self, xyz, conf=None, cell=32, min_z=0.05):
|
||
xyz = np.asarray(xyz, dtype=np.float64)
|
||
good = np.isfinite(xyz).all(axis=1) & (np.abs(xyz).sum(axis=1) > 1e-6)
|
||
self.idx = np.where(good)[0]
|
||
self.xyz = xyz[good]
|
||
self.conf = None if conf is None else np.asarray(conf)[good]
|
||
|
||
self.p_cam = calib.lidar_to_cam(self.xyz)
|
||
self.u, self.v, ok = calib.project_cam_to_pixel(self.p_cam)
|
||
self.rng = np.linalg.norm(self.p_cam, axis=1)
|
||
self.ok = ok & (self.p_cam[:, 2] > float(min_z))
|
||
|
||
self.cell = max(4, int(cell))
|
||
self.nx = int(np.ceil(calib.IMAGE_WIDTH / float(self.cell)))
|
||
self.ny = int(np.ceil(calib.IMAGE_HEIGHT / float(self.cell)))
|
||
|
||
sel = np.where(self.ok)[0]
|
||
if sel.size:
|
||
cx = np.clip((self.u[sel] // self.cell).astype(np.int64),
|
||
0, self.nx - 1)
|
||
cy = np.clip((self.v[sel] // self.cell).astype(np.int64),
|
||
0, self.ny - 1)
|
||
keys = cy * self.nx + cx
|
||
order = np.argsort(keys, kind='stable')
|
||
keys = keys[order]
|
||
sel = sel[order]
|
||
uniq, starts = np.unique(keys, return_index=True)
|
||
ends = np.append(starts[1:], keys.size)
|
||
self.buckets = {int(k): sel[s:e]
|
||
for k, s, e in zip(uniq, starts, ends)}
|
||
else:
|
||
self.buckets = {}
|
||
|
||
def _gather(self, query_u, query_v, radius):
|
||
c = self.cell
|
||
x0 = max(0, int(np.floor((query_u - radius) / c)))
|
||
x1 = min(self.nx - 1, int(np.floor((query_u + radius) / c)))
|
||
y0 = max(0, int(np.floor((query_v - radius) / c)))
|
||
y1 = min(self.ny - 1, int(np.floor((query_v + radius) / c)))
|
||
|
||
parts = []
|
||
for cy in range(y0, y1 + 1):
|
||
base = cy * self.nx
|
||
for cx in range(x0, x1 + 1):
|
||
arr = self.buckets.get(base + cx)
|
||
if arr is not None and arr.size:
|
||
parts.append(arr)
|
||
if not parts:
|
||
return np.empty(0, dtype=np.int64)
|
||
return np.concatenate(parts)
|
||
|
||
def locate(self, query_u, query_v, radius=10.0, depth_band=0.6,
|
||
fit_plane=True, max_residual=0.05, min_conf=None):
|
||
"""局部查找，返回与 locate() 相同的字段。"""
|
||
query_u = float(query_u)
|
||
query_v = float(query_v)
|
||
|
||
res = {
|
||
'found': False, 'n_valid': int(self.idx.size),
|
||
'n_cand': 0, 'radius': float(radius),
|
||
'query_u': query_u, 'query_v': query_v,
|
||
'indexed': True,
|
||
}
|
||
|
||
cand = self._gather(query_u, query_v, radius)
|
||
if cand.size == 0:
|
||
return res
|
||
|
||
dpx = np.hypot(self.u[cand] - query_u, self.v[cand] - query_v)
|
||
keep = dpx <= radius
|
||
if min_conf is not None and self.conf is not None:
|
||
keep &= self.conf[cand] >= min_conf
|
||
cand = cand[keep]
|
||
dpx = dpx[keep]
|
||
res['n_cand'] = int(cand.size)
|
||
if cand.size == 0:
|
||
return res
|
||
|
||
# 距离最近的点
|
||
order = np.argsort(self.rng[cand])
|
||
cand = cand[order]
|
||
dpx = dpx[order]
|
||
sel = int(cand[0])
|
||
|
||
ray = calib.pixel_to_ray(query_u, query_v)
|
||
r_sorted = np.sort(self.rng[cand])
|
||
multi = bool(r_sorted.size > 3 and (r_sorted[-1] - r_sorted[0]) > 2.0)
|
||
|
||
p_final = self.p_cam[sel]
|
||
plane_used = False
|
||
residual = None
|
||
if fit_plane:
|
||
band = self.p_cam[cand[self.rng[cand] <= self.rng[sel] + depth_band]]
|
||
p_fit, residual = _fit_plane_ray(band, ray, max_residual)
|
||
if p_fit is not None:
|
||
p_final = p_fit
|
||
plane_used = True
|
||
|
||
res.update({
|
||
'found': True,
|
||
'p_cam': p_final,
|
||
'p_lidar': calib.cam_to_lidar(p_final)[0],
|
||
'hit_u': float(self.u[sel]),
|
||
'hit_v': float(self.v[sel]),
|
||
'px_err': float(dpx[0]),
|
||
'lidar_index': int(self.idx[sel]),
|
||
'plane_used': plane_used,
|
||
'plane_residual': residual,
|
||
'multi_layer': multi,
|
||
'confidence': None if self.conf is None else float(self.conf[sel]),
|
||
})
|
||
return res
|