Update app.py

app.py

@@ -1,15 +1,18 @@
 import os
 import sys
-import spaces
-from typing import Iterable
-import gradio as gr
-import torch
+import io
+import json
 import requests
+from typing import Iterable, List, Tuple, Dict, Any
+
 from PIL import Image, ImageDraw, ImageFont
+import gradio as gr
+import torch
 from transformers import AutoProcessor, Florence2ForConditionalGeneration
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
+# ---------- Theme (kept from your original) ----------
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8", c100="#D3E5F0", c200="#A8CCE1", c300="#7DB3D2",
@@ -69,6 +72,7 @@ css = """
 }
 """
 
+# ---------- Models ----------
 MODEL_IDS = {
     "Florence-2-base": "florence-community/Florence-2-base",
     "Florence-2-base-ft": "florence-community/Florence-2-base-ft",
@@ -76,15 +80,15 @@ MODEL_IDS = {
     "Florence-2-large-ft": "florence-community/Florence-2-large-ft",
 }
 
-models = {}
-processors = {}
+models: Dict[str, Florence2ForConditionalGeneration] = {}
+processors: Dict[str, AutoProcessor] = {}
 
 print("Loading Florence-2 models... This may take a while.")
 for name, repo_id in MODEL_IDS.items():
     print(f"Loading {name}...")
     model = Florence2ForConditionalGeneration.from_pretrained(
         repo_id,
-        torch_dtype=torch.bfloat16,
+        dtype=torch.bfloat16,
         device_map="auto",
         trust_remote_code=True
     )
@@ -95,86 +99,201 @@ for name, repo_id in MODEL_IDS.items():
 
 print("\n🎉 All models loaded successfully!")
 
-def draw_bboxes(image, results, task):
+# ---------- Utilities ----------
+def _safe_parse_json_like(text: Any) -> Any:
+    """
+    If text is a dict already, return it. If it's a JSON-like string, try to json.loads it.
+    Otherwise return the original text.
+    """
+    if isinstance(text, dict):
+        return text
+    if isinstance(text, str):
+        text_str = text.strip()
+        # try to decode if it looks like JSON
+        if (text_str.startswith("{") and text_str.endswith("}")) or (text_str.startswith("[") and text_str.endswith("]")):
+            try:
+                return json.loads(text_str)
+            except Exception:
+                # fallback to returning original string
+                return text
+    return text
+
+def _find_bboxes_and_labels(obj: Any) -> List[Tuple[List[int], str]]:
     """
-    Draws bounding boxes on the image based on the model's results.
+    Recursively search `obj` (dict/list) for pairs of 'bboxes' and 'labels' (or region entries).
+    Returns list of tuples: (bbox, label)
+    bbox assumed as [x1,y1,x2,y2] (integers/floats)
     """
-    if task not in results:
-        return image
+    found: List[Tuple[List[int], str]] = []
+
+    def recurse(o: Any):
+        if isinstance(o, dict):
+            # direct pair case
+            if "bboxes" in o:
+                bboxes = o.get("bboxes", [])
+                labels = o.get("labels", [])
+                # if labels length mismatch, fill with empty strings
+                for i, bx in enumerate(bboxes):
+                    lbl = labels[i] if i < len(labels) else ""
+                    # sometimes bboxes come as dicts with keys or lists
+                    if isinstance(bx, dict) and {"x","y","w","h"}.issubset(bx.keys()):
+                        # convert xywh to x1,y1,x2,y2
+                        x = bx["x"]; y = bx["y"]; w = bx["w"]; h = bx["h"]
+                        found.append(([int(x), int(y), int(x + w), int(y + h)], lbl))
+                    else:
+                        # assume list-like [x1,y1,x2,y2] or [x,y,w,h]
+                        try:
+                            bx_list = list(map(int, bx))
+                            if len(bx_list) == 4:
+                                x1, y1, x2, y2 = bx_list
+                                # Heuristic: if x2>x1 and y2>y1 assume x1,y1,x2,y2 otherwise maybe xywh
+                                if x2 > x1 and y2 > y1:
+                                    found.append(([x1, y1, x2, y2], lbl))
+                                else:
+                                    # try treat as xywh
+                                    found.append(([x1, y1, x1 + x2, y1 + y2], lbl))
+                            else:
+                                # skip unexpected format
+                                pass
+                        except Exception:
+                            pass
+            # also check for region entries like {'bbox': ..., 'text': ...} or list of regions
+            if "regions" in o and isinstance(o["regions"], list):
+                for reg in o["regions"]:
+                    if isinstance(reg, dict) and "bbox" in reg:
+                        bx = reg["bbox"]
+                        lbl = reg.get("label", reg.get("text", ""))
+                        try:
+                            bx_list = list(map(int, bx))
+                            if len(bx_list) == 4:
+                                found.append(([bx_list[0], bx_list[1], bx_list[2], bx_list[3]], lbl))
+                        except Exception:
+                            pass
+            # recurse deeper
+            for v in o.values():
+                recurse(v)
+        elif isinstance(o, list):
+            for item in o:
+                recurse(item)
+        # else ignore primitives
 
-    predictions = results[task]
-    if not predictions:
-        return image
+    recurse(obj)
+    return found
 
-    draw = ImageDraw.Draw(image)
+def _draw_bboxes_on_image(img: Image.Image, boxes_and_labels: List[Tuple[List[int], str]]) -> Image.Image:
+    """
+    Draw bounding boxes and labels on a copy of `img`.
+    """
+    annotated = img.convert("RGB").copy()
+    draw = ImageDraw.Draw(annotated)
+    # try to get a default font (PIL may not have a TTF available)
     try:
-        font = ImageFont.truetype("arial.ttf", 20)
-    except IOError:
+        font = ImageFont.truetype("DejaVuSans.ttf", size=14)
+    except Exception:
         font = ImageFont.load_default()
 
-    for pred in predictions:
-        if 'bboxes' in pred and 'labels' in pred:
-            for box, label in zip(pred['bboxes'], pred['labels']):
-                draw.rectangle(box, outline="red", width=3)
-                text_position = [box[0], box[1] - 25]
-                # Add a background to the text for better visibility
-                text_bbox = draw.textbbox(tuple(text_position), label, font=font)
-                draw.rectangle(text_bbox, fill="red")
-                draw.text(tuple(text_position), label, fill="white", font=font)
-        elif 'bboxes' in pred: # For tasks like REGION_PROPOSAL without labels
-            for box in pred['bboxes']:
-                draw.rectangle(box, outline="red", width=3)
+    for bbox, label in boxes_and_labels:
+        # bbox should be [x1,y1,x2,y2]
+        x1, y1, x2, y2 = bbox
+        # keep coordinates within image bounds
+        x1 = max(0, int(x1)); y1 = max(0, int(y1))
+        x2 = min(annotated.width - 1, int(x2)); y2 = min(annotated.height - 1, int(y2))
+        # draw rectangle (thicker by drawing several offsets)
+        thickness = max(2, int(round(min(annotated.width, annotated.height) / 200)))
+        for t in range(thickness):
+            draw.rectangle([x1 - t, y1 - t, x2 + t, y2 + t], outline="red")
+        # draw label background
+        if label is None:
+            label = ""
+        label_text = str(label)
+        text_w, text_h = draw.textsize(label_text, font=font)
+        # background rectangle for label (semi-opaque)
+        label_bg = [x1, max(0, y1 - text_h - 4), x1 + text_w + 6, y1]
+        draw.rectangle(label_bg, fill="red")
+        # text
+        draw.text((x1 + 3, max(0, y1 - text_h - 2)), label_text, fill="white", font=font)
 
-    return image
+    return annotated
 
+# ---------- Inference function ----------
+# tasks for which we attempt to extract/display bboxes
+VISUAL_REGION_TASKS = {"<OD>", "<DENSE_REGION_CAPTION>", "<OCR_WITH_REGION>", "<REGION_PROPOSAL>"}
 
-@spaces.GPU
+# If you are using Spaces with GPU decorator; keep it as-is in your environment
 def run_florence2_inference(model_name: str, image: Image.Image, task_prompt: str,
                             max_new_tokens: int = 1024, num_beams: int = 3):
     """
     Runs inference using the selected Florence-2 model.
+    Returns a tuple: (parsed_answer, annotated_image_or_none)
     """
     if image is None:
-        return "Please upload an image to get started.", None
+        return {"error": "Please upload an image to get started."}, None
 
     model = models[model_name]
     processor = processors[model_name]
 
-    inputs = processor(text=task_prompt, images=image, return_tensors="pt").to(model.device, torch.bfloat16)
+    # Prepare inputs (move to model device)
+    inputs = processor(text=task_prompt, images=image, return_tensors="pt")
+    # send tensors to model device and set dtype
+    device = model.device
+    for k, v in inputs.items():
+        if isinstance(v, torch.Tensor):
+            inputs[k] = v.to(device, dtype=torch.bfloat16)
 
+    # Generate
     generated_ids = model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
+        input_ids=inputs.get("input_ids"),
+        pixel_values=inputs.get("pixel_values"),
         max_new_tokens=max_new_tokens,
         num_beams=num_beams,
         do_sample=False
     )
 
+    # Decode
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
 
+    # Post-process (the processor provided by Florence models sometimes provides
+    # a structured output such as dict with 'bboxes' etc.)
     image_size = image.size
     parsed_answer = processor.post_process_generation(
         generated_text, task=task_prompt, image_size=image_size
     )
 
-    if task_prompt in ["<OD>", "<DENSE_REGION_CAPTION>", "<OCR_WITH_REGION>", "<REGION_PROPOSAL>"]:
-        image_with_boxes = draw_bboxes(image.copy(), parsed_answer, task_prompt)
-        return image_with_boxes, parsed_answer
-    else:
-        return image, parsed_answer
+    # Try to make parsed_answer JSON-serializable and easily inspectable
+    parsed_serializable = parsed_answer
+    # If it's a string that contains JSON, attempt to parse
+    if isinstance(parsed_answer, str):
+        parsed_serializable = _safe_parse_json_like(parsed_answer)
+
+    annotated_image = None
+    # If the task is in our visual region tasks, try to find bboxes and labels
+    if task_prompt in VISUAL_REGION_TASKS:
+        # parsed_serializable may be dict/list or string; try to find bboxes
+        boxes_and_labels = _find_bboxes_and_labels(parsed_serializable)
+        if boxes_and_labels:
+            try:
+                annotated_image = _draw_bboxes_on_image(image, boxes_and_labels)
+            except Exception as e:
+                # if drawing fails, set annotated_image to None but keep parsed answer
+                print("Failed to draw boxes:", e)
+                annotated_image = None
 
+    # Return parsed answer (prefer a dict or serializable structure) and annotated image (PIL) or None
+    return parsed_serializable, annotated_image
 
+# ---------- UI ----------
 florence_tasks = [
     "<OD>", "<CAPTION>", "<DETAILED_CAPTION>", "<MORE_DETAILED_CAPTION>",
     "<DENSE_REGION_CAPTION>", "<REGION_PROPOSAL>", "<OCR>", "<OCR_WITH_REGION>"
 ]
 
+# Example image (keeps your example)
 url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
 example_image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
 
 with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
     gr.Markdown("# **Florence-2 Vision Models**", elem_id="main-title")
-    gr.Markdown("Select a model, upload an image, choose a task, and click Submit to see the results.")
+    gr.Markdown("Select a model, upload an image, choose a task, and click Submit to see the parsed output and an annotated image (when bounding boxes are present).")
 
     with gr.Row():
         with gr.Column(scale=2):
@@ -182,7 +301,7 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
             task_prompt = gr.Dropdown(
                 label="Select Task",
                 choices=florence_tasks,
-                value="<OD>"
+                value="<MORE_DETAILED_CAPTION>"
             )
             model_choice = gr.Radio(
                 choices=list(MODEL_IDS.keys()),
@@ -202,13 +321,13 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
         with gr.Column(scale=3):
             gr.Markdown("## Output", elem_id="output-title")
             parsed_output = gr.JSON(label="Parsed Answer")
-            output_image = gr.Image(label="Image with Bounding Boxes", type="pil")
+            annotated_output = gr.Image(label="Annotated Image (if available)", type="pil")
 
     image_submit.click(
         fn=run_florence2_inference,
         inputs=[model_choice, image_upload, task_prompt, max_new_tokens, num_beams],
-        outputs=[output_image, parsed_output]
+        outputs=[parsed_output, annotated_output]
     )
 
 if __name__ == "__main__":
-    demo.queue().launch(debug=True, mcp_server=True, ssr_mode=False, show_error=True)
+    demo.queue().launch(debug=True, mcp_server=True, ssr_mode=False, show_error=True)