app.py

import gradio as gr
import cv2
import tempfile, os
from ultralytics import YOLO
from datetime import datetime
import torch

# Disable gradients for CPU
torch.set_grad_enabled(False)

# Load YOLOv8n model
model = YOLO("yolov8n.pt")

def video_object_detection(video_path, progress=gr.Progress()):
    cap = cv2.VideoCapture(video_path)

    # Resize for speed
    W, H = 416, 234
    fps = int(cap.get(cv2.CAP_PROP_FPS)) or 15
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    temp_dir = tempfile.mkdtemp()
    out_path = os.path.join(
        temp_dir, f"detected_{datetime.now().strftime('%H%M%S')}.mp4"
    )

    out = cv2.VideoWriter(
        out_path,
        cv2.VideoWriter_fourcc(*"mp4v"),
        fps,
        (W, H)
    )

    frame_skip = 20
    frame_id = 0
    last_boxes = []
    last_labels = []

    progress(0, desc="Processing video...")

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

        frame = cv2.resize(frame, (W, H))

        # Run YOLO occasionally
        if frame_id % frame_skip == 0:
            results = model(frame, conf=0.4, verbose=False)
            last_boxes = []
            last_labels = []

            if results[0].boxes is not None:
                boxes = results[0].boxes.xyxy.cpu().numpy()
                classes = results[0].boxes.cls.cpu().numpy()
                names = results[0].names  # class id → label

                for i, box in enumerate(boxes):
                    last_boxes.append(box)
                    last_labels.append(names[int(classes[i])])

        # Draw boxes and labels
        for i, box in enumerate(last_boxes):
            x1, y1, x2, y2 = map(int, box)
            label = last_labels[i]

            # Draw rectangle
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # Draw label background
            ((text_w, text_h), _) = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(frame, (x1, y1 - 20), (x1 + text_w, y1), (0, 255, 0), -1)

            # Put label text
            cv2.putText(frame, label, (x1, y1 - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

        out.write(frame)

        if total > 0:
            progress(frame_id / total)

        frame_id += 1

    cap.release()
    out.release()
    progress(1.0, desc="Completed")

    return out_path


# Gradio Interface
demo = gr.Interface(
    fn=video_object_detection,
    inputs=gr.Video(label="Upload video (≤1 min, 480p recommended)"),
    outputs=gr.Video(label="Detected Video"),
    title="🎥YOLOv8 Video Object Detection with Labels",
    description="CPU-safe • Stable UI • Prints object labels"
)

demo.queue()
demo.launch()