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

def load_model(model_path="DDR.pt"):
    return YOLO(model_path)

model = load_model()

def predict(image):
    results = model(image, conf=0.1)
    pred_img = results[0].plot()  # Visualize detections
    return pred_img

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="numpy"),
    title="DDR-Detection",
    description="Upload an image, and the model will detect objects using YOLO11.",
)

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