File size: 1,819 Bytes
cf76678 1f6fcb2 cf76678 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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()
|