import torch
import cv2
import gradio as gr
from pathlib import Path

# Load the YOLOv5 model
MODEL_PATH = "best.pt"  # Path to your custom model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.hub.load('ultralytics/yolov5', 'custom', path=MODEL_PATH, force_reload=True).to(device)

# Define the prediction function
def predict(image):
    # Convert the input image (numpy) to a compatible format
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    results = model(img_bgr)  # Run inference

    # Render results
    annotated_image = results.render()[0]
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    return annotated_image_rgb

# Gradio Interface
title = "Custom YOLOv5 Model Deployment"
description = "Upload an image to get predictions using a custom YOLOv5 model trained with your dataset."

interface = gr.Interface(
    fn=predict,                  # Prediction function
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Detected Objects"),
    title=title,
    description=description,
    live=True  # Enable live processing
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch()