Improve interactivity by allowing direct bounding box drawing and tool selection

app.py

@@ -129,11 +129,9 @@ def viz_pred_mask(img,
             else:
                 cv2.circle(out,(col, row), marker_size, (255,0,0), -1)
 
-    if bbox_coords is not None:
-        for i in range(len(bbox_coords)//2):
-            cv2.rectangle(out, bbox_coords[2*i], bbox_coords[2*i+1], (255,165,0), marker_size)
-        if len(bbox_coords) % 2 == 1:
-            cv2.circle(out, tuple(bbox_coords[-1]), marker_size, (255,165,0), -1)
+    if bbox_coords:
+        for bbox in bbox_coords:
+            cv2.rectangle(out, bbox[0], bbox[1], (255,165,0), 2)
 
     return out.astype(np.uint8)
 
@@ -242,7 +240,8 @@ def refresh_predictions(predictor, input_img, output_img, click_coords, click_la
 def get_select_coords(predictor, input_img, brush_label, bbox_label, best_mask, low_res_mask, 
                       click_coords, click_labels, bbox_coords,
                       seperate_scribble_masks, last_scribble_mask, scribble_img, img_features,
-                      output_img, binary_checkbox, multimask_mode, autopredict_checkbox, evt: gr.SelectData):
+                      output_img, binary_checkbox, multimask_mode, autopredict_checkbox,
+                      tool_select, evt: gr.SelectData):
     """
     Record user click and update the prediction
     """
@@ -255,6 +254,16 @@ def get_select_coords(predictor, input_img, brush_label, bbox_label, best_mask,
     else:
         raise TypeError("Invalid brush label: {brush_label}")
 
+    if tool_select == "Rectangle" and evt.bounding_box:
+        # User drew a rectangle
+        bbox_coords.append((evt.bounding_box[0], evt.bounding_box[1]))
+    elif tool_select == "Point" and evt.index is not None:
+        # User clicked a point
+        click_coords.append(evt.index)
+        click_labels.append(1 if brush_label == 'Positive (green)' else 0)
+    else:
+        gr.Error("Invalid selection")
+        
     # Only make new prediction if not waiting for additional bounding box click
     if (len(bbox_coords) % 2 == 0) and autopredict_checkbox:
 
@@ -402,15 +411,20 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
                 )
 
             with gr.Tab("Clicks/Boxes") as click_tab:
+                # Update click_img to be interactive and use the selected tool
                 click_img = gr.Image(
                     label="Input",
                     type='numpy',
                     value=default_example, 
+                    tool="select",  # Use 'select' to capture both points and rectangles
+                    interactive=True,
                     show_download_button=True,
                     container=True,
                     height=display_height
                 )
                 with gr.Row():
+                    # Add a tool selection radio button
+                    tool_select = gr.Radio(["Point", "Rectangle"], label="Tool", value="Point")
                     undo_click_button = gr.Button("Undo Last Click")
                     clear_click_button = gr.Button("Clear Clicks/Boxes", variant="stop") 
             
@@ -546,7 +560,8 @@ with gr.Blocks(theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg)) as de
                         input_img, brush_label, bbox_label, best_mask, low_res_mask, click_coords, click_labels, bbox_coords,
                         seperate_scribble_masks, last_scribble_mask, scribble_img, img_features,
                         output_img, binary_checkbox, multimask_mode, autopredict_checkbox
-                      ], 
+                        tool_select  # Add tool_select here
+                     ], 
                      outputs=[click_img, scribble_img, output_img, best_mask, low_res_mask, img_features,
                               click_coords, click_labels, bbox_coords, seperate_scribble_masks, last_scribble_mask],
                     api_name = "get_select_coords"