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

# Load the trained YOLOv8 model
model = YOLO('best.pt')

def predict(image):
    results = model(image)
    # You might want to process results to return bounding boxes, class labels, etc.
    annotated_image = results[0].plot()  # plot the results on the image
    return annotated_image

def predict_video(video):
    # Read the video file
    cap = cv2.VideoCapture(video)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    
    # Create a temporary file to save the output video
    out_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    out_path = out_file.name
    
    # Define the codec and create VideoWriter object
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        results = model(frame)
        annotated_frame = results[0].plot()  # plot the results on the frame
        out.write(annotated_frame)
    
    cap.release()
    out.release()
    
    return out_path

# Create Gradio interface
interface = gr.Interface(
    fn=lambda img, vid: (predict(img), predict_video(vid)),
    inputs=[
        gr.inputs.Image(type="numpy", label="Input Image"),
        gr.inputs.Video(label="Input Video")
    ],
    outputs=[
        gr.outputs.Image(type="numpy", label="Output Image"),
        gr.outputs.Video(label="Output Video")
    ],
    title="YOLOv8 Object Detection",
    description="Upload an image or a video and get the object detection results using a YOLOv8 model."
)

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