jt5d/splatter_image
0
1import torch2 3import os4from omegaconf import OmegaConf5import spaces 6 7from utils.app_utils import (8 remove_background, 9 resize_foreground, 10 set_white_background,11 resize_to_128,12 to_tensor,13 get_source_camera_v2w_rmo_and_quats,14 export_to_obj)15 16 17from scene.gaussian_predictor import GaussianSplatPredictor18 19import gradio as gr20 21import rembg22 23from huggingface_hub import hf_hub_download24 25def main():26 27 if torch.cuda.is_available():28 device = "cuda:0"29 else:30 device = "cpu"31 32 model_cfg_path = hf_hub_download(repo_id="szymanowiczs/splatter-image-v1", 33 filename="config_objaverse.yaml")34 model_path = hf_hub_download(repo_id="szymanowiczs/splatter-image-v1", 35 filename="model_latest.pth")36 37 model_cfg = OmegaConf.load(model_cfg_path) 38 model = GaussianSplatPredictor(model_cfg)39 40 ckpt_loaded = torch.load(model_path, map_location="cpu")41 model.load_state_dict(ckpt_loaded["model_state_dict"])42 model.to(device)43 44 # ============= image preprocessing =============45 rembg_session = rembg.new_session()46 47 def check_input_image(input_image):48 if input_image is None:49 raise gr.Error("No image uploaded!")50 51 def preprocess(input_image, preprocess_background=True, foreground_ratio=0.65):52 # 0.7 seems to be a reasonable foreground ratio53 if preprocess_background:54 image = input_image.convert("RGB")55 image = remove_background(image, rembg_session)56 image = resize_foreground(image, foreground_ratio)57 image = set_white_background(image)58 else:59 image = input_image60 if image.mode == "RGBA":61 image = set_white_background(image)62 image = resize_to_128(image)63 return image64 65 ply_out_path = f'./mesh.ply'66 67 @spaces.GPU()68 def reconstruct_and_export(image):69 """70 Passes image through model, outputs reconstruction in form of a dict of tensors.71 """72 image = to_tensor(image).to(device)73 view_to_world_source, rot_transform_quats = get_source_camera_v2w_rmo_and_quats()74 view_to_world_source = view_to_world_source.to(device)75 rot_transform_quats = rot_transform_quats.to(device)76 77 reconstruction_unactivated = model(78 image.unsqueeze(0).unsqueeze(0),79 view_to_world_source,80 rot_transform_quats,81 None,82 activate_output=False)83 84 # export reconstruction to ply85 export_to_obj(reconstruction_unactivated, ply_out_path)86 87 return ply_out_path88 89 css = """90 h1 {91 text-align: center;92 display:block;93 }94 """95 96 with gr.Blocks(css=css) as demo:97 gr.Markdown(98 """99 # Splatter Image100 101 **Splatter Image (CVPR 2024)** [[code](https://github.com/szymanowiczs/splatter-image), [project page](https://szymanowiczs.github.io/splatter-image)] is a fast, super cheap-to-train method for object 3D reconstruction from a single image. 102 The model used in the demo was trained on **Objaverse-LVIS on 2 A6000 GPUs for 3.5 days**.103 Locally, on an NVIDIA V100 GPU, reconstruction (forward pass of the network) can be done at 38FPS and rendering (with Gaussian Splatting) at 588FPS.104 Upload an image of an object or click on one of the provided examples to see how the Splatter Image does.105 The 3D viewer will render a .ply object exported from the 3D Gaussians, which is only an approximation.106 For best results run the demo locally and render locally with Gaussian Splatting - to do so, clone the [main repository](https://github.com/szymanowiczs/splatter-image).107 """108 )109 with gr.Row(variant="panel"):110 with gr.Column():111 with gr.Row():112 input_image = gr.Image(113 label="Input Image",114 image_mode="RGBA",115 sources="upload",116 type="pil",117 elem_id="content_image",118 )119 processed_image = gr.Image(label="Processed Image", interactive=False)120 with gr.Row():121 with gr.Group():122 preprocess_background = gr.Checkbox(123 label="Remove Background", value=True124 )125 with gr.Row():126 submit = gr.Button("Generate", elem_id="generate", variant="primary")127 128 with gr.Row(variant="panel"): 129 gr.Examples(130 examples=[131 './demo_examples/01_bigmac.png',132 './demo_examples/02_hydrant.jpg',133 './demo_examples/03_spyro.png',134 './demo_examples/04_lysol.png',135 './demo_examples/05_pinapple_bottle.png',136 './demo_examples/06_unsplash_broccoli.png',137 './demo_examples/07_objaverse_backpack.png',138 './demo_examples/08_unsplash_chocolatecake.png',139 './demo_examples/09_realfusion_cherry.png',140 './demo_examples/10_triposr_teapot.png'141 ],142 inputs=[input_image],143 cache_examples=False,144 label="Examples",145 examples_per_page=20,146 )147 with gr.Column():148 with gr.Row():149 with gr.Tab("Reconstruction"):150 output_model = gr.Model3D(151 height=512,152 label="Output Model",153 interactive=False154 )155 156 gr.Markdown(157 """158 ## Comments:159 1. If you run the demo online, the first example you upload should take about 4.5 seconds (with preprocessing, saving and overhead), the following take about 1.5s.160 2. The 3D viewer shows a .ply mesh extracted from a mix of 3D Gaussians. This is only an approximations and artefacts might show.161 3. Known limitations include:162 - a black dot appearing on the model from some viewpoints163 - see-through parts of objects, especially on the back: this is due to the model performing less well on more complicated shapes164 - back of objects are blurry: this is a model limiation due to it being deterministic165 4. Our model is of comparable quality to state-of-the-art methods, and is **much** cheaper to train and run.166 167 ## How does it work?168 169 Splatter Image formulates 3D reconstruction as an image-to-image translation task. It maps the input image to another image, 170 in which every pixel represents one 3D Gaussian and the channels of the output represent parameters of these Gaussians, including their shapes, colours and locations.171 The resulting image thus represents a set of Gaussians (almost like a point cloud) which reconstruct the shape and colour of the object.172 The method is very cheap: the reconstruction amounts to a single forward pass of a neural network with only 2D operators (2D convolutions and attention).173 The rendering is also very fast, due to using Gaussian Splatting.174 Combined, this results in very cheap training and high-quality results.175 For more results see the [project page](https://szymanowiczs.github.io/splatter-image) and the [CVPR article](https://arxiv.org/abs/2312.13150).176 """177 )178 179 180 181 submit.click(fn=check_input_image, inputs=[input_image]).success(182 fn=preprocess,183 inputs=[input_image, preprocess_background],184 outputs=[processed_image],185 ).success(186 fn=reconstruct_and_export,187 inputs=[processed_image],188 outputs=[output_model],189 )190 191 demo.queue(max_size=1)192 demo.launch()193 194 195if __name__ == "__main__":196 main()197 198# gradio app interface199