open vocabulary detection with Florence2 + masks with SAM2

app.py

@@ -1,15 +1,16 @@
 from typing import Tuple, Optional
 
 import gradio as gr
-import numpy as np
 import supervision as sv
 import torch
 from PIL import Image
 
 from utils.florence import load_florence_model, run_florence_inference, \
     FLORENCE_DETAILED_CAPTION_TASK, \
-    FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK
-from utils.sam import load_sam_model
+    FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK, FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
+from utils.modes import INFERENCE_MODES, OPEN_VOCABULARY_DETECTION, \
+    CAPTION_GROUNDING_MASKS
+from utils.sam import load_sam_model, run_sam_inference
 
 MARKDOWN = """
 # Florence2 + SAM2 🔥
@@ -21,94 +22,122 @@ into masks.
 """
 
 EXAMPLES = [
-    "https://media.roboflow.com/notebooks/examples/dog-2.jpeg",
-    "https://media.roboflow.com/notebooks/examples/dog-3.jpeg",
-    "https://media.roboflow.com/notebooks/examples/dog-4.jpeg"
+    [OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 'straw'],
+    [OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 'napkin'],
+    [OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", 'tail'],
+    [CAPTION_GROUNDING_MASKS, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None],
+    [CAPTION_GROUNDING_MASKS, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", None],
 ]
 
 DEVICE = torch.device("cuda")
-
 FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE)
 SAM_MODEL = load_sam_model(device=DEVICE)
 BOX_ANNOTATOR = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX)
 LABEL_ANNOTATOR = sv.LabelAnnotator(
     color_lookup=sv.ColorLookup.INDEX,
     text_position=sv.Position.CENTER_OF_MASS,
-    text_color=sv.Color.BLACK,
+    text_color=sv.Color.from_hex("#FFFFFF"),
     border_radius=5
 )
 MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
 
 
-def process(
-    image_input,
-) -> Tuple[Optional[Image.Image], Optional[str]]:
-    if image_input is None:
-        return None, None
+def annotate_image(image, detections):
+    output_image = image.copy()
+    output_image = MASK_ANNOTATOR.annotate(output_image, detections)
+    output_image = BOX_ANNOTATOR.annotate(output_image, detections)
+    output_image = LABEL_ANNOTATOR.annotate(output_image, detections)
+    return output_image
 
-    _, result = run_florence_inference(
-        model=FLORENCE_MODEL,
-        processor=FLORENCE_PROCESSOR,
-        device=DEVICE,
-        image=image_input,
-        task=FLORENCE_DETAILED_CAPTION_TASK
-    )
-    caption = result[FLORENCE_DETAILED_CAPTION_TASK]
-    _, result = run_florence_inference(
-        model=FLORENCE_MODEL,
-        processor=FLORENCE_PROCESSOR,
-        device=DEVICE,
-        image=image_input,
-        task=FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK,
-        text=caption
-    )
-    detections = sv.Detections.from_lmm(
-        lmm=sv.LMM.FLORENCE_2,
-        result=result,
-        resolution_wh=image_input.size
-    )
-    image = np.array(image_input.convert("RGB"))
-    SAM_MODEL.set_image(image)
-    mask, score, _ = SAM_MODEL.predict(box=detections.xyxy, multimask_output=False)
 
-    # dirty fix; remove this later
-    if len(mask.shape) == 4:
-        mask = np.squeeze(mask)
+def on_mode_dropdown_change(text):
+    return [
+        gr.Textbox(visible=text == OPEN_VOCABULARY_DETECTION),
+        gr.Textbox(visible=text == CAPTION_GROUNDING_MASKS),
+    ]
 
-    detections.mask = mask.astype(bool)
 
-    output_image = image_input.copy()
-    output_image = MASK_ANNOTATOR.annotate(output_image, detections)
-    output_image = BOX_ANNOTATOR.annotate(output_image, detections)
-    output_image = LABEL_ANNOTATOR.annotate(output_image, detections)
-    return output_image, caption
+def process(
+    mode_dropdown, image_input, text_input
+) -> Tuple[Optional[Image.Image], Optional[str]]:
+    if not image_input:
+        return None, None
+
+    if mode_dropdown == OPEN_VOCABULARY_DETECTION:
+        if not text_input:
+            return None, None
+
+        _, result = run_florence_inference(
+            model=FLORENCE_MODEL,
+            processor=FLORENCE_PROCESSOR,
+            device=DEVICE,
+            image=image_input,
+            task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK,
+            text=text_input
+        )
+        detections = sv.Detections.from_lmm(
+            lmm=sv.LMM.FLORENCE_2,
+            result=result,
+            resolution_wh=image_input.size
+        )
+        detections = run_sam_inference(SAM_MODEL, image_input, detections)
+        return annotate_image(image_input, detections), None
+
+    if mode_dropdown == CAPTION_GROUNDING_MASKS:
+        _, result = run_florence_inference(
+            model=FLORENCE_MODEL,
+            processor=FLORENCE_PROCESSOR,
+            device=DEVICE,
+            image=image_input,
+            task=FLORENCE_DETAILED_CAPTION_TASK
+        )
+        caption = result[FLORENCE_DETAILED_CAPTION_TASK]
+        _, result = run_florence_inference(
+            model=FLORENCE_MODEL,
+            processor=FLORENCE_PROCESSOR,
+            device=DEVICE,
+            image=image_input,
+            task=FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK,
+            text=caption
+        )
+        detections = sv.Detections.from_lmm(
+            lmm=sv.LMM.FLORENCE_2,
+            result=result,
+            resolution_wh=image_input.size
+        )
+        detections = run_sam_inference(SAM_MODEL, image_input, detections)
+        return annotate_image(image_input, detections), caption
 
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
+    mode_dropdown_component = gr.Dropdown(
+        choices=INFERENCE_MODES,
+        value=INFERENCE_MODES[0],
+        label="Mode",
+        info="Select a mode to use.",
+        interactive=True
+    )
     with gr.Row():
         with gr.Column():
             image_input_component = gr.Image(
                 type='pil', label='Upload image')
+            text_input_component = gr.Textbox(
+                label='Text prompt')
             submit_button_component = gr.Button(value='Submit', variant='primary')
-
         with gr.Column():
             image_output_component = gr.Image(type='pil', label='Image output')
-            text_output_component = gr.Textbox(label='Caption output')
+            text_output_component = gr.Textbox(label='Caption output', visible=False)
 
-    submit_button_component.click(
-        fn=process,
-        inputs=[image_input_component],
-        outputs=[
-            image_output_component,
-            text_output_component
-        ]
-    )
     with gr.Row():
         gr.Examples(
             fn=process,
             examples=EXAMPLES,
-            inputs=[image_input_component],
+            inputs=[
+                mode_dropdown_component,
+                image_input_component,
+                text_input_component
+            ],
             outputs=[
                 image_output_component,
                 text_output_component
@@ -116,4 +145,25 @@ with gr.Blocks() as demo:
             run_on_click=True
         )
 
+    submit_button_component.click(
+        fn=process,
+        inputs=[
+            mode_dropdown_component,
+            image_input_component,
+            text_input_component
+        ],
+        outputs=[
+            image_output_component,
+            text_output_component
+        ]
+    )
+    mode_dropdown_component.change(
+        on_mode_dropdown_change,
+        inputs=[mode_dropdown_component],
+        outputs=[
+            text_input_component,
+            text_output_component
+        ]
+    )
+
 demo.launch(debug=False, show_error=True, max_threads=1)

utils/florence.py

@@ -8,8 +8,11 @@ from transformers import AutoModelForCausalLM, AutoProcessor
 from transformers.dynamic_module_utils import get_imports
 
 FLORENCE_CHECKPOINT = "microsoft/Florence-2-large"
+FLORENCE_OBJECT_DETECTION_TASK = '<OD>'
 FLORENCE_DETAILED_CAPTION_TASK = '<MORE_DETAILED_CAPTION>'
 FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK = '<CAPTION_TO_PHRASE_GROUNDING>'
+FLORENCE_OPEN_VOCABULARY_DETECTION_TASK = '<OPEN_VOCABULARY_DETECTION>'
+FLORENCE_DENSE_REGION_CAPTION_TASK = '<DENSE_REGION_CAPTION>'
 
 
 def fixed_get_imports(filename: Union[str, os.PathLike]) -> list[str]:

utils/modes.py

@@ -0,0 +1,7 @@
+OPEN_VOCABULARY_DETECTION = "open vocabulary detection + masks"
+CAPTION_GROUNDING_MASKS = "caption + grounding + masks"
+
+INFERENCE_MODES = [
+    OPEN_VOCABULARY_DETECTION,
+    CAPTION_GROUNDING_MASKS
+]

utils/sam.py

@@ -1,4 +1,9 @@
+from typing import Any
+
+import numpy as np
+import supervision as sv
 import torch
+from PIL import Image
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
@@ -13,3 +18,20 @@ def load_sam_model(
 ) -> SAM2ImagePredictor:
     model = build_sam2(config, checkpoint, device=device)
     return SAM2ImagePredictor(sam_model=model)
+
+
+def run_sam_inference(
+    model: Any,
+    image: Image,
+    detections: sv.Detections
+) -> sv.Detections:
+    image = np.array(image.convert("RGB"))
+    model.set_image(image)
+    mask, score, _ = model.predict(box=detections.xyxy, multimask_output=False)
+
+    # dirty fix; remove this later
+    if len(mask.shape) == 4:
+        mask = np.squeeze(mask)
+
+    detections.mask = mask.astype(bool)
+    return detections