File size: 4,367 Bytes
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 | """COCO -> manifest rows for stop signs.
Free stop-sign data that the pretrained weights have already seen. No depth
ground truth, so gt_distance_m is null -- these rows train the detector and are
invisible to the phase 2 scoring.
Domain caveat, worth remembering when stop-sign mAP looks strange: COCO stop
signs are photographs from arbitrary viewpoints, including close-ups and
non-road contexts. nuScenes and AV2 frames are dashcam views. A model trained
on both can do well on COCO-style stop signs and poorly on dashcam ones.
Usage:
python -m src.data.extract_coco --dataroot /data/coco --split train2017
"""
from __future__ import annotations
import argparse
import json
from collections import defaultdict
from pathlib import Path
from src.common import paths, schema
from src.common.geometry import clip_box
SOURCE = "coco"
COCO_CATEGORY = "stop sign"
UNIFIED_CLASS = "stop_sign"
# COCO images have no shared calibration. Phase 2 must not try to estimate
# distance from them, and this sensor_id has no calib/*.yaml on purpose.
SENSOR_ID = "coco_unknown"
def extract(dataroot: Path, split: str, root: Path, max_images: int | None) -> None:
annotation_path = dataroot / "annotations" / f"instances_{split}.json"
print(f"reading {annotation_path} (this takes a moment, it is a large file)")
with open(annotation_path) as handle:
coco = json.load(handle)
category_ids = [c["id"] for c in coco["categories"] if c["name"] == COCO_CATEGORY]
if not category_ids:
raise SystemExit(f"no category named {COCO_CATEGORY!r} in {annotation_path}")
category_id = category_ids[0]
annotations_by_image: dict[int, list[dict]] = defaultdict(list)
for annotation in coco["annotations"]:
if annotation["category_id"] == category_id and not annotation.get("iscrowd", 0):
annotations_by_image[annotation["image_id"]].append(annotation)
images = {image["id"]: image for image in coco["images"]}
image_ids = sorted(annotations_by_image)
if max_images is not None:
image_ids = image_ids[:max_images]
paths.link_source(root, SOURCE, dataroot)
rows: list[dict] = []
skipped = 0
for image_id in image_ids:
image = images[image_id]
image_path = f"{SOURCE}/{split}/{image['file_name']}"
for annotation in annotations_by_image[image_id]:
x, y, w, h = annotation["bbox"]
clipped, truncation = clip_box(
(x, y, x + w, y + h), image["width"], image["height"]
)
if clipped[2] - clipped[0] <= 1.0 or clipped[3] - clipped[1] <= 1.0:
skipped += 1
continue
rows.append(
schema.object_row(
image_path=image_path,
sensor_id=SENSOR_ID,
source=SOURCE,
# Each COCO photo is its own scene; there is no video to leak.
scene_id=f"{SOURCE}_{image_id}",
class_name=UNIFIED_CLASS,
box=clipped,
gt_distance_m=None,
gt_dims_hwl=None,
visibility=None,
truncation=truncation,
)
)
frame = schema.rows_to_frame(rows)
part = paths.part_path(root, SOURCE)
part.parent.mkdir(parents=True, exist_ok=True)
schema.write_manifest(frame, part)
print(f"\nwrote {part}")
print(f" frames : {frame['image_path'].nunique()}")
print(f" objects : {len(frame)}")
if skipped:
print(f" dropped : {skipped} degenerate boxes")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--dataroot", required=True, type=Path,
help="COCO root, containing annotations/ and train2017/")
parser.add_argument("--split", default="train2017")
parser.add_argument("--unified-root", type=Path, default=None)
parser.add_argument("--max-images", type=int, default=None)
args = parser.parse_args()
extract(
dataroot=args.dataroot,
split=args.split,
root=paths.unified_root(args.unified_root),
max_images=args.max_images,
)
if __name__ == "__main__":
main()
|