add quad boxes

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import torch
-from PIL import Image, ImageDraw
+import cv2
+import numpy as np
+from matplotlib import pyplot as pltfrom PIL import Image, ImageDraw
 from transformers import AutoProcessor
 from modeling_florence2 import Florence2ForConditionalGeneration
 import io
@@ -115,14 +117,15 @@ def fig_to_pil(fig):
 def plot_bbox(image, data, use_quad_boxes=False):
     fig, ax = plt.subplots()
     ax.imshow(image)
-    
+
     # Handle both 'bboxes' and 'quad_boxes'
     if use_quad_boxes:
         for quad_box, label in zip(data.get('quad_boxes', []), data.get('labels', [])):
             quad_box = np.array(quad_box).reshape(-1, 2)
-            poly = patches.Polygon(quad_box, linewidth=1, edgecolor='r', facecolor='none')
+            poly = Polygon(quad_box, linewidth=1, edgecolor='r', facecolor='none')
             ax.add_patch(poly)
-            plt.text(quad_box[0][0], quad_box[0][1], label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
+            plt.text(quad_box[0][0], quad_box[0][1], label, color='white', fontsize=8,
+                     bbox=dict(facecolor='red', alpha=0.5))
     else:
         bboxes = data.get('bboxes', [])
         labels = data.get('labels', [])
@@ -149,49 +152,60 @@ def draw_ocr_bboxes(image, prediction):
                   fill=color)
     return image
 
+def draw_bounding_boxes(image, quad_boxes, labels, color=(0, 255, 0), thickness=2):
+    """
+    Draws quadrilateral bounding boxes on the image.
+
+    Args:
+        image: The original image where the bounding boxes will be drawn.
+        quad_boxes: List of quadrilateral bounding box points. Each bounding box contains four points.
+        labels: List of labels corresponding to each bounding box.
+        color: Color of the bounding box. Default is green.
+        thickness: Thickness of the bounding box lines. Default is 2.
+    """
+    for i, quad in enumerate(quad_boxes):
+        points = np.array(quad, dtype=np.int32).reshape((-1, 1, 2))  # Reshape the quad points for drawing
+        image = cv2.polylines(image, [points], isClosed=True, color=color, thickness=thickness)
+        # Add label text near the top-left point of the bounding box
+        label_pos = (int(quad[0]), int(quad[1]) - 10)  # Positioning label slightly above the bounding box
+        cv2.putText(image, labels[i], label_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, thickness)
+
+    return image
+
 def process_image(image, task):
     prompt = TASK_PROMPTS[task]
-    
     # Print the inputs for debugging
     print(f"\n--- Processing Task: {task} ---")
     print(f"Prompt: {prompt}")
-    
     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
-    
     # Print the input tensors for debugging
     print(f"Model Input: {inputs}")
-
     generated_ids = model.generate(
         **inputs,
         max_new_tokens=1024,
         num_beams=3,
         do_sample=False
     )
-    
+
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    
     # Print the raw generated output for debugging
     print(f"Raw Model Output: {generated_text}")
-    
     parsed_answer = processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
-    
     # Print the parsed answer for debugging
     print(f"Parsed Answer: {parsed_answer}")
-
     return parsed_answer
 
+
 def main_process(image, task):
     result = process_image(image, task)
 
     if task in IMAGE_TASKS:
-        if task == "OCR with Region":
+        if task == "📸✍🏻OCR with Region":
             fig = plot_bbox(image, result.get('<OCR_WITH_REGION>', {}), use_quad_boxes=True)
             output_image = fig_to_pil(fig)
             text_output = result.get('<OCR_WITH_REGION>', {}).get('recognized_text', 'No text found')
-            
             # Debugging: Print the recognized text
             print(f"Recognized Text: {text_output}")
-            
             return output_image, gr.update(visible=True), text_output, gr.update(visible=True)
         else:
             fig = plot_bbox(image, result.get(TASK_PROMPTS[task], {}))
@@ -219,6 +233,8 @@ with gr.Blocks(title="PLeIAs/📸📈✍🏻Florence-PDF") as iface:
     output_image = gr.Image(label="PLeIAs/📸📈✍🏻Florence-PDF", visible=False)
     output_text = gr.Textbox(label="PLeIAs/📸📈✍🏻Florence-PDF", visible=True)
     
+    gr.Markdown(model_presentation)
+
     def process_and_update(image, task):
         if image is None:
             return None, gr.update(visible=False), "Please upload an image first.", gr.update(visible=True)

requirements.txt

@@ -3,4 +3,5 @@ transformers
 accelerate
 pillow
 einops
-timm
+timm
+opencv-python