backend/run_tracker.py

"""
run_tracker.py — Command-line video processing (no server needed)
Usage:
    python run_tracker.py --video traffic.mp4 --scene intersection_A --show
    python run_tracker.py --video traffic.mp4 --classes car truck bus --conf 0.4 --save
"""
import cv2
import argparse
from pathlib import Path
from tracker import TrafficTracker, DEFAULT_CLASSES


def parse_args():
    p = argparse.ArgumentParser(description="TrafficSense CLI Tracker")
    p.add_argument("--video",   required=True,             help="Path to input video")
    p.add_argument("--model",   default="best.pt",      help="YOLO model weights")
    p.add_argument("--scene",   default="scene_01",        help="Scene name for logs")
    p.add_argument("--classes", nargs="+", default=DEFAULT_CLASSES, help="Classes to track")
    p.add_argument("--conf",    type=float, default=0.35,  help="Confidence threshold")
    p.add_argument("--show",    action="store_true",       help="Display video window")
    p.add_argument("--save",    action="store_true",       help="Save annotated output video")
    p.add_argument("--logs",    default="logs",            help="Directory to save logs")
    p.add_argument("--out",     default="output",          help="Directory for output video")
    p.add_argument("--line",    type=float, default=0.55,  help="Counting line position (0-1)")
    return p.parse_args()


def main():
    args = parse_args()

    cap = cv2.VideoCapture(args.video)
    if not cap.isOpened():
        print(f"[ERROR] Cannot open video: {args.video}")
        return

    tracker = TrafficTracker(
        model_path=args.model,
        selected_classes=args.classes,
        conf_threshold=args.conf,
        scene_name=args.scene,
        output_dir=args.logs,
        counting_line_ratio=args.line,
    )
    tracker.setup_video(cap)

    out_writer = None
    if args.save:
        Path(args.out).mkdir(parents=True, exist_ok=True)
        out_path = str(Path(args.out) / f"{args.scene}_output.mp4")
        fourcc   = cv2.VideoWriter_fourcc(*"mp4v")
        out_writer = cv2.VideoWriter(
            out_path, fourcc, tracker.fps,
            (tracker.frame_width, tracker.frame_height)
        )
        print(f"[INFO] Saving to: {out_path}")

    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"[INFO] Processing {total} frames  |  Scene: {args.scene}  |  Classes: {args.classes}")

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        annotated, stats = tracker.process_frame(frame)

        if out_writer:
            out_writer.write(annotated)

        if args.show:
            cv2.imshow(f"TrafficSense — {args.scene}", annotated)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        # Progress
        if tracker.frame_index % 50 == 0:
            pct = tracker.frame_index / max(total, 1) * 100
            print(f"  [{pct:5.1f}%] frame {tracker.frame_index}/{total}  "
                  f"counts: {dict(tracker.count_per_class)}")

    cap.release()
    if out_writer:
        out_writer.release()
    if args.show:
        cv2.destroyAllWindows()

    print("\n[INFO] Processing complete.")
    log_paths = tracker.save_logs()
    summary   = tracker.get_summary()

    print("\n── SUMMARY ──────────────────────────────────")
    print(f"  Scene:          {summary['scene']}")
    print(f"  Duration:       {summary['duration_sec']}s")
    print(f"  Total frames:   {summary['total_frames']}")
    print(f"  Unique objects: {summary['total_unique_objects']}")
    print(f"  Per class:      {summary['count_per_class']}")
    print(f"\n  Logs saved to:  {log_paths}")
    print("─────────────────────────────────────────────")


if __name__ == "__main__":
    main()