bug fix

app.py

@@ -1,5 +1,5 @@
 from typing import Optional
-
+import numpy as np
 import gradio as gr
 import spaces
 import supervision as sv
@@ -8,6 +8,7 @@ from PIL import Image
 from io import BytesIO
 import PIL.Image
 import requests
+import cv2
 
 from utils.florence import load_florence_model, run_florence_inference, \
     FLORENCE_OPEN_VOCABULARY_DETECTION_TASK
@@ -29,17 +30,14 @@ SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE)
 @spaces.GPU(duration=20)
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-def process_image(image_input, image_url, task_prompt, text_input) -> Optional[Image.Image]:
+def process_image(image_input, image_url, task_prompt, text_prompt=None, dilate=0) -> Optional[Image.Image]:
     if not image_input:
         gr.Info("Please upload an image.")
         return None
     if not task_prompt:
         gr.Info("Please enter a task prompt.")
         return None
-    if not text_input:
-        gr.Info("Please enter a text prompt.")
-        return None
-        
+   
     if image_url:
         print("start to fetch image from url", image_url)
         response = requests.get(image_url)
@@ -52,8 +50,8 @@ def process_image(image_input, image_url, task_prompt, text_input) -> Optional[I
         processor=FLORENCE_PROCESSOR,
         device=DEVICE,
         image=image_input,
-        task=text_input,
-        text=prompt
+        task=task_prompt,
+        text=text_prompt
     )
     detections = sv.Detections.from_lmm(
         lmm=sv.LMM.FLORENCE_2,
@@ -66,9 +64,14 @@ def process_image(image_input, image_url, task_prompt, text_input) -> Optional[I
         return None
     images = []
     print("mask generated:", len(detections.mask))
+    kernel_size = dilate
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+
     for i in range(len(detections.mask)):
-        img = Image.fromarray(detections.mask[i].astype(np.uint8) * 255)
-        images.append(img)
+        mask = detections.mask[i].astype(np.uint8) * 255
+        if dilate > 0:
+            mask = cv2.dilate(mask, kernel, iterations=1)
+        images.append(mask)
     return images
 
 
@@ -80,6 +83,7 @@ with gr.Blocks() as demo:
             task_prompt = gr.Dropdown(
                 ["<CAPTION>", "<DETAILED_CAPTION>",  "<MORE_DETAILED_CAPTION>",  "<CAPTION_TO_PHRASE_GROUNDING>",  "<OPEN_VOCABULARY_DETECTION>", "<DENSE_REGION_CAPTION>"], value="<CAPTION_TO_PHRASE_GROUNDING>", label="Task Prompt", info="task prompts"
             )
+            dilate = gr.Slider(label="dilate mask", minimum=0, maximum=50, value=10, step=1)
             text_prompt = gr.Textbox(label='Text prompt', placeholder='Enter text prompts')
             submit_button = gr.Button(value='Submit', variant='primary')
         with gr.Column():
@@ -87,7 +91,7 @@ with gr.Blocks() as demo:
     print(image, image_url, task_prompt, text_prompt, image_gallery)
     submit_button.click(
         fn = process_image,
-        inputs = [image, image_url, task_prompt, text_prompt],
+        inputs = [image, image_url, task_prompt, text_prompt, dilate],
         outputs = [image_gallery,],
         show_api=False
     )

requirements.txt

@@ -8,4 +8,4 @@ gradio
 supervision
 opencv-python
 pytest
-requests
+requests

utils/florence.py

@@ -42,9 +42,12 @@ def run_florence_inference(
     device: torch.device,
     image: Image,
     task: str,
-    text: str = ""
+    text: str = None
 ) -> Tuple[str, Dict]:
-    prompt = task + text
+    if text:
+        prompt = task + text
+    else:
+        prompt = task
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],