Upload app.py

app.py

@@ -129,8 +129,9 @@ def reset_image(predictor, img):
       - 返回原图
     """
     predictor.set_image(img)
-    # 返回predictor，原始图像
-    return predictor, img
+    original_img = img.copy()
+    # 返回predictor，visible occlusion mask初始化, 原始图像
+    return predictor, img, img, original_img
 
 def button_clickable(selected_points):
     if len(selected_points) > 0:
@@ -372,8 +373,8 @@ 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)
+    color = (255, 0, 0) if point_type == "visible" else (0, 255, 0)
+    cv2.circle(image_with_points, (int(x), int(y)), radius=10, color=color, thickness=-1)
     return image_with_points
 
 
@@ -409,6 +410,70 @@ def delete_point(point_type, visible_points, occlusion_points):
     return visible_points, occlusion_points
 
 
+def clear_all_points(image):
+    """
+    清除所有点：返回原图、空的 visible 和 occlusion 列表，
+    以及更新后的点文本信息和空下拉菜单列表。
+    """
+    updated_image = image.copy()
+    empty_list = []
+    text = "Visible Points: []\nOcclusion Points: []"
+    return updated_image, empty_list, empty_list, text, [], []
+
+def see_visible_points(image, visible_points):
+    """
+    在图像上绘制所有 visible 点（红色）。
+    """
+    updated_image = image.copy()
+    for p in visible_points:
+        cv2.circle(updated_image, (int(p[0]), int(p[1])), radius=10, color=(255, 0, 0), thickness=-1)
+    return updated_image
+
+def see_occlusion_points(image, occlusion_points):
+    """
+    在图像上绘制所有 occlusion 点（绿色）。
+    """
+    updated_image = image.copy()
+    for p in occlusion_points:
+        cv2.circle(updated_image, (int(p[0]), int(p[1])), radius=10, color=(0, 255, 0), thickness=-1)
+    return updated_image
+
+def update_all_points(visible_points, occlusion_points):
+    """
+    返回一个文本描述当前两个点列表，以及对应的下拉菜单选项（字符串格式）。
+    """
+    text = f"Visible Points: {visible_points}\nOcclusion Points: {occlusion_points}"
+    visible_dropdown = [f"({p[0]}, {p[1]})" for p in visible_points]
+    occlusion_dropdown = [f"({p[0]}, {p[1]})" for p in occlusion_points]
+    return text, visible_dropdown, occlusion_dropdown
+
+def delete_selected_visible(image, visible_points, occlusion_points, selected_index):
+    """
+    删除 visible 点列表中指定索引的点，并更新图像和显示信息。
+    """
+    if selected_index is not None and 0 <= selected_index < len(visible_points):
+        visible_points.pop(selected_index)
+    updated_image = image.copy()
+    # 重新绘制所有 visible 点（红色）
+    for p in visible_points:
+        cv2.circle(updated_image, (int(p[0]), int(p[1])), radius=10, color=(255, 0, 0), thickness=-1)
+    updated_text, vis_dropdown, occ_dropdown = update_all_points(visible_points, occlusion_points)
+    return updated_image, visible_points, occlusion_points, updated_text, vis_dropdown, occ_dropdown
+
+def delete_selected_occlusion(image, visible_points, occlusion_points, selected_index):
+    """
+    删除 occlusion 点列表中指定索引的点，并更新图像和显示信息。
+    """
+    if selected_index is not None and 0 <= selected_index < len(occlusion_points):
+        occlusion_points.pop(selected_index)
+    updated_image = image.copy()
+    # 重新绘制所有 occlusion 点（绿色）
+    for p in occlusion_points:
+        cv2.circle(updated_image, (int(p[0]), int(p[1])), radius=10, color=(0, 255, 0), thickness=-1)
+    updated_text, vis_dropdown, occ_dropdown = update_all_points(visible_points, occlusion_points)
+    return updated_image, visible_points, occlusion_points, updated_text, vis_dropdown, occ_dropdown
+
+
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
     ## 3D Amodal Reconstruction with [Amodal3R](https://sm0kywu.github.io/Amodal3R/)
@@ -423,6 +488,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     predictor = gr.State(value=get_sam_predictor())
     visible_points_state = gr.State(value=[])
     occlusion_points_state = gr.State(value=[])
+    original_image = gr.State(value=None)
 
 
     with gr.Row():
@@ -436,38 +502,83 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             with gr.Row():
                 see_button = gr.Button("See")
                 add_button = gr.Button("Add")
-                delete_button = gr.Button("Delete")
+            with gr.Row():
+                # 新增按钮：Clear、分别查看 visible/occlusion
+                clear_button = gr.Button("Clear Points")
+                see_visible_button = gr.Button("See Visible Points")
+                see_occlusion_button = gr.Button("See Occlusion Points")
+            with gr.Row():
+                # 新增文本框实时显示点列表
+                points_text = gr.Textbox(label="Points List", interactive=False)
+            with gr.Row():
+                # 新增下拉菜单，用户可选择需要删除的点
+                visible_points_dropdown = gr.Dropdown(label="Select Visible Point to Delete", choices=[], interactive=True)
+                occlusion_points_dropdown = gr.Dropdown(label="Select Occlusion Point to Delete", choices=[], interactive=True)
+            with gr.Row():
+                delete_visible_button = gr.Button("Delete Selected Visible")
+                delete_occlusion_button = gr.Button("Delete Selected Occlusion")
         with gr.Column():
-            # 显示 SAM 分割结果（带 overlay）—— 使用 AnnotatedImage 显示更直观
-            sam_image = gr.Image(label='SAM Generated Mask', interactive=False, height=300)
-
-    
-    # 会话启动与结束
-    demo.load(start_session)
-    demo.unload(end_session)
+            # 用于显示 SAM 分割结果
+            visible_mask = gr.Image(label='Visible Mask', interactive=False, height=300)
+            occlusion_mask = gr.Image(label='Occlusion Mask', interactive=False, height=300)
     
-    # 上传图像时：重置 predictor 并将原图赋值给 original_image、preprocessed_image、selected_points 以及 output_mask
+    # ---------------------------
+    # 原有交互逻辑（略）
+    # ---------------------------
     input_image.upload(
         reset_image,
         [predictor, input_image],
-        [predictor, sam_image]
+        [predictor, visible_mask, occlusion_mask, original_image]
     )
-    # 如果点击see按钮，应该在input图片上生成对应的点，
     see_button.click(
         see_point, 
         inputs=[input_image, x_input, y_input, point_type], 
         outputs=[input_image]
     )
-    # 如果点击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]
     )
-    delete_button.click(
-        delete_point, 
-        inputs=[point_type, visible_points_state, occlusion_points_state], 
-        outputs=[visible_points_state, occlusion_points_state]
+    
+    # ---------------------------
+    # 新增的交互逻辑
+    # ---------------------------
+    clear_button.click(
+        clear_all_points,
+        inputs=[original_image],
+        outputs=[input_image, visible_points_state, occlusion_points_state, points_text, visible_points_dropdown, occlusion_points_dropdown]
+    )
+    see_visible_button.click(
+        see_visible_points,
+        inputs=[input_image, visible_points_state],
+        outputs=input_image
+    )
+    see_occlusion_button.click(
+        see_occlusion_points,
+        inputs=[input_image, occlusion_points_state],
+        outputs=input_image
+    )
+    # 当 visible_points_state 或 occlusion_points_state 变化时，更新文本框和下拉菜单
+    visible_points_state.change(
+        update_all_points,
+        inputs=[visible_points_state, occlusion_points_state],
+        outputs=[points_text, visible_points_dropdown, occlusion_points_dropdown]
+    )
+    occlusion_points_state.change(
+        update_all_points,
+        inputs=[visible_points_state, occlusion_points_state],
+        outputs=[points_text, visible_points_dropdown, occlusion_points_dropdown]
+    )
+    delete_visible_button.click(
+        delete_selected_visible,
+        inputs=[input_image, visible_points_state, occlusion_points_state, visible_points_dropdown],
+        outputs=[input_image, visible_points_state, occlusion_points_state, points_text, visible_points_dropdown, occlusion_points_dropdown]
+    )
+    delete_occlusion_button.click(
+        delete_selected_occlusion,
+        inputs=[input_image, visible_points_state, occlusion_points_state, occlusion_points_dropdown],
+        outputs=[input_image, visible_points_state, occlusion_points_state, points_text, visible_points_dropdown, occlusion_points_dropdown]
     )