detect.py

"""
detect.py — Real-Time Vehicle & Object Detection System
Author: Raghav Matta

Uses Python + OpenCV + YOLOv8 (Ultralytics) to detect and classify
objects from a webcam stream, video file, or image.

Usage:
    python detect.py                   # live webcam (default)
    python detect.py --source video.mp4
    python detect.py --source image.jpg
    python detect.py --source 0        # explicit webcam index
"""

import argparse
import sys
import time
from pathlib import Path

import cv2
import numpy as np

try:
    from ultralytics import YOLO
except ImportError:
    sys.exit(
        "[ERROR] 'ultralytics' package not found.\n"
        "Run:  pip install ultralytics opencv-python"
    )


# ── Config ──────────────────────────────────────────────────────────────────
MODEL_NAME    = "yolov8n.pt"   # nano model – fast inference, auto-downloaded
CONF_THRESH   = 0.45           # minimum confidence to display a detection
IOU_THRESH    = 0.45           # NMS IoU threshold
FRAME_WIDTH   = 1280
FRAME_HEIGHT  = 720

# Colour palette: one BGR colour per COCO class (rotates via modulo)
PALETTE = [
    (0,   200, 255), (0,   255, 128), (255, 128,   0),
    (255,   0, 128), (128,   0, 255), (0,   128, 255),
    (128, 255,   0), (255,   0,   0), (0,   255,   0),
    (0,     0, 255), (200, 200,   0), (0,   200, 200),
]


def get_color(class_id: int) -> tuple[int, int, int]:
    return PALETTE[class_id % len(PALETTE)]


def draw_detection(
    frame: np.ndarray,
    box: list[float],
    class_id: int,
    label: str,
    confidence: float,
) -> None:
    """Draw a bounding box, label, and confidence score on the frame."""
    x1, y1, x2, y2 = int(box[0]), int(box[1]), int(box[2]), int(box[3])
    color = get_color(class_id)

    # Bounding box
    cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness=2)

    # Label background
    text = f"{label}  {confidence:.0%}"
    (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
    cv2.rectangle(frame, (x1, y1 - th - 10), (x1 + tw + 6, y1), color, -1)

    # Label text
    cv2.putText(
        frame, text, (x1 + 3, y1 - 5),
        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA,
    )


def draw_hud(
    frame: np.ndarray,
    fps: float,
    n_detections: int,
    model_name: str,
) -> None:
    """Draw FPS counter and detection count in the top-left corner."""
    h, w = frame.shape[:2]
    overlay = frame.copy()
    cv2.rectangle(overlay, (0, 0), (340, 80), (0, 0, 0), -1)
    cv2.addWeighted(overlay, 0.45, frame, 0.55, 0, frame)

    cv2.putText(frame, f"FPS:    {fps:5.1f}", (12, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 255, 160), 2, cv2.LINE_AA)
    cv2.putText(frame, f"Objects: {n_detections:3d}", (12, 58),
                cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 200, 255), 2, cv2.LINE_AA)
    cv2.putText(frame, model_name, (w - 170, h - 12),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (180, 180, 180), 1, cv2.LINE_AA)


def process_frame(frame: np.ndarray, model: "YOLO") -> tuple[np.ndarray, int]:
    """Run YOLOv8 inference on a single frame and annotate it."""
    results = model(frame, conf=CONF_THRESH, iou=IOU_THRESH, verbose=False)[0]

    n = 0
    if results.boxes is not None:
        for box in results.boxes:
            cls_id = int(box.cls[0].item())
            conf   = float(box.conf[0].item())
            label  = results.names[cls_id]
            xyxy   = box.xyxy[0].tolist()
            draw_detection(frame, xyxy, cls_id, label, conf)
            n += 1

    return frame, n


def run_image(source: str, model: "YOLO") -> None:
    """Run detection on a single image file."""
    img = cv2.imread(source)
    if img is None:
        sys.exit(f"[ERROR] Could not read image: {source}")

    annotated, n = process_frame(img, model)
    print(f"[INFO] Detected {n} object(s) in '{source}'")

    cv2.imshow("Detection — press any key to exit", annotated)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    out_path = Path(source).stem + "_detected.jpg"
    cv2.imwrite(out_path, annotated)
    print(f"[INFO] Saved annotated image → {out_path}")


def run_video(source, model: "YOLO") -> None:
    """Run detection on a webcam stream or video file."""
    # Convert source to int for webcam index
    if isinstance(source, str) and source.isdigit():
        source = int(source)

    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        sys.exit(f"[ERROR] Cannot open video source: {source}")

    cap.set(cv2.CAP_PROP_FRAME_WIDTH,  FRAME_WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)

    print("[INFO] Starting detection. Press  Q  to quit.")
    fps_timer = time.perf_counter()
    fps = 0.0
    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            print("[INFO] Stream ended or no frame received.")
            break

        annotated, n_det = process_frame(frame, model)

        frame_count += 1
        elapsed = time.perf_counter() - fps_timer
        if elapsed >= 0.5:                       # update FPS every 0.5 s
            fps = frame_count / elapsed
            fps_timer   = time.perf_counter()
            frame_count = 0

        draw_hud(annotated, fps, n_det, MODEL_NAME)
        cv2.imshow("Real-Time Detection  |  Q = quit", annotated)

        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

    cap.release()
    cv2.destroyAllWindows()
    print("[INFO] Detection stopped.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Real-Time Object & Vehicle Detection (YOLOv8 + OpenCV)"
    )
    parser.add_argument(
        "--source", default="0",
        help="Webcam index (0), video file path, or image file path. Default: 0",
    )
    args = parser.parse_args()

    print(f"[INFO] Loading model: {MODEL_NAME}  (downloads on first run ~6 MB)")
    model = YOLO(MODEL_NAME)

    # Detect whether source is a still image
    image_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}
    source_path = Path(args.source)
    if source_path.suffix.lower() in image_exts:
        run_image(args.source, model)
    else:
        run_video(args.source, model)


if __name__ == "__main__":
    main()