checkpoint selection added

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import spaces
 import supervision as sv
-from rfdetr import RFDETRBase
+from rfdetr import RFDETRBase, RFDETRLarge
 from rfdetr.util.coco_classes import COCO_CLASSES
 
 MARKDOWN = """
@@ -23,17 +23,26 @@ RF-DETR is a real-time, transformer-based object detection model architecture de
 by [Roboflow](https://roboflow.com/) and released under the Apache 2.0 license.
 """
 
+IMAGE_EXAMPLES = [
+    ['https://media.roboflow.com/supervision/image-examples/people-walking.png', 0.3, "large"],
+    ['https://media.roboflow.com/supervision/image-examples/vehicles.png', 0.3, "large"],
+    ['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, "base"],
+]
+
 COLOR = sv.ColorPalette.from_hex([
     "#ffff00", "#ff9b00", "#ff8080", "#ff66b2", "#ff66ff", "#b266ff",
     "#9999ff", "#3399ff", "#66ffff", "#33ff99", "#66ff66", "#99ff00"
 ])
 
-MODEL = RFDETRBase()
+MODEL_BASE = RFDETRBase(resolution=728)
+MODEL_LARGE = RFDETRLarge(resolution=728)
 
 
 @spaces.GPU()
-def inference(image, confidence):
-    detections = MODEL.predict(image, threshold=confidence)
+def inference(image, confidence: float, checkpoint: str):
+    detections = MODEL_BASE.predict(image, threshold=confidence) \
+        if checkpoint == "base" \
+        else MODEL_LARGE.predict(image, threshold=confidence)
 
     text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
     thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size)
@@ -74,7 +83,13 @@ with gr.Blocks() as demo:
                 step=0.05,
                 value=0.5,
             )
-            submit_button = gr.Button("Submit")
+            with gr.Row():
+                checkpoint_dropdown = gr.Dropdown(
+                    label="Checkpoint",
+                    choices=["base", "large"],
+                    value="base"
+                )
+                submit_button = gr.Button("Submit")
         with gr.Column():
             output_image = gr.Image(
                 label="Input Image",
@@ -82,11 +97,17 @@ with gr.Blocks() as demo:
                 type='pil',
                 height=600
             )
+    gr.Examples(
+        fn=inference,
+        examples=IMAGE_EXAMPLES,
+        inputs=[input_image, confidence_slider, checkpoint_dropdown],
+        outputs=output_image
+    )
 
-        submit_button.click(
-            inference,
-            inputs=[input_image, confidence_slider],
-            outputs=output_image
-        )
+    submit_button.click(
+        inference,
+        inputs=[input_image, confidence_slider, checkpoint_dropdown],
+        outputs=output_image
+    )
 
 demo.launch(debug=False, show_error=True)