Update app.py

app.py

@@ -140,8 +140,8 @@ mask_adapter = None
 @torch.no_grad()
 @torch.autocast(device_type="cuda", dtype=torch.float32)
 def inference_box(input_img,  img_state,class_names_input):
-    if len(img_state.selected_bboxes) != 2:
-        return None
+    # if len(img_state.selected_bboxes) != 2:
+    #     return None
     mp.set_start_method("spawn", force=True)
     
     box_points = img_state.selected_bboxes
@@ -239,6 +239,12 @@ def get_bbox_with_draw(image, img_state, evt: gr.SelectData):
         )
     return img_state, image
 
+def check_and_infer_box(img_state_bbox, input_image, class_names_input_box):
+    if len(img_state_bbox.selected_bboxes) == 2:
+        return inference_box(input_image, img_state_bbox, class_names_input_box)
+    return None
+
+
 def initialize_models(sam_path, adapter_pth, model_cfg, cfg):
     cfg = setup_cfg(cfg)
     global sam2_model, clip_model, mask_adapter
@@ -356,7 +362,7 @@ with gr.Blocks() as demo:
                 [input_image, img_state_bbox],
                 outputs=[img_state_bbox, input_image]
             ).then(
-                inference_box,
+                check_and_infer_box,
                 inputs=[input_image, img_state_bbox,class_names_input_box],
                 outputs=[output_image_box]
             )