import gradio as gr
from transformers import pipeline
from PIL import Image
import os

classifier = pipeline("image-classification", model="gfichetdc/casting-defect-vit")

examples_dir = "examples"
example_files = sorted([f for f in os.listdir(examples_dir) if f.endswith(".jpeg")])

def ground_truth(filename):
    return "defective" if filename.startswith("def") else "ok"

def predict(name):
    image = Image.open(os.path.join(examples_dir, name))
    results = classifier(image)
    bars = {r["label"]: r["score"] for r in results}
    truth = ground_truth(name)
    return image, bars, truth

with gr.Blocks(title="Casting Defect Detection") as demo:
    gr.Markdown("### Casting Defect Detection (ViT)\nClick a test image to classify it.")
    gallery = gr.Dataset(
        components=[gr.Image(visible=False)],
        samples=[[os.path.join(examples_dir, f)] for f in example_files],
        label="Test images (5 defective, 5 ok)",
        samples_per_page=10,
    )
    with gr.Row():
        img_out = gr.Image(label="Selected image", height=224)
        with gr.Column():
            label_out = gr.Label(num_top_classes=2, label="Prediction")
            truth_out = gr.Textbox(label="Ground truth", interactive=False)

    def on_select(evt: gr.SelectData):
        name = example_files[evt.index]
        return predict(name)

    gallery.select(on_select, outputs=[img_out, label_out, truth_out])

demo.launch()