scorevision: push artifact

training/verify_dataset.py

@@ -0,0 +1,170 @@
+"""
+Quick QA over an assembled YOLO dataset for Detect-crime.
+
+Checks:
+  - Every label file has a corresponding image and vice versa.
+  - All class ids are within range (0..5).
+  - All bbox coords are normalized (0 < x < 1, 0 < y < 1, 0 < w/h <= 1).
+  - Reports per-class instance counts per split + total positive images per class.
+  - Flags samples with degenerate (near-zero) boxes.
+  - Optional: writes overlay PNGs for the first --visualize K val images.
+
+Usage:
+  python verify_dataset.py --data ../data/data.yaml
+  python verify_dataset.py --data ../data/data.yaml --visualize 20 --vis-out /tmp/crime_vis
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+TARGET_CLASS_NAMES = ["balaclava", "bat", "glove", "graffiti", "hoodie", "spray paint"]
+
+
+def parse_yaml(path: Path) -> dict:
+    try:
+        import yaml
+        return yaml.safe_load(path.read_text())
+    except Exception:
+        cfg: dict = {}
+        for line in path.read_text().splitlines():
+            line = line.strip()
+            if not line or line.startswith("#") or ":" not in line:
+                continue
+            k, v = line.split(":", 1)
+            cfg[k.strip()] = v.strip()
+        return cfg
+
+
+def check_split(root: Path, split: str, log) -> dict:
+    img_dir = root / "images" / split
+    lbl_dir = root / "labels" / split
+    if not img_dir.exists():
+        log(f"[{split}] images/{split} missing: {img_dir}")
+        return {"images": 0}
+    if not lbl_dir.exists():
+        log(f"[{split}] labels/{split} missing: {lbl_dir}")
+        return {"images": 0}
+
+    image_names = {p.stem for p in img_dir.iterdir() if p.is_file()}
+    label_names = {p.stem for p in lbl_dir.iterdir() if p.is_file() and p.suffix == ".txt"}
+    only_img = image_names - label_names
+    only_lbl = label_names - image_names
+
+    n_total = len(image_names)
+    cls_counts = [0] * len(TARGET_CLASS_NAMES)
+    cls_image_counts = [set() for _ in range(len(TARGET_CLASS_NAMES))]
+    bad_rows = 0
+    degenerate = 0
+    for lbl in lbl_dir.glob("*.txt"):
+        rows = lbl.read_text().splitlines()
+        for r in rows:
+            parts = r.strip().split()
+            if len(parts) != 5:
+                bad_rows += 1
+                continue
+            try:
+                c = int(parts[0])
+                cx, cy, bw, bh = (float(parts[1]), float(parts[2]), float(parts[3]), float(parts[4]))
+            except Exception:
+                bad_rows += 1
+                continue
+            if not (0 <= c < len(TARGET_CLASS_NAMES)):
+                bad_rows += 1
+                continue
+            if not (0 < cx < 1 and 0 < cy < 1 and 0 < bw <= 1 and 0 < bh <= 1):
+                degenerate += 1
+                continue
+            cls_counts[c] += 1
+            cls_image_counts[c].add(lbl.stem)
+
+    return {
+        "images": n_total,
+        "only_img": only_img,
+        "only_lbl": only_lbl,
+        "cls_counts": cls_counts,
+        "cls_image_counts": [len(s) for s in cls_image_counts],
+        "bad_rows": bad_rows,
+        "degenerate": degenerate,
+    }
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", required=True)
+    ap.add_argument("--visualize", type=int, default=0,
+                    help="render this many val images with boxes overlaid")
+    ap.add_argument("--vis-out", default="/tmp/crime_vis", help="dir for overlay images")
+    args = ap.parse_args()
+
+    data_yaml = Path(args.data).resolve()
+    cfg = parse_yaml(data_yaml)
+    root = Path(cfg.get("path") or data_yaml.parent).resolve()
+
+    def log(msg: str):
+        print(msg, file=sys.stderr, flush=True)
+
+    log(f"[verify] root={root}")
+    overall_ok = True
+    for split in ("train", "val"):
+        res = check_split(root, split, log)
+        if not res["images"]:
+            continue
+        log(f"[{split}] images: {res['images']}")
+        if res["only_img"]:
+            log(f"[{split}] {len(res['only_img'])} images without labels (sample: "
+                f"{list(res['only_img'])[:3]})")
+        if res["only_lbl"]:
+            log(f"[{split}] {len(res['only_lbl'])} labels without images (sample: "
+                f"{list(res['only_lbl'])[:3]})")
+            overall_ok = False
+        if res["bad_rows"]:
+            log(f"[{split}] BAD ROWS: {res['bad_rows']}"); overall_ok = False
+        if res["degenerate"]:
+            log(f"[{split}] degenerate boxes (out-of-range): {res['degenerate']}")
+            overall_ok = False
+        log(f"[{split}] per-class instance counts:")
+        for i, n in enumerate(TARGET_CLASS_NAMES):
+            log(f"    {i} {n:13s}  instances={res['cls_counts'][i]:6d}  "
+                f"images={res['cls_image_counts'][i]:5d}")
+
+    if args.visualize > 0:
+        try:
+            import cv2
+        except ImportError:
+            log("[viz] cv2 not installed; pip install opencv-python")
+            return 0 if overall_ok else 1
+        out = Path(args.vis_out).resolve()
+        out.mkdir(parents=True, exist_ok=True)
+        val_imgs = sorted((root / "images" / "val").iterdir())[: args.visualize]
+        for img_path in val_imgs:
+            lbl_path = root / "labels" / "val" / (img_path.stem + ".txt")
+            img = cv2.imread(str(img_path))
+            if img is None:
+                continue
+            h, w = img.shape[:2]
+            if lbl_path.exists():
+                for r in lbl_path.read_text().splitlines():
+                    p = r.split()
+                    if len(p) != 5:
+                        continue
+                    c = int(p[0])
+                    cx, cy, bw, bh = (float(p[1]) * w, float(p[2]) * h,
+                                      float(p[3]) * w, float(p[4]) * h)
+                    x1, y1, x2, y2 = int(cx - bw / 2), int(cy - bh / 2), int(cx + bw / 2), int(cy + bh / 2)
+                    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                    cv2.putText(img, TARGET_CLASS_NAMES[c], (x1, max(0, y1 - 6)),
+                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
+            cv2.imwrite(str(out / img_path.name), img)
+        log(f"[viz] wrote {len(val_imgs)} overlays to {out}")
+
+    return 0 if overall_ok else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())