""" detector.py ----------- Core AI module for the Traffic Scene Interpretation System. Responsibilities: - Load a YOLO model (Ultralytics) once and reuse it. - Run detection on a single image (numpy array / PIL image). - Run detection on a video file, frame by frame, and write an annotated output video. - Aggregate per-frame detections into simple traffic-scene statistics (vehicle counts, congestion level) — this is the "scene interpretation" layer on top of raw object detection. Kept deliberately simple and well-commented so it's easy to explain during a project defense. """ from __future__ import annotations import time from collections import Counter from dataclasses import dataclass, field import cv2 import numpy as np from ultralytics import YOLO # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- # Classes from the COCO dataset (what pretrained YOLO already knows) that are # relevant to a traffic scene. No custom training needed for the MVP. VEHICLE_CLASSES = { "car": "Cars", "bus": "Buses", "truck": "Trucks", "motorcycle": "Motorcycles", "bicycle": "Bicycles", "person": "Pedestrians", } # Thresholds used to translate a raw vehicle count into a human-readable # "scene interpretation" label. Tune these once you see real results. CONGESTION_THRESHOLDS = { "free": 5, # 0-5 vehicles -> Free flowing "moderate": 15, # 6-15 vehicles -> Moderate traffic # >15 vehicles -> Congested } IMAGE_MODEL = "yolov8m.pt" # Used for single-image detection. Accuracy matters more than # speed here since it only runs once per upload. VIDEO_MODEL = "yolov8m.pt" # Bumped up from yolov8s for better accuracy, per testing feedback. # This is noticeably slower (roughly 2-3x the compute of yolov8s). # Frame skipping + reduced imgsz below help offset that cost. # If processing time becomes uncomfortable, drop back to "yolov8s.pt". VIDEO_IMGSZ = 480 # Shrinking the frame before detection speeds video up further. # Lower = faster but less accurate on small/distant objects. # Try 384 if still too slow; try 640 (native) if you have room to spare. VIDEO_FRAME_SKIP = 2 # Run detection on 1 out of every N frames; reuse the previous # frame's boxes for the skipped ones. 2 = run detection on half # the frames (~2x faster). Set to 1 to disable (detect every frame). CONFIDENCE_THRESHOLD = 0.25 # Lowered from 0.35 to catch smaller/more distant vehicles. # If you start seeing false detections (boxes on things that aren't # vehicles), raise this back up toward 0.35-0.4. @dataclass class SceneStats: """Aggregated statistics for one image or one video.""" counts: Counter = field(default_factory=Counter) total_frames: int = 1 fps: float = 0.0 def per_frame_average(self) -> Counter: if self.total_frames <= 0: return self.counts return Counter({k: round(v / self.total_frames, 1) for k, v in self.counts.items()}) def congestion_label(self) -> str: vehicle_count = sum( v for k, v in self.per_frame_average().items() if k != "Pedestrians" ) if vehicle_count <= CONGESTION_THRESHOLDS["free"]: return "Free flowing" elif vehicle_count <= CONGESTION_THRESHOLDS["moderate"]: return "Moderate traffic" else: return "Congested" class TrafficDetector: """ Wraps two YOLO models: - self.image_model: larger/more accurate, used for single-image detection. - self.video_model: smaller/faster, used for video (and would be used for webcam too, if that's added later) since it runs once per frame. Both are loaded once at startup and reused. """ def __init__( self, image_model_path: str = IMAGE_MODEL, video_model_path: str = VIDEO_MODEL, conf: float = CONFIDENCE_THRESHOLD, ): self.image_model = YOLO(image_model_path) # Avoid loading the same weights twice if someone sets both to the same file. self.video_model = ( self.image_model if video_model_path == image_model_path else YOLO(video_model_path) ) self.conf = conf # -- Image ------------------------------------------------------------- def detect_image(self, image: np.ndarray) -> tuple[np.ndarray, SceneStats]: """ Run detection on a single BGR image (as read by cv2). Returns (annotated_image, stats). """ results = self.image_model.predict(image, conf=self.conf, verbose=False) result = results[0] counts = self._count_from_result(result) annotated = result.plot() # draws boxes + labels + confidence stats = SceneStats(counts=counts, total_frames=1) return annotated, stats # -- Video --------------------------------------------------------------- def detect_video(self, input_path: str, output_path: str, progress_callback=None) -> SceneStats: """ Process a video file frame-by-frame: read frame -> YOLO detect -> draw boxes -> write frame to output. Uses the faster video_model + a reduced inference size (VIDEO_IMGSZ) to keep processing time reasonable without a GPU. Frame skipping (VIDEO_FRAME_SKIP): to save time, detection only runs on every Nth frame. For the frames in between, we reuse the last detected boxes and re-draw them onto the new frame. Since consecutive frames are 1/25th-1/30th of a second apart, objects barely move between them, so this looks smooth while cutting detection calls (the expensive part) roughly in half. progress_callback(current_frame, total_frames) is called after each frame if provided, so a Streamlit progress bar can be updated. """ cap = cv2.VideoCapture(input_path) if not cap.isOpened(): raise RuntimeError(f"Could not open video: {input_path}") fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or None fourcc = cv2.VideoWriter_fourcc(*"mp4v") writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) total_counts: Counter = Counter() frame_idx = 0 detected_frame_count = 0 # frames actually run through YOLO (for stats averaging) start_time = time.time() last_result = None # cached YOLO result, reused on skipped frames while True: ok, frame = cap.read() if not ok: break run_detection = (frame_idx % VIDEO_FRAME_SKIP == 0) or (last_result is None) if run_detection: results = self.video_model.predict(frame, conf=self.conf, imgsz=VIDEO_IMGSZ, verbose=False) last_result = results[0] total_counts.update(self._count_from_result(last_result)) detected_frame_count += 1 annotated = last_result.plot() else: # Re-draw the previous frame's boxes onto the current frame. # plot(img=...) lets us reuse a YOLO result's boxes on a new image. annotated = last_result.plot(img=frame) writer.write(annotated) frame_idx += 1 if progress_callback: progress_callback(frame_idx, total_frames) cap.release() writer.release() elapsed = max(time.time() - start_time, 1e-6) processing_fps = frame_idx / elapsed return SceneStats( counts=total_counts, total_frames=max(detected_frame_count, 1), fps=round(processing_fps, 1), ) # -- Webcam (single-frame step, called repeatedly by the UI layer) ----- def detect_frame(self, frame: np.ndarray) -> tuple[np.ndarray, Counter]: """ Used for live webcam mode: process exactly one frame and return it annotated, plus its own counts. Uses the fast video_model, same reasoning as detect_video above. The caller (app.py) is responsible for the capture loop, since Streamlit needs to own that loop to keep the UI responsive. """ results = self.video_model.predict(frame, conf=self.conf, imgsz=VIDEO_IMGSZ, verbose=False) result = results[0] counts = self._count_from_result(result) return result.plot(), counts # -- Helpers ------------------------------------------------------------- def _count_from_result(self, result) -> Counter: """Turn one YOLO result into a Counter of {readable_label: count}.""" counts: Counter = Counter() names = result.names if result.boxes is None: return counts for cls_id in result.boxes.cls.tolist(): raw_name = names[int(cls_id)] label = VEHICLE_CLASSES.get(raw_name) if label: counts[label] += 1 return counts