add rf-detr seg

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 RFDETRNano, RFDETRSmall, RFDETRMedium, RFDETRBase, RFDETRLarge
+from rfdetr import RFDETRNano, RFDETRSmall, RFDETRMedium, RFDETRBase, RFDETRLarge, RFDETRSegPreview
 from rfdetr.detr import RFDETR
 from rfdetr.util.coco_classes import COCO_CLASSES
 
@@ -25,16 +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, 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"],
+    ['https://media.roboflow.com/supervision/image-examples/people-walking.png', 0.3, 1024, "medium (object detection)"],
+    ['https://media.roboflow.com/supervision/image-examples/vehicles.png', 0.3, 1024, "medium (object detection)"],
+    ['https://media.roboflow.com/supervision/image-examples/motorbike.png', 0.3, 1024, "medium (object detection)"],
+    ['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)"],
 ]
 VIDEO_PROCESSING_EXAMPLES = [
-    ["videos/people-walking.mp4", 0.3, 1024, "medium"],
-    ["videos/vehicles.mp4", 0.3, 1024, "medium"],
+    ["videos/people-walking.mp4", 0.3, 1024, "medium (object detection)"],
+    ["videos/vehicles.mp4", 0.3, 1024, "medium (object detection)"],
 ]
 
 COLOR = sv.ColorPalette.from_hex([
@@ -53,6 +53,7 @@ def detect_and_annotate(
         model: RFDETR,
         image: ImageType,
         confidence: float,
+        checkpoint: str = "medium (object detection)"
 ) -> ImageType:
     detections = model.predict(image, threshold=confidence)
 
@@ -76,27 +77,34 @@ def detect_and_annotate(
     annotated_image = image.copy()
     annotated_image = bbox_annotator.annotate(annotated_image, detections)
     annotated_image = label_annotator.annotate(annotated_image, detections, labels)
+    if checkpoint == "segmentation preview":
+        mask_annotator = sv.MaskAnnotator()
+        annotated_image = mask_annotator.annotate(annotated_image, detections)
     return annotated_image
 
 
 def load_model(resolution: int, checkpoint: str) -> RFDETR:
-    if checkpoint == "nano":
+    if checkpoint == "nano (object detection)":
         return RFDETRNano(resolution=resolution)
-    if checkpoint == "small":
+    if checkpoint == "small (object detection)":
         return RFDETRSmall(resolution=resolution)
-    if checkpoint == "medium":
+    if checkpoint == "medium (object detection)":
         return RFDETRMedium(resolution=resolution)
-    if checkpoint == "base":
+    if checkpoint == "base (object detection)":
         return RFDETRBase(resolution=resolution)
-    elif checkpoint == "large":
+    elif 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.")
 
 
 def adjust_resolution(checkpoint: str, resolution: int) -> int:
-    if checkpoint in {"nano", "small", "medium"}:
+    if checkpoint == "segmentation preview":
+        divisor = 24
+    elif checkpoint in {"nano (object detection)", "small (object detection)", "medium (object detection)"}:
         divisor = 32
-    elif checkpoint in {"base", "large"}:
+    elif checkpoint in {"base (object detection)", "large (object detection)"}:
         divisor = 56
     else:
         raise ValueError(f"Unknown checkpoint: {checkpoint}")
@@ -121,7 +129,7 @@ def image_processing_inference(
 ):
     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)
+    return detect_and_annotate(model=model, image=input_image, confidence=confidence, checkpoint=checkpoint)
 
 
 def video_processing_inference(
@@ -149,6 +157,7 @@ def video_processing_inference(
                 model=model,
                 image=frame,
                 confidence=confidence,
+                checkpoint=checkpoint
             )
             annotated_frame = sv.scale_image(annotated_frame, VIDEO_SCALE_FACTOR)
             sink.write_frame(annotated_frame)
@@ -189,8 +198,8 @@ with gr.Blocks() as demo:
                 )
                 image_processing_checkpoint_dropdown = gr.Dropdown(
                     label="Checkpoint",
-                    choices=["nano", "small", "medium"],
-                    value="medium"
+                    choices=["nano (object detection)", "small (object detection)", "medium (object detection)", "segmentation preview"],
+                    value="segmentation preview"
                 )
             with gr.Column():
                 image_processing_submit_button = gr.Button("Submit", value="primary")
@@ -245,8 +254,8 @@ with gr.Blocks() as demo:
                 )
                 video_processing_checkpoint_dropdown = gr.Dropdown(
                     label="Checkpoint",
-                    choices=["nano", "small", "medium"],
-                    value="medium"
+                    choices=["nano (object detection)", "small (object detection)", "medium (object detection)", "segmentation preview"],
+                    value="segmentation preview"
                 )
             with gr.Column():
                 video_processing_submit_button = gr.Button("Submit", value="primary")