Update app.py

app.py

@@ -1,19 +1,22 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoProcessor
-from PIL import Image, ImageDraw
+from transformers import AutoProcessor, AutoModelForCausalLM
+from PIL import Image
 import requests
+import copy
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
-import numpy as np
 import random
+import numpy as np
 
-# Load model and processor
 model_id = 'microsoft/Florence-2-large'
 model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).eval()
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 
 def run_example(task_prompt, image, text_input=None):
-    prompt = task_prompt if text_input is None else task_prompt + text_input
+    if text_input is None:
+        prompt = task_prompt
+    else:
+        prompt = task_prompt + text_input
     inputs = processor(text=prompt, images=image, return_tensors="pt")
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],
@@ -25,8 +28,8 @@ def run_example(task_prompt, image, text_input=None):
     )
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
     parsed_answer = processor.post_process_generation(
-        generated_text, 
-        task=task_prompt, 
+        generated_text,
+        task=task_prompt,
         image_size=(image.width, image.height)
     )
     return parsed_answer
@@ -39,43 +42,147 @@ def plot_bbox(image, data):
         rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none')
         ax.add_patch(rect)
         plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
-    plt.axis('off')
-    plt.show()
+    ax.axis('off')
+    return fig
 
 def draw_polygons(image, prediction, fill_mask=False):
     draw = ImageDraw.Draw(image)
-    colormap = ['blue', 'orange', 'green', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan', 'red']
+    scale = 1
     for polygons, label in zip(prediction['polygons'], prediction['labels']):
         color = random.choice(colormap)
-        fill_color = color if fill_mask else None
-        for polygon in polygons:
-            draw.polygon(polygon, outline=color, fill=fill_color)
-            draw.text((polygon[0][0], polygon[0][1]), label, fill=color)
-    image.show()
+        fill_color = random.choice(colormap) if fill_mask else None
+        for _polygon in polygons:
+            _polygon = np.array(_polygon).reshape(-1, 2)
+            if len(_polygon) < 3:
+                print('Invalid polygon:', _polygon)
+                continue
+            _polygon = (_polygon * scale).reshape(-1).tolist()
+            if fill_mask:
+                draw.polygon(_polygon, outline=color, fill=fill_color)
+            else:
+                draw.polygon(_polygon, outline=color)
+            draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color)
+    return image
+
+def convert_to_od_format(data):
+    bboxes = data.get('bboxes', [])
+    labels = data.get('bboxes_labels', [])
+    od_results = {
+        'bboxes': bboxes,
+        'labels': labels
+    }
+    return od_results
+
+def draw_ocr_bboxes(image, prediction):
+    scale = 1
+    draw = ImageDraw.Draw(image)
+    bboxes, labels = prediction['quad_boxes'], prediction['labels']
+    for box, label in zip(bboxes, labels):
+        color = random.choice(colormap)
+        new_box = (np.array(box) * scale).tolist()
+        draw.polygon(new_box, width=3, outline=color)
+        draw.text((new_box[0]+8, new_box[1]+2),
+                  "{}".format(label),
+                  align="right",
+                  fill=color)
+    return image
+
+def process_image(image, task_prompt, text_input=None):
+    if task_prompt == '<CAPTION>':
+        result = run_example(task_prompt, image)
+        return result
+    elif task_prompt == '<DETAILED_CAPTION>':
+        result = run_example(task_prompt, image)
+        return result
+    elif task_prompt == '<MORE_DETAILED_CAPTION>':
+        result = run_example(task_prompt, image)
+        return result
+    elif task_prompt == '<OD>':
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<OD>'])
+        return fig
+    elif task_prompt == '<DENSE_REGION_CAPTION>':
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<DENSE_REGION_CAPTION>'])
+        return fig
+    elif task_prompt == '<REGION_PROPOSAL>':
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<REGION_PROPOSAL>'])
+        return fig
+    elif task_prompt == '<CAPTION_TO_PHRASE_GROUNDING>':
+        results = run_example(task_prompt, image, text_input)
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return fig
+    elif task_prompt == '<REFERRING_EXPRESSION_SEGMENTATION>':
+        results = run_example(task_prompt, image, text_input)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True)
+        return output_image
+    elif task_prompt == '<REGION_TO_SEGMENTATION>':
+        results = run_example(task_prompt, image, text_input)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results['<REGION_TO_SEGMENTATION>'], fill_mask=True)
+        return output_image
+    elif task_prompt == '<OPEN_VOCABULARY_DETECTION>':
+        results = run_example(task_prompt, image, text_input)
+        bbox_results = convert_to_od_format(results['<OPEN_VOCABULARY_DETECTION>'])
+        fig = plot_bbox(image, bbox_results)
+        return fig
+    elif task_prompt == '<REGION_TO_CATEGORY>':
+        results = run_example(task_prompt, image, text_input)
+        return results
+    elif task_prompt == '<REGION_TO_DESCRIPTION>':
+        results = run_example(task_prompt, image, text_input)
+        return results
+    elif task_prompt == '<OCR>':
+        result = run_example(task_prompt, image)
+        return result
+    elif task_prompt == '<OCR_WITH_REGION>':
+        results = run_example(task_prompt, image)
+        output_image = copy.deepcopy(image)
+        output_image = draw_ocr_bboxes(output_image, results['<OCR_WITH_REGION>'])
+        return output_image
+
+css = """
+  #output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
+"""
+
+with gr.Blocks(css=css) as demo:
+    gr.HTML("<h1><center>Florence-2 Demo<center><h1>")
+    with gr.Tab(label="Florence-2 Image Captioning"):
+        with gr.Row():
+            with gr.Column():
+                input_img = gr.Image(label="Input Picture")
+                task_prompt = gr.Dropdown(choices=[
+                    '<CAPTION>', '<DETAILED_CAPTION>', '<MORE_DETAILED_CAPTION>', '<OD>',
+                    '<DENSE_REGION_CAPTION>', '<REGION_PROPOSAL>', '<CAPTION_TO_PHRASE_GROUNDING>',
+                    '<REFERRING_EXPRESSION_SEGMENTATION>', '<REGION_TO_SEGMENTATION>',
+                    '<OPEN_VOCABULARY_DETECTION>', '<REGION_TO_CATEGORY>', '<REGION_TO_DESCRIPTION>',
+                    '<OCR>', '<OCR_WITH_REGION>'
+                ], label="Task Prompt")
+                text_input = gr.Textbox(label="Text Input (optional)")
+                submit_btn = gr.Button(value="Submit")
+            with gr.Column():
+                output_text = gr.Textbox(label="Output Text")
+                output_img = gr.Image(label="Output Image")
 
-def gradio_interface(image, task_prompt, text_input):
-    result = run_example(task_prompt, image, text_input)
-    if task_prompt in ['<OD>', '<OPEN_VOCABULARY_DETECTION>']:
-        plot_bbox(image, result)
-    elif task_prompt in ['<REFERRING_EXPRESSION_SEGMENTATION>', '<REGION_TO_SEGMENTATION>']:
-        draw_polygons(image, result, fill_mask=True)
-    return result
+        gr.Examples(
+            examples=[
+                ["image1.jpg", '<CAPTION>'],
+                ["image1.jpg", '<OD>'],
+                ["image1.jpg", '<OCR_WITH_REGION>']
+            ],
+            inputs=[input_img, task_prompt],
+            outputs=[output_text, output_img],
+            fn=process_image,
+            cache_examples=True,
+            label='Try examples'
+        )
 
-with gr.Blocks() as demo:
-    gr.Markdown("## Florence Model Advanced Tasks")
-    with gr.Row():
-        image_input = gr.Image(type="pil")
-        task_input = gr.Dropdown(label="Select Task", choices=[
-            '<CAPTION>', '<DETAILED_CAPTION>', '<MORE_DETAILED_CAPTION>', 
-            '<OD>', '<DENSE_REGION_CAPTION>', '<REGION_PROPOSAL>', 
-            '<CAPTION_TO_PHRASE_GROUNDING>', '<REFERRING_EXPRESSION_SEGMENTATION>',
-            '<REGION_TO_SEGMENTATION>', '<OPEN_VOCABULARY_DETECTION>',
-            '<REGION_TO_CATEGORY>', '<REGION_TO_DESCRIPTION>', '<OCR>', '<OCR_WITH_REGION>'
-        ])
-        text_input = gr.Textbox(label="Optional Text Input", placeholder="Enter text here if required by the task")
-        submit_btn = gr.Button("Run Task")
-        output = gr.Textbox(label="Output")
-    
-    submit_btn.click(fn=gradio_interface, inputs=[image_input, task_input, text_input], outputs=output)
+        submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
 
-demo.launch()
+demo.launch(debug=True)