Update app.py

app.py

@@ -1,12 +1,12 @@
 import os
-import re
 import time
 from threading import Thread
+import re
+from PIL import Image, ImageDraw
 
 import gradio as gr
 import spaces
 import torch
-from PIL import Image, ImageDraw
 
 from transformers import (
     Qwen2_5_VLForConditionalGeneration,
@@ -30,62 +30,76 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-
-def extract_coordinates(text: str):
-    """Extract all (x, y) coordinates from model output text."""
-    pattern = r"\((\d+),\s*(\d+)\)"
-    coords = re.findall(pattern, text)
-    return [(int(x), int(y)) for x, y in coords]
-
-
-def draw_boxes(image: Image.Image, coords, box_type="solid"):
-    """Draw bounding boxes on the image."""
-    img_copy = image.copy()
-    draw = ImageDraw.Draw(img_copy)
-
-    side = 50  # square size
-    for (x, y) in coords:
-        box = [(x - side//2, y - side//2), (x + side//2, y + side//2)]
-        if box_type == "solid":
-            draw.rectangle(box, outline="red", width=3)
-        elif box_type == "dotted":
-            # Draw dotted (dashed) rectangle
-            dash_len = 5
-            x1, y1 = box[0]
-            x2, y2 = box[1]
-
-            # Top edge
-            for i in range(x1, x2, dash_len*2):
-                draw.line([(i, y1), (min(i+dash_len, x2), y1)], fill="blue", width=2)
-            # Bottom edge
-            for i in range(x1, x2, dash_len*2):
-                draw.line([(i, y2), (min(i+dash_len, x2), y2)], fill="blue", width=2)
-            # Left edge
-            for i in range(y1, y2, dash_len*2):
-                draw.line([(x1, i), (x1, min(i+dash_len, y2))], fill="blue", width=2)
-            # Right edge
-            for i in range(y1, y2, dash_len*2):
-                draw.line([(x2, i), (x2, min(i+dash_len, y2))], fill="blue", width=2)
-
-    return img_copy
-
+def parse_model_output(text: str):
+    """
+    Parses the model output to extract the answer and bounding box coordinates.
+    """
+    # Extract coordinates from the <think> block
+    think_match = re.search(r"<think>(.*?)</think>", text, re.DOTALL)
+    coordinates = []
+    if think_match:
+        think_content = think_match.group(1)
+        # Find all occurrences of (x, y) coordinates
+        coords_raw = re.findall(r'\((\d+),\s*(\d+)\)', think_content)
+        coordinates = [(int(x), int(y)) for x, y in coords_raw]
+
+    # Extract the answer from the <answer> block
+    answer_match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
+    answer = answer_match.group(1).strip() if answer_match else text
+
+    return answer, coordinates
+
+def draw_bounding_boxes(image: Image.Image, coordinates: list, box_size: int = 60, use_dotted_style: bool = False):
+    """
+    Draws square bounding boxes on the image at the given coordinates.
+    """
+    if not coordinates:
+        return image
+
+    img_with_boxes = image.copy()
+    draw = ImageDraw.Draw(img_with_boxes, "RGBA")
+    
+    half_box = box_size // 2
+
+    for (x, y) in coordinates:
+        # Define the bounding box corners
+        x1 = x - half_box
+        y1 = y - half_box
+        x2 = x + half_box
+        y2 = y + half_box
+        
+        if use_dotted_style:
+            # "Dotted like seaborn" - a semi-transparent fill with a solid outline
+            fill_color = (0, 100, 255, 60)  # Light blue, semi-transparent
+            outline_color = (0, 0, 255) # Solid blue
+            draw.rectangle([x1, y1, x2, y2], fill=fill_color, outline=outline_color, width=2)
+        else:
+            # Default solid box
+            outline_color = (255, 0, 0) # Red
+            draw.rectangle([x1, y1, x2, y2], outline=outline_color, width=3)
+            
+    return img_with_boxes
 
 @spaces.GPU
 def generate_image(text: str, image: Image.Image,
-                   max_new_tokens: int = 1024,
-                   temperature: float = 0.6,
-                   top_p: float = 0.9,
-                   top_k: int = 50,
-                   repetition_penalty: float = 1.2,
-                   box_type: str = "solid"):
+                   max_new_tokens: int,
+                   temperature: float,
+                   top_p: float,
+                   top_k: int,
+                   repetition_penalty: float,
+                   draw_boxes: bool,
+                   use_dotted_style: bool):
     """
-    Generates responses using the Lumian2-VLR-7B-Thinking model for image input.
-    Yields raw text, Markdown-formatted text, and annotated image with bounding boxes.
+    Generates responses and draws bounding boxes based on model output.
+    Yields raw text, markdown-formatted text, and the processed image.
     """
     if image is None:
         yield "Please upload an image.", "Please upload an image.", None
         return
 
+    # Yield the original image immediately for the output display
+    yield "", "", image
+
     messages = [{
         "role": "user",
         "content": [
@@ -115,15 +129,23 @@ def generate_image(text: str, image: Image.Image,
     }
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
+    
     buffer = ""
-    coords = []
     for new_text in streamer:
         buffer += new_text
-        coords = extract_coordinates(buffer)
-        annotated_image = draw_boxes(image, coords, box_type) if coords else None
         time.sleep(0.01)
-        yield buffer, buffer, annotated_image
-
+        # During generation, yield text updates but keep the original image
+        yield buffer, buffer, image
+
+    # After generation is complete, parse the output and draw boxes
+    final_answer, coordinates = parse_model_output(buffer)
+    
+    output_image = image
+    if draw_boxes and coordinates:
+        output_image = draw_bounding_boxes(image, coordinates, use_dotted_style=use_dotted_style)
+    
+    # Yield the final result with the processed image
+    yield buffer, final_answer, output_image
 
 # Define examples for image inference
 image_examples = [
@@ -154,52 +176,50 @@ css = """
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown("# **Lumian2-VLR-7B-Thinking Image Inference**")
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             gr.Markdown("## Image Inference")
             image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
             image_upload = gr.Image(type="pil", label="Image")
             image_submit = gr.Button("Submit", elem_classes="submit-btn")
-            gr.Examples(
-                examples=image_examples,
-                inputs=[image_query, image_upload]
-            )
+            
             with gr.Accordion("Advanced options", open=False):
                 max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
                 temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
                 top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
                 top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
                 repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
-
-            # New options for bounding box visualization
-            box_type = gr.Radio(
-                choices=["solid", "dotted"],
-                value="solid",
-                label="Bounding Box Style"
+                
+            gr.Examples(
+                examples=image_examples,
+                inputs=[image_query, image_upload]
             )
 
-        with gr.Column():
-            with gr.Column(elem_classes="canvas-output"):
-                gr.Markdown("## Output")
-                output = gr.Textbox(label="Raw Output", interactive=False, lines=3, scale=2)
-
-                with gr.Accordion("(Result.md)", open=False):
-                    markdown_output = gr.Markdown()
-
-                annotated_output = gr.Image(label="Annotated Image with Bounding Boxes")
+        with gr.Column(scale=2):
+            gr.Markdown("## Output")
+            with gr.Tabs():
+                with gr.TabItem("Image with Bounding Box"):
+                    image_output = gr.Image(label="Processed Image")
+                with gr.TabItem("Raw Text"):
+                    output = gr.Textbox(label="Raw Model Output", interactive=False, lines=10)
+                with gr.TabItem("Parsed Answer"):
+                    markdown_output = gr.Markdown(label="Parsed Answer")
 
             gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/Qwen2.5-VL/discussions)")
 
             gr.Markdown(
-                """
-                > [Lumian2-VLR-7B-Thinking](https://huggingface.co/prithivMLmods/Lumian2-VLR-7B-Thinking): The Lumian2-VLR-7B-Thinking model is a high-fidelity vision-language reasoning (experimental model) system designed for fine-grained multimodal understanding. Built on Qwen2.5-VL-7B-Instruct, this model enhances image captioning, and document comprehension through explicit grounded reasoning. It produces structured reasoning traces aligned with visual coordinates, enabling explainable multimodal reasoning.
-                """
+                """> [Lumian2-VLR-7B-Thinking](https://huggingface.co/prithivMLmods/Lumian2-VLR-7B-Thinking): The Lumian2-VLR-7B-Thinking model is a high-fidelity vision-language reasoning (experimental model) system designed for fine-grained multimodal understanding. Built on Qwen2.5-VL-7B-Instruct, this model enhances image captioning, and document comprehension through explicit grounded reasoning. It produces structured reasoning traces aligned with visual coordinates, enabling explainable multimodal reasoning."""
             )
 
+            with gr.Row():
+                draw_boxes_checkbox = gr.Checkbox(label="Draw Bounding Boxes", value=True)
+                dotted_style_checkbox = gr.Checkbox(label="Use Dotted Style for Boxes", value=False)
+
+
     image_submit.click(
         fn=generate_image,
-        inputs=[image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty, box_type],
-        outputs=[output, markdown_output, annotated_output]
+        inputs=[image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty, draw_boxes_checkbox, dotted_style_checkbox],
+        outputs=[output, markdown_output, image_output]
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=50).launch(share=True)
+    demo.queue(max_size=50).launch(share=True)```