File size: 6,341 Bytes
9274a8d 5d449ff 9274a8d 5d449ff 9274a8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | """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()
|