add bounding boxes

app.py

@@ -5,6 +5,12 @@ import requests
 from transformers import AutoProcessor
 from modeling_florence2 import Florence2ForConditionalGeneration
 from configuration_florence2 import Florence2Config
+import io
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import random
+import copy
 
 # Initialize model and processor
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -25,6 +31,40 @@ TASK_PROMPTS = {
     "Region Proposal": "<REGION_PROPOSAL>"
 }
 
+colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
+            'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']
+
+def fig_to_pil(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png')
+    buf.seek(0)
+    return Image.open(buf)
+
+def plot_bbox(image, data):
+    fig, ax = plt.subplots()
+    ax.imshow(image)
+    for bbox, label in zip(data['bboxes'], data['labels']):
+        x1, y1, x2, y2 = bbox
+        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))
+    ax.axis('off')
+    return fig
+
+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 = TASK_PROMPTS[task]
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
@@ -39,19 +79,40 @@ def process_image(image, task):
     
     parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
     
-    return str(parsed_answer)
+    return parsed_answer
+
+def main_process(image, task):
+    result = process_image(image, task)
+    
+    if task in ["Object Detection", "Dense Region Caption", "Region Proposal"]:
+        fig = plot_bbox(image, result[TASK_PROMPTS[task]])
+        output_image = fig_to_pil(fig)
+    elif task == "OCR with Region":
+        output_image = draw_ocr_bboxes(image.copy(), result[TASK_PROMPTS[task]])
+    else:
+        output_image = None
+    
+    return {task: str(result)}, output_image
 
 # Define Gradio interface
-iface = gr.Interface(
-    fn=process_image,
-    inputs=[
-        gr.Image(type="pil"),
-        gr.Dropdown(list(TASK_PROMPTS.keys()), label="Task")
-    ],
-    outputs=gr.Textbox(label="Result"),
-    title="Florence-2 Demo",
-    description="Upload an image and select a task to process with Florence-2."
-)
+with gr.Blocks(title="Florence-2 Demo") as iface:
+    gr.Markdown("# Florence-2 Demo")
+    gr.Markdown("Upload an image and select a task to process with Florence-2.")
+    
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Input Image")
+        task_dropdown = gr.Dropdown(list(TASK_PROMPTS.keys()), label="Task")
+    
+    submit_button = gr.Button("Process")
+    
+    output_text = gr.JSON(label="Output")
+    output_image = gr.Image(label="Processed Image")
+    
+    submit_button.click(
+        fn=main_process,
+        inputs=[image_input, task_dropdown],
+        outputs=[output_text, output_image]
+    )
 
 # Launch the interface
 iface.launch()