from __future__ import annotations from dataclasses import dataclass from math import acos import numpy as np from scipy.ndimage import gaussian_filter1d from .structures import CandidateBand, SliceResult, TrackSolution def _contiguous_ranges(mask: np.ndarray) -> list[tuple[int, int]]: ranges: list[tuple[int, int]] = [] start = None for idx, value in enumerate(mask): if value and start is None: start = idx elif not value and start is not None: ranges.append((start, idx - 1)) start = None if start is not None: ranges.append((start, len(mask) - 1)) return ranges def _candidate_score( z_low: float, slice_bottom_ref: float, z_scale: float, flatness: float, flat_thresh: float, normal_alignment: float, point_count: int, width: float, preferred_width: float, edge_ratio: float, edge_preference_weight: float, ) -> float: low_term = np.clip((slice_bottom_ref - z_low) / max(z_scale, 1.0) + 1.0, 0.0, 3.0) flat_term = np.clip(1.0 - flatness / max(flat_thresh, 1e-6), 0.0, 1.5) width_term = np.clip( 1.0 - abs(width - preferred_width) / max(preferred_width, 1.0), -0.5, 1.0, ) density_term = np.clip(np.log1p(point_count) / 6.0, 0.0, 1.2) edge_term = np.clip(edge_ratio, 0.0, 1.0) return float( 1.2 * low_term + 0.9 * flat_term + 0.8 * normal_alignment + 0.5 * width_term + 0.4 * density_term + edge_preference_weight * edge_term ) def extract_beam_slices( points_vehicle: np.ndarray, normals_vehicle: np.ndarray, config: dict, ) -> list[SliceResult]: beam_cfg = config["beam"] x_min = float(points_vehicle[:, 0].min()) x_max = float(points_vehicle[:, 0].max()) slice_edges = np.arange(x_min, x_max + beam_cfg["slice_step_mm"], beam_cfg["slice_step_mm"]) if slice_edges[-1] < x_max: slice_edges = np.append(slice_edges, x_max) y_global_lo, y_global_hi = np.quantile( points_vehicle[:, 1], [beam_cfg["midline_trim_quantile"], 1.0 - beam_cfg["midline_trim_quantile"]], ) flat_thresh = beam_cfg["flatness_mm"] results: list[SliceResult] = [] for slice_index, x_start in enumerate(slice_edges[:-1]): x_center = x_start + beam_cfg["slice_thickness_mm"] * 0.5 x_end = x_start + beam_cfg["slice_thickness_mm"] mask = (points_vehicle[:, 0] >= x_start) & (points_vehicle[:, 0] < x_end) slice_points = points_vehicle[mask] if len(slice_points) < beam_cfg["slice_min_points"]: results.append( SliceResult( slice_index=slice_index, x_start=float(x_start), x_end=float(x_end), x_center=float(x_center), y_mid=float(0.5 * (y_global_lo + y_global_hi)), point_count=int(len(slice_points)), ) ) continue slice_normals = normals_vehicle[mask] y_lo, y_hi = np.quantile( slice_points[:, 1], [beam_cfg["midline_trim_quantile"], 1.0 - beam_cfg["midline_trim_quantile"]], ) y_mid = float(0.5 * (y_lo + y_hi)) y_edges = np.arange(y_lo, y_hi + beam_cfg["y_bin_width_mm"], beam_cfg["y_bin_width_mm"]) if len(y_edges) < 3: y_edges = np.linspace(y_lo, y_hi, 4) bin_count = len(y_edges) - 1 low_z = np.full(bin_count, np.nan) flatness = np.full(bin_count, np.nan) normal_alignment = np.full(bin_count, np.nan) counts = np.zeros(bin_count, dtype=np.int32) low_counts = np.zeros(bin_count, dtype=np.int32) for bin_index in range(bin_count): bin_mask = (slice_points[:, 1] >= y_edges[bin_index]) & ( slice_points[:, 1] < y_edges[bin_index + 1] ) counts[bin_index] = int(np.count_nonzero(bin_mask)) if counts[bin_index] < beam_cfg["min_bin_points"]: continue bin_points = slice_points[bin_mask] z_threshold = np.quantile(bin_points[:, 2], beam_cfg["low_quantile"]) low_mask = bin_points[:, 2] <= z_threshold low_points = bin_points[low_mask] low_counts[bin_index] = int(np.count_nonzero(low_mask)) if len(low_points) == 0: continue low_z[bin_index] = float(np.median(low_points[:, 2])) flatness[bin_index] = float( np.median(np.abs(low_points[:, 2] - low_z[bin_index])) ) normal_alignment[bin_index] = float( np.mean(np.abs(slice_normals[bin_mask][low_mask, 2])) ) valid = ~np.isnan(low_z) if valid.sum() < 5: results.append( SliceResult( slice_index=slice_index, x_start=float(x_start), x_end=float(x_end), x_center=float(x_center), y_mid=y_mid, point_count=int(len(slice_points)), ) ) continue smoothed_low_z = low_z.copy() smoothed_low_z[valid] = gaussian_filter1d(low_z[valid], sigma=1.0) z_cut = float( np.quantile(smoothed_low_z[valid], 0.25) + beam_cfg["low_band_margin_mm"] ) qualified_mask = ( valid & (smoothed_low_z <= z_cut) & (flatness <= flat_thresh) & (normal_alignment >= beam_cfg["normal_alignment_min"]) & (counts >= beam_cfg["min_bin_points"]) ) slice_result = SliceResult( slice_index=slice_index, x_start=float(x_start), x_end=float(x_end), x_center=float(x_center), y_mid=y_mid, point_count=int(len(slice_points)), ) z_scale = float(np.std(slice_points[:, 2]) + 1e-6) slice_bottom_ref = float(np.quantile(slice_points[:, 2], 0.15)) candidate_counter = 0 def append_candidate(start: int, end: int, side: str) -> None: nonlocal candidate_counter y_min = float(y_edges[start]) y_max = float(y_edges[end + 1]) width = float(y_max - y_min) if width < beam_cfg["width_min_mm"] or width > beam_cfg["width_max_mm"]: return weighted_counts = counts[start : end + 1].astype(np.float64) weighted_y = 0.5 * (y_edges[start : end + 1] + y_edges[start + 1 : end + 2]) y_center = float(np.average(weighted_y, weights=weighted_counts)) half_span = max(0.5 * (y_hi - y_lo), 1.0) edge_ratio = float(np.clip(abs(y_center - y_mid) / half_span, 0.0, 1.0)) z_low_group = float(np.nanmedian(smoothed_low_z[start : end + 1])) flatness_group = float(np.nanmean(flatness[start : end + 1])) normal_group = float(np.nanmean(normal_alignment[start : end + 1])) point_count = int(np.sum(counts[start : end + 1])) low_point_count = int(np.sum(low_counts[start : end + 1])) score = _candidate_score( z_low=z_low_group, slice_bottom_ref=slice_bottom_ref, z_scale=z_scale, flatness=flatness_group, flat_thresh=flat_thresh, normal_alignment=normal_group, point_count=point_count, width=width, preferred_width=beam_cfg["width_preferred_mm"], edge_ratio=edge_ratio, edge_preference_weight=beam_cfg["edge_preference_weight"], ) candidate = CandidateBand( candidate_id=f"s{slice_index:03d}_c{candidate_counter:02d}", slice_index=slice_index, x_start=float(x_start), x_end=float(x_end), x_center=float(x_center), y_min=y_min, y_max=y_max, y_center=y_center, z_low=z_low_group, width=width, point_count=point_count, low_point_count=low_point_count, flatness=flatness_group, normal_alignment=normal_group, score=score, edge_ratio=edge_ratio, side=side, y_mid=y_mid, bin_start=start, bin_end=end, ) candidate_counter += 1 slice_result.candidate_bands.append(candidate) if side == "left": slice_result.left_candidates.append(candidate) else: slice_result.right_candidates.append(candidate) def side_candidates(side_mask: np.ndarray, side: str) -> None: indices = np.where(qualified_mask & side_mask)[0] if len(indices) == 0: return index_set = set(indices.tolist()) minima: list[int] = [] for idx in indices: left_neighbor = idx - 1 if (idx - 1) in index_set else None right_neighbor = idx + 1 if (idx + 1) in index_set else None left_ok = left_neighbor is None or smoothed_low_z[idx] <= smoothed_low_z[left_neighbor] right_ok = right_neighbor is None or smoothed_low_z[idx] <= smoothed_low_z[right_neighbor] if left_ok and right_ok: minima.append(idx) if not minima: minima = [int(indices[np.argmin(smoothed_low_z[indices])])] minima = sorted(minima, key=lambda idx: smoothed_low_z[idx]) used_bins: set[int] = set() rise_limit_extra = max(beam_cfg["flatness_mm"] * 2.0, beam_cfg["low_band_margin_mm"] * 0.6) for min_idx in minima[:3]: if min_idx in used_bins: continue start = end = min_idx rise_limit = smoothed_low_z[min_idx] + rise_limit_extra while ( (start - 1) in index_set and smoothed_low_z[start - 1] <= rise_limit and (y_edges[end + 1] - y_edges[start - 1]) <= beam_cfg["width_max_mm"] ): start -= 1 while ( (end + 1) in index_set and smoothed_low_z[end + 1] <= rise_limit and (y_edges[end + 2] - y_edges[start]) <= beam_cfg["width_max_mm"] ): end += 1 append_candidate(start, end, side) used_bins.update(range(start, end + 1)) # The beam usually sits on the outer bottom edge, so seed one # boundary-driven candidate even if that edge is not a local minimum. outer_idx = int(indices[0] if side == "left" else indices[-1]) if outer_idx not in used_bins: start = end = outer_idx rise_limit = smoothed_low_z[outer_idx] + rise_limit_extra if side == "left": while ( (end + 1) in index_set and smoothed_low_z[end + 1] <= rise_limit and (y_edges[end + 2] - y_edges[start]) <= beam_cfg["width_max_mm"] ): end += 1 else: while ( (start - 1) in index_set and smoothed_low_z[start - 1] <= rise_limit and (y_edges[end + 1] - y_edges[start - 1]) <= beam_cfg["width_max_mm"] ): start -= 1 append_candidate(start, end, side) used_bins.update(range(start, end + 1)) side_candidates(np.asarray([0.5 * (y_edges[i] + y_edges[i + 1]) < y_mid for i in range(bin_count)]), "left") side_candidates(np.asarray([0.5 * (y_edges[i] + y_edges[i + 1]) >= y_mid for i in range(bin_count)]), "right") for attr_name in ("left_candidates", "right_candidates"): side_list = getattr(slice_result, attr_name) if not side_list: continue strongest_edge = max(candidate.edge_ratio for candidate in side_list) edge_floor = max( beam_cfg["inner_reject_ratio"], beam_cfg["edge_keep_min_ratio"], ) if strongest_edge >= beam_cfg["strong_edge_ratio"]: edge_floor = max(edge_floor, strongest_edge - beam_cfg["relative_edge_window"]) filtered = [ candidate for candidate in side_list if candidate.edge_ratio >= edge_floor ] if filtered or beam_cfg.get("drop_weak_edge_slices", False): setattr(slice_result, attr_name, filtered) slice_result.candidate_bands = slice_result.left_candidates + slice_result.right_candidates results.append(slice_result) return results @dataclass(slots=True) class _PathState: score: float candidate_ids: tuple[str | None, ...] selected: tuple[CandidateBand, ...] coverage_count: int last_y: float | None last_z: float | None last_width: float | None last_normal: np.ndarray | None prev_dy: float | None prev_dz: float | None gap_run: int def _normal_to_vector(candidate: CandidateBand) -> np.ndarray: lateral = 1.0 if candidate.side == "right" else -1.0 return np.array([0.0, lateral * np.sqrt(max(1.0 - candidate.normal_alignment**2, 0.0)), candidate.normal_alignment]) def _transition_penalty( prev_state: _PathState, candidate: CandidateBand, config: dict, ) -> float: weights = config["tracking"]["weights"] beam_cfg = config["beam"] penalty = 0.0 if prev_state.last_y is not None: penalty += weights["y_penalty"] * abs(candidate.y_center - prev_state.last_y) / max( beam_cfg["y_bin_width_mm"], 1.0, ) if prev_state.last_z is not None: penalty += weights["z_penalty"] * abs(candidate.z_low - prev_state.last_z) / max( beam_cfg["flatness_mm"], 1.0, ) if prev_state.last_width is not None: penalty += weights["width_penalty"] * abs(candidate.width - prev_state.last_width) / max( beam_cfg["width_preferred_mm"], 1.0, ) current_normal = _normal_to_vector(candidate) if prev_state.last_normal is not None: cosine = float( np.clip( np.dot(prev_state.last_normal, current_normal) / max(np.linalg.norm(prev_state.last_normal) * np.linalg.norm(current_normal), 1e-6), -1.0, 1.0, ) ) penalty += weights["normal_penalty"] * acos(cosine) if prev_state.prev_dy is not None and prev_state.last_y is not None: dy = candidate.y_center - prev_state.last_y penalty += weights["jitter_penalty"] * abs(dy - prev_state.prev_dy) / max( beam_cfg["y_bin_width_mm"], 1.0, ) if prev_state.prev_dz is not None and prev_state.last_z is not None: dz = candidate.z_low - prev_state.last_z penalty += weights["jitter_penalty"] * abs(dz - prev_state.prev_dz) / max( beam_cfg["flatness_mm"], 1.0, ) return penalty def _insert_topk(states: list[_PathState], state: _PathState, k: int) -> list[_PathState]: states.append(state) states.sort(key=lambda item: item.score, reverse=True) deduped: list[_PathState] = [] seen = set() for candidate in states: key = candidate.candidate_ids if key in seen: continue deduped.append(candidate) seen.add(key) if len(deduped) >= k: break return deduped def track_side_candidates( slices: list[SliceResult], side: str, config: dict, ) -> list[TrackSolution]: top_k = config["tracking"]["top_k"] beam_width = config["tracking"]["beam_width"] weights = config["tracking"]["weights"] states: dict[str, list[_PathState]] = { "gap": [ _PathState( score=0.0, candidate_ids=tuple(), selected=tuple(), coverage_count=0, last_y=None, last_z=None, last_width=None, last_normal=None, prev_dy=None, prev_dz=None, gap_run=0, ) ] } for slice_result in slices: candidates = ( slice_result.left_candidates if side == "left" else slice_result.right_candidates ) current_states: dict[str, list[_PathState]] = {"gap": []} for candidate in candidates: current_states[candidate.candidate_id] = [] for previous_group in states.values(): for prev_state in previous_group: gap_score = prev_state.score - weights["gap_penalty"] - weights["gap_growth_penalty"] * prev_state.gap_run gap_state = _PathState( score=gap_score, candidate_ids=prev_state.candidate_ids + (None,), selected=prev_state.selected, coverage_count=prev_state.coverage_count, last_y=prev_state.last_y, last_z=prev_state.last_z, last_width=prev_state.last_width, last_normal=prev_state.last_normal, prev_dy=prev_state.prev_dy, prev_dz=prev_state.prev_dz, gap_run=prev_state.gap_run + 1, ) current_states["gap"] = _insert_topk(current_states["gap"], gap_state, top_k) for candidate in candidates: penalty = _transition_penalty(prev_state, candidate, config) dy = None if prev_state.last_y is None else candidate.y_center - prev_state.last_y dz = None if prev_state.last_z is None else candidate.z_low - prev_state.last_z score = ( prev_state.score + weights["coverage_reward"] + weights["node_reward"] * candidate.score - penalty ) candidate_state = _PathState( score=score, candidate_ids=prev_state.candidate_ids + (candidate.candidate_id,), selected=prev_state.selected + (candidate,), coverage_count=prev_state.coverage_count + 1, last_y=candidate.y_center, last_z=candidate.z_low, last_width=candidate.width, last_normal=_normal_to_vector(candidate), prev_dy=dy, prev_dz=dz, gap_run=0, ) current_states[candidate.candidate_id] = _insert_topk( current_states[candidate.candidate_id], candidate_state, top_k, ) flat_states = [item for group in current_states.values() for item in group] flat_states.sort(key=lambda item: item.score, reverse=True) states = {"gap": [], **{key: [] for key in current_states if key != "gap"}} for state in flat_states[:beam_width]: state_key = "gap" if state.candidate_ids[-1] is None else state.candidate_ids[-1] states[state_key] = _insert_topk(states[state_key], state, top_k) solutions: list[_PathState] = [item for group in states.values() for item in group] solutions.sort(key=lambda item: item.score, reverse=True) deduped: list[_PathState] = [] seen = set() for solution in solutions: if solution.candidate_ids in seen: continue seen.add(solution.candidate_ids) deduped.append(solution) if len(deduped) >= top_k: break final_solutions: list[TrackSolution] = [] min_coverage = config["tracking"]["min_coverage_ratio"] * max(len(slices), 1) if not deduped: return final_solutions score_baseline = deduped[0].score for rank, state in enumerate(deduped, start=1): if state.coverage_count < min_coverage: continue margin = None if rank < len(deduped): margin = float(state.score - deduped[rank].score) confidence = float( np.clip( 0.5 + 0.25 * state.coverage_count / max(len(slices), 1) + 0.25 * (state.score / max(abs(score_baseline), 1e-6)), 0.0, 1.0, ) ) final_solutions.append( TrackSolution( side=side, rank=rank, score=float(state.score), coverage_ratio=float(state.coverage_count / max(len(slices), 1)), coverage_count=int(state.coverage_count), candidate_ids=list(state.candidate_ids), candidates=list(state.selected), score_margin=margin, confidence=confidence, ) ) return final_solutions def candidate_rows(slices: list[SliceResult]) -> list[dict]: rows: list[dict] = [] for slice_result in slices: for candidate in slice_result.candidate_bands: rows.append( { "slice_index": candidate.slice_index, "candidate_id": candidate.candidate_id, "side": candidate.side, "x_start": candidate.x_start, "x_end": candidate.x_end, "x_center": candidate.x_center, "y_min": candidate.y_min, "y_max": candidate.y_max, "y_center": candidate.y_center, "z_low": candidate.z_low, "width": candidate.width, "point_count": candidate.point_count, "low_point_count": candidate.low_point_count, "flatness": candidate.flatness, "normal_alignment": candidate.normal_alignment, "score": candidate.score, "edge_ratio": candidate.edge_ratio, "y_mid": candidate.y_mid, } ) return rows