Update app.py

app.py

@@ -1,7 +1,7 @@
 import spaces
 import multiprocessing as mp
 import numpy as np
-from PIL import Image
+from PIL import Image,ImageDraw
 import torch
 try:
     import detectron2
@@ -32,6 +32,49 @@ def setup_cfg(config_file):
     cfg.freeze()
     return cfg
 
+class IMGState:
+    def __init__(self):
+        self.img = None
+        self.img_feat = None
+        self.selected_points = []
+        self.selected_points_labels = []
+        self.selected_bboxes = []
+
+        self.available_to_set = True
+
+    def set_img(self, img, img_feat):
+        self.img = img
+        self.img_feat = img_feat
+
+        self.available_to_set = False
+
+    def clear(self):
+        self.img = None
+        self.img_feat = None
+        self.selected_points = []
+        self.selected_points_labels = []
+        self.selected_bboxes = []
+
+        self.available_to_set = True
+
+    def clean(self):
+        self.selected_points = []
+        self.selected_points_labels = []
+        self.selected_bboxes = []
+
+    def to_device(self, device=torch.device("cuda")):
+        if self.img_feat is not None:
+            for k in self.img_feat:
+                if isinstance(self.img_feat[k], torch.Tensor):
+                    self.img_feat[k] = self.img_feat[k].to(device)
+                elif isinstance(self.img_feat[k], tuple):
+                    self.img_feat[k] = tuple(v.to(device) for v in self.img_feat[k])
+
+    @property
+    def available(self):
+        return self.available_to_set
+    
+    
 @spaces.GPU
 @torch.no_grad()
 @torch.autocast(device_type="cuda", dtype=torch.float32)
@@ -60,13 +103,12 @@ def inference_automatic(input_img, class_names):
 @spaces.GPU
 @torch.no_grad()
 @torch.autocast(device_type="cuda", dtype=torch.float32)
-def inference_point(input_img, evt: gr.SelectData,):
+def inference_point(input_img,  img_state,):
 
 
     mp.set_start_method("spawn", force=True)
     
-    x, y = evt.index[0], evt.index[1]
-    points = [[x, y]]
+    points = img_state.selected_points
     print(f"Selected point: {points}")
     
     config_file = './configs/ground-truth-warmup/mask-adapter/mask_adapter_convnext_large_cocopan_eval_ade20k.yaml'
@@ -74,10 +116,8 @@ def inference_point(input_img, evt: gr.SelectData,):
     
     demo = SAMPointVisualizationDemo(cfg, 0.8, sam2_model, clip_model,mask_adapter)
     
-    img = read_image(input_img, format="BGR")
-    
 
-    _, visualized_output = demo.run_on_image_with_points(img, points)
+    _, visualized_output = demo.run_on_image_with_points(img_state.img, points)
     return visualized_output
 
 
@@ -85,6 +125,21 @@ sam2_model = None
 clip_model = None
 mask_adapter = None
 
+def get_points_with_draw(image, img_state, evt: gr.SelectData):
+    label = 'Add Mask'
+
+    x, y = evt.index[0], evt.index[1]
+    point_radius, point_color = 10, (97, 217, 54) if label == "Add Mask" else (237, 34, 13)
+
+    img_state.selected_points.append([x, y])
+    img_state.selected_points_labels.append(1 if label == "Add Mask" else 0)
+    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,
+    )
+    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
@@ -107,6 +162,16 @@ def initialize_models(sam_path, adapter_pth, model_cfg, cfg):
         mask_adapter.load_state_dict(adapter_state_dict)
         print("Mask Adapter model initialized.")
 
+def clear_everything(img_state):
+    img_state.clear()
+    return img_state, None, None
+
+
+def clean_prompts(img_state):
+    img_state.clean()
+    return img_state, Image.fromarray(img_state.img), None
+
+# 初始化配置和模型
 model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
 sam_path = './sam2.1_hiera_large.pt'
 adapter_pth = './model_0279999_with_sem_new.pth'
@@ -162,21 +227,35 @@ with gr.Blocks() as demo:
                 gr.Examples(examples=examples, inputs=[input_image, class_names], outputs=output_image)
 
         with gr.TabItem("Point Mode"):
-            with gr.Row():  
-                with gr.Column(): 
-                    def init_state():
-                        return []
-                    input_image = gr.Image(type='filepath', label="Upload Image", interactive=True)
-                    points_input = gr.State(value=init_state()) 
-
-                with gr.Column(): 
-                    output_image_point = gr.Image(type="pil", label='Segmentation Map')
-
-            input_image.select(inference_point, inputs=[input_image], outputs=output_image_point)
-
-            clear_button_point = gr.Button("Clear Segmentation Map")
-            clear_button_point.click(lambda: None, inputs=None, outputs=output_image_point)
-        
+            img_state_points = 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_point = gr.Image(type="pil", label='Segmentation Map',interactive=False)  # 输出分割图
+
+            input_image.select(
+                get_points_with_draw,
+                [input_image, img_state_points],
+                outputs=[img_state_points, input_image]
+            ).then(
+                inference_point,
+                inputs=[input_image, img_state_points],
+                outputs=[output_image_point]
+            )
+            clear_prompt_button_point = gr.Button("Clean Prompt")
+            clear_prompt_button_point.click(
+                clean_prompts,
+                inputs=[img_state_points],
+                outputs=[img_state_points, input_image, output_image_point]
+            )
+            clear_button_point = gr.Button("Restart")
+            clear_button_point.click(
+                clear_everything,
+                inputs=[img_state_points],
+                outputs=[img_state_points, input_image, output_image_point]
+            )
+