"""The unified schema: class names, manifest columns, and the filter thresholds. Every extractor writes rows in this shape, merge.py concatenates them, and to_yolo.py is a pure function of the result. If you change a column here, change it here only. A manifest row is one of two things: * an annotated object -- `class` is set, the box columns are set * a negative image -- `class` is null, the box columns are null Negatives exist so that images with no objects (BDD100K road scenes) still reach the YOLO export as empty label files. """ from __future__ import annotations import pandas as pd # --------------------------------------------------------------------------- # Classes # --------------------------------------------------------------------------- # Order matters: this is the YOLO class index order, written into data.yaml. CLASSES: tuple[str, ...] = ("cone", "barrier", "stop_sign") CLASS_TO_INDEX: dict[str, int] = {name: i for i, name in enumerate(CLASSES)} # Classes whose 2D box bottom is a ground contact point. Phase 2 routes these # to the ground-plane estimator; everything else (stop_sign, pole-mounted) goes # to the known-size estimator. GROUND_CONTACT_CLASSES: tuple[str, ...] = ("cone", "barrier") # --------------------------------------------------------------------------- # Manifest columns # --------------------------------------------------------------------------- # gt_distance_m is the forward depth of the object centre along the camera # optical axis (the Z component in the camera frame), in metres. It is NOT # Euclidean range. Phase 2 must compare like with like. MANIFEST_COLUMNS: tuple[str, ...] = ( "image_path", # str relative to data/unified/images/ "sensor_id", # str joins to calib/.yaml "source", # str nuscenes / av2 / coco / bdd "scene_id", # str split unit -- never split on frames "class", # str one of CLASSES, or null for a negative image "x1", # float pixels, absolute, clipped to the image "y1", "x2", "y2", "gt_distance_m", # float null for COCO/BDD "gt_dims_hwl", # list [height, width, length] metres, null where absent "visibility", # float 0-1, null where unavailable "truncation", # float 0-1, fraction of the projected box lost to clipping "split", # str train / val, assigned by merge.py ) # --------------------------------------------------------------------------- # Filter thresholds (applied in merge.py, documented in the README) # --------------------------------------------------------------------------- # A box shorter than this is not a learnable label, it is false-positive # pressure. 15 px at 640 input is roughly 10 px after the resize. MIN_BOX_HEIGHT_PX = 15.0 # nuScenes visibility buckets are v0-40 / v40-60 / v60-80 / v80-100. This # threshold drops the first bucket only. MIN_VISIBILITY = 0.4 # Fraction of the projected box that may fall outside the frame. Above this the # box is mostly a guess about what is off-screen. MAX_TRUNCATION = 0.5 # Sanity range for the median cone distance, asserted in merge.py. A broken # camera-frame transform lands far outside this. PLAUSIBLE_MEDIAN_CONE_DISTANCE_M = (3.0, 80.0) # --------------------------------------------------------------------------- # Row builders -- the only places a manifest row is constructed # --------------------------------------------------------------------------- def object_row( *, image_path: str, sensor_id: str, source: str, scene_id: str, class_name: str, box: tuple[float, float, float, float], gt_distance_m: float | None = None, gt_dims_hwl: list[float] | None = None, visibility: float | None = None, truncation: float = 0.0, ) -> dict: """One annotated object.""" if class_name not in CLASS_TO_INDEX: raise ValueError(f"unknown class {class_name!r}, expected one of {CLASSES}") x1, y1, x2, y2 = box return { "image_path": image_path, "sensor_id": sensor_id, "source": source, "scene_id": scene_id, "class": class_name, "x1": float(x1), "y1": float(y1), "x2": float(x2), "y2": float(y2), "gt_distance_m": None if gt_distance_m is None else float(gt_distance_m), "gt_dims_hwl": gt_dims_hwl, "visibility": None if visibility is None else float(visibility), "truncation": float(truncation), "split": None, } def negative_row( *, image_path: str, sensor_id: str, source: str, scene_id: str, ) -> dict: """One image that contains none of our classes.""" return { "image_path": image_path, "sensor_id": sensor_id, "source": source, "scene_id": scene_id, "class": None, "x1": None, "y1": None, "x2": None, "y2": None, "gt_distance_m": None, "gt_dims_hwl": None, "visibility": None, "truncation": None, "split": None, } # --------------------------------------------------------------------------- # Manifest I/O # --------------------------------------------------------------------------- def rows_to_frame(rows: list[dict]) -> pd.DataFrame: """Build a DataFrame with exactly MANIFEST_COLUMNS, in order.""" frame = pd.DataFrame(rows, columns=list(MANIFEST_COLUMNS)) for column in ("x1", "y1", "x2", "y2", "gt_distance_m", "visibility", "truncation"): frame[column] = pd.to_numeric(frame[column], errors="coerce") return frame def write_manifest(frame: pd.DataFrame, path) -> None: missing = set(MANIFEST_COLUMNS) - set(frame.columns) if missing: raise ValueError(f"manifest is missing columns: {sorted(missing)}") frame = frame[list(MANIFEST_COLUMNS)] frame.to_parquet(path, index=False) def read_manifest(path) -> pd.DataFrame: frame = pd.read_parquet(path) missing = set(MANIFEST_COLUMNS) - set(frame.columns) if missing: raise ValueError(f"manifest at {path} is missing columns: {sorted(missing)}") return frame def objects_only(frame: pd.DataFrame) -> pd.DataFrame: """Drop negative rows, keeping only rows that describe an object.""" return frame[frame["class"].notna()].copy()