src/common/schema.py

"""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/<sensor_id>.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()