Upload app.py

app.py

@@ -33,16 +33,9 @@ def end_session(req: gr.Request):
     shutil.rmtree(user_dir)
 
 def reset_image(predictor, img):
-    """
-    上传图像后调用：
-      - 重置 predictor，
-      - 设置 predictor 的输入图像，
-      - 返回原图
-    """
     predictor.set_image(img)
     original_img = img.copy()
-    # 返回predictor，visible occlusion mask初始化, 原始图像
-    return predictor, original_img, "The models are ready."
+    return predictor, original_img, "The models are ready.", []
 
 def button_clickable(selected_points):
     if len(selected_points) > 0:
@@ -51,10 +44,6 @@ def button_clickable(selected_points):
         return gr.Button.update(interactive=False)
 
 def run_sam(predictor, selected_points):
-    """
-    调用 SAM 模型进行分割。
-    """
-    # predictor.set_image(image)
     if len(selected_points) == 0:
         return [], None
     input_points = [p for p in selected_points]
@@ -62,7 +51,7 @@ def run_sam(predictor, selected_points):
     masks, _, _ = predictor.predict(
         point_coords=np.array(input_points),
         point_labels=np.array(input_labels),
-        multimask_output=False,  # 单对象输出
+        multimask_output=False,
     )
     best_mask = masks[0].astype(np.uint8)
     # dilate
@@ -73,9 +62,6 @@ def run_sam(predictor, selected_points):
     return best_mask
 
 def apply_mask_overlay(image, mask, color=(255, 0, 0)):
-    """
-    在原图上叠加 mask：使用红色绘制 mask 的轮廓，非 mask 区域叠加浅灰色半透明遮罩。
-    """
     img_arr = image
     overlay = img_arr.copy()
     gray_color = np.array([200, 200, 200], dtype=np.uint8)
@@ -86,9 +72,6 @@ def apply_mask_overlay(image, mask, color=(255, 0, 0)):
     return overlay
 
 def segment_and_overlay(image, points, sam_predictor):
-    """
-    调用 run_sam 获得 mask，然后叠加显示分割结果。
-    """
     visible_mask = run_sam(sam_predictor, points)
     overlaid = apply_mask_overlay(image, visible_mask * 255)
     return overlaid, visible_mask
@@ -106,22 +89,6 @@ def image_to_3d(
     erode_kernel_size: int,
     req: gr.Request,
 ) -> Tuple[dict, str]:
-    """
-    Convert an image to a 3D model.
-    Args:
-        image (Image.Image): The input image.
-        multiimages (List[Tuple[Image.Image, str]]): The input images in multi-image mode.
-        is_multiimage (bool): Whether is in multi-image mode.
-        seed (int): The random seed.
-        ss_guidance_strength (float): The guidance strength for sparse structure generation.
-        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
-        slat_guidance_strength (float): The guidance strength for structured latent generation.
-        slat_sampling_steps (int): The number of sampling steps for structured latent generation.
-        multiimage_algo (Literal["multidiffusion", "stochastic"]): The algorithm for multi-image generation.
-    Returns:
-        dict: The information of the generated 3D model.
-        str: The path to the video of the 3D model.
-    """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     outputs = pipeline.run_multi_image(
         [image],
@@ -156,9 +123,6 @@ def extract_glb(
     texture_size: int,
     req: gr.Request,
 ) -> tuple:
-    """
-    从生成的 3D 模型中提取 GLB 文件。
-    """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh = unpack_state(state)
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
@@ -170,9 +134,6 @@ def extract_glb(
 
 @spaces.GPU
 def extract_gaussian(state: dict, req: gr.Request) -> tuple:
-    """
-    从生成的 3D 模型中提取 Gaussian 文件。
-    """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, _ = unpack_state(state)
     gaussian_path = os.path.join(user_dir, 'sample.ply')
@@ -229,7 +190,6 @@ def get_sam_predictor():
 
 
 def draw_points_on_image(image, point):
-    """在图像上绘制所有点，points 为 [(x, y, point_type), ...]"""
     image_with_points = image.copy()
     x, y = point
     color = (255, 0, 0)
@@ -238,44 +198,24 @@ def draw_points_on_image(image, point):
 
 
 def see_point(image, x, y):
-    """
-    see操作：不修改 points 列表，仅在图像上临时显示这个点，
-    并返回更新后的图像和当前列表（不更新）。
-    """
-    # 复制当前列表，并在副本中加上新点（仅用于显示）
     updated_image = draw_points_on_image(image, [x,y])
     return updated_image
 
 def add_point(x, y, visible_points):
-    """
-    add操作：将新点添加到 points 列表中，
-    并返回更新后的图像和新的点列表。
-    """
     if [x, y] not in visible_points:
         visible_points.append([x, y])
     return visible_points
 
 def delete_point(visible_points):
-    """
-    delete操作：删除 points 列表中的最后一个点，
-    并��回更新后的图像和新的点列表。
-    """
     visible_points.pop()
     return visible_points
 
 
 def clear_all_points(image):
-    """
-    清除所有点：返回原图、空的 visible 和 occlusion 列表，
-    以及更新后的点文本信息和空下拉菜单列表。
-    """
     updated_image = image.copy()
     return updated_image
 
 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)
@@ -284,11 +224,9 @@ def see_visible_points(image, visible_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]
-    # 返回更新字典来明确设置 choices 和 value
     return text, gr.Dropdown(label="Select Point to Delete", choices=visible_dropdown_choices, value=None, interactive=True)
 
 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)
     except ValueError:
@@ -296,14 +234,12 @@ def delete_selected_visible(image, visible_points, selected_value):
     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 = update_all_points(visible_points)
     return updated_image, visible_points, updated_text, vis_dropdown
 
-def add_mask(mask, mask_list):
-    # check if the mask if same as the last mask in the list 
+def add_mask(mask, mask_list): 
     if len(mask_list) > 0:
         if np.array_equal(mask, mask_list[-1]):
             return mask_list
@@ -312,11 +248,9 @@ def add_mask(mask, mask_list):
 
 def vis_mask(image, mask_list):
     updated_image = image.copy()
-    # combine all the mask:
     combined_mask = np.zeros_like(updated_image[:, :, 0])
     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
 
@@ -327,7 +261,6 @@ def delete_mask(mask_list):
 
 def check_combined_mask(image, visibility_mask, mask_list, scale=0.65):
     updated_image = image.copy()
-    # combine all the mask:
     combined_mask = np.zeros_like(updated_image[:, :, 0])
     occluded_mask = np.zeros_like(updated_image[:, :, 0])
     if len(mask_list) == 0:
@@ -345,7 +278,6 @@ def check_combined_mask(image, visibility_mask, mask_list, scale=0.65):
     masked_img = updated_image * combined_mask[:, :, None]
     occluded_mask[combined_mask == 1] = 127
 
-    # move the visible part to the center of the image
     x, y, w, h = cv2.boundingRect(combined_mask.astype(np.uint8))
     cropped_occluded_mask = (occluded_mask[y:y+h, x:x+w]).astype(np.uint8)
     cropped_img = masked_img[y:y+h, x:x+w]
@@ -383,7 +315,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     ## 3D Amodal Reconstruction with [Amodal3R](https://sm0kywu.github.io/Amodal3R/)          
     """)
 
-     # 定义各状态变量
     predictor = gr.State(value=get_sam_predictor())
     visible_points_state = gr.State(value=[])
     occlusion_points_state = gr.State(value=[])
@@ -466,9 +397,9 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
                 with gr.Row():
                     with gr.Column():
                         seed = gr.Slider(0, MAX_SEED, label="Seed", value=1, step=1)
-                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
                     with gr.Column():
-                        erode_kernel_size = gr.Slider(0, 5, label="Erode Kernel Size", value=0, step=1)
+                        erode_kernel_size = gr.Slider(0, 5, label="Erode Kernel Size", value=3, step=1)
                 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)
@@ -500,18 +431,15 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     demo.load(start_session)
     demo.unload(end_session)
 
-    # ---------------------------
-    # 原有交互逻辑（略）
-    # ---------------------------
     input_image.upload(
         reset_image,
         [predictor, input_image],
-        [predictor, original_image, message],
+        [predictor, original_image, message, visible_points_state],
     )
     apply_example_btn.click(
         reset_image,
         inputs=[predictor, input_image],
-        outputs=[predictor, original_image, message]
+        outputs=[predictor, original_image, message, visible_points_state]
     )
     see_button.click(
         see_point, 
@@ -524,9 +452,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[visible_points_state]
     )
     
-    # ---------------------------
-    # 新增的交互逻辑
-    # ---------------------------
     clear_button.click(
         clear_all_points,
         inputs=[original_image],
@@ -537,7 +462,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         inputs=[input_image, visible_points_state],
         outputs=input_image
     )
-    # 当 visible_points_state 或 occlusion_points_state 变化时，更新文本框和下拉菜单
+
     visible_points_state.change(
         update_all_points,
         inputs=[visible_points_state],
@@ -549,7 +474,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[input_image, visible_points_state, points_text, visible_points_dropdown]
     )
 
-    # 生成mask的逻辑
     gen_vis_mask.click(
         segment_and_overlay,
         inputs=[original_image, visible_points_state, predictor],
@@ -622,8 +546,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 
 
     
-
-# 启动 Gradio App
 if __name__ == "__main__":
     pipeline = Amodal3RImageTo3DPipeline.from_pretrained("Sm0kyWu/Amodal3R")
     pipeline.cuda()