Update app.py

app.py

@@ -1,33 +1,53 @@
 import random
 import requests
 import json
+import ast
 
 import matplotlib.pyplot as plt
 from PIL import Image, ImageDraw, ImageFont
 
 import gradio as gr
 import torch
-from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration, AutoModelForCausalLM
 from qwen_vl_utils import process_vision_info
 from spaces import GPU
 from gradio.themes.ocean import Ocean
 
 # --- Config ---
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2.5-VL-3B-Instruct", torch_dtype="auto", device_map="auto"
+model_qwen_id = "Qwen/Qwen2.5-VL-3B-Instruct"
+model_moondream_id = "vikhyatk/moondream2"
+
+model_qwen = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    model_qwen_id, torch_dtype="auto", device_map="auto"
+)
+model_moondream = AutoModelForCausalLM.from_pretrained(
+    model_moondream_id,
+    revision="2025-06-21",
+    trust_remote_code=True,
+    device_map={"": "cuda"}
 )
 
+def extract_model_short_name(model_id):
+    return model_id.split("/")[-1].replace("-", " ").replace("_", " ")
+
+model_qwen_name = extract_model_short_name(model_qwen_id)  # → "Qwen2.5 VL 3B Instruct"
+model_moondream_name = extract_model_short_name(model_moondream_id)  # → "moondream2"
+
+
 min_pixels = 224 * 224
-max_pixels = 512 * 512
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
+max_pixels = 1024 * 1024
+processor_qwen = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
+#processor_moondream = AutoProcessor.from_pretrained("vikhyatk/moondream2", trust_remote_code=True)
 
 label2color = {}
+vivid_colors = ["#e6194b", "#3cb44b", "#0082c8", "#f58231", "#911eb4", "#46f0f0", "#f032e6", "#d2f53c", "#fabebe", "#008080", "#e6beff", "#aa6e28", "#fffac8", "#800000", "#aaffc3", "#808000", "#ffd8b1", "#000080", "#808080", "#000000"]
 
 def get_color(label, explicit_color=None):
     if explicit_color:
         return explicit_color
     if label not in label2color:
-        label2color[label] = "#" + ''.join(random.choices('0123456789ABCDEF', k=6))
+        index = len(label2color) % len(vivid_colors)
+        label2color[label] = vivid_colors[index]
     return label2color[label]
 
 def create_annotated_image(image, json_data, height, width):
@@ -47,11 +67,8 @@ def create_annotated_image(image, json_data, height, width):
     draw = ImageDraw.Draw(draw_image)
 
     try:
-        print(1)
-        print('int(12 * scale_factor)', int(12 * scale_factor))
-        font = ImageFont.truetype("arial.ttf", int(12 * scale_factor))
+        font = ImageFont.truetype("DejaVuSans-Bold.ttf", int(12 * scale_factor))
     except:
-        print(2)
         font = ImageFont.load_default()
 
     for item in bbox_data:
@@ -84,10 +101,51 @@ def create_annotated_image(image, json_data, height, width):
 
     return draw_image
 
+def create_annotated_image_normalized(image, json_data, label="object", explicit_color=None):
+    if not isinstance(json_data, dict):
+        return image
+
+    original_width, original_height = image.size
+    scale_factor = max(original_width, original_height) / 512
+    draw_image = image.copy()
+    draw = ImageDraw.Draw(draw_image)
+
+    try:
+        font = ImageFont.truetype("DejaVuSans-Bold.ttf", int(12 * scale_factor))
+    except:
+        font = ImageFont.load_default()
+
+    color = get_color(label, explicit_color)
+
+    for point in json_data.get("points", []):
+        x = int(point["x"] * original_width)
+        y = int(point["y"] * original_height)
+        radius = int(4 * scale_factor)
+        draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=color, outline=color)
+
+    for item in json_data.get("objects", []):
+        x_min = int(item["x_min"] * original_width)
+        y_min = int(item["y_min"] * original_height)
+        x_max = int(item["x_max"] * original_width)
+        y_max = int(item["y_max"] * original_height)
+        draw.rectangle([x_min, y_min, x_max, y_max], outline=color, width=int(2 * scale_factor))
+        draw.text((x_min, max(0, y_min - int(15 * scale_factor))), label, fill=color, font=font)
+
+    if "reasoning" in json_data:
+        for grounding in json_data["reasoning"].get("grounding", []):
+            for x_norm, y_norm in grounding.get("points", []):
+                x = int(x_norm * original_width)
+                y = int(y_norm * original_height)
+                radius = int(4 * scale_factor)
+                draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=color, outline=color)
+
+    return draw_image
+
+
+
 @GPU
-def detect(image, prompt):
-    STANDARD_SIZE = (512, 512)
-    image.thumbnail(STANDARD_SIZE)
+def detect_qwen(image, prompt):
+
     messages = [
         {
             "role": "user",
@@ -98,21 +156,21 @@ def detect(image, prompt):
         }
     ]
 
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    text = processor_qwen.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
+    inputs = processor_qwen(
         text=[text],
         images=image_inputs,
         videos=video_inputs,
         padding=True,
         return_tensors="pt",
-    ).to(model.device)
+    ).to(model_qwen.device)
 
-    generated_ids = model.generate(**inputs, max_new_tokens=1024)
+    generated_ids = model_qwen.generate(**inputs, max_new_tokens=1024)
     generated_ids_trimmed = [
         out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
     ]
-    output_text = processor.batch_decode(
+    output_text = processor_qwen.batch_decode(
         generated_ids_trimmed, do_sample=True, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )[0]
 
@@ -123,6 +181,30 @@ def detect(image, prompt):
 
     return annotated_image, output_text
 
+
+@GPU
+def detect_moondream(image, prompt, category_input):
+    if category_input in ["Object Detection", "Visual Grounding + Object Detection"]:
+        output_text = model_moondream.detect(image=image, object=prompt)
+    elif category_input == "Visual Grounding + Keypoint Detection":
+        output_text = model_moondream.point(image=image, object=prompt)
+    else:
+        output_text = model_moondream.query(image=image, question=prompt, reasoning=True)
+
+    annotated_image = create_annotated_image_normalized(image=image, json_data=output_text, label="object", explicit_color=None)
+    
+    return annotated_image, output_text
+
+@GPU
+def detect(image, prompt_model_1, prompt_model_2, category_input):
+    STANDARD_SIZE = (1024, 1024)
+    image.thumbnail(STANDARD_SIZE)
+    
+    annotated_image_model_1, output_text_model_1 = detect_qwen(image, prompt_model_1)
+    annotated_image_model_2, output_text_model_2 = detect_moondream(image, prompt_model_2, category_input)
+
+    return annotated_image_model_1, output_text_model_1, annotated_image_model_2, output_text_model_2
+
 css_hide_share = """
 button#gradio-share-link-button-0 {
     display: none !important;
@@ -132,43 +214,70 @@ button#gradio-share-link-button-0 {
 # --- Gradio Interface ---
 with gr.Blocks(theme=Ocean(), css=css_hide_share) as demo:
 
-    gr.Markdown("# Object Understanding with Vision Language Models")
+    gr.Markdown("# 👓 Object Understanding with Vision Language Models")
     gr.Markdown("### Explore object detection, visual grounding, keypoint detection, and/or object counting through natural language prompts.")
     gr.Markdown("""
-    *Powered by Qwen2.5-VL*  
-    *Inspired by the tutorial [Object Detection and Visual Grounding with Qwen 2.5](https://pyimagesearch.com/2025/06/09/object-detection-and-visual-grounding-with-qwen-2-5/) on PyImageSearch.*
+    *Powered by [Qwen2.5-VL 3B](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) and [Moondream 2B (revision="2025-06-21")](https://huggingface.co/vikhyatk/moondream2). Inspired by the tutorial [Object Detection and Visual Grounding with Qwen 2.5](https://pyimagesearch.com/2025/06/09/object-detection-and-visual-grounding-with-qwen-2-5/) on PyImageSearch.*  
+    *Moondream 2B uses the [moondream.py API](https://huggingface.co/vikhyatk/moondream2/blob/main/moondream.py), selecting `detect` for categories with "Object Detection" `point` for the ones with "Keypoint Detection", and reasoning-based querying for all others.*
     """)
 
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column(scale=2):
             image_input = gr.Image(label="Upload an image", type="pil", height=400)
-            prompt_input = gr.Textbox(label="Enter your prompt", placeholder="e.g., Detect all red cars in the image")
-            category_input = gr.Textbox(label="Category", interactive=False)
+            prompt_input_model_1 = gr.Textbox(
+                label=f"Enter your prompt for {model_qwen_name}",
+                placeholder="e.g., Detect all red cars in the image"
+            )
+
+            prompt_input_model_2 = gr.Textbox(
+                label=f"Enter your prompt for {model_moondream_name}",
+                placeholder="e.g., Detect all blue cars in the image"
+            )
+
+            
+            categories = [
+                "Object Detection",
+                "Object Counting",
+                "Visual Grounding + Keypoint Detection",
+                "Visual Grounding + Object Detection",
+                "General query"
+            ]
+
+            category_input = gr.Dropdown(
+                choices=categories,
+                label="Category",
+                interactive=True
+            )
             generate_btn = gr.Button(value="Generate")
 
         with gr.Column(scale=1):
-            output_image = gr.Image(type="pil", label="Annotated image", height=400)
-            output_textbox = gr.Textbox(label="Model response", lines=10)
+            output_image_model_1 = gr.Image(type="pil", label=f"Annotated image for {model_qwen_name}", height=400)
+            output_textbox_model_1 = gr.Textbox(label=f"Model response for {model_qwen_name}", lines=10)
+        
+        with gr.Column(scale=1):
+            output_image_model_2 = gr.Image(type="pil", label=f"Annotated image for {model_moondream_name}", height=400)
+            output_textbox_model_2 = gr.Textbox(label=f"Model response for {model_moondream_name}", lines=10)
 
     gr.Markdown("### Examples")
     example_prompts = [
-        ["examples/example_1.jpg", "Detect all objects in the image and return their locations and labels.", "Object Detection"],
-        ["examples/example_2.JPG", "Detect all the individual candies in the image and return their locations and labels.", "Object Detection"],
-        ["examples/example_1.jpg", "Count the number of red cars in the image.", "Object Counting"],
-        ["examples/example_2.JPG", "Count the number of blue candies in the image.", "Object Counting"],
-        ["examples/example_1.jpg", "Identify the red cars in this image, detect their key points and return their positions in the form of points.", "Visual Grounding + Keypoint Detection"],
-        ["examples/example_2.JPG", "Identify the blue candies in this image, detect their key points and return their positions in the form of points.", "Visual Grounding + Keypoint Detection"],
-        ["examples/example_1.jpg", "Detect the red car that is leading in this image and return its location and label.", "Visual Grounding + Object Detection"],
-        ["examples/example_2.JPG", "Detect the blue candy located at the top of the group in this image and return its location and label.", "Visual Grounding + Object Detection"],
+        ["examples/example_1.jpg", "Detect all objects in the image and return their locations and labels.", "Detect all objects in the image and return their locations and labels.", "Object Detection"],
+        ["examples/example_2.JPG", "Detect all the individual candies in the image and return their locations and labels.", "Detect all the individual candies in the image and return their locations and labels.", "Object Detection"],
+        ["examples/example_1.jpg", "Count the number of red cars in the image.", "Count the number of red cars in the image.", "Object Counting"],
+        ["examples/example_2.JPG", "Count the number of blue candies in the image.", "Count the number of blue candies in the image.", "Object Counting"],
+        ["examples/example_1.jpg", "Identify the red cars in this image, detect their key points and return their positions in the form of points.", "Identify the red cars in this image, detect their key points and return their positions in the form of points.", "Visual Grounding + Keypoint Detection"],
+        ["examples/example_2.JPG", "Identify the blue candies in this image, detect their key points and return their positions in the form of points.", "Identify the blue candies in this image, detect their key points and return their positions in the form of points.", "Visual Grounding + Keypoint Detection"],
+        ["examples/example_1.jpg", "Detect the red car that is leading in this image and return its location and label.", "Detect the red car that is leading in this image and return its location and label.", "Visual Grounding + Object Detection"],
+        ["examples/example_2.JPG", "Detect the blue candy located at the top of the group in this image and return its location and label.", "Detect the blue candy located at the top of the group in this image and return its location and label.", "Visual Grounding + Object Detection"],
     ]
 
     gr.Examples(
         examples=example_prompts,
-        inputs=[image_input, prompt_input, category_input],
+        inputs=[image_input, prompt_input_model_1, prompt_input_model_2, category_input],
         label="Click an example to populate the input"
     )
 
-    generate_btn.click(fn=detect, inputs=[image_input, prompt_input], outputs=[output_image, output_textbox])
+    generate_btn.click(fn=detect, inputs=[image_input, prompt_input_model_1, prompt_input_model_2, category_input], outputs=[output_image_model_1, output_textbox_model_1, output_image_model_2, output_textbox_model_2])
 
 if __name__ == "__main__":
     demo.launch()
+