Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoPro
 from qwen_vl_utils import process_vision_info
 import torch
 import base64
-from PIL import Image
+from PIL import Image, ImageDraw
 from io import BytesIO
 import re
 
@@ -27,6 +27,14 @@ def image_to_base64(image):
     return img_str
 
 
+def draw_bounding_boxes(image, bounding_boxes, outline_color="red", line_width=2):
+    draw = ImageDraw.Draw(image)
+    for box in bounding_boxes:
+        xmin, xmax, ymin, ymax = box
+        draw.rectangle([xmin, ymin, xmax, ymax], outline=outline_color, width=line_width)
+    return image
+
+
 @spaces.GPU
 def run_example(image, text_input, model_id="Qwen/Qwen2-VL-7B-Instruct"):
     model = models[model_id].eval()
@@ -67,7 +75,7 @@ def run_example(image, text_input, model_id="Qwen/Qwen2-VL-7B-Instruct"):
     pattern = r'\[\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*\]'
     matches = re.findall(pattern, str(output_text))
     parsed_boxes = [[int(num) for num in match] for match in matches]
-    return output_text, parsed_boxes
+    return output_text, parsed_boxes, draw_bounding_boxes(image, parsed_boxes)
 
 css = """
   #output {
@@ -89,7 +97,8 @@ with gr.Blocks(css=css) as demo:
             with gr.Column():
                 model_output_text = gr.Textbox(label="Model Output Text")
                 parsed_boxes = gr.Textbox(label="Parsed Boxes")
+                annotated_image = gr.Image(label="Annotated Picture")
 
-        submit_btn.click(run_example, [input_img, text_input, model_selector], [model_output_text, parsed_boxes])
+        submit_btn.click(run_example, [input_img, text_input, model_selector], [model_output_text, parsed_boxes, annotated_image])
 
 demo.launch(debug=True)