Update app.py

app.py

@@ -125,6 +125,33 @@ sam2_model = None
 clip_model = None
 mask_adapter = None
 
+@spaces.GPU
+@torch.no_grad()
+@torch.autocast(device_type="cuda", dtype=torch.float32)
+def inference_box(input_img,  img_state,):
+
+
+    mp.set_start_method("spawn", force=True)
+    
+    box_points = img_state.selected_bboxes
+    bbox = (
+        min(box_points[0][0], box_points[1][0]),
+        min(box_points[0][1], box_points[1][1]),
+        max(box_points[0][0], box_points[1][0]),
+        max(box_points[0][1], box_points[1][1]),
+    )
+    bbox = np.array(bbox)
+    config_file = './configs/ground-truth-warmup/mask-adapter/mask_adapter_convnext_large_cocopan_eval_ade20k.yaml'
+    cfg = setup_cfg(config_file)
+    
+    demo = SAMPointVisualizationDemo(cfg, 0.8, sam2_model, clip_model,mask_adapter)
+    
+    text_features = torch.from_numpy(np.load("./text_embedding/lvis_coco_text_embedding.npy")).cuda()
+    _, visualized_output = demo.run_on_image_with_boxes(img_state.img, bbox,text_features)
+    return visualized_output
+
+
+
 def get_points_with_draw(image, img_state, evt: gr.SelectData):
     label = 'Add Mask'
 
@@ -140,6 +167,41 @@ def get_points_with_draw(image, img_state, evt: gr.SelectData):
         fill=point_color,
     )
     return img_state, image
+
+def get_bbox_with_draw(image, img_state, evt: gr.SelectData):
+    x, y = evt.index[0], evt.index[1]
+    point_radius, point_color, box_outline = 5, (237, 34, 13), 2
+    box_color = (237, 34, 13)
+
+    if len(img_state.selected_bboxes) in [0, 1]:
+        img_state.selected_bboxes.append([x, y])
+    elif len(img_state.selected_bboxes) == 2:
+        img_state.selected_bboxes = [[x, y]]
+        image = Image.fromarray(img_state.img)
+    else:
+        raise ValueError(f"Cannot be {len(img_state.selected_bboxes)}")
+    img_state.set_img(np.array(image), None)
+
+    draw = ImageDraw.Draw(image)
+    draw.ellipse(
+        [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
+        fill=point_color,
+    )
+
+    if len(img_state.selected_bboxes) == 2:
+        box_points = img_state.selected_bboxes
+        bbox = (min(box_points[0][0], box_points[1][0]),
+                min(box_points[0][1], box_points[1][1]),
+                max(box_points[0][0], box_points[1][0]),
+                max(box_points[0][1], box_points[1][1]),
+                )
+        draw.rectangle(
+            bbox,
+            outline=box_color,
+            width=box_outline
+        )
+    return img_state, image
+
 def initialize_models(sam_path, adapter_pth, model_cfg, cfg):
     cfg = setup_cfg(cfg)
     global sam2_model, clip_model, mask_adapter
@@ -234,6 +296,39 @@ with gr.Blocks() as demo:
             with gr.Row():
                 gr.Examples(examples=examples, inputs=[input_image, class_names], outputs=output_image)
 
+        with gr.TabItem("Box Mode"):
+            img_state_bbox = gr.State(value=IMGState())
+            with gr.Row():  # 水平排列
+                with gr.Column(scale=1): 
+                    input_image = gr.Image( label="Input Image", type="pil") 
+                with gr.Column(scale=1):  # 第二列：分割图输出
+                    output_image_box = gr.Image(type="pil", label='Segmentation Map',interactive=False)  # 输出分割图
+
+            input_image.select(
+                get_bbox_with_draw,
+                [input_image, img_state_bbox],
+                outputs=[img_state_bbox, input_image]
+            ).then(
+                inference_box,
+                inputs=[input_image, img_state_bbox],
+                outputs=[output_image_box]
+            )
+            clear_prompt_button_box = gr.Button("Clean Prompt")
+            clear_prompt_button_box.click(
+                clean_prompts,
+                inputs=[img_state_bbox],
+                outputs=[img_state_bbox, input_image, output_image_box]
+            )
+            clear_button_box = gr.Button("Restart")
+            clear_button_box.click(
+                clear_everything,
+                inputs=[img_state_bbox],
+                outputs=[img_state_bbox, input_image, output_image_box]
+            )
+
+            with gr.Row():
+                gr.Examples(examples=examples_point, inputs=input_image, outputs=output_image_box,examples_per_page=5)
+                
         with gr.TabItem("Point Mode"):
             img_state_points = gr.State(value=IMGState())
             with gr.Row():  # 水平排列