Spaces:
Configuration error
Configuration error
| """ | |
| RCLane relay-chain decoder -- ported from src/lane_codec.py (decode / decode_branch / | |
| thresh_line / iou_nms) and src/lane_geometry.py (PointSelf / FloatLengthLine._iou). | |
| Turns the 5 predicted maps back into lane polylines. This is the inference-time | |
| counterpart of `encode`. Pure NumPy + OpenCV, framework-agnostic. | |
| Pipeline: | |
| 1. seeds = point-NMS on the seg foreground map (keep peaks >= min_dist apart). | |
| 2. for each seed: crawl a forward (down_arrow) and backward (up_arrow) relay chain, | |
| each step = normalized transfer vector * step_length; stop via the distance | |
| (bound) heuristic once the walk leaves the foreground. | |
| 3. merge reversed(up) + down into one lane. | |
| 4. drop low-score lanes (thresh_line), then IoU-NMS to remove duplicates. | |
| Note: the stopping rule inside `decode_branch` (RMS of remaining-step estimates, | |
| the 0.75 factor, "keep going while on foreground") is NOT described in the paper -- | |
| it is ported verbatim from the MindSpore code. A `norm == 0` guard is added so a | |
| zero transfer vector breaks the walk instead of producing NaNs. | |
| Map layout (matches encode.py / rclane.py, channel-first): | |
| seg_prob: (H, W) foreground probability in [0, 1] | |
| up/down_arrow/bound: (2, H, W) float32 | |
| """ | |
| import numpy as np | |
| import cv2 | |
| try: | |
| from numba import njit, prange | |
| except ImportError: # portable fallback for environments without the JIT | |
| njit = None | |
| prange = range | |
| if njit is not None: | |
| def _greedy_seed_select_numba(sorted_x, sorted_y, height, width, | |
| radius, max_seeds): | |
| taken = np.zeros((height, width), dtype=np.uint8) | |
| limit = len(sorted_x) if max_seeds < 0 else min( | |
| len(sorted_x), max_seeds | |
| ) | |
| seeds = np.empty((limit, 2), dtype=np.int32) | |
| seed_count = 0 | |
| for candidate in range(len(sorted_x)): | |
| x = int(sorted_x[candidate]) | |
| y = int(sorted_y[candidate]) | |
| if taken[y, x] != 0: | |
| continue | |
| seeds[seed_count, 0] = x | |
| seeds[seed_count, 1] = y | |
| seed_count += 1 | |
| y_start = max(0, y - radius) | |
| y_stop = min(height, y + radius + 1) | |
| x_start = max(0, x - radius) | |
| x_stop = min(width, x + radius + 1) | |
| for yy in range(y_start, y_stop): | |
| for xx in range(x_start, x_stop): | |
| taken[yy, xx] = 1 | |
| if max_seeds >= 0 and seed_count >= max_seeds: | |
| break | |
| return seeds[:seed_count] | |
| else: | |
| _greedy_seed_select_numba = None | |
| # --------------------------------------------------------------------------- # | |
| # Lane container (replaces FloatLengthLine + PointSelf) | |
| # --------------------------------------------------------------------------- # | |
| class Lane: | |
| def __init__(self, width, height): | |
| self.width = width | |
| self.height = height | |
| self.points = [] # list of (x, y, score) | |
| self._score_sum = 0.0 | |
| self.lane_id = None | |
| self.lane_role = None | |
| self.is_ego_boundary = False | |
| self.lateral_rank = None | |
| def append(self, x, y, score): | |
| if isinstance(self.points, np.ndarray): | |
| self.points = [tuple(map(float, point)) for point in self.points] | |
| self.points.append((float(x), float(y), float(score))) | |
| self._score_sum += float(score) | |
| def reverse(self): | |
| if isinstance(self.points, np.ndarray): | |
| self.points = self.points[::-1].copy() | |
| else: | |
| self.points.reverse() | |
| def __len__(self): | |
| return len(self.points) | |
| def score(self): | |
| if len(self.points) == 0: | |
| return 0.0 | |
| return self._score_sum / len(self.points) | |
| def concat(self, other): | |
| out = Lane(self.width, self.height) | |
| if isinstance(self.points, np.ndarray) or isinstance( | |
| other.points, np.ndarray | |
| ): | |
| out.points = np.concatenate(( | |
| np.asarray(self.points, dtype=np.float32), | |
| np.asarray(other.points, dtype=np.float32), | |
| )) | |
| else: | |
| out.points = self.points + other.points | |
| out._score_sum = self._score_sum + other._score_sum | |
| return out | |
| def xy(self): | |
| points = np.asarray(self.points, dtype=np.float32) | |
| if points.size == 0: | |
| return np.empty((0, 2), dtype=np.float32) | |
| return points[:, :2] | |
| def iou(self, other, lane_width=15): | |
| """Rasterize both lanes (cv2.line, width 15) and return mask IoU.""" | |
| im1 = np.zeros((self.height, self.width), np.uint8) | |
| im2 = np.zeros((self.height, self.width), np.uint8) | |
| p1 = self.xy().astype(np.int32) | |
| p2 = other.xy().astype(np.int32) | |
| for i in range(len(p1) - 1): | |
| cv2.line(im1, tuple(p1[i]), tuple(p1[i + 1]), 255, lane_width) | |
| for i in range(len(p2) - 1): | |
| cv2.line(im2, tuple(p2[i]), tuple(p2[i + 1]), 255, lane_width) | |
| union = int((cv2.bitwise_or(im1, im2) > 0).sum()) | |
| if union == 0: | |
| return 0.0 | |
| inter = int((im1 > 0).sum()) + int((im2 > 0).sum()) - union | |
| return inter / float(union) | |
| # --------------------------------------------------------------------------- # | |
| # seeding | |
| # --------------------------------------------------------------------------- # | |
| def point_nms(prob, thr=0.5, min_dist=2, max_seeds=1024, | |
| backend="auto"): | |
| """Greedy point-NMS: keep highest-prob foreground pixels >= min_dist apart.""" | |
| H, W = prob.shape | |
| ys, xs = np.where(prob > thr) | |
| if len(ys) == 0: | |
| return [] | |
| order = np.argsort(-prob[ys, xs]) | |
| selected_backend = backend | |
| if backend == "auto": | |
| selected_backend = ( | |
| "numba" if _greedy_seed_select_numba is not None else "python" | |
| ) | |
| if selected_backend == "numba": | |
| limit = -1 if max_seeds is None else int(max_seeds) | |
| return _greedy_seed_select_numba( | |
| np.ascontiguousarray(xs[order], dtype=np.int32), | |
| np.ascontiguousarray(ys[order], dtype=np.int32), | |
| H, W, int(min_dist), limit, | |
| ) | |
| if selected_backend != "python": | |
| raise ValueError("point-NMS backend must be auto, numba, or python") | |
| taken = np.zeros((H, W), dtype=bool) | |
| seeds = [] | |
| r = min_dist | |
| for idx in order: | |
| y, x = int(ys[idx]), int(xs[idx]) | |
| if taken[y, x]: | |
| continue | |
| seeds.append((x, y)) | |
| taken[max(0, y - r):y + r + 1, max(0, x - r):x + r + 1] = True | |
| if max_seeds is not None and len(seeds) >= max_seeds: | |
| break | |
| return seeds | |
| # --------------------------------------------------------------------------- # | |
| # relay-chain crawl (port of decode_branch) | |
| # --------------------------------------------------------------------------- # | |
| def decode_branch(cx, cy, semantic_fine, arrow, bound, step_length, seg_threshold): | |
| H, W = semantic_fine.shape | |
| arrow_dx, arrow_dy = arrow[0], arrow[1] | |
| lane = Lane(W, H) | |
| remain_sq_sum = 0.0 | |
| remain_count = 0 | |
| cx, cy = int(cx), int(cy) | |
| for index in range(H): | |
| if semantic_fine[cy, cx] > seg_threshold: | |
| remain = bound[cy, cx] * 100 / step_length + index | |
| remain_sq_sum += remain * remain | |
| remain_count += 1 | |
| dx = arrow_dx[cy, cx] | |
| dy = arrow_dy[cy, cx] | |
| norm = np.sqrt(dx * dx + dy * dy) | |
| if norm == 0: # guard (not in original): dead transfer vector -> stop | |
| break | |
| cx = int(np.floor(cx + dx / norm * step_length)) | |
| cy = int(np.floor(cy + dy / norm * step_length)) | |
| if not (0 <= cx < W and 0 <= cy < H): | |
| break | |
| lane.append(cx, cy, semantic_fine[cy, cx]) | |
| if remain_count: | |
| ret = np.sqrt(remain_sq_sum / remain_count) | |
| else: | |
| ret = 1 | |
| if semantic_fine[cy, cx] > seg_threshold: | |
| continue | |
| if index > ret * 0.75: | |
| break | |
| return lane | |
| def decode_branches_batch(seeds, semantic_fine, arrow, bound, step_length, | |
| seg_threshold): | |
| """Vectorized equivalent of :func:`decode_branch` for many seeds. | |
| The relay walk is inherently sequential along each lane, but every seed at | |
| a given crawl step is independent. Advancing all active seeds with NumPy | |
| gathers removes the expensive Python ``seeds x steps`` nested loop while | |
| retaining the original stopping rule and integer pixel trajectory. | |
| Returns: | |
| ``(points, lengths)`` where ``points`` has shape ``(N, H, 3)`` and each | |
| valid prefix stores ``(x, y, score)`` exactly like ``Lane.points``. | |
| """ | |
| H, W = semantic_fine.shape | |
| seed_array = np.asarray(seeds, dtype=np.int32) | |
| if seed_array.size == 0: | |
| return np.empty((0, H, 3), dtype=np.float32), np.zeros(0, np.int32) | |
| if seed_array.ndim != 2 or seed_array.shape[1] != 2: | |
| raise ValueError("seeds must have shape (N, 2)") | |
| count = len(seed_array) | |
| cx = seed_array[:, 0].copy() | |
| cy = seed_array[:, 1].copy() | |
| active = np.ones(count, dtype=bool) | |
| lengths = np.zeros(count, dtype=np.int32) | |
| remain_sq_sum = np.zeros(count, dtype=np.float64) | |
| remain_count = np.zeros(count, dtype=np.int32) | |
| points = np.empty((count, H, 3), dtype=np.float32) | |
| arrow_dx, arrow_dy = arrow[0], arrow[1] | |
| for index in range(H): | |
| active_ids = np.flatnonzero(active) | |
| if len(active_ids) == 0: | |
| break | |
| current_x = cx[active_ids] | |
| current_y = cy[active_ids] | |
| current_score = semantic_fine[current_y, current_x] | |
| foreground = current_score > seg_threshold | |
| if np.any(foreground): | |
| foreground_ids = active_ids[foreground] | |
| remain = ( | |
| bound[current_y[foreground], current_x[foreground]] | |
| * 100.0 / step_length + index | |
| ) | |
| remain_sq_sum[foreground_ids] += remain * remain | |
| remain_count[foreground_ids] += 1 | |
| dx = arrow_dx[current_y, current_x] | |
| dy = arrow_dy[current_y, current_x] | |
| norm = np.sqrt(dx * dx + dy * dy) | |
| movable = np.isfinite(norm) & (norm != 0.0) | |
| if np.any(~movable): | |
| active[active_ids[~movable]] = False | |
| active_ids = active_ids[movable] | |
| if len(active_ids) == 0: | |
| continue | |
| next_x = np.floor( | |
| cx[active_ids] + dx[movable] / norm[movable] * step_length | |
| ).astype(np.int32) | |
| next_y = np.floor( | |
| cy[active_ids] + dy[movable] / norm[movable] * step_length | |
| ).astype(np.int32) | |
| in_bounds = ( | |
| (next_x >= 0) & (next_x < W) & (next_y >= 0) & (next_y < H) | |
| ) | |
| if np.any(~in_bounds): | |
| active[active_ids[~in_bounds]] = False | |
| active_ids = active_ids[in_bounds] | |
| if len(active_ids) == 0: | |
| continue | |
| next_x = next_x[in_bounds] | |
| next_y = next_y[in_bounds] | |
| cx[active_ids] = next_x | |
| cy[active_ids] = next_y | |
| next_score = semantic_fine[next_y, next_x] | |
| points[active_ids, index, 0] = next_x | |
| points[active_ids, index, 1] = next_y | |
| points[active_ids, index, 2] = next_score | |
| lengths[active_ids] = index + 1 | |
| has_remaining = remain_count[active_ids] > 0 | |
| remaining = np.ones(len(active_ids), dtype=np.float64) | |
| remaining[has_remaining] = np.sqrt( | |
| remain_sq_sum[active_ids[has_remaining]] | |
| / remain_count[active_ids[has_remaining]] | |
| ) | |
| stop = ( | |
| (next_score <= seg_threshold) | |
| & (index > remaining * 0.75) | |
| ) | |
| if np.any(stop): | |
| active[active_ids[stop]] = False | |
| return points, lengths | |
| if njit is not None: | |
| def _decode_branches_numba_impl(seeds, semantic_fine, arrow, bound, | |
| step_length, seg_threshold): | |
| """Parallel scalar relay walks compiled to native CPU code.""" | |
| height, width = semantic_fine.shape | |
| seed_count = len(seeds) | |
| points = np.empty((seed_count, height, 3), dtype=np.float32) | |
| lengths = np.zeros(seed_count, dtype=np.int32) | |
| for seed_index in prange(seed_count): | |
| cx = int(seeds[seed_index, 0]) | |
| cy = int(seeds[seed_index, 1]) | |
| remain_sq_sum = 0.0 | |
| remain_count = 0 | |
| for index in range(height): | |
| if semantic_fine[cy, cx] > seg_threshold: | |
| remain = ( | |
| bound[cy, cx] * 100.0 / step_length + index | |
| ) | |
| remain_sq_sum += remain * remain | |
| remain_count += 1 | |
| dx = arrow[0, cy, cx] | |
| dy = arrow[1, cy, cx] | |
| norm = np.sqrt(dx * dx + dy * dy) | |
| if norm == 0.0 or not np.isfinite(norm): | |
| break | |
| cx = int(np.floor(cx + dx / norm * step_length)) | |
| cy = int(np.floor(cy + dy / norm * step_length)) | |
| if not (0 <= cx < width and 0 <= cy < height): | |
| break | |
| score = semantic_fine[cy, cx] | |
| points[seed_index, index, 0] = cx | |
| points[seed_index, index, 1] = cy | |
| points[seed_index, index, 2] = score | |
| lengths[seed_index] = index + 1 | |
| ret = ( | |
| np.sqrt(remain_sq_sum / remain_count) | |
| if remain_count else 1.0 | |
| ) | |
| if score > seg_threshold: | |
| continue | |
| if index > ret * 0.75: | |
| break | |
| return points, lengths | |
| def _candidate_metadata_numba_impl(up_points, up_lengths, | |
| down_points, down_lengths, bin_px): | |
| seed_count = len(up_lengths) | |
| scores = np.zeros(seed_count, dtype=np.float64) | |
| bins = np.zeros(seed_count, dtype=np.int32) | |
| total_lengths = up_lengths + down_lengths | |
| for seed_index in prange(seed_count): | |
| up_length = int(up_lengths[seed_index]) | |
| down_length = int(down_lengths[seed_index]) | |
| total_length = up_length + down_length | |
| if total_length <= 1: | |
| continue | |
| y_values = np.empty(total_length, dtype=np.float32) | |
| score_sum = 0.0 | |
| position = 0 | |
| for point_index in range(up_length): | |
| score_sum += up_points[seed_index, point_index, 2] | |
| y_values[position] = up_points[seed_index, point_index, 1] | |
| position += 1 | |
| for point_index in range(down_length): | |
| score_sum += down_points[seed_index, point_index, 2] | |
| y_values[position] = down_points[seed_index, point_index, 1] | |
| position += 1 | |
| scores[seed_index] = score_sum / total_length | |
| median_y = np.median(y_values) | |
| lower_x = np.empty(total_length, dtype=np.float32) | |
| lower_count = 0 | |
| for point_index in range(up_length): | |
| if up_points[seed_index, point_index, 1] >= median_y: | |
| lower_x[lower_count] = up_points[ | |
| seed_index, point_index, 0 | |
| ] | |
| lower_count += 1 | |
| for point_index in range(down_length): | |
| if down_points[seed_index, point_index, 1] >= median_y: | |
| lower_x[lower_count] = down_points[ | |
| seed_index, point_index, 0 | |
| ] | |
| lower_count += 1 | |
| bins[seed_index] = int( | |
| np.median(lower_x[:lower_count]) // bin_px | |
| ) | |
| return total_lengths, scores, bins | |
| else: | |
| _decode_branches_numba_impl = None | |
| _candidate_metadata_numba_impl = None | |
| def decode_branches_numba(seeds, semantic_fine, arrow, bound, step_length, | |
| seg_threshold): | |
| if _decode_branches_numba_impl is None: | |
| raise RuntimeError( | |
| "Numba crawl requested but numba is not installed; " | |
| "install requirements.txt or use crawl_backend='numpy'" | |
| ) | |
| seed_array = np.asarray(seeds, dtype=np.int32) | |
| if seed_array.size == 0: | |
| height = semantic_fine.shape[0] | |
| return ( | |
| np.empty((0, height, 3), dtype=np.float32), | |
| np.zeros(0, dtype=np.int32), | |
| ) | |
| return _decode_branches_numba_impl( | |
| seed_array, | |
| np.ascontiguousarray(semantic_fine, dtype=np.float32), | |
| np.ascontiguousarray(arrow, dtype=np.float32), | |
| np.ascontiguousarray(bound, dtype=np.float32), | |
| float(step_length), | |
| float(seg_threshold), | |
| ) | |
| def warmup_decode_backend(crawl_backend="auto"): | |
| """Compile the optional Numba backend before latency measurements.""" | |
| backend = crawl_backend | |
| if crawl_backend == "auto": | |
| backend = "numba" if njit is not None else "numpy" | |
| if backend != "numba": | |
| return backend | |
| semantic = np.zeros((8, 8), dtype=np.float32) | |
| arrow = np.zeros((2, 8, 8), dtype=np.float32) | |
| arrow[1] = 1.0 | |
| bound = np.zeros((8, 8), dtype=np.float32) | |
| points, lengths = decode_branches_numba( | |
| [(4, 4)], semantic, arrow, bound, 1, 0.5 | |
| ) | |
| _candidate_metadata_numba_impl( | |
| points, lengths, points, lengths, 16 | |
| ) | |
| _greedy_seed_select_numba( | |
| np.array((4,), dtype=np.int32), | |
| np.array((4,), dtype=np.int32), | |
| 8, 8, 2, 1, | |
| ) | |
| return backend | |
| def configure_decode_threads(thread_count=8): | |
| """Set Numba's relay-crawl worker count and return the applied value.""" | |
| if njit is None: | |
| return 1 | |
| import numba | |
| maximum = int(numba.config.NUMBA_NUM_THREADS) | |
| if not 1 <= int(thread_count) <= maximum: | |
| raise ValueError(f"decode threads must be in [1, {maximum}]") | |
| numba.set_num_threads(int(thread_count)) | |
| return numba.get_num_threads() | |
| def _candidate_metadata_numpy(up_points, up_lengths, | |
| down_points, down_lengths, bin_px): | |
| total_lengths = up_lengths + down_lengths | |
| scores = np.zeros(len(total_lengths), dtype=np.float64) | |
| bins = np.zeros(len(total_lengths), dtype=np.int32) | |
| for seed_index, total_length in enumerate(total_lengths): | |
| if total_length <= 1: | |
| continue | |
| up = up_points[seed_index, :up_lengths[seed_index]] | |
| down = down_points[seed_index, :down_lengths[seed_index]] | |
| merged = np.concatenate((up, down)) | |
| scores[seed_index] = merged[:, 2].sum(dtype=np.float64) / len(merged) | |
| median_y = np.median(merged[:, 1]) | |
| lower = merged[merged[:, 1] >= median_y] | |
| bins[seed_index] = int(np.median(lower[:, 0]) // bin_px) | |
| return total_lengths, scores, bins | |
| def _preselect_batched_candidates(up_points, up_lengths, | |
| down_points, down_lengths, score_threshold, | |
| max_lanes, backend, bin_px=16): | |
| metadata = ( | |
| _candidate_metadata_numba_impl | |
| if backend == "numba" else _candidate_metadata_numpy | |
| ) | |
| total_lengths, scores, bins = metadata( | |
| up_points, up_lengths, down_points, down_lengths, bin_px | |
| ) | |
| valid = np.flatnonzero( | |
| (total_lengths > 1) & (scores >= score_threshold) | |
| ) | |
| order = valid[np.argsort(-scores[valid], kind="stable")] | |
| if max_lanes is None or len(order) <= max_lanes: | |
| return order | |
| buckets = {} | |
| for candidate in order: | |
| buckets.setdefault(int(bins[candidate]), []).append(int(candidate)) | |
| keys = list(buckets) | |
| positions = {key: 0 for key in keys} | |
| selected = [] | |
| while len(selected) < max_lanes: | |
| progressed = False | |
| for key in keys: | |
| position = positions[key] | |
| if position < len(buckets[key]): | |
| selected.append(buckets[key][position]) | |
| positions[key] += 1 | |
| progressed = True | |
| if len(selected) >= max_lanes: | |
| break | |
| if not progressed: | |
| break | |
| return np.asarray( | |
| sorted(selected, key=lambda index: scores[index], reverse=True), | |
| dtype=np.int32, | |
| ) | |
| def _lines_from_batched_crawls(up_points, up_lengths, | |
| down_points, down_lengths, width, height, | |
| candidate_indices=None): | |
| lines = [] | |
| if candidate_indices is None: | |
| candidate_indices = range(len(up_lengths)) | |
| for seed_index in candidate_indices: | |
| up_length = int(up_lengths[seed_index]) | |
| down_length = int(down_lengths[seed_index]) | |
| if up_length + down_length <= 1: | |
| continue | |
| merged = np.concatenate(( | |
| up_points[seed_index, :up_length][::-1], | |
| down_points[seed_index, :down_length], | |
| )) | |
| lane = Lane(width, height) | |
| lane.points = merged | |
| lane._score_sum = float(merged[:, 2].sum(dtype=np.float64)) | |
| lines.append(lane) | |
| return lines | |
| # --------------------------------------------------------------------------- # | |
| # post-processing | |
| # --------------------------------------------------------------------------- # | |
| def thresh_line(lines, thr=0.10): | |
| return [ln for ln in lines if ln.score >= thr] | |
| def _diverse_preselect(order, max_lanes, bin_px=16): | |
| """Cap the candidate list while keeping spatial diversity. | |
| `order` is already sorted by score (desc). Taking the top `max_lanes` | |
| outright lets the single strongest lane monopolize every slot -- its copies | |
| all score highest -- so genuine but slightly weaker lanes get dropped before | |
| IoU-NMS ever compares them (this collapsed multi-lane curves to one lane). | |
| Instead, bucket candidates by their horizontal position on the lower half of | |
| the line (where lanes are well separated) and pick round-robin across | |
| buckets, so every lane keeps representatives within the cap. | |
| """ | |
| if max_lanes is None: | |
| return order | |
| if max_lanes <= 0: | |
| raise ValueError("max_lanes must be positive or None") | |
| if bin_px <= 0: | |
| raise ValueError("bin_px must be positive") | |
| if len(order) <= max_lanes: | |
| return order | |
| buckets = {} | |
| for ln in order: # already score-sorted | |
| xy = ln.xy() | |
| if len(xy) == 0: | |
| continue | |
| low = xy[xy[:, 1] >= np.median(xy[:, 1])] # lower (near) half | |
| ref = low if len(low) else xy | |
| key = int(np.median(ref[:, 0]) // bin_px) | |
| buckets.setdefault(key, []).append(ln) | |
| keys = list(buckets.keys()) | |
| idx = {k: 0 for k in keys} | |
| selected = [] | |
| while len(selected) < max_lanes: | |
| progressed = False | |
| for k in keys: | |
| if idx[k] < len(buckets[k]): | |
| selected.append(buckets[k][idx[k]]) | |
| idx[k] += 1 | |
| progressed = True | |
| if len(selected) >= max_lanes: | |
| break | |
| if not progressed: | |
| break | |
| # NMS is greedy, so restore global score priority after choosing a spatially | |
| # diverse candidate set. | |
| return sorted(selected, key=lambda ln: ln.score, reverse=True) | |
| def iou_nms(lines, thr=0.5, max_lanes=128, scale=0.25, | |
| lane_width=15): | |
| """Lane IoU NMS with one cached, downscaled mask per candidate. | |
| Rasterizes each candidate once (downscaled) instead of re-rasterizing both | |
| masks per pair. The candidate list is capped with `_diverse_preselect` rather | |
| than a plain top-score cut, so the strongest lane cannot crowd out the others | |
| before NMS runs. | |
| """ | |
| if not lines: | |
| return [] | |
| order = sorted(lines, key=lambda ln: ln.score, reverse=True) | |
| order = _diverse_preselect(order, max_lanes) | |
| height = max(1, int(round(order[0].height * scale))) | |
| width = max(1, int(round(order[0].width * scale))) | |
| scaled_width = max(1, int(round(lane_width * scale))) | |
| masks = [] | |
| areas = [] | |
| for line in order: | |
| mask = np.zeros((height, width), np.uint8) | |
| points = line.xy() * scale | |
| if len(points) >= 2: | |
| points[:, 0] = np.clip(points[:, 0], 0, width - 1) | |
| points[:, 1] = np.clip(points[:, 1], 0, height - 1) | |
| cv2.polylines(mask, [points.astype(np.int32)], False, 1, | |
| scaled_width) | |
| masks.append(mask) | |
| areas.append(int(np.count_nonzero(mask))) | |
| suppressed = [False] * len(order) | |
| keep = [] | |
| for i in range(len(order)): | |
| if suppressed[i]: | |
| continue | |
| keep.append(order[i]) | |
| for j in range(i + 1, len(order)): | |
| if suppressed[j]: | |
| continue | |
| inter = int(np.count_nonzero(masks[i] & masks[j])) | |
| union = areas[i] + areas[j] - inter | |
| if union > 0 and inter / union >= thr: | |
| suppressed[j] = True | |
| return keep | |
| # --------------------------------------------------------------------------- # | |
| # lane identity (left-to-right ordering) | |
| # --------------------------------------------------------------------------- # | |
| def _bottom_x(lane): | |
| """x where a lane meets its nearest row (largest y). Lanes fan out near the | |
| camera, so this is the most reliable place to order them left-to-right.""" | |
| xy = lane.xy() | |
| if len(xy) == 0: | |
| return float("inf") | |
| return float(xy[int(np.argmax(xy[:, 1])), 0]) | |
| def _ego_reference_x(lane, target_y=None): | |
| """Estimate where a boundary meets the near-camera reference row. | |
| On a sharp bend, multiple boundaries can leave through the same image side. | |
| Comparing their last visible x then becomes ambiguous (both are about 0 or | |
| ``width - 1``), and a short outer boundary can be mistaken for the ego-lane | |
| boundary. Extrapolating the lower 40% of each polyline to a common row keeps | |
| their lateral order after they leave the image. | |
| """ | |
| xy = lane.xy() | |
| if len(xy) < 2: | |
| return _bottom_x(lane) | |
| xy = xy[np.isfinite(xy).all(axis=1)] | |
| if len(xy) < 2: | |
| return _bottom_x(lane) | |
| if target_y is None: | |
| target_y = lane.height - 1.0 | |
| cutoff = np.quantile(xy[:, 1], 0.6) | |
| lower = xy[xy[:, 1] >= cutoff] | |
| if len(lower) < 2 or np.ptp(lower[:, 1]) < 1.0: | |
| return _bottom_x(lane) | |
| ys = lower[:, 1] | |
| xs = lower[:, 0] | |
| centered_y = ys - ys.mean() | |
| denominator = float(np.dot(centered_y, centered_y)) | |
| if denominator <= 1e-6: | |
| return _bottom_x(lane) | |
| slope = float(np.dot(centered_y, xs - xs.mean()) / denominator) | |
| return float(xs.mean() + slope * (float(target_y) - ys.mean())) | |
| def order_lanes(lanes): | |
| """Sort lanes left-to-right and assign a frame-local index. | |
| RCLane is anchor-free: `decode` emits lane instances in score order with no | |
| inherent identity. This helper only establishes spatial order inside one | |
| frame; it must not be used as a persistent video identity because a missing | |
| outer lane would shift every following index. | |
| """ | |
| ordered = sorted(lanes, key=_bottom_x) | |
| for i, ln in enumerate(ordered): | |
| ln.lane_id = i | |
| return ordered | |
| def assign_ego_lane_roles(lanes, ego_x=None): | |
| """Assign stable semantic IDs relative to the ego vehicle. | |
| IDs describe a lane boundary's role rather than its position in a variable | |
| length list: | |
| * P1: nearest boundary left of ego (current-lane left boundary) | |
| * P2: nearest boundary right of ego (current-lane right boundary) | |
| * P0: next boundary to the left | |
| * P3: next boundary to the right | |
| Consequently P1/P2 do not become P0/P1 merely because an outer boundary is | |
| temporarily missing. With the default four-lane cap, returned IDs are in | |
| ``[0, 3]``. More uncapped lanes continue outward with negative IDs on the | |
| left and IDs greater than three on the right. | |
| """ | |
| if not lanes: | |
| return [] | |
| if ego_x is None: | |
| ego_x = lanes[0].width / 2.0 | |
| ego_x = float(ego_x) | |
| for lane in lanes: | |
| lane.lane_id = None | |
| lane.lane_role = None | |
| lane.is_ego_boundary = False | |
| lane.lateral_rank = None | |
| reference_x = {lane: _ego_reference_x(lane) for lane in lanes} | |
| left = sorted( | |
| (lane for lane in lanes if reference_x[lane] < ego_x), | |
| key=lambda lane: (abs(reference_x[lane] - ego_x), -lane.score), | |
| ) | |
| right = sorted( | |
| (lane for lane in lanes if reference_x[lane] >= ego_x), | |
| key=lambda lane: (abs(reference_x[lane] - ego_x), -lane.score), | |
| ) | |
| for rank, lane in enumerate(left, 1): | |
| lane.lane_id = 2 - rank # nearest left=P1, next=P0 | |
| lane.lateral_rank = -rank | |
| lane.is_ego_boundary = rank == 1 | |
| lane.lane_role = "ego_left" if rank == 1 else f"left_{rank}" | |
| for rank, lane in enumerate(right, 1): | |
| lane.lane_id = 1 + rank # nearest right=P2, next=P3 | |
| lane.lateral_rank = rank | |
| lane.is_ego_boundary = rank == 1 | |
| lane.lane_role = "ego_right" if rank == 1 else f"right_{rank}" | |
| return sorted(lanes, key=lambda lane: lane.lane_id) | |
| def ego_lane_boundaries(lanes): | |
| """Return ``(left, right)`` boundaries of the lane containing ego. | |
| Either value can be ``None`` when that side was not detected. | |
| """ | |
| left = next( | |
| (lane for lane in lanes if lane.lane_role == "ego_left"), None | |
| ) | |
| right = next( | |
| (lane for lane in lanes if lane.lane_role == "ego_right"), None | |
| ) | |
| return left, right | |
| def select_ego_lanes(lanes, max_lanes=4, ego_x=None, | |
| min_score_ratio=0.5, balance_sides=True): | |
| """Keep the closest reliable lane boundaries around the ego vehicle. | |
| The decoder can occasionally return an extra low-confidence crawl in | |
| addition to the real road boundaries. When more than ``max_lanes`` are | |
| present, first prefer candidates whose score is at least | |
| ``min_score_ratio`` of the best candidate (provided that still leaves enough | |
| lanes), then select the nearest boundaries using their near-camera x. | |
| For the usual four-lane output, ``balance_sides`` reserves two slots on | |
| either side of the camera centre when possible. Any unfilled slots are | |
| taken from the remaining closest candidates. The returned IDs are semantic: | |
| P1/P2 are the current-lane boundaries, while P0/P3 are the adjacent outer | |
| boundaries. Missing outer lanes therefore do not shift the ego-lane IDs. | |
| """ | |
| if max_lanes is None: | |
| return assign_ego_lane_roles(lanes, ego_x) | |
| if max_lanes <= 0: | |
| raise ValueError("max_lanes must be positive or None") | |
| if not 0.0 <= min_score_ratio <= 1.0: | |
| raise ValueError("min_score_ratio must be in [0, 1]") | |
| ordered = order_lanes(lanes) | |
| if not ordered: | |
| return [] | |
| if ego_x is None: | |
| ego_x = ordered[0].width / 2.0 | |
| ego_x = float(ego_x) | |
| # The four-lane semantic contract has exactly two possible boundaries per | |
| # side: P0/P1 on the left and P2/P3 on the right. If one side is missing, | |
| # return fewer lanes instead of filling the gap with P4/P-1 farther out. | |
| if len(ordered) <= max_lanes: | |
| if balance_sides and max_lanes == 4: | |
| reference_x = {lane: _ego_reference_x(lane) for lane in ordered} | |
| def near_ego(lane): | |
| return (abs(reference_x[lane] - ego_x), -lane.score) | |
| left = sorted( | |
| (lane for lane in ordered if reference_x[lane] < ego_x), | |
| key=near_ego, | |
| ) | |
| right = sorted( | |
| (lane for lane in ordered if reference_x[lane] >= ego_x), | |
| key=near_ego, | |
| ) | |
| ordered = left[:2] + right[:2] | |
| return assign_ego_lane_roles(ordered, ego_x) | |
| best_score = max(lane.score for lane in ordered) | |
| reliable = [ | |
| lane for lane in ordered | |
| if lane.score >= best_score * min_score_ratio | |
| ] | |
| # Never let the reliability gate force the output below the requested cap. | |
| pool = reliable if len(reliable) >= max_lanes else ordered | |
| reference_x = {lane: _ego_reference_x(lane) for lane in pool} | |
| def proximity_key(lane): | |
| return (abs(reference_x[lane] - ego_x), -lane.score) | |
| ranked = sorted(pool, key=proximity_key) | |
| selected = [] | |
| if balance_sides and max_lanes >= 2: | |
| left = sorted( | |
| (lane for lane in pool if reference_x[lane] < ego_x), | |
| key=proximity_key, | |
| ) | |
| right = sorted( | |
| (lane for lane in pool if reference_x[lane] >= ego_x), | |
| key=proximity_key, | |
| ) | |
| left_slots = max_lanes // 2 | |
| right_slots = max_lanes - left_slots | |
| selected.extend(left[:left_slots]) | |
| selected.extend(right[:right_slots]) | |
| if not (balance_sides and max_lanes == 4): | |
| for lane in ranked: | |
| if lane not in selected: | |
| selected.append(lane) | |
| if len(selected) == max_lanes: | |
| break | |
| return assign_ego_lane_roles(selected[:max_lanes], ego_x) | |
| # --------------------------------------------------------------------------- # | |
| # full decode | |
| # --------------------------------------------------------------------------- # | |
| def decode(seg_prob, up_arrow, down_arrow, up_bound, down_bound, | |
| step_length=10, seg_threshold=0.5, seed_min_dist=2, | |
| score_thresh=0.10, iou_thresh=0.5, seed_threshold=None, | |
| max_seeds=1024, nms_max_lanes=128, nms_scale=0.25, | |
| sort_lanes=True, max_output_lanes=4, ego_x=None, | |
| ego_min_score_ratio=0.5, balance_ego_sides=True, | |
| batch_crawl=True, crawl_backend="auto", | |
| point_nms_backend="auto"): | |
| """ | |
| Args: | |
| seg_prob: (H, W) foreground probability. | |
| up_arrow, down_arrow, up_bound, down_bound: (2, H, W). | |
| seg_threshold: foreground test used while crawling a chain. | |
| seed_threshold: threshold for picking seeds (defaults to seg_threshold). | |
| RCLane seg maps are low-magnitude (OHEM 15:1), so seeds often sit below | |
| 0.5 -- set this lower (e.g. 0.3) for under-trained models. | |
| max_output_lanes: final ego-centric lane cap. Defaults to four; pass | |
| ``None`` to preserve every lane surviving NMS. | |
| Returns: | |
| list of Lane. Use `lane.xy()` for the (N, 2) point array and `lane.score`. | |
| """ | |
| H, W = seg_prob.shape | |
| if seed_threshold is None: | |
| seed_threshold = seg_threshold | |
| seeds = point_nms( | |
| seg_prob, seed_threshold, seed_min_dist, max_seeds, | |
| backend=point_nms_backend, | |
| ) | |
| ub0, db0 = up_bound[0], down_bound[0] # bound channel 0 (both channels equal) | |
| if batch_crawl: | |
| backend = crawl_backend | |
| if crawl_backend == "auto": | |
| backend = "numba" if njit is not None else "numpy" | |
| if backend not in ("numba", "numpy"): | |
| raise ValueError("crawl_backend must be auto, numba, or numpy") | |
| crawl = ( | |
| decode_branches_numba if backend == "numba" | |
| else decode_branches_batch | |
| ) | |
| up_points, up_lengths = crawl( | |
| seeds, seg_prob, up_arrow, ub0, step_length, seg_threshold | |
| ) | |
| down_points, down_lengths = crawl( | |
| seeds, seg_prob, down_arrow, db0, step_length, seg_threshold | |
| ) | |
| candidate_indices = _preselect_batched_candidates( | |
| up_points, up_lengths, down_points, down_lengths, | |
| score_thresh, nms_max_lanes, backend, | |
| ) | |
| lines = _lines_from_batched_crawls( | |
| up_points, up_lengths, down_points, down_lengths, W, H, | |
| candidate_indices, | |
| ) | |
| else: | |
| lines = [] | |
| for (x, y) in seeds: | |
| up = decode_branch( | |
| x, y, seg_prob, up_arrow, ub0, step_length, seg_threshold | |
| ) | |
| down = decode_branch( | |
| x, y, seg_prob, down_arrow, db0, step_length, seg_threshold | |
| ) | |
| up.reverse() | |
| full = up.concat(down) | |
| if len(full) > 1: | |
| lines.append(full) | |
| lines = thresh_line(lines, score_thresh) | |
| lines = iou_nms(lines, iou_thresh, max_lanes=nms_max_lanes, | |
| scale=nms_scale) | |
| if max_output_lanes is not None: | |
| lines = select_ego_lanes( | |
| lines, | |
| max_lanes=max_output_lanes, | |
| ego_x=ego_x, | |
| min_score_ratio=ego_min_score_ratio, | |
| balance_sides=balance_ego_sides, | |
| ) | |
| elif sort_lanes: | |
| lines = order_lanes(lines) | |
| return lines | |
| def decode_predictions(pred_dict, **kwargs): | |
| """Convenience: decode a batch of network outputs (torch tensors, B,2,H,W). | |
| Returns a list (len B) of lists of Lane. Requires torch only for the input. | |
| """ | |
| import torch # local import so the module stays torch-free otherwise | |
| seg = torch.softmax(pred_dict["seg_map"], dim=1)[:, 1] # (B,H,W) fg prob | |
| seg = seg.detach().cpu().numpy() | |
| ua = pred_dict["up_arrow"].detach().cpu().numpy() | |
| da = pred_dict["down_arrow"].detach().cpu().numpy() | |
| ub = pred_dict["up_bound"].detach().cpu().numpy() | |
| db = pred_dict["down_bound"].detach().cpu().numpy() | |
| return [decode(seg[b], ua[b], da[b], ub[b], db[b], **kwargs) for b in range(seg.shape[0])] | |
| # --------------------------------------------------------------------------- # | |
| # smoke test -- encode/decode ROUND TRIP | |
| # --------------------------------------------------------------------------- # | |
| if __name__ == "__main__": | |
| from encode import encode, IMG_SIZE | |
| H, W = IMG_SIZE | |
| # ground-truth lane (an S-curve), points in 800x320 space | |
| ys = np.linspace(20, 300, 40) | |
| xs = 400 + 150 * np.sin(ys / 300 * np.pi) | |
| gt_lane = np.stack([xs, ys], axis=1) | |
| gt = encode([list(map(tuple, gt_lane))]) | |
| # decode straight from the GT maps -> should reconstruct the lane. | |
| # use a coarser seed spacing to keep the test fast. | |
| lanes = decode( | |
| gt["seg_map"], gt["up_arrow"], gt["down_arrow"], gt["up_bound"], gt["down_bound"], | |
| seed_min_dist=12, | |
| ) | |
| print(f"decoded {len(lanes)} lane(s) from GT maps") | |
| assert len(lanes) >= 1, "round trip produced no lanes!" | |
| # take the longest recovered lane, measure how far its points sit from the GT curve | |
| best = max(lanes, key=len) | |
| pred_xy = best.xy() | |
| print(f"longest lane: {len(best)} points, score={best.score:.2f}") | |
| def dist_to_gt(pt): | |
| d = np.hypot(gt_lane[:, 0] - pt[0], gt_lane[:, 1] - pt[1]) | |
| return d.min() | |
| errs = np.array([dist_to_gt(p) for p in pred_xy]) | |
| print(f"mean dist to GT curve = {errs.mean():.2f}px, max = {errs.max():.2f}px") | |
| assert errs.mean() < 8.0, "reconstructed lane strays too far from GT!" | |
| print("OK -- encode/decode round trip reconstructs the lane.") | |
| # Regression: a high-scoring lane may have hundreds of near-duplicate | |
| # crawls. The cap must still retain weaker candidates from other lanes, | |
| # while greedy NMS must receive candidates in descending score order. | |
| def vertical_lane(x, score): | |
| lane = Lane(W, H) | |
| lane.append(x, 200, score) | |
| lane.append(x, 300, score) | |
| return lane | |
| candidates = [ | |
| vertical_lane(100 + index % 2, 0.99 - index * 0.001) | |
| for index in range(24) | |
| ] | |
| candidates += [vertical_lane(350, 0.90), vertical_lane(650, 0.89)] | |
| candidates.sort(key=lambda lane: lane.score, reverse=True) | |
| selected = _diverse_preselect(candidates, max_lanes=8, bin_px=16) | |
| selected_bins = { | |
| int(np.median(lane.xy()[:, 0]) // 16) for lane in selected | |
| } | |
| expected_bins = {100 // 16, 350 // 16, 650 // 16} | |
| assert expected_bins <= selected_bins, "spatial preselection dropped a lane" | |
| selected_scores = [lane.score for lane in selected] | |
| assert selected_scores == sorted(selected_scores, reverse=True), ( | |
| "spatial preselection changed greedy NMS score priority" | |
| ) | |
| print("OK -- diverse NMS preselection retains spatially distinct lanes.") | |
| # Regression: cap the final output around the ego vehicle without keeping a | |
| # weak extra crawl merely because its endpoint is slightly closer laterally. | |
| ego_candidates = [ | |
| vertical_lane(8, 0.93), | |
| vertical_lane(20, 0.24), # spurious fifth crawl | |
| vertical_lane(95, 0.94), | |
| vertical_lane(748, 0.89), | |
| vertical_lane(796, 0.81), | |
| ] | |
| ego_lanes = select_ego_lanes(ego_candidates, max_lanes=4) | |
| ego_xs = [int(_bottom_x(lane)) for lane in ego_lanes] | |
| assert ego_xs == [8, 95, 748, 796], ( | |
| f"ego selector kept the wrong lanes: {ego_xs}" | |
| ) | |
| assert [lane.lane_id for lane in ego_lanes] == [0, 1, 2, 3] | |
| ego_left, ego_right = ego_lane_boundaries(ego_lanes) | |
| assert ego_left is not None and int(_bottom_x(ego_left)) == 95 | |
| assert ego_right is not None and int(_bottom_x(ego_right)) == 748 | |
| # Semantic IDs must not shift when an outer lane disappears. P1/P2 remain | |
| # the current-lane boundaries and P3 remains the next boundary on the right. | |
| missing_outer_left = select_ego_lanes( | |
| [vertical_lane(95, 0.94), vertical_lane(748, 0.89), | |
| vertical_lane(796, 0.81)], | |
| max_lanes=4, | |
| ) | |
| assert [lane.lane_id for lane in missing_outer_left] == [1, 2, 3] | |
| missing_outer_right = select_ego_lanes( | |
| [vertical_lane(8, 0.93), vertical_lane(95, 0.94), | |
| vertical_lane(748, 0.89)], | |
| max_lanes=4, | |
| ) | |
| assert [lane.lane_id for lane in missing_outer_right] == [0, 1, 2] | |
| right_only = select_ego_lanes( | |
| [vertical_lane(500, 0.95), vertical_lane(600, 0.90), | |
| vertical_lane(700, 0.80)], | |
| max_lanes=4, | |
| ) | |
| assert [lane.lane_id for lane in right_only] == [2, 3] | |
| # Two right boundaries can both leave through x=width on a sharp curve. | |
| # The longer/nearer curve must remain P2 even if its last visible x is a | |
| # little farther right than the short outer curve's last x. | |
| near_right = Lane(800, 320) | |
| outer_right = Lane(800, 320) | |
| for x, y in ((700, 200), (740, 225), (780, 250)): | |
| near_right.append(x, y, 0.9) | |
| for x, y in ((700, 125), (750, 137.5), (790, 147.5)): | |
| outer_right.append(x, y, 0.8) | |
| curved_right = assign_ego_lane_roles([outer_right, near_right], ego_x=400) | |
| assert near_right.lane_id == 2 and near_right.lane_role == "ego_right" | |
| assert outer_right.lane_id == 3 and outer_right.lane_role == "right_2" | |
| assert [lane.lane_id for lane in curved_right] == [2, 3] | |
| print("OK -- ego post-processing keeps four reliable nearby lanes.") | |