import gradio as gr
import supervision as sv
from rfdetr import RFDETRBase

models = {
    "GC10-DET": "gc10-det.pth",
    "NEU-DET": "neu-det.pth"
}

def predict(model_name, img, conf, iou):
    model_path = models[model_name]
    model = RFDETRBase(pretrain_weights=model_path)
    model.optimize_for_inference()

    results = model.predict(img, threshold=conf)
    results = results.with_nmm(threshold=iou)
    
    labels = [f"defect {confidence:.2f}" for confidence in results.confidence]
    
    annotated_image = img.copy()
    annotated_image = sv.BoxAnnotator().annotate(annotated_image, results)
    annotated_image = sv.LabelAnnotator().annotate(annotated_image, results, labels)

    return annotated_image

base_conf, base_iou = 0.25, 0.5
title = "Detection with RF-DETR trained on NEU-DET and GC10-DET"
des = """
        **Instructions:**
        1. Choose a model:  
        \- **GC10-DET** - for images from the GC10 dataset (10 defect types, steel surface images)  
        \- **NEU-DET** - for images from the NEU dataset (6 defect types, steel strip images)
        2. Upload Image.
        3. Adjust thresholds
        4. Click the **Submit** button. Wait while predictions are generated.
        """
interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Dropdown(list(models.keys()), label="Select Model"),
        'image', 
        gr.Slider(maximum=1, minimum=0, value=base_conf, label="Confidence Threshold"), 
        gr.Slider(maximum=1, minimum=0, value=base_iou, label="NMM IoU Threshold")
    ],
    outputs=["image"],
    title=title,
    description=des,
    examples=[
        ["GC10-DET", "example1.jpg", base_conf, base_iou],
        ["GC10-DET", "example2.jpg", base_conf, base_iou],
        ["NEU-DET", "example3.jpg", base_conf, base_iou]
    ]
)

interface.launch()