import gradio as gr
from ultralytics import YOLO
import os

# Load YOLO model
model = YOLO("best.pt")

def detect_defects(image):
    results = model.predict(image, conf=0.25)
    return results[0].plot()

# ---------------------------
# Defect Classes (10)
# ---------------------------
classes = [
    "punching_hole",
    "welding_line",
    "crescent_gap",
    "water_spot",
    "oil_spot",
    "silk_spot",
    "inclusion",
    "rolled_pit",
    "crease",
    "waist_folding"
]

# ---------------------------
# Load category-wise examples
# ---------------------------
example_dir = "examples"
category_examples = {}

for cls in classes:
    class_folder = os.path.join(example_dir, cls)
    if os.path.exists(class_folder):
        imgs = [
            os.path.join(class_folder, f)
            for f in sorted(os.listdir(class_folder))
            if f.lower().endswith((".jpg", ".png", ".jpeg"))
        ]
        category_examples[cls] = imgs
    else:
        category_examples[cls] = []

# ---------------------------
# Build Gradio Interface
# ---------------------------
with gr.Blocks(title="Metal Surface Defect Detection (YOLOv8)") as demo:

    gr.Markdown("""
    # 🔍 Metal Surface Defect Detection (YOLOv8)
    Upload an image or choose an example from the defect categories below.
    """)

    with gr.Row():
        input_img = gr.Image(type="numpy", label="Input Image")
        output_img = gr.Image(type="numpy", label="Detection Result")

    detect_btn = gr.Button("Run Detection")
    detect_btn.click(detect_defects, inputs=input_img, outputs=output_img)

    gr.Markdown("## 📂 Choose Example Images by Category")

    with gr.Tabs():
        for cls in classes:
            with gr.Tab(cls):
                if len(category_examples[cls]) == 0:
                    gr.Markdown("_No example images found for this category._")
                else:
                    gr.Examples(
                        examples=category_examples[cls],
                        inputs=input_img,
                        outputs=output_img,
                        fn=detect_defects,
                        cache_examples=False
                    )

demo.launch()