add quad boxes

app.py

@@ -118,13 +118,15 @@ def plot_bbox(image, data, use_quad_boxes=False):
     
     # Handle both 'bboxes' and 'quad_boxes'
     if use_quad_boxes:
-        for quad_box, label in zip(data['quad_boxes'], data['labels']):
+        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')
             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))
     else:
-        for bbox, label in zip(data['bboxes'], data['labels']):
+        bboxes = data.get('bboxes', [])
+        labels = data.get('labels', [])
+        for bbox, label in zip(bboxes, labels):
             x1, y1, x2, y2 = bbox
             rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=1, edgecolor='r', facecolor='none')
             ax.add_patch(rect)
@@ -183,7 +185,7 @@ def main_process(image, task):
 
     if task in IMAGE_TASKS:
         if task == "OCR with Region":
-            fig = plot_bbox(image, result['<OCR_WITH_REGION>'], use_quad_boxes=True)
+            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')
             
@@ -192,12 +194,13 @@ def main_process(image, task):
             
             return output_image, gr.update(visible=True), text_output, gr.update(visible=True)
         else:
-            fig = plot_bbox(image, result[TASK_PROMPTS[task]])
+            fig = plot_bbox(image, result.get(TASK_PROMPTS[task], {}))
             output_image = fig_to_pil(fig)
             return output_image, gr.update(visible=True), None, gr.update(visible=False)
     else:
         return None, gr.update(visible=False), str(result), gr.update(visible=True)
 
+
 def reset_outputs():
     return None, gr.update(visible=False), None, gr.update(visible=True)
 
@@ -207,7 +210,7 @@ with gr.Blocks(title="PLeIAs/📸📈✍🏻Florence-PDF") as iface:
     
     with gr.Column():
         image_input = gr.Image(type="pil", label="Input Image")
-        task_dropdown = gr.Dropdown(list(TASK_PROMPTS.keys()), label="Task", value="Caption")
+        task_dropdown = gr.Dropdown(list(TASK_PROMPTS.keys()), label="Task", value="✍🏻Caption")
     
     with gr.Row():
         submit_button = gr.Button("Process")