Upload app.py

app.py

@@ -51,7 +51,7 @@ def button_clickable(selected_points):
     else:
         return gr.Button.update(interactive=False)
 
-def run_sam(image, predictor, selected_points):
+def run_sam(predictor, selected_points):
     """
     调用 SAM 模型进行分割。
     """
@@ -70,11 +70,10 @@ def run_sam(image, predictor, selected_points):
     best_mask = masks[0].astype(np.uint8)
     # dilate
     if len(input_points) > 1:
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
         best_mask = cv2.dilate(best_mask, kernel, iterations=1)
         best_mask = cv2.erode(best_mask, kernel, iterations=1)
-    visible_mask = 255 * best_mask
-    return visible_mask
+    return best_mask
 
 def apply_mask_overlay(image, mask):
     """
@@ -93,8 +92,8 @@ def segment_and_overlay(image, points, sam_predictor):
     """
     调用 run_sam 获得 mask，然后叠加显示分割结果。
     """
-    visible_mask = run_sam(image, sam_predictor, points)
-    overlaid = apply_mask_overlay(image, visible_mask)
+    visible_mask = run_sam(sam_predictor, points)
+    overlaid = apply_mask_overlay(image, visible_mask * 255)
     return overlaid, visible_mask
 
 
@@ -375,15 +374,39 @@ def delete_selected_occlusion(image, visible_points, occlusion_points, selected_
     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 add_mask(mask, mask_list):
+    mask_list.append(mask)
+    return mask_list
+
+def vis_mask(image, mask_list):
+    updated_image = image.copy()
+    # combine all the mask:
+    combined_mask = np.zeros_like(updated_image)
+    for mask in mask_list:
+        combined_mask = cv2.bitwise_or(combined_mask, mask)
+    # overlay the mask on the image
+    updated_image = apply_mask_overlay(updated_image, combined_mask)
+    return updated_image
+
 
 with gr.Blocks(delete_cache=(600, 600)) as demo:
+    # gr.Markdown("""
+    # ## 3D Amodal Reconstruction with [Amodal3R](https://sm0kywu.github.io/Amodal3R/)
+    # * Upload an image and click "Generate" to create a 3D asset.
+    # * Target object selection. Multiple point prompts are supported until you get the ideal visible area.
+    # * Occluders selection, this can be done by squential point prompts. You can choose "all occ", then all the other areas except the target object will be treated as occluders.
+    # * Different random seeds can be tried in "Generation Settings", if you think the results are not ideal. 
+    # * If the reconstruction 3D asset is satisfactory, you can extract the GLB file and download it.
+    # """)
     gr.Markdown("""
     ## 3D Amodal Reconstruction with [Amodal3R](https://sm0kywu.github.io/Amodal3R/)
-    * Upload an image and click "Generate" to create a 3D asset.
-    * Target object selection. Multiple point prompts are supported until you get the ideal visible area.
-    * Occluders selection, this can be done by squential point prompts. You can choose "all occ", then all the other areas except the target object will be treated as occluders.
-    * Different random seeds can be tried in "Generation Settings", if you think the results are not ideal. 
-    * If the reconstruction 3D asset is satisfactory, you can extract the GLB file and download it.
+    * Step 1 - Generate Visibility Mask and Occlusion Mask.
+    * Please wait for a few seconds after uploading the image. The 2D segmenter is getting ready.
+    * Add the point prompts to indicate the target object and occluders separately.
+    * "Render Point", see the position of the point to be added.
+    * "Add Point", the point will be added to the list.
+    * "Generate mask", see the segmented area corresponding to current point list.
+    * "Add mask", current mask will be added for 3D amodal completion.          
     """)
 
      # 定义各状态变量
@@ -393,6 +416,8 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     original_image = gr.State(value=None)
     visibility_mask = gr.State(value=None)
     occlusion_mask = gr.State(value=None)
+    visibility_mask_list = gr.State(value=[])
+    occlusion_mask_list = gr.State(value=[])
 
 
     with gr.Row():
@@ -430,14 +455,14 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
                 add_vis_mask = gr.Button("Add Mask")
             with gr.Row():
                 render_vis_mask = gr.Button("Render Mask")
-                undo_vis_mask = gr.Button("Undo Occlusion Mask")
-            occlusion_mask = gr.Image(label='Occlusion Mask', interactive=False, height=300)
+                undo_vis_mask = gr.Button("Undo Last Mask")
+            occluded_mask = gr.Image(label='Occlusion Mask', interactive=False, height=300)
             with gr.Row():
                 gen_occ_mask = gr.Button("Generate Mask")
                 add_occ_mask = gr.Button("Add Mask")
             with gr.Row():
                 render_occ_mask = gr.Button("Render Mask")
-                undo_occ_mask = gr.Button("Undo Occlusion Mask")
+                undo_occ_mask = gr.Button("Undo Last Mask")
             # 
     
     # ---------------------------
@@ -505,6 +530,31 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         inputs=[original_image, visible_points_state, predictor],
         outputs=[visible_mask, visibility_mask]
     )
+    add_vis_mask.click(
+        add_mask,
+        inputs=[visibility_mask, visibility_mask_list],
+        outputs=[visibility_mask_list]
+    )
+    render_vis_mask.click(
+        vis_mask,
+        inputs=[original_image, visibility_mask_list],
+        outputs=[input_image]
+    )
+    gen_occ_mask.click(
+        segment_and_overlay,
+        inputs=[original_image, occlusion_points_state, predictor],
+        outputs=[occluded_mask, occlusion_mask]
+    )
+    add_occ_mask.click(
+        add_mask,
+        inputs=[occlusion_mask, occlusion_mask_list],
+        outputs=[occlusion_mask_list]
+    )
+    render_occ_mask.click(
+        vis_mask,
+        inputs=[original_image, occlusion_mask_list],
+        outputs=[input_image]
+    )
     
 
 # 启动 Gradio App