| 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() |
|
|
|
|
