import gradio as gr from gradio_litmodel3d import LitModel3D import spaces 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.set_image(img) original_img = img.copy() 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): if len(selected_points) == 0: return [], None input_points = [p for p in selected_points] input_labels = [1 for _ in range(len(selected_points))] 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 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)): 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): visible_mask = run_sam(sam_predictor, points) overlaid = apply_mask_overlay(image, visible_mask * 255) return overlaid, visible_mask @spaces.GPU def image_to_3d( image: np.ndarray, mask: np.ndarray, seed: int, ss_guidance_strength: float, ss_sampling_steps: int, slat_guidance_strength: float, slat_sampling_steps: int, erode_kernel_size: int, req: gr.Request, ) -> Tuple[dict, str]: user_dir = os.path.join(TMP_DIR, str(req.session_hash)) outputs = pipeline.run_multi_image( [image], [mask], seed=seed, formats=["gaussian", "mesh"], 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="stochastic", erode_kernel_size=erode_kernel_size, ) video = render_utils.render_video(outputs['gaussian'][0], num_frames=120, bg_color=(1,1,1))['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: 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: 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 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_predictor = SamPredictor(sam) return sam_predictor def draw_points_on_image(image, point): image_with_points = image.copy() x, y = point 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): updated_image = draw_points_on_image(image, [x,y]) return updated_image def add_point(x, y, visible_points): if [x, y] not in visible_points: visible_points.append([x, y]) return visible_points def delete_point(visible_points): visible_points.pop() return visible_points def clear_all_points(image): updated_image = image.copy() return updated_image def see_visible_points(image, visible_points): 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 update_all_points(visible_points): text = f"Points: {visible_points}" visible_dropdown_choices = [f"({p[0]}, {p[1]})" for p in visible_points] 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): 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() 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): 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() combined_mask = np.zeros_like(updated_image[:, :, 0]) for mask in mask_list: combined_mask = cv2.bitwise_or(combined_mask, mask) 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 check_combined_mask(image, visibility_mask, mask_list, scale=0.65): updated_image = image.copy() combined_mask = np.zeros_like(updated_image[:, :, 0]) occluded_mask = np.zeros_like(updated_image[:, :, 0]) if len(mask_list) == 0: combined_mask = visibility_mask else: for mask in mask_list: combined_mask = cv2.bitwise_or(combined_mask, mask) if len(mask_list) > 1: kernel = np.ones((5, 5), np.uint8) dilate_iterations = 1 combined_mask = cv2.dilate(combined_mask, kernel, iterations=dilate_iterations) combined_mask = cv2.erode(combined_mask, kernel, iterations=dilate_iterations) masked_img = updated_image * combined_mask[:, :, None] occluded_mask[combined_mask == 1] = 127 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] target_size = 512 scale_factor = target_size / max(w, h) new_w = int(round(w * scale_factor * scale)) new_h = int(round(h * scale_factor * scale)) resized_occluded_mask = cv2.resize(cropped_occluded_mask.astype(np.uint8), (new_w, new_h), cv2.INTER_NEAREST) resized_img = cv2.resize(cropped_img, (new_w, new_h), cv2.INTER_NEAREST) final_img = np.zeros((target_size, target_size, 3), dtype=updated_image.dtype) final_occluded_mask = np.zeros((target_size, target_size), dtype=np.uint8) x_offset = (target_size - new_w) // 2 y_offset = (target_size - new_h) // 2 final_img[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized_img final_occluded_mask[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized_occluded_mask return final_img, final_occluded_mask 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) visibility_mask_list = gr.State(value=[]) occluded_mask = gr.State(value=None) output_buf = gr.State() with gr.Row(): gr.Markdown(""" ### Step 1 - Generate Visibility 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. "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. * The target object need to be put in the center of the image and the scale can be adjusted for better reconstruction. * Please click "Load Example Image" when using the provided example images. """) with gr.Row(): with gr.Column(): input_image = gr.Image(type="numpy", label='Input Occlusion Image', sources="upload", height=300) with gr.Row(): apply_example_btn = gr.Button("Load Example Image") 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) with gr.Row(): see_button = gr.Button("Render Point") add_button = gr.Button("Add Point") with gr.Row(): clear_button = gr.Button("Clear 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 Point to Delete", choices=[], value=None, interactive=True) delete_visible_button = gr.Button("Delete Selected Visible") 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") vis_input = gr.Image(label='Visible Input', interactive=False, height=300) with gr.Row(): zoom_scale = gr.Slider(0.3, 1.0, label="Target Object Scale", value=0.68, step=0.1) check_visible_input = gr.Button("Generate Occluded Input") with gr.Row(): examples = gr.Examples( examples=[ f'assets/example_image/{image}' for image in os.listdir("assets/example_image") ], inputs=[input_image], fn=lambda x: x, outputs=[input_image], run_on_click=True, examples_per_page=12, ) 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. * The boundary of the segmentation may not be accurate, so here we provide the option to erode the visible area. * 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): 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=False) with gr.Column(): 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) 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("Amodal 3D Reconstruction") with gr.Accordion(label="GLB Extraction Settings", open=False): mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01) texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512) with gr.Row(): extract_glb_btn = gr.Button("Extract GLB") extract_gs_btn = gr.Button("Extract Gaussian") gr.Markdown(""" *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.* """) with gr.Column(): video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300) model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=10.0, height=300) with gr.Row(): download_glb = gr.DownloadButton(label="Download GLB", interactive=False) download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False) # # Handlers demo.load(start_session) demo.unload(end_session) input_image.upload( reset_image, [predictor, input_image], [predictor, original_image, message, visible_points_state], ) apply_example_btn.click( reset_image, inputs=[predictor, input_image], outputs=[predictor, original_image, message, visible_points_state] ) see_button.click( see_point, inputs=[original_image, x_input, y_input], outputs=[input_image] ) add_button.click( add_point, inputs=[x_input, y_input, visible_points_state], outputs=[visible_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 ) visible_points_state.change( update_all_points, inputs=[visible_points_state], outputs=[points_text, visible_points_dropdown] ) delete_visible_button.click( delete_selected_visible, inputs=[input_image, visible_points_state, visible_points_dropdown], outputs=[input_image, visible_points_state, points_text, visible_points_dropdown] ) 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] ) check_visible_input.click( check_combined_mask, inputs=[original_image, visibility_mask, visibility_mask_list, zoom_scale], outputs=[vis_input, occluded_mask] ) # 3D Amodal Reconstruction generate_btn.click( get_seed, inputs=[randomize_seed, seed], outputs=[seed], ).then( image_to_3d, inputs=[vis_input, occluded_mask, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, erode_kernel_size], outputs=[output_buf, video_output], ).then( lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]), outputs=[extract_glb_btn, extract_gs_btn], ) video_output.clear( lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]), outputs=[extract_glb_btn, extract_gs_btn], ) extract_glb_btn.click( extract_glb, inputs=[output_buf, mesh_simplify, texture_size], outputs=[model_output, download_glb], ).then( lambda: gr.Button(interactive=True), outputs=[download_glb], ) extract_gs_btn.click( extract_gaussian, inputs=[output_buf], outputs=[model_output, download_gs], ).then( lambda: gr.Button(interactive=True), outputs=[download_gs], ) model_output.clear( lambda: gr.Button(interactive=False), outputs=[download_glb], ) 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()