Update app to support confidence adjustment and mask options

Updated the app.py and requirements.txt to introduce a slide bar for adjusting confidence threshold and adding a solid mask annotator in addition to the semi-transparent mask annotator. Also, modifications have been made to include the confidence parameter in the inference function call. The gradio version has been pinned to 3.50.2 in requirements.txt for consistency. These changes were made to offer users a better interaction experience and greater control over the image processing results.

app.py

@@ -14,11 +14,12 @@ This is the demo for a Open Vocabulary Image Segmentation using
 [MetaCLIP](https://github.com/facebookresearch/MetaCLIP) combo.
 """
 EXAMPLES = [
-    ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "dog"],
-    ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "building"],
-    ["https://media.roboflow.com/notebooks/examples/dog-3.jpeg", "jacket"],
+    ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "dog", 0.5],
+    ["https://media.roboflow.com/notebooks/examples/dog.jpeg", "building", 0.5],
+    ["https://media.roboflow.com/notebooks/examples/dog-3.jpeg", "jacket", 0.5],
+    ["https://media.roboflow.com/notebooks/examples/dog-3.jpeg", "coffee", 0.6],
 ]
-
+MIN_AREA_THRESHOLD = 0.01
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 SAM_GENERATOR = pipeline(
     task="mask-generation",
@@ -26,9 +27,13 @@ SAM_GENERATOR = pipeline(
     device=DEVICE)
 CLIP_MODEL = CLIPModel.from_pretrained("facebook/metaclip-b32-400m").to(DEVICE)
 CLIP_PROCESSOR = CLIPProcessor.from_pretrained("facebook/metaclip-b32-400m")
-MASK_ANNOTATOR = sv.MaskAnnotator(
+SEMITRANSPARENT_MASK_ANNOTATOR = sv.MaskAnnotator(
     color=sv.Color.red(),
     color_lookup=sv.ColorLookup.INDEX)
+SOLID_MASK_ANNOTATOR = sv.MaskAnnotator(
+    color=sv.Color.red(),
+    color_lookup=sv.ColorLookup.INDEX,
+    opacity=1)
 
 
 def run_sam(image_rgb_pil: Image.Image) -> sv.Detections:
@@ -54,9 +59,13 @@ def reverse_mask_image(image: np.ndarray, mask: np.ndarray, gray_value=128):
     return np.where(mask[..., None], image, gray_color)
 
 
-def annotate(image_rgb_pil: Image.Image, detections: sv.Detections) -> Image.Image:
+def annotate(
+    image_rgb_pil: Image.Image,
+    detections: sv.Detections,
+    annotator: sv.MaskAnnotator
+) -> Image.Image:
     img_bgr_numpy = np.array(image_rgb_pil)[:, :, ::-1]
-    annotated_bgr_image = MASK_ANNOTATOR.annotate(
+    annotated_bgr_image = annotator.annotate(
         scene=img_bgr_numpy, detections=detections)
     return Image.fromarray(annotated_bgr_image[:, :, ::-1])
 
@@ -64,7 +73,8 @@ def annotate(image_rgb_pil: Image.Image, detections: sv.Detections) -> Image.Ima
 def filter_detections(
     image_rgb_pil: Image.Image,
     detections: sv.Detections,
-    prompt: str
+    prompt: str,
+    confidence: float
 ) -> sv.Detections:
     img_rgb_numpy = np.array(image_rgb_pil)
     text = [f"a picture of {prompt}", "a picture of background"]
@@ -76,27 +86,38 @@ def filter_detections(
         masked_crop = reverse_mask_image(image=crop, mask=mask_crop)
         masked_crop_pil = Image.fromarray(masked_crop)
         probs = run_clip(image_rgb_pil=masked_crop_pil, text=text)
-        lass_index = np.argmax(probs)
-        filtering_mask.append(lass_index == 0)
+        filtering_mask.append(probs[0][0] > confidence)
 
     filtering_mask = np.array(filtering_mask)
     return detections[filtering_mask]
 
 
-def inference(image_rgb_pil: Image.Image, prompt: str) -> List[Image.Image]:
+def inference(
+    image_rgb_pil: Image.Image,
+    prompt: str,
+    confidence: float
+) -> List[Image.Image]:
     width, height = image_rgb_pil.size
     area = width * height
 
     detections = run_sam(image_rgb_pil)
-    detections = detections[detections.area / area > 0.01]
+    detections = detections[detections.area / area > MIN_AREA_THRESHOLD]
     detections = filter_detections(
         image_rgb_pil=image_rgb_pil,
         detections=detections,
-        prompt=prompt)
+        prompt=prompt,
+        confidence=confidence)
 
+    blank_image = Image.new("RGB", (width, height), "black")
     return [
-        annotate(image_rgb_pil=image_rgb_pil, detections=detections),
-        annotate(image_rgb_pil=Image.new("RGB", (width, height), "black"), detections=detections)
+        annotate(
+            image_rgb_pil=image_rgb_pil,
+            detections=detections,
+            annotator=SEMITRANSPARENT_MASK_ANNOTATOR),
+        annotate(
+            image_rgb_pil=blank_image,
+            detections=detections,
+            annotator=SOLID_MASK_ANNOTATOR)
     ]
 
 
@@ -104,15 +125,19 @@ with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(image_mode='RGB', type='pil', height=500)
-            prompt_text = gr.Textbox(label="Prompt", value="dog")
+            input_image = gr.Image(
+                image_mode='RGB', type='pil', height=500)
+            prompt_text = gr.Textbox(
+                label="Prompt", value="dog")
+            confidence_slider = gr.Slider(
+                label="Confidence", minimum=0.5, maximum=1.0, step=0.05, value=0.6)
             submit_button = gr.Button("Submit")
         gallery = gr.Gallery(label="Result", object_fit="scale-down", preview=True)
     with gr.Row():
         gr.Examples(
             examples=EXAMPLES,
             fn=inference,
-            inputs=[input_image, prompt_text],
+            inputs=[input_image, prompt_text, confidence_slider],
             outputs=[gallery],
             cache_examples=True,
             run_on_click=True
@@ -120,7 +145,7 @@ with gr.Blocks() as demo:
 
     submit_button.click(
         inference,
-        inputs=[input_image, prompt_text],
+        inputs=[input_image, prompt_text, confidence_slider],
         outputs=gallery)
 
 demo.launch(debug=False, show_error=True)

requirements.txt

@@ -4,6 +4,6 @@ torchvision
 
 numpy
 pillow
-gradio
+gradio==gradio==3.50.2
 transformers
 supervision