Update app.py

app.py

@@ -1,8 +1,16 @@
+import os
 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"
 
@@ -11,48 +19,41 @@ 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
+    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):
     if image is None:
-        return None
-
+        raise gr.Error("Please upload an image first.")
     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
+    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)
 
@@ -62,7 +63,7 @@ def run_inference(image: Image.Image, task: str, confidence: float):
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-        # CIAT RF-DETR Demo
+        # CIAT RF-DETR Demo 
         Upload an image and choose **Object Detection** or **Segmentation**.
         """
     )