52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
"""Extract per-lane polylines from binary (JPEG-compressed) lane-segmentation masks.
|
|
|
|
The dataset at hand stores lanes as a single-class binary mask (white=lane,
|
|
black=background) rather than TuSimple's official per-lane (x,y) JSON annotations, so
|
|
lane *instances* have to be recovered from connected components of the thresholded
|
|
mask. Masks are lossy JPEG, so we clean small compression-noise blobs before labeling.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
MIN_BLOB_AREA = 150 # drop connected components smaller than this (JPEG noise)
|
|
BINARY_THRESHOLD = 127
|
|
MORPH_KERNEL = np.ones((3, 3), np.uint8)
|
|
MIN_POINTS_PER_LANE = 5 # discard components too short to be a real lane
|
|
|
|
|
|
def mask_to_lanes(mask: np.ndarray) -> list[list[tuple[float, float]]]:
|
|
"""Convert an (H,W) or (H,W,3) binary-ish mask into a list of lane polylines.
|
|
|
|
Each polyline is a list of (x, y) pixel points, one point per mask row that the
|
|
lane instance covers, ordered top-to-bottom (increasing y).
|
|
"""
|
|
if mask.ndim == 3:
|
|
mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
|
|
|
|
binary = (mask > BINARY_THRESHOLD).astype(np.uint8)
|
|
binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, MORPH_KERNEL)
|
|
|
|
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
|
|
|
|
lanes: list[list[tuple[float, float]]] = []
|
|
for label_id in range(1, num_labels): # 0 is background
|
|
area = stats[label_id, cv2.CC_STAT_AREA]
|
|
if area < MIN_BLOB_AREA:
|
|
continue
|
|
|
|
ys, xs = np.where(labels == label_id)
|
|
if ys.size == 0:
|
|
continue
|
|
|
|
row_to_xs: dict[int, list[int]] = {}
|
|
for y, x in zip(ys, xs):
|
|
row_to_xs.setdefault(int(y), []).append(int(x))
|
|
|
|
points = [(float(np.mean(row_to_xs[y])), float(y)) for y in sorted(row_to_xs)]
|
|
if len(points) >= MIN_POINTS_PER_LANE:
|
|
lanes.append(points)
|
|
|
|
return lanes
|