batch_infer.py

"""
batch_infer.py — Batch inference using the new ParallelDetectionPipeline.
==========================================================================
Processes all images in ./test/ and saves annotated outputs to ./output/.
"""

import sys
import json
import time
import logging
from pathlib import Path

import cv2

from config import SUPPORTED_EXTENSIONS, OUTPUT_DIR
from pipeline import ParallelDetectionPipeline
from annotations import annotate_from_pipeline_result
from storage import init_db, store_result

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
)
logger = logging.getLogger("gridlock.batch")


def main():
    test_dir = Path("./test").resolve()
    output_dir = OUTPUT_DIR.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    # Collect images
    images = sorted([
        p for p in test_dir.iterdir()
        if p.is_file() and p.suffix.lower() in SUPPORTED_EXTENSIONS
    ])

    if not images:
        print(f"No images found in {test_dir}")
        sys.exit(1)

    print(f"\n{'='*60}")
    print(f"  GRIDLOCK — Parallel Pipeline Batch Inference")
    print(f"  Images  : {len(images)}")
    print(f"  Output  : {output_dir}")
    print(f"{'='*60}\n")

    # Initialize
    print("Initializing database...")
    init_db()

    print("Loading models (this may take a minute)...")
    t_load = time.time()
    pipeline = ParallelDetectionPipeline()
    print(f"Models loaded in {time.time() - t_load:.1f}s\n")

    results_summary = []
    t_total = time.time()

    for idx, img_path in enumerate(images, 1):
        print(f"[{idx}/{len(images)}] Processing: {img_path.name}")

        # Run the parallel pipeline with annotation data
        result = pipeline.process(str(img_path), return_annotations=True)

        # Read image and draw annotations
        img = cv2.imread(str(img_path))
        if img is not None:
            annotated = annotate_from_pipeline_result(img, result)
            out_name = img_path.stem + "_annotated.jpg"
            out_path = output_dir / out_name
            cv2.imwrite(str(out_path), annotated, [cv2.IMWRITE_JPEG_QUALITY, 95])
            print(f"  → Saved: {out_name}")
        else:
            print(f"  → Skipped (could not read image)")

        # Store in database
        store_result(result, img_path.name, str(out_path) if img is not None else None)

        # Print violations summary
        violations = result.get("violations", [])
        proc_ms = result.get("processing_time_ms", 0)

        if violations:
            for v in violations:
                vtype = " | ".join(v.get("violation_types", []))
                if v["type"] == "motorcycle":
                    print(f"  ⚠  [{v['type']}] {vtype} | Riders: {v['num_riders']} "
                          f"Helmet violations: {v['helmet_violations']} "
                          f"Plate: {v['license_plate']}")
                else:
                    print(f"  ⚠  [{v['type']}] {vtype} | No Seatbelt: {v['no_seatbelt']} "
                          f"Plate: {v['license_plate']}")
        if not violations:
            print(f"  ✓  No violations detected")

        print(f"  ⏱  {proc_ms}ms")

        results_summary.append({
            "image": img_path.name,
            "violations": len(violations),
            "time_ms": proc_ms,
        })
        print()

    # Final summary
    total_time = time.time() - t_total
    total_v = sum(r["violations"] for r in results_summary)

    print(f"{'='*60}")
    print(f"  DONE — {len(images)} image(s) processed in {total_time:.1f}s")
    print(f"  Total violations : {total_v}")
    print(f"  Avg time/image   : {total_time/len(images)*1000:.0f}ms")
    print(f"  Output saved to  : {output_dir}")
    print(f"{'='*60}\n")

    # Save JSON summary
    summary_path = output_dir / "results_summary.json"
    with open(summary_path, "w") as f:
        json.dump(results_summary, f, indent=2)
    print(f"Summary saved to: {summary_path}")


if __name__ == "__main__":
    main()