Update app.py

app.py

@@ -1,102 +1,84 @@
-import os
+import io
 import gradio as gr
 import supervision as sv
-from inference import get_model
 from PIL import Image
-
-ROBOFLOW_API_KEY = os.getenv("ROBOFLOW_API_KEY")
-
-if ROBOFLOW_API_KEY is None:
-    raise RuntimeError(
-        "ROBOFLOW_API_KEY ERROR. "
-    )
-
-DET_MODEL_ID = "rfdetr-base"
-SEG_MODEL_ID = "rfdetr-seg-preview"
+from rfdetr import RFDETRBase, RFDETRSegPreview
+from rfdetr.util.coco_classes import COCO_CLASSES
 
 det_model = None
 seg_model = None
 
-
-def load_model(task: str):
-    global det_model, seg_model
-    try:
-        if task == "Object Detection":
-            if det_model is None:
-                det_model = get_model(DET_MODEL_ID, api_key=ROBOFLOW_API_KEY)
-            return det_model
-        else:
-            if seg_model is None:
-                seg_model = get_model(SEG_MODEL_ID, api_key=ROBOFLOW_API_KEY)
-            return seg_model
-    except Exception as e:
-        raise gr.Error(
-            "Failed to load model from Roboflow. "
-            "Check that your API key is correct and has access to this model.\n\n"
-            f"Details: {type(e).__name__}: {e}"
-        )
-
-
-def run_inference(image: Image.Image, task: str, confidence: float):
+def load_det_model():
+    global det_model
+    if det_model is None:
+        det_model = RFDETRBase()
+        det_model.optimize_for_inference()
+    return det_model
+
+def load_seg_model():
+    global seg_model
+    if seg_model is None:
+        seg_model = RFDETRSegPreview()
+        seg_model.optimize_for_inference()
+    return seg_model
+
+def run_inference(image, task, threshold):
     if image is None:
-        raise gr.Error("Please upload an image first.")
-    model = load_model(task)
-    predictions = model.infer(image, confidence=confidence)[0]
-    detections = sv.Detections.from_inference(predictions)
-    labels = [prediction.class_name for prediction in predictions.predictions]
-    annotated_image = image.copy()
-
+        return None
+    if isinstance(image, str):
+        image = Image.open(io.BytesIO(image)).convert("RGB")
+    else:
+        image = image.convert("RGB")
     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)
+        model = load_det_model()
+        detections = model.predict(image, threshold=threshold)
+        labels = [
+            f"{COCO_CLASSES[int(class_id)]} {confidence:.2f}"
+            for class_id, confidence in zip(detections.class_id, detections.confidence)
+        ]
+        annotated = image.copy()
+        annotated = sv.BoxAnnotator().annotate(annotated, detections)
+        annotated = sv.LabelAnnotator().annotate(annotated, detections, labels)
+        return annotated
     else:
-        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
-
+        model = load_seg_model()
+        detections = model.predict(image, threshold=threshold)
+        labels = [
+            f"{COCO_CLASSES[int(class_id)]} {confidence:.2f}"
+            for class_id, confidence in zip(detections.class_id, detections.confidence)
+        ]
+        annotated = image.copy()
+        try:
+            annotated = sv.MaskAnnotator().annotate(annotated, detections)
+        except Exception:
+            annotated = sv.BoxAnnotator().annotate(annotated, detections)
+        annotated = sv.LabelAnnotator().annotate(annotated, detections, labels)
+        return annotated
 
 with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # CIAT RF-DETR Demo 
-        Upload an image and choose **Object Detection** or **Segmentation**.
-        """
-    )
-
+    gr.Markdown("# RF-DETR: Detection and Segmentation Preview")
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(
-                label="Input image",
-                type="pil"
-            )
-            task_input = gr.Radio(
-                choices=["Object Detection", "Segmentation"],
+            inp_image = gr.Image(type="pil", label="Input image")
+            task = gr.Radio(
+                ["Object Detection", "Segmentation"],
                 value="Object Detection",
                 label="Task"
             )
-            conf_input = gr.Slider(
+            threshold = gr.Slider(
                 minimum=0.1,
-                maximum=1.0,
+                maximum=0.9,
                 value=0.5,
                 step=0.05,
                 label="Confidence threshold"
             )
-            run_button = gr.Button("Run RF-DETR")
-
+            run_btn = gr.Button("Run")
         with gr.Column():
-            image_output = gr.Image(
-                label="Annotated output"
-            )
-
-    run_button.click(
+            out_image = gr.Image(type="pil", label="Output", interactive=False)
+    run_btn.click(
         fn=run_inference,
-        inputs=[image_input, task_input, conf_input],
-        outputs=image_output
+        inputs=[inp_image, task, threshold],
+        outputs=out_image
     )
 
 if __name__ == "__main__":