Upload app.py

app.py

@@ -75,7 +75,7 @@ def run_sam(predictor, selected_points):
         best_mask = cv2.erode(best_mask, kernel, iterations=1)
     return best_mask
 
-def apply_mask_overlay(image, mask):
+def apply_mask_overlay(image, mask, color=(255, 0, 0)):
     """
     在原图上叠加 mask：使用红色绘制 mask 的轮廓，非 mask 区域叠加浅灰色半透明遮罩。
     """
@@ -85,7 +85,7 @@ def apply_mask_overlay(image, mask):
     non_mask = mask == 0
     overlay[non_mask] = (0.5 * overlay[non_mask] + 0.5 * gray_color).astype(np.uint8)
     contours, _ = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    cv2.drawContours(overlay, contours, -1, (255, 0, 0), 2)
+    cv2.drawContours(overlay, contours, -1, color, 2)
     return overlay
 
 def segment_and_overlay(image, points, sam_predictor):
@@ -401,6 +401,42 @@ def delete_mask(mask_list):
     return mask_list
 
 
+def apply_combined_mask_overlay(image, vis_mask, occ_mask):
+    """
+    在原图上叠加 mask：使用红色绘制 mask 的轮廓，非 mask 区域叠加浅灰色半透明遮罩。
+    """
+    img_arr = image
+    overlay = img_arr.copy()
+    gray_color = np.array([200, 200, 200], dtype=np.uint8)
+    non_mask = (vis_mask == 0) & (occ_mask == 0)
+    overlay[non_mask] = (0.5 * overlay[non_mask] + 0.5 * gray_color).astype(np.uint8)
+    contours_occ, _ = cv2.findContours(occ_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    cv2.drawContours(overlay, contours_occ, -1, (255,0,0), 2)
+    contours_vis, _ = cv2.findContours(vis_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    cv2.drawContours(overlay, contours_vis, -1, (255,0,0), 2)
+    return overlay
+
+
+def combine_mask(image, visible_mask_list, occlusion_mask_list):
+    combined_vis_mask = np.zeros_like(image[:, :, 0])
+    combined_occ_mask = np.zeros_like(image[:, :, 0])
+    combined_mask = np.zeros_like(image[:, :, 0])
+    for mask in visible_mask_list:
+        combined_vis_mask = cv2.bitwise_or(combined_mask, mask)
+    for mask in occlusion_mask_list:
+        combined_occ_mask = cv2.bitwise_or(combined_mask, mask)
+    overlay = apply_combined_mask_overlay(image, combined_vis_mask, combined_occ_mask)
+    # 5*5 kernel dilate for occlusion mask
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    combined_occ_mask = cv2.dilate(combined_occ_mask, kernel, iterations=1)
+    combined_mask[combined_occ_mask > 0] = 128
+    combined_mask[combined_vis_mask > 0] = 255
+    # concat the mask and overlay to be a single image
+    print(overlay.shape, combined_mask.shape)
+    result = cv2.hconcat([overlay, combined_mask[..., None].repeat(3, axis=-1)])
+    return result
+
+
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
     ## 3D Amodal Reconstruction with [Amodal3R](https://sm0kywu.github.io/Amodal3R/)          
@@ -557,6 +593,11 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         inputs=[original_image, visibility_mask_list],
         outputs=[visible_mask]
     )
+    undo_vis_mask.click(
+        delete_mask,
+        inputs=[visibility_mask_list],
+        outputs=[visibility_mask_list]
+    )
     gen_occ_mask.click(
         segment_and_overlay,
         inputs=[original_image, occlusion_points_state, predictor],
@@ -572,6 +613,18 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         inputs=[original_image, occlusion_mask_list],
         outputs=[occluded_mask]
     )
+    undo_occ_mask.click(
+        delete_mask,
+        inputs=[occlusion_mask_list],
+        outputs=[occlusion_mask_list]
+    )
+
+    # check combined mask
+    check_combine_button.click(
+        combine_mask,
+        inputs=[original_image, visibility_mask_list, occlusion_mask_list],
+        outputs=[combined_mask]
+    )
     
 
 # 启动 Gradio App