app.py

import cv2
import torch
import gradio as gr
import numpy as np
from ultralytics import YOLO

# Load YOLOv8 model and set device (GPU if available)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLO('./data/best.pt')  # Path to your model
model.to(device)

# Define the function that processes the uploaded video
def process_video(video):
    # video is now the file path string, not a file object
    input_video = cv2.VideoCapture(video)  # Directly pass the path to cv2.VideoCapture

    # Get frame width, height, and fps from input video
    frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = input_video.get(cv2.CAP_PROP_FPS)

    # Resize to reduce computation (optional)
    new_width, new_height = 640, 480  # Resize to 640x480 resolution
    frame_width, frame_height = new_width, new_height

    while True:
        # Read a frame from the video
        ret, frame = input_video.read()
        if not ret:
            break  # End of video

        # Resize the frame to reduce computational load
        frame = cv2.resize(frame, (new_width, new_height))

        # Perform inference on the frame
        results = model(frame)  # Automatically uses GPU if available

        # Check if any object was detected
        if len(results[0].boxes) > 0:  # If there are detected objects
            # Annotate the frame with bounding boxes
            annotated_frame = results[0].plot()  # Plot the frame with bounding boxes

            # Convert the annotated frame to RGB format for displaying
            annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)

            # Display the frame with detections
            cv2.imshow("Detected Frame", annotated_frame_rgb)

            # Wait for a key press (optional: press 'q' to quit early)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    # Release resources
    input_video.release()
    cv2.destroyAllWindows()

    return "Video processing complete!"

# Create a Gradio interface for video upload
iface = gr.Interface(fn=process_video, 
                     inputs=gr.Video(label="Upload Video"),  # Updated line
                     outputs=gr.Textbox(label="Processing Status"),  # Output text showing processing status
                     title="YOLOv8 Object Detection - Real-Time Display",
                     description="Upload a video for object detection using YOLOv8. The frames with detections will be shown in real-time.")

# Launch the interface
iface.launch()