Update app.py

app.py

@@ -22,16 +22,40 @@ processor = PaliGemmaProcessor.from_pretrained(model_id)
 def infer(
         image: PIL.Image.Image,
         text: str,
-        max_new_tokens: int
-) -> str:
+        max_new_tokens: int = 2048
+) -> tuple:
     inputs = processor(text=text, images=image, return_tensors="pt", padding="longest", do_convert_rgb=True).to(device).to(dtype=model.dtype)
     with torch.no_grad():
         generated_ids = model.generate(
             **inputs,
-            max_length=2048
+            max_length=max_new_tokens
         )
     result = processor.decode(generated_ids[0], skip_special_tokens=True)
-    return result
+
+    # Placeholder to extract bounding box info from the result (you should replace this with actual bounding box extraction)
+    bounding_boxes = extract_bounding_boxes(result)
+    
+    # Draw bounding boxes on the image
+    annotated_image = image.copy()
+    draw = ImageDraw.Draw(annotated_image)
+    
+    # Example of drawing bounding boxes (replace with actual coordinates)
+    for idx, (box, label) in enumerate(bounding_boxes):
+        color = COLORS[idx % len(COLORS)]
+        draw.rectangle(box, outline=color, width=3)
+        draw.text((box[0], box[1]), label, fill=color)
+    
+    return result, annotated_image
+
+def extract_bounding_boxes(result):
+    """
+    Extract bounding boxes and labels from the model result.
+    Placeholder logic - replace this with actual parsing logic from model output.
+    
+    Example return: [((x1, y1, x2, y2), "Label")]
+    """
+    # Example static bounding box and label
+    return [((50, 50, 200, 200), "Damage"), ((300, 300, 400, 400), "Dent")]
 
 ######## Demo
 
@@ -47,15 +71,12 @@ with gr.Blocks(css="style.css") as demo:
             text_input = gr.Text(label="Input Text")
 
             text_output = gr.Text(label="Text Output")
+            output_image = gr.Image(label="Annotated Image")
             chat_btn = gr.Button()
 
-        chat_inputs = [
-            image,
-            text_input
-        ]
-        chat_outputs = [
-            text_output
-        ]
+        chat_inputs = [image, text_input]
+        chat_outputs = [text_output, output_image]
+
         chat_btn.click(
             fn=infer,
             inputs=chat_inputs,