|
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) |
|
|
|
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: |
|
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, |
|
) |
|
|
|
|
|
|
|
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] |
|
) |
|
|
|
|
|
|
|
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() |
