Create app.py

app.py

@@ -0,0 +1,103 @@
+import gradio as gr
+import supervision as sv
+from inference import get_model
+from PIL import Image
+
+# Model IDs from RF-DETR docs
+DET_MODEL_ID = "rfdetr-base"
+SEG_MODEL_ID = "rfdetr-seg-preview"
+
+det_model = None
+seg_model = None
+
+
+def load_model(task: str):
+    """
+    Lazily load the selected model once, then reuse it.
+    """
+    global det_model, seg_model
+
+    if task == "Object Detection":
+        if det_model is None:
+            det_model = get_model(DET_MODEL_ID)
+        return det_model
+    else:  # "Segmentation"
+        if seg_model is None:
+            seg_model = get_model(SEG_MODEL_ID)
+        return seg_model
+
+
+def run_inference(image: Image.Image, task: str, confidence: float):
+    if image is None:
+        return None
+
+    model = load_model(task)
+    # Run inference
+    predictions = model.infer(image, confidence=confidence)[0]
+
+    # Convert to supervision.Detections
+    detections = sv.Detections.from_inference(predictions)
+
+    # Labels (class names)
+    labels = [prediction.class_name for prediction in predictions.predictions]
+
+    annotated_image = image.copy()
+
+    if task == "Object Detection":
+        box_annotator = sv.BoxAnnotator(color=sv.ColorPalette.ROBOFLOW)
+        label_annotator = sv.LabelAnnotator(color=sv.ColorPalette.ROBOFLOW)
+
+        annotated_image = box_annotator.annotate(annotated_image, detections)
+        annotated_image = label_annotator.annotate(annotated_image, detections, labels)
+
+    else:  # Segmentation
+        mask_annotator = sv.MaskAnnotator(color=sv.ColorPalette.ROBOFLOW)
+        label_annotator = sv.LabelAnnotator(color=sv.ColorPalette.ROBOFLOW)
+
+        annotated_image = mask_annotator.annotate(annotated_image, detections)
+        annotated_image = label_annotator.annotate(annotated_image, detections, labels)
+
+    return annotated_image
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # CIAT RF-DETR Demo
+        Upload an image and choose **Object Detection** or **Segmentation**.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(
+                label="Input image",
+                type="pil"
+            )
+            task_input = gr.Radio(
+                choices=["Object Detection", "Segmentation"],
+                value="Object Detection",
+                label="Task"
+            )
+            conf_input = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.5,
+                step=0.05,
+                label="Confidence threshold"
+            )
+            run_button = gr.Button("Run RF-DETR")
+
+        with gr.Column():
+            image_output = gr.Image(
+                label="Annotated output"
+            )
+
+    run_button.click(
+        fn=run_inference,
+        inputs=[image_input, task_input, conf_input],
+        outputs=image_output
+    )
+
+if __name__ == "__main__":
+    demo.launch()