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

device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLO("best.pt").to(device)

def predict_images(image):
    results = model.predict(image, save=False, conf=0.5, iou=0.6)
    # draws the bounding boxes predicted by the model onto the images
    annotated_image = Image.fromarray(results[0].plot()[..., ::-1])

    return annotated_image

def predict_videos(video):
    cap = cv2.VideoCapture(video)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        # predict the objects for each frame
        results = model.predict(frame, save=False, conf=0.5, iou=0.6)
        annotated_frame = results[0].plot()
                                           
        out.write(annotated_frame)

    cap.release()
    out.release()

    return output_path

# gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Drone Type Detection w YOLO11s")
    gr.Markdown("Upload an **image** or **video**, to detect different types of drones (Rotary-Winged vs Fixed-Winged).")

    device_info = gr.Markdown(f"Device: {device}")

    with gr.Tab("Image"):
        image_input = gr.Image(type="pil", label="Upload an Image")
        image_output = gr.Image(type="pil", label="Detected Objects")
        image_button = gr.Button("Run Detection")
        image_button.click(predict_images, inputs=image_input, outputs=image_output)

    with gr.Tab("Video"):
        video_input = gr.Video(label="Upload a Video")
        video_output = gr.Video(label="Detected Objects")
        video_button = gr.Button("Run Detection")
        video_button.click(predict_videos, inputs=video_input, outputs=video_output)

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