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

# Load YOLOv8 model
model = YOLO('./data/model.pt')  # Path to your model

# 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)

    # Define output video writer
    output_video_path = "/mnt/data/output_video.mp4"  # Path to save the output video
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    output_video = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

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

        # Perform inference on the frame
        results = model(frame)

        # The results object contains annotations for the frame
        annotated_frame = results[0].plot()  # Plot the frame with bounding boxes

        # Write the annotated frame to the output video
        output_video.write(annotated_frame)

    # Release resources
    input_video.release()
    output_video.release()

    return output_video_path

# Create a Gradio interface for video upload
iface = gr.Interface(fn=process_video, 
                     inputs=gr.Video(label="Upload Video"),  # Updated line
                     outputs="file", 
                     title="YOLOv8 Object Detection on Video",
                     description="Upload a video for object detection using YOLOv8")

# Launch the interface
iface.launch()