import gradio as gr import spaces from gradio_litmodel3d import LitModel3D import os import shutil os.environ['SPCONV_ALGO'] = 'native' from typing import * import torch import numpy as np import imageio from easydict import EasyDict as edict from PIL import Image from Amodal3R.pipelines import Amodal3RImageTo3DPipeline from Amodal3R.representations import Gaussian, MeshExtractResult from Amodal3R.utils import render_utils, postprocessing_utils from segment_anything import sam_model_registry, SamPredictor from huggingface_hub import hf_hub_download import cv2 MAX_SEED = np.iinfo(np.int32).max TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp') os.makedirs(TMP_DIR, exist_ok=True) def start_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) os.makedirs(user_dir, exist_ok=True) def end_session(req: gr.Request): user_dir = os.path.join(TMP_DIR, str(req.session_hash)) 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." def button_clickable(selected_points): if len(selected_points) > 0: return gr.Button.update(interactive=True) else: 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] input_labels = [1 for _ in range(len(selected_points))] # input_points = np.array([[210, 300]]) # input_labels = np.array([1]) masks, _, _ = predictor.predict( point_coords=np.array(input_points), point_labels=np.array(input_labels), multimask_output=False, # 单对象输出 ) 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) 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) 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, color, 2) 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 def reset_points(): """ 清空点击点提示。 """ return [], "" @spaces.GPU def image_to_3d( image: List[tuple], masks: List[np.ndarray], seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, multiimage_algo: str, req: gr.Request, ) -> tuple: """ 将图像转换为 3D 模型。 """ user_dir = os.path.join(TMP_DIR, str(req.session_hash)) 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))] video_path = os.path.join(user_dir, 'sample.mp4') imageio.mimsave(video_path, video, fps=15) state = pack_state(outputs['gaussian'][0], outputs['mesh'][0]) torch.cuda.empty_cache() return state, video_path @spaces.GPU(duration=90) def extract_glb( state: dict, mesh_simplify: float, 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) glb_path = os.path.join(user_dir, 'sample.glb') glb.export(glb_path) torch.cuda.empty_cache() return glb_path, glb_path @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') gs.save_ply(gaussian_path) torch.cuda.empty_cache() return gaussian_path, gaussian_path def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict: return { 'gaussian': { **gs.init_params, '_xyz': gs._xyz.cpu().numpy(), '_features_dc': gs._features_dc.cpu().numpy(), '_scaling': gs._scaling.cpu().numpy(), '_rotation': gs._rotation.cpu().numpy(), '_opacity': gs._opacity.cpu().numpy(), }, 'mesh': { 'vertices': mesh.vertices.cpu().numpy(), 'faces': mesh.faces.cpu().numpy(), }, } def unpack_state(state: dict) -> tuple: gs = Gaussian( aabb=state['gaussian']['aabb'], sh_degree=state['gaussian']['sh_degree'], mininum_kernel_size=state['gaussian']['mininum_kernel_size'], scaling_bias=state['gaussian']['scaling_bias'], opacity_bias=state['gaussian']['opacity_bias'], scaling_activation=state['gaussian']['scaling_activation'], ) gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda') gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda') gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda') gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda') gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda') mesh = edict( vertices=torch.tensor(state['mesh']['vertices'], device='cuda'), faces=torch.tensor(state['mesh']['faces'], device='cuda'), ) return gs, mesh def prepare_multi_example() -> list: multi_case = list(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")])) images = [] for case in multi_case: _images = [] for i in range(1, 4): img = Image.open(f'assets/example_multi_image/{case}_{i}.png') W, H = img.size img = img.resize((int(W / H * 512), 512)) _images.append(np.array(img)) images.append(Image.fromarray(np.concatenate(_images, axis=1))) return images def split_image(image: Image.Image) -> list: """ 将图像拆分为多个视图（不进行预处理）。 """ image = np.array(image) alpha = image[..., 3] alpha = np.any(alpha > 0, axis=0) start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist() end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist() images = [] for s, e in zip(start_pos, end_pos): images.append(Image.fromarray(image[:, s:e+1])) return [image for image in images] def get_sam_predictor(): sam_checkpoint = hf_hub_download("ybelkada/segment-anything", "checkpoints/sam_vit_h_4b8939.pth") model_type = "vit_h" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) # sam.cuda() sam_predictor = SamPredictor(sam) return sam_predictor def draw_points_on_image(image, point, point_type): """在图像上绘制所有点，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) 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): """ see操作：不修改 points 列表，仅在图像上临时显示这个点， 并返回更新后的图像和当前列表（不更新）。 """ # 复制当前列表，并在副本中加上新点（仅用于显示） updated_image = draw_points_on_image(image, [x,y], point_type) return updated_image def add_point(x, y, point_type, visible_points, occlusion_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 def delete_point(point_type, visible_points, occlusion_points): """ delete操作：删除 points 列表中的最后一个点， 并返回更新后的图像和新的点列表。 """ if point_type == "vis": visible_points.pop() else: occlusion_points.pop() return visible_points, occlusion_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) 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_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) def delete_selected_visible(image, visible_points, occlusion_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: selected_index = None 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_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 def add_mask(mask, mask_list): # check if the mask if same as the last mask in the list if len(mask_list) > 0: if np.array_equal(mask, mask_list[-1]): return 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[:, :, 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 def delete_mask(mask_list): if len(mask_list) > 0: mask_list.pop() 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, (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 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, combined_mask, occluded_image 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: gr.Markdown(""" ## 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=[]) 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) occluded_image = gr.State(value=None) with gr.Row(): gr.Markdown("""* 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. """) with gr.Row(): with gr.Column(): input_image = gr.Image(type="numpy", label='Input Occlusion Image', sources="upload", height=300) with gr.Row(): message = gr.Markdown("Please wait a few seconds after uploading the image.", label="Message") # 用于显示提示信息 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") 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) with gr.Row(): gen_vis_mask = gr.Button("Generate Mask") add_vis_mask = gr.Button("Add Mask") 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) 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") with gr.Row(): with gr.Column(): mask_check = gr.Image(label='Combined Mask', interactive=False, height=300) with gr.Row(): check_combine_button = 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. * 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(): 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) # --------------------------- # 原有交互逻辑（略） # --------------------------- input_image.upload( reset_image, [predictor, input_image], [predictor, original_image, message], ) see_button.click( see_point, inputs=[original_image, x_input, y_input, point_type], 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] ) # --------------------------- # 新增的交互逻辑 # --------------------------- clear_button.click( clear_all_points, inputs=[original_image], outputs=[input_image] ) 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] ) # 生成mask的逻辑 gen_vis_mask.click( segment_and_overlay, 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=[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], 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=[mask_check, combined_mask] ) # 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 if __name__ == "__main__": pipeline = Amodal3RImageTo3DPipeline.from_pretrained("Sm0kyWu/Amodal3R") pipeline.cuda() try: pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))) except: pass demo.launch()