from ultralytics import YOLO
from PIL import Image
import gradio as gr
import os

# Define the model path
model_path = "best.onnx"


# Function to load the YOLO model
def load_model(model_path):
    print(f"Loading model from: {model_path}")
    detection_model = YOLO(model_path, task='detect')
    return detection_model

# Load the YOLO model
detection_model = load_model(model_path)

# Function to perform prediction on an image
def predict_image(Input_Image):
    try:
        result = detection_model.predict(Input_Image, conf=0.5, iou=0.6)
        img_bgr = result[0].plot()
        out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB for PIL image
        return out_pilimg
    except Exception as e:
        return f"Error processing image: {str(e)}"

# Function to perform prediction on a video
def predict_video(Input_Video):
    try:
        result_video = detection_model.predict(Input_Video, conf=0.5, iou=0.6, save=True)
        
        # Find the processed video output path
        processed_video_path = str(result_video[0].save_dir / result_video[0].path.name)
        return processed_video_path
    except Exception as e:
        return f"Error processing video: {str(e)}"

# Define Gradio interface for image and video separately
image_interface = gr.Interface(
    fn=predict_image, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Image(type="pil", label="Processed_Image"),
    title="Object Detection - Image",
    description="Upload an image and the model will detect objects."
)

video_interface = gr.Interface(
    fn=predict_video, 
    inputs=gr.Video(), 
    outputs=gr.File(label="Processed_Video"),
    title="Object Detection - Video",
    description="Upload a video and the model will detect objects."
)

# Combine both into a tabbed interface
app = gr.TabbedInterface(
    [image_interface, video_interface], 
    ["Image Detection", "Video Detection"]
)

# Launch the app
app.launch(share=True)