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

# Load the optimal trained YOLOv8s weights (CutPaste-Poisson n=100)
model = YOLO("best.pt")

def detect_defects(image):
    # Run YOLOv8 inference on the uploaded image
    results = model(image)
    # Generate the annotated image with bounding boxes and confidence scores
    annotated_img = results[0].plot()
    # Convert BGR (OpenCV format) to RGB (PIL/Gradio format)
    return Image.fromarray(annotated_img[..., ::-1])

# Define the Gradio web interface
iface = gr.Interface(
    fn=detect_defects,
    inputs=gr.Image(type="pil", label="Upload Steel Surface Image"),
    outputs=gr.Image(type="pil", label="Detected Defects & Confidence"),
    title="Real-Time Steel Surface Defect Detection",
    description="Upload a grayscale steel surface image (NEU-DET) to detect 6 types of industrial defects: **crazing, inclusion, patches, pitted_surface, rolled-in_scale, and scratches**.\n\nThis model is powered by **YOLOv8s** trained with an optimal volume of Cut-Paste Poisson blending synthetic data augmentation.",
    examples=[["sample_scratch.jpg"], ["sample_patch.jpg"]]
)

iface.launch()