""" 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") 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, ) 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()