cone-distance / src /data /merge.py
Aryan Sethi
Claude Opus 5 (1M context)
Drop barrel, add held-out stop-sign distance, add the GPU runbook
5d449ff
Raw History Blame Contribute Delete
13.7 kB
"""Concatenate the per-source manifests, filter, split by scene, assert sanity.
This is where the thresholds from schema.py are applied and where the
cross-dataset agreement checks live. The checks are real asserts, not things to
eyeball, because a coordinate-convention mistake between nuScenes and AV2
produces a manifest that trains YOLO perfectly well and makes the depth numbers
nonsense a day later.
Also writes class_priors.yaml -- per-class mean physical dimensions derived from
gt_dims_hwl. Phase 2's known-size estimator consumes it. Data-driven beats
looked-up, and it is one groupby.
Usage:
python -m src.data.merge
python -m src.data.merge --val-fraction 0.2
"""
from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
import pandas as pd
import yaml
from src.common import paths, schema
# The US MUTCD R1-1 stop sign is 0.75 m across the flats. Stop signs have no
# cuboid dimensions in COCO and are pole-mounted, so phase 2 uses width, not
# height, for them. This is the one prior we cannot derive from the data.
STOP_SIGN_WIDTH_M = 0.75
# ---------------------------------------------------------------------------
# Load and filter
# ---------------------------------------------------------------------------
def load_parts(root: Path) -> pd.DataFrame:
part_paths = sorted(paths.parts_dir(root).glob("*.parquet"))
if not part_paths:
raise SystemExit(f"no per-source manifests under {paths.parts_dir(root)}; "
f"run an extractor first")
frames = []
for path in part_paths:
frame = schema.read_manifest(path)
print(f" {path.name:28s} {len(frame):7d} rows")
frames.append(frame)
return pd.concat(frames, ignore_index=True)
def apply_filters(frame: pd.DataFrame) -> pd.DataFrame:
"""Drop objects we do not intend to detect. Thresholds live in schema.py."""
objects = frame["class"].notna()
height = frame["y2"] - frame["y1"]
too_small = objects & (height < schema.MIN_BOX_HEIGHT_PX)
too_occluded = objects & frame["visibility"].notna() & (
frame["visibility"] < schema.MIN_VISIBILITY
)
too_truncated = objects & frame["truncation"].notna() & (
frame["truncation"] > schema.MAX_TRUNCATION
)
print(f"\nfilters (thresholds in src/common/schema.py)")
print(f" box height < {schema.MIN_BOX_HEIGHT_PX:g} px : {int(too_small.sum())}")
print(f" visibility < {schema.MIN_VISIBILITY:g} : {int(too_occluded.sum())}")
print(f" truncation > {schema.MAX_TRUNCATION:g} : {int(too_truncated.sum())}")
return frame[~(too_small | too_occluded | too_truncated)].copy()
def restore_frames_as_negatives(filtered: pd.DataFrame,
original: pd.DataFrame) -> pd.DataFrame:
"""Keep images that lost every object to the filters, as negatives.
The alternative is dropping those images entirely. We keep them: everything
the filters removed is below the size we intend to detect at all, so
teaching the model not to fire there matches the deployment target. Use
--drop-filtered-frames to choose the other behaviour.
"""
surviving = set(filtered["image_path"])
lost = original[~original["image_path"].isin(surviving)]
if lost.empty:
return filtered
negatives = lost.drop_duplicates("image_path").apply(
lambda row: schema.negative_row(
image_path=row["image_path"],
sensor_id=row["sensor_id"],
source=row["source"],
scene_id=row["scene_id"],
),
axis=1,
)
print(f" {len(negatives)} images lost every object -> kept as negatives")
return pd.concat(
[filtered, schema.rows_to_frame(list(negatives))], ignore_index=True
)
# ---------------------------------------------------------------------------
# Split
# ---------------------------------------------------------------------------
def assign_split(frame: pd.DataFrame, val_fraction: float, seed: int) -> pd.DataFrame:
"""Split by scene, never by frame, and do it independently per source.
Adjacent video frames are near-duplicates, so a frame-level random split
leaks the val set into training and inflates val mAP a long way. Splitting
per source also guarantees every source is represented in val -- otherwise
a random scene draw can leave val with no nuScenes cones, which is the one
thing phase 2 needs.
Within a source, val is seeded with one scene per class before being filled
to the frame target. With few scenes this can overshoot --val-fraction; that
is the better failure, since a class with no val instances cannot be scored
at all.
"""
frame = frame.copy()
frame["split"] = None
rng = np.random.default_rng(seed)
for source in sorted(frame["source"].unique()):
mask = frame["source"] == source
source_frame = frame[mask]
frames_per_scene = source_frame.groupby("scene_id")["image_path"].nunique()
scene_ids = frames_per_scene.index.to_numpy()
rng.shuffle(scene_ids)
# Which scenes contain each class. Instances cluster hard by scene -- in
# AV2 nearly all cones live in a few logs -- so a purely random draw over
# a handful of scenes can leave a class with zero val instances and no
# way to score it. Seed val with one scene per class before filling.
objects = source_frame[source_frame["class"].notna()]
scenes_with = {name: set(group["scene_id"])
for name, group in objects.groupby("class")}
target = val_fraction * frames_per_scene.sum()
val_scenes: list[str] = []
accumulated = 0
for class_name in sorted(scenes_with):
if any(scene in scenes_with[class_name] for scene in val_scenes):
continue
# scene_ids is already shuffled, so the first match is a random one.
for scene_id in scene_ids:
if scene_id in scenes_with[class_name]:
val_scenes.append(scene_id)
accumulated += int(frames_per_scene[scene_id])
break
for scene_id in scene_ids:
if accumulated >= target:
break
if scene_id not in val_scenes:
val_scenes.append(scene_id)
accumulated += int(frames_per_scene[scene_id])
val_scenes = set(val_scenes)
frame.loc[mask, "split"] = np.where(
frame.loc[mask, "scene_id"].isin(val_scenes), "val", "train"
)
print(f" {source:10s} {len(val_scenes):5d}/{len(scene_ids):5d} scenes -> val "
f"({accumulated} frames)")
return frame
# ---------------------------------------------------------------------------
# Sanity checks -- asserts, not eyeballing
# ---------------------------------------------------------------------------
def check(frame: pd.DataFrame) -> None:
objects = schema.objects_only(frame)
# 1. Boxes are non-degenerate.
assert (objects["x2"] > objects["x1"]).all(), "boxes with x2 <= x1 survived merge"
assert (objects["y2"] > objects["y1"]).all(), "boxes with y2 <= y1 survived merge"
# 2. Distance is always positive where it exists. A negative distance means
# an object behind the camera was projected as if it were in front.
with_distance = objects[objects["gt_distance_m"].notna()]
assert (with_distance["gt_distance_m"] > 0).all(), (
f"{int((with_distance['gt_distance_m'] <= 0).sum())} objects with "
f"non-positive gt_distance_m"
)
# 3. Distance grows toward the horizon. Within one camera, an object whose
# box bottom is lower in the frame is nearer. A sign error or a swapped
# axis in the camera-frame transform flips or destroys this.
ground = with_distance[with_distance["class"].isin(schema.GROUND_CONTACT_CLASSES)]
for sensor_id, group in ground.groupby("sensor_id"):
if len(group) < 50:
continue
correlation = float(np.corrcoef(group["y2"], group["gt_distance_m"])[0, 1])
assert correlation < 0.0, (
f"{sensor_id}: box bottom row and distance correlate at {correlation:+.2f}. "
f"Objects lower in the frame should be NEARER. The camera-frame "
f"transform is wrong."
)
if correlation > -0.3:
print(f" !! {sensor_id}: weak horizon correlation {correlation:+.2f}")
else:
print(f" {sensor_id}: horizon correlation {correlation:+.2f}")
# 4. Median cone distance is plausible. Getting this wrong by a factor of
# 100 (millimetres, or a distance that is really an index) is the classic
# unit mistake and it survives every check above.
cones = with_distance[with_distance["class"] == "cone"]
if len(cones):
median = float(cones["gt_distance_m"].median())
low, high = schema.PLAUSIBLE_MEDIAN_CONE_DISTANCE_M
assert low <= median <= high, (
f"median cone distance is {median:.2f} m, expected {low}-{high} m. "
f"Check the units and the camera-frame transform."
)
print(f" median cone distance {median:.1f} m")
# 5. The thing phase 2 actually runs on. Report every class, because a
# scene-level split on a small number of scenes can easily land zero
# instances of a class in val, and a silent zero here means phase 2 has
# nothing to score that class against.
print(" val instances with gt_distance_m, per class:")
val = with_distance[with_distance["split"] == "val"]
train_classes = set(objects[objects["split"] == "train"]["class"])
for class_name in sorted(train_classes):
count = int((val["class"] == class_name).sum())
flag = " <-- ZERO, change --seed or --val-fraction" if count == 0 else ""
print(f" {class_name:10s} {count:6d}{flag}")
if len(val[val["class"] == "cone"]) < 200:
print(" !! fewer than 200 val cones -- phase 2 wants at least a few hundred")
# ---------------------------------------------------------------------------
# Class priors
# ---------------------------------------------------------------------------
def write_class_priors(frame: pd.DataFrame, root: Path) -> Path:
"""Per-class mean physical size, for phase 2's known-size estimator."""
objects = schema.objects_only(frame)
with_dims = objects[objects["gt_dims_hwl"].notna()]
priors: dict[str, dict] = {}
for class_name, group in with_dims.groupby("class"):
dims = np.array(group["gt_dims_hwl"].tolist(), dtype=float)
priors[class_name] = {
"height_m": float(np.mean(dims[:, 0])),
"width_m": float(np.mean(dims[:, 1])),
"length_m": float(np.mean(dims[:, 2])),
"height_std_m": float(np.std(dims[:, 0])),
"n_samples": int(len(dims)),
"source": "derived from manifest gt_dims_hwl",
}
if "stop_sign" not in priors:
priors["stop_sign"] = {
"width_m": STOP_SIGN_WIDTH_M,
"source": "US MUTCD R1-1, across the flats (no cuboid dims available)",
}
path = Path(root) / "class_priors.yaml"
with open(path, "w") as handle:
yaml.safe_dump(priors, handle, sort_keys=True)
print(f"\nwrote {path}")
for class_name, prior in sorted(priors.items()):
if "height_m" in prior:
print(f" {class_name:10s} height {prior['height_m']:.2f} m "
f"(n={prior['n_samples']})")
else:
print(f" {class_name:10s} width {prior['width_m']:.2f} m ({prior['source']})")
return path
# ---------------------------------------------------------------------------
def merge(root: Path, val_fraction: float, seed: int, drop_filtered_frames: bool) -> None:
print("loading per-source manifests")
frame = load_parts(root)
print(f" {'total':28s} {len(frame):7d} rows")
filtered = apply_filters(frame)
if not drop_filtered_frames:
filtered = restore_frames_as_negatives(filtered, frame)
print("\nsplitting by scene")
filtered = assign_split(filtered, val_fraction, seed)
print("\nsanity checks")
check(filtered)
manifest_path = Path(root) / "manifest.parquet"
schema.write_manifest(filtered, manifest_path)
objects = schema.objects_only(filtered)
print(f"\nwrote {manifest_path}")
print(f" frames : {filtered['image_path'].nunique()}")
print(f" objects: {len(objects)}")
print("\ninstances per source x class")
print(pd.crosstab(objects["source"], objects["class"]).to_string())
print("\nframes per split")
print(filtered.groupby("split")["image_path"].nunique().to_string())
write_class_priors(filtered, root)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--unified-root", type=Path, default=None)
parser.add_argument("--val-fraction", type=float, default=0.15)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--drop-filtered-frames", action="store_true",
help="drop images that lost every object to the filters, "
"instead of keeping them as negatives")
args = parser.parse_args()
merge(
root=paths.unified_root(args.unified_root),
val_fraction=args.val_fraction,
seed=args.seed,
drop_filtered_frames=args.drop_filtered_frames,
)
if __name__ == "__main__":
main()