scorevision: push artifact

training/train.py

@@ -0,0 +1,140 @@
+"""
+Two-stage training recipe for Detect-crime.
+
+Stage A (silver pretrain): warm a YOLOv11s/m on the full silver corpus
+                           (Roboflow + COCO bat + king-distilled manako frames).
+Stage B (clean fine-tune):  short fine-tune on a manually-verified set of manako frames.
+
+Defaults are tuned for a single GPU with >=24 GB VRAM. Pro_6000 should run YOLOv11s at
+imgsz=1536 with batch 12 comfortably; YOLOv11m at imgsz=1280 needs batch ~8.
+
+Usage:
+  # Stage A
+  python train.py --data ../data/data.yaml --weights yolo11s.pt --imgsz 1280 --stage A --epochs 200
+
+  # Stage B (clean fine-tune)
+  python train.py --data ../data_clean/data.yaml \
+                  --weights runs/detect/crime_a/weights/best.pt \
+                  --imgsz 1280 --stage B --epochs 50
+
+Notes:
+  - Heavy mosaic+copy_paste is essential for crime — class imbalance is brutal (hoodie
+    dominates GT count; balaclava/glove/spray-paint are rare). Copy-paste pastes rare-class
+    crops into otherwise-normal frames; this is the most reliable fix for low recall on the
+    rare classes.
+  - HSV jitter cranked on V channel: CCTV crime footage spans daylight, IR night, and
+    mixed indoor lighting. Augment hard.
+  - `cls=1.0` (raised from default 0.5) because mAP@50 is per-class-averaged; pushing the
+    classification head matters more than localization here.
+  - Always save with ema=True so export uses smoothed weights.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+from ultralytics import YOLO
+
+
+COMMON_AUG = dict(
+    hsv_h=0.02,
+    hsv_s=0.7,
+    hsv_v=0.5,         # CCTV brightness varies massively (IR night, daylight, mixed indoor)
+    degrees=0.0,       # CCTV cameras are mounted level
+    translate=0.10,
+    scale=0.5,
+    shear=0.0,
+    perspective=0.0,
+    flipud=0.0,
+    fliplr=0.5,
+    mosaic=1.0,
+    mixup=0.20,
+    copy_paste=0.50,   # high: rebalance toward rare classes (balaclava, glove, spray paint)
+    erasing=0.0,
+)
+
+STAGE_A = dict(
+    epochs=200,
+    optimizer="AdamW",
+    lr0=2e-3,
+    lrf=0.01,
+    momentum=0.937,
+    weight_decay=5e-4,
+    warmup_epochs=3,
+    cos_lr=True,
+    label_smoothing=0.05,
+    cls=1.0,           # higher cls weight: per-class AP matters
+    box=7.5,
+    dfl=1.5,
+    close_mosaic=15,
+    patience=40,
+    amp=True,
+    **COMMON_AUG,
+)
+
+STAGE_B = dict(
+    epochs=50,
+    optimizer="AdamW",
+    lr0=5e-4,
+    lrf=0.01,
+    momentum=0.9,
+    weight_decay=5e-4,
+    warmup_epochs=2,
+    cos_lr=True,
+    label_smoothing=0.05,
+    cls=0.8,
+    box=7.5,
+    dfl=1.5,
+    close_mosaic=10,
+    patience=15,
+    amp=True,
+    **{**COMMON_AUG, "mosaic": 0.5, "mixup": 0.0, "copy_paste": 0.10},
+)
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", required=True, help="path to data.yaml (YOLO format)")
+    ap.add_argument("--weights", default="yolo11s.pt",
+                    help="starting weights (yolo11n.pt / yolo11s.pt / yolo11m.pt / your stage-A best)")
+    ap.add_argument("--imgsz", type=int, default=1280)
+    ap.add_argument("--batch", type=int, default=16)
+    ap.add_argument("--stage", choices=["A", "B"], default="A")
+    ap.add_argument("--epochs", type=int, default=None)
+    ap.add_argument("--device", default=0)
+    ap.add_argument("--project", default="runs/detect")
+    ap.add_argument("--name", default=None)
+    ap.add_argument("--resume", action="store_true")
+    args = ap.parse_args()
+
+    cfg = STAGE_A.copy() if args.stage == "A" else STAGE_B.copy()
+    if args.epochs is not None:
+        cfg["epochs"] = args.epochs
+
+    model = YOLO(args.weights)
+    name = args.name or f"crime_{args.stage.lower()}"
+
+    model.train(
+        data=args.data,
+        imgsz=args.imgsz,
+        batch=args.batch,
+        device=args.device,
+        project=args.project,
+        name=name,
+        resume=args.resume,
+        save=True,
+        plots=True,
+        seed=42,
+        rect=False,
+        single_cls=False,
+        **cfg,
+    )
+
+    metrics = model.val(data=args.data, imgsz=args.imgsz, batch=args.batch,
+                        device=args.device, plots=True, save_json=True)
+    print(f"[train] {args.stage} done. map50={metrics.box.map50:.4f} map={metrics.box.map:.4f}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())