Update app.py to allow new nano, small, and medium checkpoints

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 import numpy as np
 import supervision as sv
 from PIL import Image
-from rfdetr import RFDETRBase, RFDETRLarge
+from rfdetr import RFDETRNano, RFDETRSmall, RFDETRMedium, RFDETRBase, RFDETRLarge
 from rfdetr.detr import RFDETR
 from rfdetr.util.coco_classes import COCO_CLASSES
 
@@ -25,13 +25,16 @@ by [Roboflow](https://roboflow.com/) and released under the Apache 2.0 license.
 """
 
 IMAGE_PROCESSING_EXAMPLES = [
-    ['https://media.roboflow.com/supervision/image-examples/people-walking.png', 0.3, 728, "large"],
-    ['https://media.roboflow.com/supervision/image-examples/vehicles.png', 0.3, 728, "large"],
-    ['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, 560, "base"],
+    ['https://media.roboflow.com/supervision/image-examples/people-walking.png', 0.3, 1024, "medium"],
+    ['https://media.roboflow.com/supervision/image-examples/vehicles.png', 0.3, 1024, "medium"],
+    ['https://media.roboflow.com/supervision/image-examples/motorbike.png', 0.3, 1024, "medium"],
+    ['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, 512, "nano"],
+    ['https://media.roboflow.com/notebooks/examples/dog-3.jpeg', 0.5, 512, "nano"],
+    ['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 0.5, 512, "nano"],
 ]
 VIDEO_PROCESSING_EXAMPLES = [
-    ["videos/people-walking.mp4", 0.3, 728, "large"],
-    ["videos/vehicles.mp4", 0.3, 728, "large"],
+    ["videos/people-walking.mp4", 0.3, 1024, "medium"],
+    ["videos/vehicles.mp4", 0.3, 1024, "medium"],
 ]
 
 COLOR = sv.ColorPalette.from_hex([
@@ -77,6 +80,12 @@ def detect_and_annotate(
 
 
 def load_model(resolution: int, checkpoint: str) -> RFDETR:
+    if checkpoint == "nano":
+        return RFDETRNano(resolution=resolution)
+    if checkpoint == "small":
+        return RFDETRSmall(resolution=resolution)
+    if checkpoint == "medium":
+        return RFDETRMedium(resolution=resolution)
     if checkpoint == "base":
         return RFDETRBase(resolution=resolution)
     elif checkpoint == "large":
@@ -84,12 +93,33 @@ def load_model(resolution: int, checkpoint: str) -> RFDETR:
     raise TypeError("Checkpoint must be a base or large.")
 
 
+def adjust_resolution(checkpoint: str, resolution: int) -> int:
+    if checkpoint in {"nano", "small", "medium"}:
+        divisor = 32
+    elif checkpoint in {"base", "large"}:
+        divisor = 56
+    else:
+        raise ValueError(f"Unknown checkpoint: {checkpoint}")
+
+    remainder = resolution % divisor
+    if remainder == 0:
+        return resolution
+    lower = resolution - remainder
+    upper = lower + divisor
+
+    if resolution - lower < upper - resolution:
+        return lower
+    else:
+        return upper
+
+
 def image_processing_inference(
         input_image: Image.Image,
         confidence: float,
         resolution: int,
         checkpoint: str
 ):
+    resolution = adjust_resolution(checkpoint=checkpoint, resolution=resolution)
     model = load_model(resolution=resolution, checkpoint=checkpoint)
     return detect_and_annotate(model=model, image=input_image, confidence=confidence)
 
@@ -100,6 +130,7 @@ def video_processing_inference(
         resolution: int,
         checkpoint: str,
 ):
+    resolution = adjust_resolution(checkpoint=checkpoint, resolution=resolution)
     model = load_model(resolution=resolution, checkpoint=checkpoint)
 
     name = generate_unique_name()
@@ -151,14 +182,14 @@ with gr.Blocks() as demo:
                 )
                 image_processing_resolution_slider = gr.Slider(
                     label="Inference resolution",
-                    minimum=560,
-                    maximum=1120,
-                    step=56,
-                    value=728,
+                    minimum=224,
+                    maximum=2240,
+                    step=1,
+                    value=896,
                 )
                 image_processing_checkpoint_dropdown = gr.Dropdown(
                     label="Checkpoint",
-                    choices=["base", "large"],
+                    choices=["nano", "small", "medium", "base", "large"],
                     value="base"
                 )
             with gr.Column():