app.py

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 select_point(predictor: SamPredictor,
                 annotated_img: np.ndarray,
                 orig_img: np.ndarray,
                 sel_pix: list,
                 point_type: str, 
                 evt: gr.SelectData):
    """
    当用户在标注图像上点击时：
      - 将点击坐标添加到 sel_pix（正/负 prompt 根据 point_type），
      - 根据 sel_pix 调用 SAM 得到 mask，
      - 在 annotated_img 上绘制所有已选点的标记，
      - 返回更新后的标注图像、SAM 输出（用于显示）及生成的 visible_mask（用于后续 pix2gestalt）。
    """
    # 拷贝原图（用于标注）
    img = annotated_img.copy()
    h_original, w_original, _ = orig_img.shape
    h_new, w_new = 256, 256
    scale_x = w_new / w_original
    scale_y = h_new / h_original

    # 根据 prompt 类型添加点击点（evt.index 格式为 (x, y)）
    if point_type == 'positive_prompt':
        sel_pix.append((evt.index, 1))
    elif point_type == 'negative_prompt':
        sel_pix.append((evt.index, 0))
    else:
        sel_pix.append((evt.index, 1))

    # 将原始尺寸的点转换到 256x256 尺寸（SAM 输入要求）
    processed_sel_pix = []
    for point, label in sel_pix:
        x, y = point
        new_x = int(x * scale_x)
        new_y = int(y * scale_y)
        processed_sel_pix.append(([new_x, new_y], label))

    visible_mask, overlay_mask = run_sam(predictor, processed_sel_pix)
    # overlay_mask 是 SAM 输出的 mask（256x256），调整尺寸到原图尺寸以便显示
    mask = np.squeeze(overlay_mask[0][0])  # (256, 256)
    resized_mask = cv2.resize(mask.astype(np.uint8) * 255, (w_original, h_original), interpolation=cv2.INTER_AREA)
    resized_mask = resized_mask > 127
    # 制作 overlay 信息（供 output_mask 使用）
    resized_overlay_mask = [(resized_mask, 'visible_mask')]

    # 绘制所有点的标记
    COLORS = [(255, 0, 0), (0, 255, 0)]
    MARKERS = [1, 4]
    scaling_factor = min(h_original / 256, w_original / 256)
    marker_size = int(6 * scaling_factor)
    marker_thickness = int(2 * scaling_factor)
    for point, label in sel_pix:
        cv2.drawMarker(img, tuple(point), COLORS[label], markerType=MARKERS[label],
                       markerSize=marker_size, thickness=marker_thickness)

    return img, (orig_img, resized_overlay_mask), visible_mask

def undo_points(predictor, orig_img, sel_pix):
    """
    撤销最后一次点击：
      - 从 sel_pix 中 pop 出最后一个点，
      - 根据剩余点重新调用 SAM 得到 mask，
      - 返回更新后的图像和 mask。
    """
    temp = orig_img.copy()
    h_original, w_original, _ = orig_img.shape
    COLORS = [(255, 0, 0), (0, 255, 0)]
    MARKERS = [0, 5]
    scaling_factor = min(h_original / 256, w_original / 256)
    marker_size = int(6 * scaling_factor)
    marker_thickness = int(2 * scaling_factor)
    if len(sel_pix) > 0:
        sel_pix.pop()
        # 重新绘制剩余点
        for point, label in sel_pix:
            cv2.drawMarker(temp, tuple(point), COLORS[label],
                           markerType=MARKERS[label], markerSize=marker_size, thickness=marker_thickness)
    else:
        dummy_overlay_mask = [(np.zeros((h_original, w_original), dtype=np.uint8), 'visible_mask')]
        return orig_img, (orig_img, dummy_overlay_mask), []

    visible_mask, overlay_mask = run_sam(predictor, sel_pix)
    mask = np.squeeze(overlay_mask[0][0])
    resized_mask = cv2.resize(mask.astype(np.uint8) * 255, (w_original, h_original), interpolation=cv2.INTER_AREA)
    resized_mask = resized_mask > 127
    resized_overlay_mask = [(resized_mask, 'visible_mask')]
    return temp, (orig_img, resized_overlay_mask), visible_mask

def reset_image(predictor, img):
    """
    上传图像后调用：
      - 重置 predictor，
      - 设置 predictor 的输入图像，
      - 返回原图
    """
    predictor.set_image(img)
    # 返回predictor，原始图像
    return predictor, img

def button_clickable(selected_points):
    if len(selected_points) > 0:
        return gr.Button.update(interactive=True)
    else:
        return gr.Button.update(interactive=False)


@spaces.GPU
def run_sam(predictor: SamPredictor, image, selected_points):
    """
    调用 SAM 模型进行分割。
    """
    # 确保图像为 RGB 模式
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    if image.mode != 'RGB':
        image = image.convert("RGB")
    if len(selected_points) == 0:
        return [], None
    input_points = [p for p, _ in selected_points]
    input_labels = [int(l) for _, l in selected_points]
    masks, _, _ = predictor.predict(
        point_coords=np.array(input_points),
        point_labels=input_labels,
        multimask_output=False,  # 单对象输出
    )
    visible_mask = 255 * np.squeeze(masks).astype(np.uint8)
    return visible_mask, None


def apply_mask_overlay(image: Image.Image, mask: np.ndarray) -> Image.Image:
    """
    在原图上叠加 mask：使用红色绘制 mask 的轮廓，非 mask 区域叠加浅灰色半透明遮罩。
    """
    img_arr = np.array(image)
    if mask.ndim == 3:
        mask = mask[:, :, 0]
    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, (255, 0, 0), 2)
    return Image.fromarray(overlay)


def segment_and_overlay(image: np.ndarray, points):
    """
    调用 run_sam 获得 mask，然后叠加显示分割结果。
    """
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    if image.mode != "RGB":
        image = image.convert("RGB")
    mask, _ = run_sam(sam_predictor, image, points)
    if mask == []:
        return image
    overlaid = apply_mask_overlay(image, mask)
    return overlaid


def reset_points():
    """
    清空点击点提示。
    """
    return [], ""


@spaces.GPU
def image_to_3d(
    image: Image.Image,
    multiimages: List[tuple],
    is_multiimage: bool,
    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))
    if not is_multiimage:
        outputs = pipeline.run(
            image,
            seed=seed,
            formats=["gaussian", "mesh"],
            preprocess_image=False,
            sparse_structure_sampler_params={
                "steps": ss_sampling_steps,
                "cfg_strength": ss_guidance_strength,
            },
            slat_sampler_params={
                "steps": slat_sampling_steps,
                "cfg_strength": slat_guidance_strength,
            },
        )
    else:
        outputs = pipeline.run_multi_image(
            [img[0] for img in multiimages],
            seed=seed,
            formats=["gaussian", "mesh"],
            preprocess_image=False,
            sparse_structure_sampler_params={
                "steps": ss_sampling_steps,
                "cfg_strength": ss_guidance_strength,
            },
            slat_sampler_params={
                "steps": slat_sampling_steps,
                "cfg_strength": slat_guidance_strength,
            },
            mode=multiimage_algo,
        )
    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 = (0, 0, 255) if point_type == "visible" else (0, 255, 0)
    cv2.circle(image_with_points, (int(x), int(y)), radius=5, 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 == "visible":
        visible_points.append([x, y])
    else:
        occlusion_points.append([x, y])
    return visible_points, occlusion_points

def delete_point(point_type, visible_points, occlusion_points):
    """
    delete操作：删除 points 列表中的最后一个点，
    并返回更新后的图像和新的点列表。
    """
    if point_type == "visible":
        visible_points.pop()
    else:
        occlusion_points.pop()
    return visible_points, occlusion_points


with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    ## 3D Amodal Reconstruction with [Amodal3R](https://sm0kywu.github.io/Amodal3R/)
    * Upload an image and click "Generate" to create a 3D asset.
    * Target object selection. Multiple point prompts are supported until you get the ideal visible area.
    * Occluders selection, this can be done by squential point prompts. You can choose "all occ", then all the other areas except the target object will be treated as occluders.
    * 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.
    """)

     # 定义各状态变量
    predictor = gr.State(value=get_sam_predictor())
    visible_points_state = gr.State(value=[])
    occlusion_points_state = gr.State(value=[])


    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="numpy", label='Input Occlusion Image', height=300)
            fg_bg_radio = gr.Radio(['positive_prompt', 'negative_prompt'], label='Point Prompt Type')
            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(choices=["visible", "occlusion"], label="Point Type", value="visible")
            with gr.Row():
                see_button = gr.Button("See")
                add_button = gr.Button("Add")
                delete_button = gr.Button("Delete")
        with gr.Column():
            # 显示 SAM 分割结果（带 overlay）—— 使用 AnnotatedImage 显示更直观
            sam_image = gr.Image(label='SAM Generated Mask', interactive=False, height=300)

    
    # 会话启动与结束
    demo.load(start_session)
    demo.unload(end_session)
    
    # 上传图像时：重置 predictor 并将原图赋值给 original_image、preprocessed_image、selected_points 以及 output_mask
    input_image.upload(
        reset_image,
        [predictor, input_image],
        [predictor, sam_image]
    )
    # 如果点击see按钮，应该在input图片上生成对应的点，
    see_button.click(
        see_point, 
        inputs=[input_image, x_input, y_input, point_type], 
        outputs=[input_image]
    )
    # 如果点击add按钮，应该将对应的点添加到visible_points_state中
    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]
    )
    delete_button.click(
        delete_point, 
        inputs=[point_type, visible_points_state, occlusion_points_state], 
        outputs=[visible_points_state, occlusion_points_state]
    )
    

# 启动 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()