Update app.py

app.py

@@ -35,22 +35,19 @@ def setup_cfg(config_file):
 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):
+    def set_img(self, img):
         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 = []
@@ -62,13 +59,6 @@ class IMGState:
         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):
@@ -160,15 +150,27 @@ def get_points_with_draw(image, img_state, evt: gr.SelectData):
 
     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)
+    if img_state.img is None:
+        img_state.set_img(np.array(image))
+        
     draw = ImageDraw.Draw(image)
-    draw.ellipse(
-        [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
+    draw.polygon(
+        [
+            (x, y - point_radius),
+            (x + point_radius * 0.25, y - point_radius * 0.25),
+            (x + point_radius, y),
+            (x + point_radius * 0.25, y + point_radius * 0.25),
+            (x, y + point_radius),
+            (x - point_radius * 0.25, y + point_radius * 0.25),
+            (x - point_radius, y),
+            (x - point_radius * 0.25, y - point_radius * 0.25)
+        ],
         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)
@@ -180,7 +182,9 @@ def get_bbox_with_draw(image, img_state, evt: gr.SelectData):
         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)
+    
+    if img_state.img is None:
+        img_state.set_img(np.array(image))
 
     draw = ImageDraw.Draw(image)
     draw.ellipse(
@@ -250,9 +254,10 @@ examples = [
 
 examples_point = [
     ['./demo/images/ADE_val_00000739.jpg'],
-    ['./demo/images/000000290833.jpg'],
-    ['./demo/images/2010_001315.jpg'],
+    ['./demo/images/000000052462.jpg'],
+    ['./demo/images/000000081766.jpg'],
     ['./demo/images/ADE_val_00000001.jpg'],
+    ['./demo/images/000000033707.jpg'],
     ['./demo/images/000000000785.jpg'],
 ]
 
@@ -274,8 +279,8 @@ Mask-Adapter effectively extends to SAM or SAM-2 without additional training, ac
 """
 
 with gr.Blocks() as demo:
-    gr.Markdown(title)  # Title
-    gr.Markdown(description)  # Description
+    gr.Markdown(title)  
+    gr.Markdown(description) 
 
     with gr.Tabs():
         with gr.TabItem("Automatic Mode"):
@@ -303,6 +308,8 @@ with gr.Blocks() as demo:
                     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)  # 输出分割图
+            
+            gr.Markdown("Click the top-left and bottom-right corners of the image to select a rectangular area") 
 
             input_image.select(
                 get_bbox_with_draw,
@@ -325,9 +332,24 @@ with gr.Blocks() as demo:
                 inputs=[img_state_bbox],
                 outputs=[img_state_bbox, input_image, output_image_box]
             )
+            input_image.clear(
+                clear_everything,
+                inputs=[img_state_bbox],
+                outputs=[img_state_bbox, input_image, output_image_box]
+            )
+            output_image_box.clear(
+                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)
+            
+            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())
@@ -358,9 +380,22 @@ with gr.Blocks() as demo:
                 inputs=[img_state_points],
                 outputs=[img_state_points, input_image, output_image_point]
             )
-
+            input_image.clear(
+                clear_everything,
+                inputs=[img_state_points],
+                outputs=[img_state_points, input_image, output_image_point]
+            )
+            output_image_point.clear(
+                clear_everything,
+                inputs=[img_state_points],
+                outputs=[img_state_points, input_image, output_image_point]
+            )
+            def clear_and_set_example_point(example):
+                clear_everything(img_state_points)
+                return example
             with gr.Row():
-                gr.Examples(examples=examples_point, inputs=input_image, outputs=output_image_point,examples_per_page=5)
+                gr.Examples(examples=examples_point, inputs=input_image, outputs=output_image_point,examples_per_page=5, preprocess=clear_and_set_example_point)
+