Aryan Sethi
Claude Opus 5 (1M context)
Drop barrel, add held-out stop-sign distance, add the GPU runbook
5d449ff Download src/data/merge.py from Aryan006/cone-distance: direct link, hf CLI and curl.
- Browser
- Download file 13.7 kB
-
https://huggingface.co/Aryan006/cone-distance/resolve/main/src/data/merge.py
- Command line
-
hf download hf://Aryan006/cone-distance/src/data/merge.py
-
curl -L -o merge.py https://huggingface.co/Aryan006/cone-distance/resolve/main/src/data/merge.py
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() | |