Update app.py

app.py

@@ -1,641 +1,643 @@
-import gradio as gr
-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 change_message():
-    return "Please wait for a few seconds after uploading the image."
-
-def reset_image(predictor, img):
-    img = np.array(img)
-    predictor.set_image(img)
-    original_img = img.copy()
-    return predictor, original_img, "The models are ready.", [], [], [], original_img
-
-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(img, predictor, selected_points):
-    if len(selected_points) == 0:
-        return np.zeros(img.shape[:2], dtype=np.uint8)
-    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
-
-
-@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 see_occlusion_points(image, occlusion_points):
-    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(points):
-    text = f"Points: {points}"
-    dropdown_choices = [f"({p[0]}, {p[1]})" for p in points]
-    return text, gr.Dropdown(show_label=False, choices=dropdown_choices, value=None, interactive=True)
-
-def delete_selected(image, visible_points, occlusion_points, occlusion_mask_list, selected_value, point_type):
-    if point_type == "visibility":
-        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)
-    else:
-        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)
-            occlusion_mask_list.pop(selected_index)
-    updated_image = image.copy()
-    updated_image = see_visible_points(updated_image, visible_points)
-    updated_image = see_occlusion_points(updated_image, occlusion_points)
-    if point_type == "visibility":
-        updated_text, dropdown = update_all_points(visible_points)
-    else:
-        updated_text, dropdown = update_all_points(occlusion_points)
-    return updated_image, visible_points, occlusion_points, updated_text, dropdown
-
-def add_current_mask(visibility_mask, visibilty_mask_list, point_type): 
-    if point_type == "visibility":
-        if len(visibilty_mask_list) > 0:
-            if np.array_equal(visibility_mask, visibilty_mask_list[-1]):
-                return visibilty_mask_list
-        visibilty_mask_list.append(visibility_mask)
-        return visibilty_mask_list
-    else: # the occlusion mask will be automatically added, so do nothing here
-        return visibilty_mask_list
-
-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 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 segment_and_overlay(image, points, sam_predictor, mask_list, point_type):
-    if point_type == "visibility":
-        visible_mask = run_sam(image, sam_predictor, points)
-        for mask in mask_list:
-            visible_mask = cv2.bitwise_or(visible_mask, mask)
-        overlaid = apply_mask_overlay(image, visible_mask * 255)
-        return overlaid, visible_mask, mask_list
-    else:
-        combined_occlusion_mask = np.zeros_like(image[:, :, 0])
-        mask_list = []
-        if len(points) != 0:
-            for point in points:
-                mask = run_sam(image, sam_predictor, [point])
-                mask_list.append(mask)
-                combined_occlusion_mask = cv2.bitwise_or(combined_occlusion_mask, mask)
-        overlaid = apply_mask_overlay(image, combined_occlusion_mask * 255, color=(0, 255, 0))
-        return overlaid, combined_occlusion_mask, mask_list
-
-def delete_mask(visibility_mask_list, occlusion_mask_list, occlusion_points_state, point_type):
-    if point_type == "visibility":
-        if len(visibility_mask_list) > 0:
-            visibility_mask_list.pop()
-    else:
-        if len(occlusion_mask_list) > 0:
-            occlusion_mask_list.pop()
-            occlusion_points_state.pop()
-    return visibility_mask_list, occlusion_mask_list, occlusion_points_state
-
-def check_combined_mask(image, visibility_mask, visibility_mask_list, occlusion_mask_list, scale=0.68):
-    if visibility_mask.sum() == 0:
-        return np.zeros_like(image), np.zeros_like(image[:, :, 0])
-    updated_image = image.copy()
-    combined_mask = np.zeros_like(updated_image[:, :, 0])
-    occluded_mask = np.zeros_like(updated_image[:, :, 0])
-    binary_visibility_masks = [(m > 0).astype(np.uint8) for m in visibility_mask_list]
-    combined_mask = np.zeros_like(binary_visibility_masks[0]) if binary_visibility_masks else (visibility_mask > 0).astype(np.uint8)
-    for m in binary_visibility_masks:
-        combined_mask = cv2.bitwise_or(combined_mask, m)
-
-    if len(binary_visibility_masks) > 1:
-        kernel = np.ones((5, 5), np.uint8)
-        combined_mask = cv2.dilate(combined_mask, kernel, iterations=1)
-
-    binary_occlusion_masks = [(m > 0).astype(np.uint8) for m in occlusion_mask_list]
-    occluded_mask = np.zeros_like(binary_occlusion_masks[0]) if binary_occlusion_masks else np.zeros_like(combined_mask)
-    for m in binary_occlusion_masks:
-        occluded_mask = cv2.bitwise_or(occluded_mask, m)
-
-    kernel_small = np.ones((3, 3), np.uint8)
-    if len(binary_occlusion_masks) > 0:
-        dilated = cv2.dilate(combined_mask, kernel_small, iterations=1)
-        boundary_mask = dilated - combined_mask
-        occluded_mask = cv2.bitwise_or(occluded_mask, boundary_mask)
-        occluded_mask = (occluded_mask > 0).astype(np.uint8)
-        occluded_mask = cv2.dilate(occluded_mask, kernel_small, iterations=1)
-        occluded_mask = (occluded_mask > 0).astype(np.uint8)
-    else:
-        occluded_mask = 1 - combined_mask
-
-    combined_mask[occluded_mask == 1] = 0
-
-    occluded_mask = (1-occluded_mask) * 255
-
-    masked_img = updated_image * combined_mask[:, :, None]
-    occluded_mask[combined_mask == 1] = 127
-
-    x, y, w, h = cv2.boundingRect(combined_mask.astype(np.uint8))
-
-    ori_h, ori_w = masked_img.shape[:2]
-    target_size = 512
-    scale_factor = target_size / max(w, h)
-    final_scale = scale_factor * scale
-    new_w = int(round(ori_w * final_scale))
-    new_h = int(round(ori_h * final_scale))
-    
-    resized_occluded_mask = cv2.resize(occluded_mask.astype(np.uint8), (new_w, new_h), interpolation=cv2.INTER_NEAREST)
-    resized_img = cv2.resize(masked_img, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
-
-    final_img = np.zeros((target_size, target_size, 3), dtype=updated_image.dtype)
-    final_occluded_mask = np.ones((target_size, target_size), dtype=np.uint8) * 255
-
-    new_x = int(round(x * final_scale))
-    new_y = int(round(y * final_scale))
-    new_w_box = int(round(w * final_scale))
-    new_h_box = int(round(h * final_scale))
-
-    new_cx = new_x + new_w_box // 2
-    new_cy = new_y + new_h_box // 2
-
-    final_cx, final_cy = target_size // 2, target_size // 2
-    x_offset = final_cx - new_cx
-    y_offset = final_cy - new_cy
-
-    final_x_start = max(0, x_offset)
-    final_y_start = max(0, y_offset)
-    final_x_end = min(target_size, x_offset + new_w)
-    final_y_end = min(target_size, y_offset + new_h)
-
-    img_x_start = max(0, -x_offset)
-    img_y_start = max(0, -y_offset)
-    img_x_end = min(new_w, target_size - x_offset)
-    img_y_end = min(new_h, target_size - y_offset)
-
-    final_img[final_y_start:final_y_end, final_x_start:final_x_end] = resized_img[img_y_start:img_y_end, img_x_start:img_x_end]
-    final_occluded_mask[final_y_start:final_y_end, final_x_start:final_x_end] = resized_occluded_mask[img_y_start:img_y_end, img_x_start:img_x_end]
-
-    return final_img, final_occluded_mask
-
-
-def get_point(img, point_type, visible_points_state, occlusion_points_state, evt: gr.SelectData):
-    updated_img = np.array(img).copy()
-    if point_type == "visibility":
-        visible_points_state = add_point(evt.index[0], evt.index[1], visible_points_state)
-    else:
-        occlusion_points_state = add_point(evt.index[0], evt.index[1], occlusion_points_state)
-    updated_img = see_visible_points(updated_img, visible_points_state)
-    updated_img = see_occlusion_points(updated_img, occlusion_points_state)
-    return updated_img, visible_points_state, occlusion_points_state
-
-
-def change_point_type(point_type, visible_points_state, occlusion_points_state):
-    if point_type == "visibility":
-        text = f"Points: {visible_points_state}"
-        dropdown_choices = [f"({p[0]}, {p[1]})" for p in visible_points_state]
-    else:
-        text = f"Points: {occlusion_points_state}"
-        dropdown_choices = [f"({p[0]}, {p[1]})" for p in occlusion_points_state]
-    return text, gr.Dropdown(show_label=False, choices=dropdown_choices, value=None, interactive=True)
-
-
-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=[])
-    occlusion_mask = gr.State(value=None)
-    occlusion_mask_list = 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():
-        with gr.Column():
-            gr.Markdown("""
-            ### Step 1 - Generate Visibility and Occlusion Mask.
-            * Please click "Load Example Image" first when using the provided example images (bottom).
-            * Please wait for a few seconds after uploading the image. Segment Anything is getting ready.
-            * **Click to add the point prompts** to indicate the target object (multiple points supported) and occluders (one point for an occluder for better usability).
-            * "Add mask", current mask will be saved if the input needs to be added sequentially.
-            * The scale of target object can be adjusted for better reconstruction, we suggest 0.4 to 0.7 for most cases.
-            """)
-            with gr.Row():
-                input_image = gr.Image(interactive=True, type='pil', label='Input Occlusion Image', show_label=True, sources="upload", height=300)
-                input_with_prompt = gr.Image(type="numpy", label='Input with Prompt', interactive=False, 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():
-                point_type = gr.Radio(["visibility", "occlusion"], label="Point Prompt Type", value="visibility")
-            with gr.Row():
-                with gr.Column():
-                    points_text = gr.Textbox(show_label=False, interactive=False)
-                with gr.Column():
-                    points_dropdown = gr.Dropdown(show_label=False, choices=[], value=None, interactive=True)
-                    delete_button = gr.Button("Delete Selected Point")
-            with gr.Row():
-                with gr.Column():
-                    render_mask = gr.Image(label='Render Mask', interactive=False, height=300)
-                    with gr.Row():
-                        add_mask = gr.Button("Add Mask")
-                        undo_mask = gr.Button("Undo Last Mask")
-                with gr.Column():
-                    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)
-                    with gr.Row():
-                        check_visible_input = gr.Button("Generate Occluded Input")
-
-        with gr.Column():
-            gr.Markdown("""
-            ### Step 2 - 3D Amodal Reconstruction. (Thanks to [TRELLIS](https://huggingface.co/spaces/JeffreyXiang/TRELLIS) for the 3D rendering component!)
-            * 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 (try 0, 3 or 5).
-            * If the reconstructed 3D asset is satisfactory, interactive GLB file can be extracted (may look dull due to the absence of light source) and downloaded.
-            """)
-            with gr.Row():
-                video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
-            with gr.Row():
-                with gr.Accordion(label="Generation Settings", open=False):
-                    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)
-            with gr.Row():
-                generate_btn = gr.Button("Amodal 3D Reconstruction")
-            with gr.Row():
-                model_output = gr.Model3D(label="Extracted GLB", pan_speed=0.5, height=300, clear_color=(0.9,0.9,0.9,1))
-            with gr.Row():
-                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")
-                download_glb = gr.DownloadButton(label="Download GLB", interactive=False) 
-    
-    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,
-        )
-    
-
-    # # Handlers
-    demo.load(start_session)
-    demo.unload(end_session)
-
-    input_image.upload(
-        change_message,
-        [],
-        [message]
-    ).then(
-        reset_image,
-        [predictor, input_image],
-        [predictor, original_image, message, visible_points_state, occlusion_points_state, occlusion_mask_list, input_with_prompt],
-    )
-    
-    apply_example_btn.click(
-        change_message,
-        [],
-        [message]
-    ).then(
-        reset_image,
-        inputs=[predictor, input_image],
-        outputs=[predictor, original_image, message, visible_points_state, occlusion_points_state, occlusion_mask_list, input_with_prompt]
-    )
-    input_image.select(
-        get_point,
-        inputs=[input_image, point_type, visible_points_state, occlusion_points_state],
-        outputs=[input_with_prompt, visible_points_state, occlusion_points_state]
-    )
-
-    point_type.change(
-        change_point_type,
-        inputs=[point_type, visible_points_state, occlusion_points_state],
-        outputs=[points_text, points_dropdown]
-    )
-
-    visible_points_state.change(
-        update_all_points,
-        inputs=[visible_points_state],
-        outputs=[points_text, points_dropdown]
-    ).then(
-        segment_and_overlay,
-        inputs=[original_image, visible_points_state, predictor, visibility_mask_list, point_type],
-        outputs=[render_mask, visibility_mask, visibility_mask_list]
-    ).then(
-        check_combined_mask,
-        inputs=[original_image, visibility_mask, visibility_mask_list, occlusion_mask_list, zoom_scale],
-        outputs=[vis_input, occluded_mask]
-    )
-
-    occlusion_points_state.change(
-        update_all_points,
-        inputs=[occlusion_points_state],
-        outputs=[points_text, points_dropdown]
-    ).then(
-        segment_and_overlay,
-        inputs=[original_image, occlusion_points_state, predictor, occlusion_mask_list, point_type],
-        outputs=[render_mask, occlusion_mask, occlusion_mask_list]
-    ).then(
-        check_combined_mask,
-        inputs=[original_image, visibility_mask, visibility_mask_list, occlusion_mask_list, zoom_scale],
-        outputs=[vis_input, occluded_mask]
-    )
-
-    delete_button.click(
-        delete_selected,
-        inputs=[original_image, visible_points_state, occlusion_points_state, occlusion_mask_list, points_dropdown, point_type],
-        outputs=[input_with_prompt, visible_points_state, occlusion_points_state, points_text, points_dropdown]
-    )
-
-    add_mask.click(
-        add_current_mask,
-        inputs=[visibility_mask, visibility_mask_list, point_type],
-        outputs=[visibility_mask_list]
-    )
-
-    undo_mask.click(
-        delete_mask,
-        inputs=[visibility_mask_list, occlusion_mask_list, occlusion_points_state, point_type],
-        outputs=[visibility_mask_list, occlusion_mask_list, occlusion_points_state]
-    )
-
-    check_visible_input.click(
-        check_combined_mask,
-        inputs=[original_image, visibility_mask, visibility_mask_list, occlusion_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],
-    )
-
-    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],
-    )
-
-    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
+import gradio as gr
+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)
+os.environ['MASTER_ADDR'] = 'localhost'
+os.environ['MASTER_PORT'] = '12355'
+
+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 change_message():
+    return "Please wait for a few seconds after uploading the image."
+
+def reset_image(predictor, img):
+    img = np.array(img)
+    predictor.set_image(img)
+    original_img = img.copy()
+    return predictor, original_img, "The models are ready.", [], [], [], original_img
+
+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(img, predictor, selected_points):
+    if len(selected_points) == 0:
+        return np.zeros(img.shape[:2], dtype=np.uint8)
+    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
+
+
+@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 see_occlusion_points(image, occlusion_points):
+    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(points):
+    text = f"Points: {points}"
+    dropdown_choices = [f"({p[0]}, {p[1]})" for p in points]
+    return text, gr.Dropdown(show_label=False, choices=dropdown_choices, value=None, interactive=True)
+
+def delete_selected(image, visible_points, occlusion_points, occlusion_mask_list, selected_value, point_type):
+    if point_type == "visibility":
+        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)
+    else:
+        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)
+            occlusion_mask_list.pop(selected_index)
+    updated_image = image.copy()
+    updated_image = see_visible_points(updated_image, visible_points)
+    updated_image = see_occlusion_points(updated_image, occlusion_points)
+    if point_type == "visibility":
+        updated_text, dropdown = update_all_points(visible_points)
+    else:
+        updated_text, dropdown = update_all_points(occlusion_points)
+    return updated_image, visible_points, occlusion_points, updated_text, dropdown
+
+def add_current_mask(visibility_mask, visibilty_mask_list, point_type): 
+    if point_type == "visibility":
+        if len(visibilty_mask_list) > 0:
+            if np.array_equal(visibility_mask, visibilty_mask_list[-1]):
+                return visibilty_mask_list
+        visibilty_mask_list.append(visibility_mask)
+        return visibilty_mask_list
+    else: # the occlusion mask will be automatically added, so do nothing here
+        return visibilty_mask_list
+
+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 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 segment_and_overlay(image, points, sam_predictor, mask_list, point_type):
+    if point_type == "visibility":
+        visible_mask = run_sam(image, sam_predictor, points)
+        for mask in mask_list:
+            visible_mask = cv2.bitwise_or(visible_mask, mask)
+        overlaid = apply_mask_overlay(image, visible_mask * 255)
+        return overlaid, visible_mask, mask_list
+    else:
+        combined_occlusion_mask = np.zeros_like(image[:, :, 0])
+        mask_list = []
+        if len(points) != 0:
+            for point in points:
+                mask = run_sam(image, sam_predictor, [point])
+                mask_list.append(mask)
+                combined_occlusion_mask = cv2.bitwise_or(combined_occlusion_mask, mask)
+        overlaid = apply_mask_overlay(image, combined_occlusion_mask * 255, color=(0, 255, 0))
+        return overlaid, combined_occlusion_mask, mask_list
+
+def delete_mask(visibility_mask_list, occlusion_mask_list, occlusion_points_state, point_type):
+    if point_type == "visibility":
+        if len(visibility_mask_list) > 0:
+            visibility_mask_list.pop()
+    else:
+        if len(occlusion_mask_list) > 0:
+            occlusion_mask_list.pop()
+            occlusion_points_state.pop()
+    return visibility_mask_list, occlusion_mask_list, occlusion_points_state
+
+def check_combined_mask(image, visibility_mask, visibility_mask_list, occlusion_mask_list, scale=0.68):
+    if visibility_mask.sum() == 0:
+        return np.zeros_like(image), np.zeros_like(image[:, :, 0])
+    updated_image = image.copy()
+    combined_mask = np.zeros_like(updated_image[:, :, 0])
+    occluded_mask = np.zeros_like(updated_image[:, :, 0])
+    binary_visibility_masks = [(m > 0).astype(np.uint8) for m in visibility_mask_list]
+    combined_mask = np.zeros_like(binary_visibility_masks[0]) if binary_visibility_masks else (visibility_mask > 0).astype(np.uint8)
+    for m in binary_visibility_masks:
+        combined_mask = cv2.bitwise_or(combined_mask, m)
+
+    if len(binary_visibility_masks) > 1:
+        kernel = np.ones((5, 5), np.uint8)
+        combined_mask = cv2.dilate(combined_mask, kernel, iterations=1)
+
+    binary_occlusion_masks = [(m > 0).astype(np.uint8) for m in occlusion_mask_list]
+    occluded_mask = np.zeros_like(binary_occlusion_masks[0]) if binary_occlusion_masks else np.zeros_like(combined_mask)
+    for m in binary_occlusion_masks:
+        occluded_mask = cv2.bitwise_or(occluded_mask, m)
+
+    kernel_small = np.ones((3, 3), np.uint8)
+    if len(binary_occlusion_masks) > 0:
+        dilated = cv2.dilate(combined_mask, kernel_small, iterations=1)
+        boundary_mask = dilated - combined_mask
+        occluded_mask = cv2.bitwise_or(occluded_mask, boundary_mask)
+        occluded_mask = (occluded_mask > 0).astype(np.uint8)
+        occluded_mask = cv2.dilate(occluded_mask, kernel_small, iterations=1)
+        occluded_mask = (occluded_mask > 0).astype(np.uint8)
+    else:
+        occluded_mask = 1 - combined_mask
+
+    combined_mask[occluded_mask == 1] = 0
+
+    occluded_mask = (1-occluded_mask) * 255
+
+    masked_img = updated_image * combined_mask[:, :, None]
+    occluded_mask[combined_mask == 1] = 127
+
+    x, y, w, h = cv2.boundingRect(combined_mask.astype(np.uint8))
+
+    ori_h, ori_w = masked_img.shape[:2]
+    target_size = 512
+    scale_factor = target_size / max(w, h)
+    final_scale = scale_factor * scale
+    new_w = int(round(ori_w * final_scale))
+    new_h = int(round(ori_h * final_scale))
+    
+    resized_occluded_mask = cv2.resize(occluded_mask.astype(np.uint8), (new_w, new_h), interpolation=cv2.INTER_NEAREST)
+    resized_img = cv2.resize(masked_img, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
+
+    final_img = np.zeros((target_size, target_size, 3), dtype=updated_image.dtype)
+    final_occluded_mask = np.ones((target_size, target_size), dtype=np.uint8) * 255
+
+    new_x = int(round(x * final_scale))
+    new_y = int(round(y * final_scale))
+    new_w_box = int(round(w * final_scale))
+    new_h_box = int(round(h * final_scale))
+
+    new_cx = new_x + new_w_box // 2
+    new_cy = new_y + new_h_box // 2
+
+    final_cx, final_cy = target_size // 2, target_size // 2
+    x_offset = final_cx - new_cx
+    y_offset = final_cy - new_cy
+
+    final_x_start = max(0, x_offset)
+    final_y_start = max(0, y_offset)
+    final_x_end = min(target_size, x_offset + new_w)
+    final_y_end = min(target_size, y_offset + new_h)
+
+    img_x_start = max(0, -x_offset)
+    img_y_start = max(0, -y_offset)
+    img_x_end = min(new_w, target_size - x_offset)
+    img_y_end = min(new_h, target_size - y_offset)
+
+    final_img[final_y_start:final_y_end, final_x_start:final_x_end] = resized_img[img_y_start:img_y_end, img_x_start:img_x_end]
+    final_occluded_mask[final_y_start:final_y_end, final_x_start:final_x_end] = resized_occluded_mask[img_y_start:img_y_end, img_x_start:img_x_end]
+
+    return final_img, final_occluded_mask
+
+
+def get_point(img, point_type, visible_points_state, occlusion_points_state, evt: gr.SelectData):
+    updated_img = np.array(img).copy()
+    if point_type == "visibility":
+        visible_points_state = add_point(evt.index[0], evt.index[1], visible_points_state)
+    else:
+        occlusion_points_state = add_point(evt.index[0], evt.index[1], occlusion_points_state)
+    updated_img = see_visible_points(updated_img, visible_points_state)
+    updated_img = see_occlusion_points(updated_img, occlusion_points_state)
+    return updated_img, visible_points_state, occlusion_points_state
+
+
+def change_point_type(point_type, visible_points_state, occlusion_points_state):
+    if point_type == "visibility":
+        text = f"Points: {visible_points_state}"
+        dropdown_choices = [f"({p[0]}, {p[1]})" for p in visible_points_state]
+    else:
+        text = f"Points: {occlusion_points_state}"
+        dropdown_choices = [f"({p[0]}, {p[1]})" for p in occlusion_points_state]
+    return text, gr.Dropdown(show_label=False, choices=dropdown_choices, value=None, interactive=True)
+
+
+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=[])
+    occlusion_mask = gr.State(value=None)
+    occlusion_mask_list = 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():
+        with gr.Column():
+            gr.Markdown("""
+            ### Step 1 - Generate Visibility and Occlusion Mask.
+            * Please click "Load Example Image" when using the provided example images (bottom).
+            * Please wait for a few seconds after uploading the image. Segment Anything is getting ready.
+            * **Click to add the point prompts** to indicate the target object (multiple points supported) and occluders (one point for an occluder for better usability).
+            * "Add mask", current mask will be saved if the input needs to be added sequentially.
+            * The scale of target object can be adjusted for better reconstruction, we suggest 0.4 to 0.7 for most cases.
+            """)
+            with gr.Row():
+                input_image = gr.Image(interactive=True, type='pil', label='Input Occlusion Image', show_label=True, sources="upload", height=300)
+                input_with_prompt = gr.Image(type="numpy", label='Input with Prompt', interactive=False, 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():
+                point_type = gr.Radio(["visibility", "occlusion"], label="Point Prompt Type", value="visibility")
+            with gr.Row():
+                with gr.Column():
+                    points_text = gr.Textbox(show_label=False, interactive=False)
+                with gr.Column():
+                    points_dropdown = gr.Dropdown(show_label=False, choices=[], value=None, interactive=True)
+                    delete_button = gr.Button("Delete Selected Point")
+            with gr.Row():
+                with gr.Column():
+                    render_mask = gr.Image(label='Render Mask', interactive=False, height=300)
+                    with gr.Row():
+                        add_mask = gr.Button("Add Mask")
+                        undo_mask = gr.Button("Undo Last Mask")
+                with gr.Column():
+                    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)
+                    with gr.Row():
+                        check_visible_input = gr.Button("Generate Occluded Input")
+
+        with gr.Column():
+            gr.Markdown("""
+            ### Step 2 - 3D Amodal Reconstruction. (Thanks to [TRELLIS](https://huggingface.co/spaces/JeffreyXiang/TRELLIS) for the 3D rendering component!)
+            * 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 (try 0, 3 or 5).
+            * If the reconstructed 3D asset is satisfactory, interactive GLB file can be extracted (may look dull due to the absence of light source) and downloaded.
+            """)
+            with gr.Row():
+                video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
+            with gr.Row():
+                with gr.Accordion(label="Generation Settings", open=False):
+                    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)
+            with gr.Row():
+                generate_btn = gr.Button("Amodal 3D Reconstruction")
+            with gr.Row():
+                model_output = gr.Model3D(label="Extracted GLB", pan_speed=0.5, height=300, clear_color=(0.9,0.9,0.9,1))
+            with gr.Row():
+                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")
+                download_glb = gr.DownloadButton(label="Download GLB", interactive=False) 
+    
+    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,
+        )
+    
+
+    # # Handlers
+    demo.load(start_session)
+    demo.unload(end_session)
+
+    input_image.upload(
+        change_message,
+        [],
+        [message]
+    ).then(
+        reset_image,
+        [predictor, input_image],
+        [predictor, original_image, message, visible_points_state, occlusion_points_state, occlusion_mask_list, input_with_prompt],
+    )
+    
+    apply_example_btn.click(
+        change_message,
+        [],
+        [message]
+    ).then(
+        reset_image,
+        inputs=[predictor, input_image],
+        outputs=[predictor, original_image, message, visible_points_state, occlusion_points_state, occlusion_mask_list, input_with_prompt]
+    )
+    input_image.select(
+        get_point,
+        inputs=[input_image, point_type, visible_points_state, occlusion_points_state],
+        outputs=[input_with_prompt, visible_points_state, occlusion_points_state]
+    )
+
+    point_type.change(
+        change_point_type,
+        inputs=[point_type, visible_points_state, occlusion_points_state],
+        outputs=[points_text, points_dropdown]
+    )
+
+    visible_points_state.change(
+        update_all_points,
+        inputs=[visible_points_state],
+        outputs=[points_text, points_dropdown]
+    ).then(
+        segment_and_overlay,
+        inputs=[original_image, visible_points_state, predictor, visibility_mask_list, point_type],
+        outputs=[render_mask, visibility_mask, visibility_mask_list]
+    ).then(
+        check_combined_mask,
+        inputs=[original_image, visibility_mask, visibility_mask_list, occlusion_mask_list, zoom_scale],
+        outputs=[vis_input, occluded_mask]
+    )
+
+    occlusion_points_state.change(
+        update_all_points,
+        inputs=[occlusion_points_state],
+        outputs=[points_text, points_dropdown]
+    ).then(
+        segment_and_overlay,
+        inputs=[original_image, occlusion_points_state, predictor, occlusion_mask_list, point_type],
+        outputs=[render_mask, occlusion_mask, occlusion_mask_list]
+    ).then(
+        check_combined_mask,
+        inputs=[original_image, visibility_mask, visibility_mask_list, occlusion_mask_list, zoom_scale],
+        outputs=[vis_input, occluded_mask]
+    )
+
+    delete_button.click(
+        delete_selected,
+        inputs=[original_image, visible_points_state, occlusion_points_state, occlusion_mask_list, points_dropdown, point_type],
+        outputs=[input_with_prompt, visible_points_state, occlusion_points_state, points_text, points_dropdown]
+    )
+
+    add_mask.click(
+        add_current_mask,
+        inputs=[visibility_mask, visibility_mask_list, point_type],
+        outputs=[visibility_mask_list]
+    )
+
+    undo_mask.click(
+        delete_mask,
+        inputs=[visibility_mask_list, occlusion_mask_list, occlusion_points_state, point_type],
+        outputs=[visibility_mask_list, occlusion_mask_list, occlusion_points_state]
+    )
+
+    check_visible_input.click(
+        check_combined_mask,
+        inputs=[original_image, visibility_mask, visibility_mask_list, occlusion_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],
+    )
+
+    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],
+    )
+
+    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()