Upgrade rfdetr to v1.6.5.post2 and migrate to proper segmentation checkpoints

app.py

@@ -6,7 +6,11 @@ import gradio as gr
 import numpy as np
 import supervision as sv
 from PIL import Image
-from rfdetr import RFDETRNano, RFDETRSmall, RFDETRMedium, RFDETRBase, RFDETRLarge, RFDETRSegPreview
+from rfdetr import (
+    RFDETRNano, RFDETRSmall, RFDETRMedium, RFDETRBase, RFDETRLarge,
+    RFDETRSegNano, RFDETRSegSmall, RFDETRSegMedium,
+    RFDETRSegLarge, RFDETRSegXLarge, RFDETRSeg2XLarge,
+)
 from rfdetr.detr import RFDETR
 from rfdetr.util.coco_classes import COCO_CLASSES
 
@@ -31,6 +35,9 @@ IMAGE_PROCESSING_EXAMPLES = [
     ['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, 512, "nano (object detection)"],
     ['https://media.roboflow.com/notebooks/examples/dog-3.jpeg', 0.5, 512, "nano (object detection)"],
     ['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 0.5, 512, "nano (object detection)"],
+    ['https://media.roboflow.com/notebooks/examples/dog-2.jpeg', 0.5, 512, "medium (instance segmentation)"],
+    ['https://media.roboflow.com/notebooks/examples/dog-3.jpeg', 0.5, 512, "medium (instance segmentation)"],
+    ['https://media.roboflow.com/supervision/image-examples/basketball-1.png', 0.5, 512, "medium (instance segmentation)"],
 ]
 VIDEO_PROCESSING_EXAMPLES = [
     ["videos/people-walking.mp4", 0.3, 1024, "medium (object detection)"],
@@ -76,10 +83,10 @@ def detect_and_annotate(
     ]
     print(detections)
     annotated_image = image.copy()
+    if checkpoint in SEGMENTATION_CHECKPOINTS:
+        annotated_image = mask_annotator.annotate(annotated_image, detections)
     annotated_image = bbox_annotator.annotate(annotated_image, detections)
     annotated_image = label_annotator.annotate(annotated_image, detections, labels)
-    if checkpoint == "segmentation preview":
-        annotated_image = mask_annotator.annotate(annotated_image, detections)
     return annotated_image
 
 
@@ -92,15 +99,35 @@ def load_model(resolution: int, checkpoint: str) -> RFDETR:
         return RFDETRMedium(resolution=resolution)
     if checkpoint == "base (object detection)":
         return RFDETRBase(resolution=resolution)
-    elif checkpoint == "large (object detection)":
+    if checkpoint == "large (object detection)":
         return RFDETRLarge(resolution=resolution)
-    elif checkpoint == "segmentation preview":
-        return RFDETRSegPreview(resolution=resolution)
-    raise TypeError("Checkpoint must be a base or large.")
+    if checkpoint == "nano (instance segmentation)":
+        return RFDETRSegNano(resolution=resolution)
+    if checkpoint == "small (instance segmentation)":
+        return RFDETRSegSmall(resolution=resolution)
+    if checkpoint == "medium (instance segmentation)":
+        return RFDETRSegMedium(resolution=resolution)
+    if checkpoint == "large (instance segmentation)":
+        return RFDETRSegLarge(resolution=resolution)
+    if checkpoint == "xlarge (instance segmentation)":
+        return RFDETRSegXLarge(resolution=resolution)
+    if checkpoint == "2xlarge (instance segmentation)":
+        return RFDETRSeg2XLarge(resolution=resolution)
+    raise TypeError(f"Unknown checkpoint: {checkpoint}")
+
+
+SEGMENTATION_CHECKPOINTS = {
+    "nano (instance segmentation)",
+    "small (instance segmentation)",
+    "medium (instance segmentation)",
+    "large (instance segmentation)",
+    "xlarge (instance segmentation)",
+    "2xlarge (instance segmentation)",
+}
 
 
 def adjust_resolution(checkpoint: str, resolution: int) -> int:
-    if checkpoint == "segmentation preview":
+    if checkpoint in SEGMENTATION_CHECKPOINTS:
         divisor = 24
     elif checkpoint in {"nano (object detection)", "small (object detection)", "medium (object detection)"}:
         divisor = 32
@@ -198,8 +225,18 @@ with gr.Blocks() as demo:
                 )
                 image_processing_checkpoint_dropdown = gr.Dropdown(
                     label="Checkpoint",
-                    choices=["nano (object detection)", "small (object detection)", "medium (object detection)", "segmentation preview"],
-                    value="segmentation preview"
+                    choices=[
+                        "nano (object detection)",
+                        "small (object detection)",
+                        "medium (object detection)",
+                        "nano (instance segmentation)",
+                        "small (instance segmentation)",
+                        "medium (instance segmentation)",
+                        "large (instance segmentation)",
+                        "xlarge (instance segmentation)",
+                        "2xlarge (instance segmentation)",
+                    ],
+                    value="medium (object detection)"
                 )
             with gr.Column():
                 image_processing_submit_button = gr.Button("Submit", value="primary")
@@ -254,8 +291,18 @@ with gr.Blocks() as demo:
                 )
                 video_processing_checkpoint_dropdown = gr.Dropdown(
                     label="Checkpoint",
-                    choices=["nano (object detection)", "small (object detection)", "medium (object detection)", "segmentation preview"],
-                    value="segmentation preview"
+                    choices=[
+                        "nano (object detection)",
+                        "small (object detection)",
+                        "medium (object detection)",
+                        "nano (instance segmentation)",
+                        "small (instance segmentation)",
+                        "medium (instance segmentation)",
+                        "large (instance segmentation)",
+                        "xlarge (instance segmentation)",
+                        "2xlarge (instance segmentation)",
+                    ],
+                    value="medium (object detection)"
                 )
             with gr.Column():
                 video_processing_submit_button = gr.Button("Submit", value="primary")