import gradio as gr
from ultralytics import YOLO
import spaces


@spaces.GPU
def predict_aerial_targets(input_img):
    model = YOLO("best.pt")
# Load our custom-trained aerial model
    # Run inference on the uploaded image
    results = model(input_img, imgsz=640)
    
    # Plot the bounding boxes directly onto the image
    annotated_img = results[0].plot()
    return annotated_img

# Define a simple drag-and-drop web interface
demo = gr.Interface(
    fn=predict_aerial_targets,
    inputs=gr.Image(type="pil", label="Upload Aerial/Satellite Image"),
    outputs=gr.Image(type="numpy", label="Detected Targets"),
    title="Nadir-Perspective Object Detection",
    description="Drop a satellite or drone image here to detect planes, ships, and infrastructure."
)

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