Upload app.py

app.py

@@ -126,12 +126,11 @@ def reset_image(predictor, img):
     上传图像后调用：
       - 重置 predictor，
       - 设置 predictor 的输入图像，
-      - 返回原图、预处理图像、清空 sel_pix、以及初始输出（无 mask）。
+      - 返回原图
     """
-    preprocessed_image = img
-    predictor.set_image(preprocessed_image)
-    # 返回原始图像、预处理图像、清空点列表、初始输出（作为 SAM mask 显示，初始为原图复制）
-    return img, preprocessed_image, [], (img.copy(), [(np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8), 'visible_mask')])
+    predictor.set_image(img)
+    # 返回predictor，原始图像
+    return predictor, img
 
 def button_clickable(selected_points):
     if len(selected_points) > 0:
@@ -369,6 +368,47 @@ def get_sam_predictor():
     return sam_predictor
 
 
+def draw_points_on_image(image, point, point_type):
+    """在图像上绘制所有点，points 为 [(x, y, point_type), ...]"""
+    image_with_points = image.copy()
+    x, y = point
+    color = (0, 0, 255) if point_type == "visible" else (0, 255, 0)
+    cv2.circle(image_with_points, (int(x), int(y)), radius=5, color=color, thickness=-1)
+    return image_with_points
+
+
+def see_point(image, x, y, point_type):
+    """
+    see操作：不修改 points 列表，仅在图像上临时显示这个点，
+    并返回更新后的图像和当前列表（不更新）。
+    """
+    # 复制当前列表，并在副本中加上新点（仅用于显示）
+    updated_image = draw_points_on_image(image, [x,y], point_type)
+    return updated_image, points
+
+def add_point(x, y, point_type, visible_points, occlusion_points):
+    """
+    add操作：将新点添加到 points 列表中，
+    并返回更新后的图像和新的点列表。
+    """
+    if point_type == "visible":
+        visible_points.append([x, y])
+    else:
+        occlusion_points.append([x, y])
+    return visible_points, occlusion_points
+
+def delete_point(point_type, visible_points, occlusion_points):
+    """
+    delete操作：删除 points 列表中的最后一个点，
+    并返回更新后的图像和新的点列表。
+    """
+    if point_type == "visible":
+        visible_points.pop()
+    else:
+        occlusion_points.pop()
+    return visible_points, occlusion_points
+
+
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
     ## 3D Amodal Reconstruction with [Amodal3R](https://sm0kywu.github.io/Amodal3R/)
@@ -381,51 +421,25 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 
      # 定义各状态变量
     predictor = gr.State(value=get_sam_predictor())
-    selected_points = gr.State(value=[])
-    original_image = gr.State(value=None)
-    preprocessed_image = gr.State(value=None)
-    visible_mask = gr.State(value=None)
+    visible_points_state = gr.State(value=[])
+    occlusion_points_state = gr.State(value=[])
 
 
     with gr.Row():
         with gr.Column():
-        #     # 上传的图像不经过预处理，直接展示原始图像
-        #     image_prompt = gr.Image(type="numpy", label="Input Occlusion Image", interactive=True, height=512)
-        #     # 用于交互标注的图像，点击时更新显示标记
-        #     # image_annotation = gr.Image(type="numpy", label="Select Point Prompts for Target Object", interactive=True, height=512)
-        #     # 存储点击点状态以及显示点击点坐标
-        #     points_state = gr.State([])
-        #     segment_button = gr.Button("Run Segmentation")
-        #     # points_output = gr.Textbox(label="Target Object Prompts", interactive=False)
-        #     # 展示 SAM 分割结果（只用于显示，不允许上传）
-        #     segmented_output = gr.Image(label="Segmented Result", height=512, interactive=False)
-            
-        #     with gr.Accordion(label="Generation Settings", open=False):
-        #         seed = gr.Slider(0, MAX_SEED, label="Seed", value=1, step=1)
-        #         randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-        #         gr.Markdown("Stage 1: Sparse Structure Generation")
-        #         with gr.Row():
-        #             ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
-        #             ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-        #         gr.Markdown("Stage 2: Structured Latent Generation")
-        #         with gr.Row():
-        #             slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
-        #             slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-        # # 其他组件（如生成按钮、视频展示、GLB 提取等）可根据需要添加\
-            input_image = gr.Image(type="numpy", label='Input Occlusion Image', height=500)
-            annotation_image = gr.Image(type="numpy", label='Annotate Image', interactive=True, height=500)
-            undo_button = gr.Button('Undo Prompt')
+            input_image = gr.Image(type="numpy", label='Input Occlusion Image', height=300)
             fg_bg_radio = gr.Radio(['positive_prompt', 'negative_prompt'], label='Point Prompt Type')
-            gr.Markdown('''
-            ### Instructions:
-            - First, upload an image.
-            - Then, click on the "Annotate Image" to select visible regions.
-            - Use "Undo Prompt" to remove the last point.
-            - Once the SAM mask is satisfactory, click "Run pix2gestalt" to perform amodal completion.
-            ''')
+            with gr.Row():
+                x_input = gr.Number(label="X Coordinate", value=0)
+                y_input = gr.Number(label="Y Coordinate", value=0)
+                point_type = gr.Radio(choices=["visible", "occlusion"], label="Point Type", value="visible")
+            with gr.Row():
+                see_button = gr.Button("See")
+                add_button = gr.Button("Add")
+                delete_button = gr.Button("Delete")
         with gr.Column():
             # 显示 SAM 分割结果（带 overlay）—— 使用 AnnotatedImage 显示更直观
-            output_mask = gr.AnnotatedImage(label='SAM Generated Visible (Modal) Mask', height=500)
+            sam_image = gr.Image(label='SAM Generated Mask', interactive=False, height=300)
 
     
     # 会话启动与结束
@@ -436,25 +450,24 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     input_image.upload(
         reset_image,
         [predictor, input_image],
-        [original_image, preprocessed_image, selected_points, output_mask]
+        [predictor, sam_image]
     )
-    # 同时更新 annotation_image（使其与上传图像保持一致）
-    input_image.upload(
-        lambda x: x, 
-        inputs=[input_image], 
-        outputs=[annotation_image]
+    # 如果点击see按钮，应该在input图片上生成对应的点，
+    see_button.click(
+        see_point, 
+        inputs=[input_image, x_input, y_input, point_type], 
+        outputs=[input_image]
     )
-    # 撤销按钮：撤销最近一次点击
-    undo_button.click(
-        undo_points,
-        [predictor, original_image, selected_points],
-        [annotation_image, output_mask, visible_mask]
+    # 如果点击add按钮，应该将对应的点添加到visible_points_state中
+    add_button.click(
+        add_point, 
+        inputs=[x_input, y_input, point_type, visible_points_state, occlusion_points_state], 
+        outputs=[visible_points_state, occlusion_points_state]
     )
-    # 在 annotation_image 上点击：调用 select_point 更新标注图像和 SAM 分割结果
-    annotation_image.select(
-        select_point,
-        [predictor, annotation_image, original_image, selected_points, fg_bg_radio],
-        [annotation_image, output_mask, visible_mask]
+    delete_button.click(
+        delete_point, 
+        inputs=[point_type, visible_points_state, occlusion_points_state], 
+        outputs=[visible_points_state, occlusion_points_state]
     )
     
 
@@ -466,4 +479,4 @@ if __name__ == "__main__":
         pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
     except:
         pass
-    demo.launch()
+    demo.launch()