Upload app.py

app.py

@@ -69,9 +69,10 @@ def run_sam(predictor, selected_points):
     )
     best_mask = masks[0].astype(np.uint8)
     # dilate
-    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)
+    if len(selected_points) > 1:
+        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)
     return best_mask
 
 def apply_mask_overlay(image, mask, color=(255, 0, 0)):
@@ -105,9 +106,8 @@ def reset_points():
 
 @spaces.GPU
 def image_to_3d(
-    image: Image.Image,
-    multiimages: List[tuple],
-    is_multiimage: bool,
+    image: List[tuple],
+    masks: List[np.ndarray],
     seed: int,
     ss_guidance_strength: float,
     ss_sampling_steps: int,
@@ -120,37 +120,22 @@ def image_to_3d(
     将图像转换为 3D 模型。
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    if not is_multiimage:
-        outputs = pipeline.run(
-            image,
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": ss_sampling_steps,
-                "cfg_strength": ss_guidance_strength,
-            },
-            slat_sampler_params={
-                "steps": slat_sampling_steps,
-                "cfg_strength": slat_guidance_strength,
-            },
-        )
-    else:
-        outputs = pipeline.run_multi_image(
-            [img[0] for img in multiimages],
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": ss_sampling_steps,
-                "cfg_strength": ss_guidance_strength,
-            },
-            slat_sampler_params={
-                "steps": slat_sampling_steps,
-                "cfg_strength": slat_guidance_strength,
-            },
-            mode=multiimage_algo,
-        )
+    outputs = pipeline.run_multi_image(
+        [img[0] for img in image],
+        [mask[0] for mask in masks],
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+        mode=multiimage_algo,
+    )
     video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
     video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
     video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
@@ -270,48 +255,40 @@ def get_sam_predictor():
     return sam_predictor
 
 
-def draw_points_on_image(image, point, point_type):
+def draw_points_on_image(image, point):
     """在图像上绘制所有点，points 为 [(x, y, point_type), ...]"""
     image_with_points = image.copy()
     x, y = point
-    color = (255, 0, 0) if point_type == "vis" else (0, 255, 0)
+    color = (255, 0, 0)
     cv2.circle(image_with_points, (int(x), int(y)), radius=10, color=color, thickness=-1)
     return image_with_points
 
 
-def see_point(image, x, y, point_type):
+def see_point(image, x, y):
     """
     see操作：不修改 points 列表，仅在图像上临时显示这个点，
     并返回更新后的图像和当前列表（不更新）。
     """
     # 复制当前列表，并在副本中加上新点（仅用于显示）
-    updated_image = draw_points_on_image(image, [x,y], point_type)
+    updated_image = draw_points_on_image(image, [x,y])
     return updated_image
 
-def add_point(x, y, point_type, visible_points, occlusion_points):
+def add_point(x, y, visible_points):
     """
     add操作：将新点添加到 points 列表中，
     并返回更新后的图像和新的点列表。
     """
-    if point_type == "vis":
-        # check duplicate
-        if [x, y] not in visible_points:
-            visible_points.append([x, y])
-    else:
-        if [x, y] not in occlusion_points:
-            occlusion_points.append([x, y])
-    return visible_points, occlusion_points
+    if [x, y] not in visible_points:
+        visible_points.append([x, y])
+    return visible_points
 
-def delete_point(point_type, visible_points, occlusion_points):
+def delete_point(visible_points):
     """
     delete操作：删除 points 列表中的最后一个点，
     并返回更新后的图像和新的点列表。
     """
-    if point_type == "vis":
-        visible_points.pop()
-    else:
-        occlusion_points.pop()
-    return visible_points, occlusion_points
+    visible_points.pop()
+    return visible_points
 
 
 def clear_all_points(image):
@@ -331,23 +308,13 @@ def see_visible_points(image, 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}"
+def update_all_points(visible_points):
+    text = f"Points: {visible_points}"
     visible_dropdown_choices = [f"({p[0]}, {p[1]})" for p in visible_points]
-    occlusion_dropdown_choices = [f"({p[0]}, {p[1]})" for p in occlusion_points]
     # 返回更新字典来明确设置 choices 和 value
-    return text, gr.Dropdown(label="Select Visible Point to Delete", choices=visible_dropdown_choices, value=None, interactive=True), gr.Dropdown(label="Select Occlusion Point to Delete", choices=occlusion_dropdown_choices, value=None, interactive=True)
+    return text, gr.Dropdown(label="Select Visible Point to Delete", choices=visible_dropdown_choices, value=None, interactive=True)
 
-def delete_selected_visible(image, visible_points, occlusion_points, selected_value):
+def delete_selected_visible(image, visible_points, selected_value):
     # selected_value 是类似 "(x, y)" 的字符串
     try:
         selected_index = [f"({p[0]}, {p[1]})" for p in visible_points].index(selected_value)
@@ -359,22 +326,8 @@ def delete_selected_visible(image, visible_points, occlusion_points, selected_va
     # 重新绘制所有 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_value):
-    try:
-        selected_index = [f"({p[0]}, {p[1]})" for p in occlusion_points].index(selected_value)
-    except ValueError:
-        selected_index = None
-    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
+    updated_text, vis_dropdown, occ_dropdown = update_all_points(visible_points)
+    return updated_image, visible_points, updated_text, vis_dropdown, occ_dropdown
 
 def add_mask(mask, mask_list):
     # check if the mask if same as the last mask in the list 
@@ -399,43 +352,25 @@ def delete_mask(mask_list):
         mask_list.pop()
     return mask_list
 
+def check_combined_mask(image, mask_list, scale=1.7):
+    updated_image = image.copy()
+    # combine all the mask:
+    combined_mask = np.zeros_like(updated_image[:, :, 0])
+    occluded_mask = np.zeros_like(updated_image[:, :, 0])
+    for mask in mask_list:
+        combined_mask = cv2.bitwise_or(combined_mask, mask)
+    masked_image = updated_image * combined_mask[:, :, None]
+    occluded_mask[combined_mask == 1] = 127
 
-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, (0,0,255), 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
+    return masked_image, occluded_mask
 
 
-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)
-    # 添加 visible mask 边缘作为 occlusion 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
+
+def get_seed(randomize_seed: bool, seed: int) -> int:
+    """
+    Get the random seed.
+    """
+    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
 
 with gr.Blocks(delete_cache=(600, 600)) as demo:
@@ -449,9 +384,9 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     occlusion_points_state = gr.State(value=[])
     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=[])
+
+    combined_mask = gr.State(value=None)
 
 
     with gr.Row():
@@ -471,25 +406,19 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             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(["vis", "occ"], label="Point Prompt Type", value="vis")
             with gr.Row():
                 see_button = gr.Button("Render Point")
                 add_button = gr.Button("Add Point")
             with gr.Row():
-                # 新增按钮：Clear、分别查看 visible/occlusion
                 clear_button = gr.Button("Clear Points")
-                see_visible_button = gr.Button("Visible Points")
-                see_occlusion_button = gr.Button("Occluded Points")
+                see_visible_button = gr.Button("Render Added 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=[], value=None, interactive=True)
-                occlusion_points_dropdown = gr.Dropdown(label="Select Occlusion Point to Delete", choices=[], value=None, 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 分割结果
             visible_mask = gr.Image(label='Visible Mask', interactive=False, height=300)
@@ -499,25 +428,31 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             with gr.Row():
                 render_vis_mask = gr.Button("Render Mask")
                 undo_vis_mask = gr.Button("Undo Last Mask")
-            occluded_mask = gr.Image(label='Occlusion Mask', interactive=False, height=300)
+            mask_check = gr.Image(label='Visible Input', 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 Last Mask")
-            # 
+                check_visible_input = gr.Button("Check Combined Mask, make sure there is no GAP between the visible area (white) and occluded area (gray)")
     with gr.Row():
         gr.Markdown("""* Step 2 - 3D Amodal Completion.
-        * Please first check the obtained mask, and make sure there is no "GAP" between the visible area (white) and occluded area (gray).
         * 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.
         """)
     with gr.Row():
         with gr.Column():
-            combined_mask = gr.Image(label='Combined Mask', interactive=False, height=300)
-            with gr.Row():
-                check_combine_button = gr.Button("Check Combined Mask")
+            with gr.Accordion(label="Generation Settings", open=True):
+                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)
+            generate_btn = gr.Button("Generate")
+        with gr.Column():
+            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
+    
 
     # ---------------------------
     # 原有交互逻辑（略）
@@ -529,13 +464,13 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     )
     see_button.click(
         see_point, 
-        inputs=[original_image, x_input, y_input, point_type], 
+        inputs=[original_image, x_input, y_input], 
         outputs=[input_image]
     )
     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]
+        inputs=[x_input, y_input, visible_points_state], 
+        outputs=[visible_points_state]
     )
     
     # ---------------------------
@@ -551,31 +486,16 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         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]
+        inputs=[visible_points_state],
+        outputs=[points_text, visible_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]
+        inputs=[input_image, visible_points_state, visible_points_dropdown],
+        outputs=[input_image, visible_points_state, points_text, visible_points_dropdown]
     )
 
     # 生成mask的逻辑
@@ -599,33 +519,20 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         inputs=[visibility_mask_list],
         outputs=[visibility_mask_list]
     )
-    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=[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]
-    )
+    check_visible_input.click(
+    
+
+    # 3D Amodal Reconstruction
+    # generate_btn.click(
+    #     get_seed,
+    #     inputs=[randomize_seed, seed],
+    #     outputs=[seed],
+    # ).then(
+    #     image_to_3d,
+    #     inputs=[original_image, [combined_mask], seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, "multiimage"],
+    #     outputs=[visibility_mask]
+    # )
     
 
 # 启动 Gradio App