app.py

import cv2
import torch
import gradio as gr
import numpy as np
from ultralytics import YOLO, __version__ as ultralytics_version

# Debug: Check environment
print(f"Torch version: {torch.__version__}")
print(f"Gradio version: {gr.__version__}")
print(f"Ultralytics version: {ultralytics_version}")
print(f"CUDA available: {torch.cuda.is_available()}")

# Load YOLOv8 model
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
model = YOLO('./data/best.pt').to(device)

def process_video(video):
    if video is None:
        return "Error: No video uploaded"
    
    # Open the input video
    cap = cv2.VideoCapture(video)
    if not cap.isOpened():
        return "Error: Could not open video file"
    
    # Get input video properties
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # Use original resolution to avoid resizing issues (optional: keep 320x240 if needed)
    # frame_width, frame_height = 320, 240
    print(f"Input video: {frame_width}x{frame_height}, {fps} FPS, {total_frames} frames")
    
    # Set up video writer
    output_path = "processed_output.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
    
    frame_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        frame_count += 1
        print(f"Processing frame {frame_count}/{total_frames}")
        
        # Optional: Resize if needed (remove if using original resolution)
        # frame = cv2.resize(frame, (frame_width, frame_height))
        
        # Run YOLOv8 inference
        results = model(frame)
        annotated_frame = results[0].plot()
        
        # Write the annotated frame to the output video
        out.write(annotated_frame)
    
    # Release resources
    cap.release()
    out.release()
    print(f"Output video saved as {output_path}")
    return output_path

# Gradio interface
iface = gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Upload Video"),
    outputs=gr.Video(label="Processed Video"),
    title="YOLOv8 Object Detection",
    description="Upload a short video for object detection"
)

if __name__ == "__main__":
    iface.launch()