import gradio as gr
from ultralytics import YOLO
import torch

# Set device to CPU
device = "cpu"

# Load your trained YOLOv8 model
model = YOLO('best.pt')
model.to(device)

def detect_corrosion(input_image):
    """
    Performs corrosion detection on the input image and returns the image with bounding boxes.
    """
    # Run inference on the image
    results = model(input_image)

    # The 'plot' method returns a NumPy array with the bounding boxes and labels drawn
    plotted_image = results[0].plot() # This returns a BGR numpy array
    
    # Convert BGR to RGB for web display
    plotted_image_rgb = plotted_image[..., ::-1]

    return plotted_image_rgb

# --- Gradio Interface ---
# Create the interface without the examples
iface = gr.Interface(
    fn=detect_corrosion,
    inputs=gr.Image(type="numpy", label="Upload Ship Hull Image"),
    outputs=gr.Image(type="numpy", label="Detection Result"),
    title="🚢 Corrosion Detection in Ship Hulls",
    description="An AI-powered tool to detect corrosion patches on ship hulls. Upload an image, and the YOLOv8 model will highlight any detected corrosion areas.",
    article="Model: YOLOv8 | Developed for maritime maintenance.",
    allow_flagging="never"
)

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