import spaces import os import numpy as np from PIL import Image from omegaconf import OmegaConf from functools import partial import gradio as gr from huggingface_hub import hf_hub_download import torch from torchvision import transforms import rembg import cv2 from pytorch_lightning import seed_everything from src.visualizer import CameraVisualizer from src.pose_estimation import load_model_from_config, estimate_poses, estimate_elevs from src.pose_funcs import find_optimal_poses from src.utils import spherical_to_cartesian, elu_to_c2w if torch.cuda.is_available(): _device_ = 'cuda:0' else: _device_ = 'cpu' _config_path_ = 'src/configs/sd-objaverse-finetune-c_concat-256.yaml' _ckpt_path_ = hf_hub_download(repo_id='tokenid/ID-Pose', filename='ckpts/zero123-xl.ckpt', repo_type='model') _matcher_ckpt_path_ = hf_hub_download(repo_id='tokenid/ID-Pose', filename='ckpts/indoor_ds_new.ckpt', repo_type='model') _config_ = OmegaConf.load(_config_path_) _model_ = load_model_from_config(_config_, _ckpt_path_, device='cpu') _model_ = _model_.to(_device_) _model_.eval() def rgba_to_rgb(img): assert img.mode == 'RGBA' img = np.asarray(img, dtype=np.float32) img[:, :, :3] = img[:, :, :3] * (img[..., 3:]/255.) + (255-img[..., 3:]) img = img.clip(0, 255).astype(np.uint8) return Image.fromarray(img[:, :, :3]) def remove_background(image, rembg_session = None, force = False, **rembg_kwargs): do_remove = True if image.mode == "RGBA" and image.getextrema()[3][0] < 255: do_remove = False do_remove = do_remove or force if do_remove: image = rembg.remove(image, session=rembg_session, **rembg_kwargs) return image def group_recenter(images, ratio=1.5, mask_thres=127, bkg_color=[255, 255, 255, 255]): ws = [] hs = [] images = [ np.asarray(img) for img in images ] for img in images: alpha = img[:, :, 3] yy, xx = np.where(alpha > mask_thres) y0, y1 = yy.min(), yy.max() x0, x1 = xx.min(), xx.max() ws.append(x1 - x0) hs.append(y1 - y0) sz_w = np.max(ws) sz_h = np.max(hs) sz = int( max(ratio*sz_w, ratio*sz_h) ) out_rgbs = [] for rgba in images: rgb = rgba[:, :, :3] alpha = rgba[:, :, 3] yy, xx = np.where(alpha > mask_thres) y0, y1 = yy.min(), yy.max() x0, x1 = xx.min(), xx.max() height, width, chn = rgb.shape cy = (y0 + y1) // 2 cx = (x0 + x1) // 2 y0 = cy - int(np.floor(sz / 2)) y1 = cy + int(np.ceil(sz / 2)) x0 = cx - int(np.floor(sz / 2)) x1 = cx + int(np.ceil(sz / 2)) out = rgba[ max(y0, 0) : min(y1, height) , max(x0, 0) : min(x1, width), : ].copy() pads = [(max(0-y0, 0), max(y1-height, 0)), (max(0-x0, 0), max(x1-width, 0)), (0, 0)] out = np.pad(out, pads, mode='constant', constant_values=0) assert(out.shape[:2] == (sz, sz)) out[:, :, :3] = out[:, :, :3] * (out[..., 3:]/255.) + np.array(bkg_color)[None, None, :3] * (1-out[..., 3:]/255.) out[:, :, -1] = bkg_color[-1] out = cv2.resize(out.astype(np.uint8), (256, 256)) out = out[:, :, :3] out_rgbs.append(out) return out_rgbs def run_preprocess(image1, image2, preprocess_chk, seed_value): seed_everything(seed_value) if preprocess_chk: rembg_session = rembg.new_session() image1 = remove_background(image1, force=True, rembg_session = rembg_session) image2 = remove_background(image2, force=True, rembg_session = rembg_session) rgbs = group_recenter([image1, image2]) image1 = Image.fromarray(rgbs[0]) image2 = Image.fromarray(rgbs[1]) return image1, image2 def image_to_tensor(img, width=256, height=256): img = transforms.ToTensor()(img).unsqueeze(0) img = img * 2 - 1 img = transforms.functional.resize(img, [height, width]) return img @spaces.GPU def run_pose_exploration_a(image1, image2, seed_value): seed_everything(seed_value) image1 = image_to_tensor(image1).to(_device_) image2 = image_to_tensor(image2).to(_device_) images = [image1, image2] elevs, elev_ranges = estimate_elevs( _model_, images, est_type='all', matcher_ckpt_path=_matcher_ckpt_path_ ) fig = None return elevs, elev_ranges, fig @spaces.GPU def run_pose_exploration_b(cam_vis, image1, image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters, seed_value): seed_everything(seed_value) noise = np.random.randn(probe_bsz, 4, 32, 32) cam_vis.set_images([np.asarray(image1, dtype=np.uint8), np.asarray(image2, dtype=np.uint8)]) image1 = image_to_tensor(image1).to(_device_) image2 = image_to_tensor(image2).to(_device_) images = [image1, image2] result_poses, aux_data = estimate_poses( _model_, images, seed_cand_num=8, explore_type='triangular', refine_type='triangular', probe_ts_range=[0.2, 0.21], ts_range=[0.2, 0.21], probe_bsz=probe_bsz, adjust_factor=10.0, adjust_iters=adj_iters, adjust_bsz=adj_bsz, refine_factor=1.0, refine_iters=0, refine_bsz=4, noise=noise, elevs=elevs, elev_ranges=elev_ranges ) theta, azimuth, radius = result_poses[0] anchor_polar = aux_data['elev'][0] if anchor_polar is None: anchor_polar = np.pi/2 xyz0 = spherical_to_cartesian((anchor_polar, 0., 4.)) c2w0 = elu_to_c2w(xyz0, np.zeros(3), np.array([0., 0., 1.])) xyz1 = spherical_to_cartesian((theta + anchor_polar, 0. + azimuth, 4. + radius)) c2w1 = elu_to_c2w(xyz1, np.zeros(3), np.array([0., 0., 1.])) cam_vis._poses = [c2w0, c2w1] fig = cam_vis.update_figure(5, base_radius=-1.2, font_size=16, show_background=True, show_grid=True, show_ticklabels=True) explored_sph = (theta, azimuth, radius) return anchor_polar, explored_sph, fig, gr.update(interactive=True) @spaces.GPU def run_pose_refinement(cam_vis, image1, image2, anchor_polar, explored_sph, refine_iters, seed_value): seed_everything(seed_value) cam_vis.set_images([np.asarray(image1, dtype=np.uint8), np.asarray(image2, dtype=np.uint8)]) image1 = image_to_tensor(image1).to(_device_) image2 = image_to_tensor(image2).to(_device_) images = [image1, image2] images = [ img.permute(0, 2, 3, 1) for img in images ] out_poses, _, loss = find_optimal_poses( _model_, images, 1.0, bsz=1, n_iter=refine_iters, init_poses={1: explored_sph}, ts_range=[0.2, 0.21], combinations=[(0, 1), (1, 0)], avg_last_n=20, print_n=100 ) final_sph = out_poses[0] theta, azimuth, radius = final_sph xyz0 = spherical_to_cartesian((anchor_polar, 0., 4.)) c2w0 = elu_to_c2w(xyz0, np.zeros(3), np.array([0., 0., 1.])) xyz1 = spherical_to_cartesian((theta + anchor_polar, 0. + azimuth, 4. + radius)) c2w1 = elu_to_c2w(xyz1, np.zeros(3), np.array([0., 0., 1.])) cam_vis._poses = [c2w0, c2w1] fig = cam_vis.update_figure(5, base_radius=-1.2, font_size=16, show_background=True, show_grid=True, show_ticklabels=True) return final_sph, fig _HEADER_ = ''' # Official 🤗 Gradio Demo for [ID-Pose: Sparse-view Camera Pose Estimation By Inverting Diffusion Models](https://github.com/xt4d/id-pose) - ID-Pose accepts input images with NO overlapping appearance. - The estimation takes about 1 minute. ZeroGPU may be halted during processing due to quota restrictions. ''' _FOOTER_ = ''' - Project Page: [https://xt4d.github.io/id-pose-web/](https://xt4d.github.io/id-pose-web/) - Github: [https://github.com/xt4d/id-pose](https://github.com/xt4d/id-pose) ''' _CITE_ = r""" ```bibtex @article{cheng2023id, title={ID-Pose: Sparse-view Camera Pose Estimation by Inverting Diffusion Models}, author={Cheng, Weihao and Cao, Yan-Pei and Shan, Ying}, journal={arXiv preprint arXiv:2306.17140}, year={2023} } ``` """ def run_demo(): demo = gr.Blocks(title='ID-Pose: Sparse-view Camera Pose Estimation By Inverting Diffusion Models') with demo: gr.Markdown(_HEADER_) with gr.Row(variant='panel'): with gr.Column(scale=1): with gr.Row(): with gr.Column(min_width=280): input_image1 = gr.Image(type='pil', image_mode='RGBA', label='Input Image 1', width=280) with gr.Column(min_width=280): input_image2 = gr.Image(type='pil', image_mode='RGBA', label='Input Image 2', width=280) with gr.Row(): with gr.Column(min_width=280): processed_image1 = gr.Image(type='numpy', image_mode='RGB', label='Processed Image 1', width=280, interactive=False) with gr.Column(min_width=280): processed_image2 = gr.Image(type='numpy', image_mode='RGB', label='Processed Image 2', width=280, interactive=False) with gr.Row(): preprocess_chk = gr.Checkbox(True, label='Remove background and recenter object') with gr.Accordion('Advanced options', open=False): probe_bsz = gr.Slider(4, 32, value=16, step=4, label='Probe Batch Size') adj_bsz = gr.Slider(1, 8, value=4, step=1, label='Adjust Batch Size') adj_iters = gr.Slider(1, 20, value=8, step=1, label='Adjust Iterations') seed_value = gr.Number(value=0, label="Seed Value", precision=0) with gr.Row(): run_btn = gr.Button('Estimate', variant='primary', interactive=True) with gr.Row(): refine_iters = gr.Slider(0, 1000, value=0, step=50, label='Refinement Iterations') with gr.Row(): refine_btn = gr.Button('Refine', variant='primary', interactive=False) with gr.Row(): gr.Markdown(_FOOTER_) with gr.Row(): gr.Markdown(_CITE_) with gr.Column(scale=1.4): with gr.Row(): vis_output = gr.Plot(label='Camera Pose Results: anchor (red) and target (blue)') with gr.Row(): with gr.Column(min_width=200): gr.Examples( examples = [ ['data/gradio_demo/duck_0.png', 'data/gradio_demo/duck_1.png'], ['data/gradio_demo/chair_0.png', 'data/gradio_demo/chair_1.png'], ['data/gradio_demo/foosball_0.png', 'data/gradio_demo/foosball_1.png'], ], inputs=[input_image1, input_image2], label='Examples (Self-captured)', cache_examples=False, examples_per_page=3 ) with gr.Column(min_width=200): gr.Examples( examples = [ ['data/gradio_demo/bunny_0.png', 'data/gradio_demo/bunny_1.png'], ['data/gradio_demo/bus_0.png', 'data/gradio_demo/bus_1.png'], ['data/gradio_demo/circo_0.png', 'data/gradio_demo/circo_1.png'], ], inputs=[input_image1, input_image2], label='Examples (Images from NAVI)', cache_examples=False, examples_per_page=3 ) with gr.Column(min_width=200): gr.Examples( examples = [ ['data/gradio_demo/status_0.png', 'data/gradio_demo/status_1.png'], ['data/gradio_demo/bag_0.png', 'data/gradio_demo/bag_1.png'], ['data/gradio_demo/cat_0.png', 'data/gradio_demo/cat_1.png'], ], inputs=[input_image1, input_image2], label='Examples (Generated)', cache_examples=False, examples_per_page=3 ) cam_vis = CameraVisualizer([np.eye(4), np.eye(4)], ['Image 1', 'Image 2'], ['red', 'blue']) explored_sph = gr.State() anchor_polar = gr.State() refined_sph = gr.State() elevs = gr.State() elev_ranges = gr.State() run_btn.click( fn=run_preprocess, inputs=[input_image1, input_image2, preprocess_chk, seed_value], outputs=[processed_image1, processed_image2], ).success( fn=run_pose_exploration_a, inputs=[processed_image1, processed_image2, seed_value], outputs=[elevs, elev_ranges, vis_output] ).success( fn=partial(run_pose_exploration_b, cam_vis), inputs=[processed_image1, processed_image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters, seed_value], outputs=[anchor_polar, explored_sph, vis_output, refine_btn] ) refine_btn.click( fn=partial(run_pose_refinement, cam_vis), inputs=[processed_image1, processed_image2, anchor_polar, explored_sph, refine_iters, seed_value], outputs=[refined_sph, vis_output] ) demo.launch() if __name__ == '__main__': run_demo()