import sys import spaces sys.path.append("flash3d") # Add the flash3d directory to the system path for importing local modules from omegaconf import OmegaConf import gradio as gr import torch import torchvision.transforms as TT import torchvision.transforms.functional as TTF from huggingface_hub import hf_hub_download import numpy as np from networks.gaussian_predictor import GaussianPredictor from util.vis3d import save_ply def main(): print("[INFO] Starting main function...") # Determine if CUDA (GPU) is available and set the device accordingly if torch.cuda.is_available(): device = "cuda:0" print("[INFO] CUDA is available. Using GPU device.") else: device = "cpu" print("[INFO] CUDA is not available. Using CPU device.") # Download model configuration and weights from Hugging Face Hub print("[INFO] Downloading model configuration...") model_cfg_path = hf_hub_download(repo_id="einsafutdinov/flash3d", filename="config_re10k_v1.yaml") print("[INFO] Downloading model weights...") model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", filename="model_re10k_v1.pth") # Load model configuration using OmegaConf print("[INFO] Loading model configuration...") cfg = OmegaConf.load(model_cfg_path) # Initialize the GaussianPredictor model with the loaded configuration print("[INFO] Initializing GaussianPredictor model...") model = GaussianPredictor(cfg) device = torch.device(device) model.to(device) # Move the model to the specified device (CPU or GPU) # Load the pre-trained model weights print("[INFO] Loading model weights...") model.load_model(model_path) # Define transformation functions for image preprocessing pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug)) # Padding to augment the image borders to_tensor = TT.ToTensor() # Convert image to tensor # Function to check if an image is uploaded by the user def check_input_image(input_image): print("[DEBUG] Checking input image...") if input_image is None: print("[ERROR] No image uploaded!") raise gr.Error("No image uploaded!") print("[INFO] Input image is valid.") # Function to preprocess the input image before passing it to the model def preprocess(image): print("[DEBUG] Preprocessing image...") # Resize the image to the desired height and width specified in the configuration image = TTF.resize( image, (cfg.dataset.height, cfg.dataset.width), interpolation=TT.InterpolationMode.BICUBIC ) # Apply padding to the image image = pad_border_fn(image) print("[INFO] Image preprocessing complete.") return image # Function to reconstruct the 3D model from the input image and export it as a PLY file import sys import spaces sys.path.append("flash3d") # Add the flash3d directory to the system path for importing local modules from omegaconf import OmegaConf import gradio as gr import torch import torchvision.transforms as TT import torchvision.transforms.functional as TTF from huggingface_hub import hf_hub_download import numpy as np from networks.gaussian_predictor import GaussianPredictor from util.vis3d import save_ply def main(): print("[INFO] Starting main function...") # Determine if CUDA (GPU) is available and set the device accordingly if torch.cuda.is_available(): device = "cuda:0" print("[INFO] CUDA is available. Using GPU device.") else: device = "cpu" print("[INFO] CUDA is not available. Using CPU device.") # Download model configuration and weights from Hugging Face Hub print("[INFO] Downloading model configuration...") model_cfg_path = hf_hub_download(repo_id="einsafutdinov/flash3d", filename="config_re10k_v1.yaml") print("[INFO] Downloading model weights...") model_path = hf_hub_download(repo_id="einsafutdinov/flash3d", filename="model_re10k_v1.pth") # Load model configuration using OmegaConf print("[INFO] Loading model configuration...") cfg = OmegaConf.load(model_cfg_path) # Initialize the GaussianPredictor model with the loaded configuration print("[INFO] Initializing GaussianPredictor model...") model = GaussianPredictor(cfg) device = torch.device(device) model.to(device) # Move the model to the specified device (CPU or GPU) # Load the pre-trained model weights print("[INFO] Loading model weights...") model.load_model(model_path) # Define transformation functions for image preprocessing pad_border_fn = TT.Pad((cfg.dataset.pad_border_aug, cfg.dataset.pad_border_aug)) # Padding to augment the image borders to_tensor = TT.ToTensor() # Convert image to tensor # Function to check if an image is uploaded by the user def check_input_image(input_image): print("[DEBUG] Checking input image...") if input_image is None: print("[ERROR] No image uploaded!") raise gr.Error("No image uploaded!") print("[INFO] Input image is valid.") # Function to preprocess the input image before passing it to the model def preprocess(image): print("[DEBUG] Preprocessing image...") # Resize the image to the desired height and width specified in the configuration image = TTF.resize( image, (cfg.dataset.height, cfg.dataset.width), interpolation=TT.InterpolationMode.BICUBIC ) # Apply padding to the image image = pad_border_fn(image) print("[INFO] Image preprocessing complete.") return image # Function to reconstruct the 3D model from the input image and export it as a PLY file @spaces.GPU(duration=120) # Decorator to allocate a GPU for this function during execution def reconstruct_and_export(image): """ Passes image through model, outputs reconstruction in form of a dict of tensors. """ print("[DEBUG] Starting reconstruction and export...") # Convert the preprocessed image to a tensor and move it to the specified device image = to_tensor(image).to(device).unsqueeze(0) inputs = { ("color_aug", 0, 0): image, } # Pass the image through the model to get the output print("[INFO] Passing image through the model...") outputs = model(inputs) # Export the reconstruction to a PLY file print(f"[INFO] Saving output to {ply_out_path}...") save_ply(outputs, ply_out_path, num_gauss=2) print("[INFO] Reconstruction and export complete.") return ply_out_path # Path to save the output PLY file ply_out_path = f'./mesh.ply' # CSS styling for the Gradio interface css = """ h1 { text-align: center; display:block; } """ # Create the Gradio user interface with gr.Blocks(css=css) as demo: gr.Markdown( """ # Flash3D """ ) # Comments about the app's behavior and known limitations gr.Markdown( """ ## Comments: 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. 2. The 3D viewer shows a .ply mesh extracted from a mix of 3D Gaussians. This is only an approximation and artifacts might show. 3. Known limitations include: - A black dot appearing on the model from some viewpoints. - See-through parts of objects, especially on the back: this is due to the model performing less well on more complicated shapes. - Back of objects are blurry: this is a model limitation due to it being deterministic. 4. Our model is of comparable quality to state-of-the-art methods, and is **much** cheaper to train and run. ## How does it work? Splatter Image formulates 3D reconstruction as an image-to-image translation task. It maps the input image to another image, 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. The resulting image thus represents a set of Gaussians (almost like a point cloud) which reconstruct the shape and colour of the object. 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). The rendering is also very fast, due to using Gaussian Splatting. Combined, this results in very cheap training and high-quality results. For more results see the [project page](https://szymanowiczs.github.io/splatter-image) and the [CVPR article](https://arxiv.org/abs/2312.13150). """ ) with gr.Row(variant="panel"): with gr.Column(scale=1): with gr.Row(): # Input image component for the user to upload an image input_image = gr.Image( label="Input Image", image_mode="RGBA", sources="upload", type="pil", elem_id="content_image", ) with gr.Row(): # Button to trigger the generation process submit = gr.Button("Generate", elem_id="generate", variant="primary") with gr.Row(variant="panel"): # Examples panel to provide sample images for users gr.Examples( examples=[ './demo_examples/bedroom_01.png', './demo_examples/kitti_02.png', './demo_examples/kitti_03.png', './demo_examples/re10k_04.jpg', './demo_examples/re10k_05.jpg', './demo_examples/re10k_06.jpg', ], inputs=[input_image], cache_examples=False, label="Examples", examples_per_page=20, ) with gr.Row(): # Display the preprocessed image (after resizing and padding) processed_image = gr.Image(label="Processed Image", interactive=False) with gr.Column(scale=2): with gr.Row(): with gr.Tab("Reconstruction"): # 3D model viewer to display the reconstructed model output_model = gr.Model3D( height=512, label="Output Model", interactive=False ) # Define the workflow for the Generate button submit.click(fn=check_input_image, inputs=[input_image]).success( fn=preprocess, inputs=[input_image], outputs=[processed_image], ).success( fn=reconstruct_and_export, inputs=[processed_image], outputs=[output_model], ) # Queue the requests to handle them sequentially (to avoid GPU resource conflicts) demo.queue(max_size=1) print("[INFO] Launching Gradio demo...") demo.launch(share=True) # Launch the Gradio interface and allow public sharing if __name__ == "__main__": print("[INFO] Running application...") main() # Decorator to allocate a GPU for this function during execution def reconstruct_and_export(image): """ Passes image through model, outputs reconstruction in form of a dict of tensors. """ print("[DEBUG] Starting reconstruction and export...") # Convert the preprocessed image to a tensor and move it to the specified device image = to_tensor(image).to(device).unsqueeze(0) inputs = { ("color_aug", 0, 0): image, } # Pass the image through the model to get the output print("[INFO] Passing image through the model...") outputs = model(inputs) # Export the reconstruction to a PLY file print(f"[INFO] Saving output to {ply_out_path}...") save_ply(outputs, ply_out_path, num_gauss=2) print("[INFO] Reconstruction and export complete.") return ply_out_path # Path to save the output PLY file ply_out_path = f'./mesh.ply' # CSS styling for the Gradio interface css = """ h1 { text-align: center; display:block; } """ # Create the Gradio user interface with gr.Blocks(css=css) as demo: gr.Markdown( """ # Flash3D """ ) # Comments about the app's behavior and known limitations gr.Markdown( """ ## Comments: 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. 2. The 3D viewer shows a .ply mesh extracted from a mix of 3D Gaussians. This is only an approximation and artifacts might show. 3. Known limitations include: - A black dot appearing on the model from some viewpoints. - See-through parts of objects, especially on the back: this is due to the model performing less well on more complicated shapes. - Back of objects are blurry: this is a model limitation due to it being deterministic. 4. Our model is of comparable quality to state-of-the-art methods, and is **much** cheaper to train and run. ## How does it work? Splatter Image formulates 3D reconstruction as an image-to-image translation task. It maps the input image to another image, 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. The resulting image thus represents a set of Gaussians (almost like a point cloud) which reconstruct the shape and colour of the object. 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). The rendering is also very fast, due to using Gaussian Splatting. Combined, this results in very cheap training and high-quality results. For more results see the [project page](https://szymanowiczs.github.io/splatter-image) and the [CVPR article](https://arxiv.org/abs/2312.13150). """ ) with gr.Row(variant="panel"): with gr.Column(scale=1): with gr.Row(): # Input image component for the user to upload an image input_image = gr.Image( label="Input Image", image_mode="RGBA", sources="upload", type="pil", elem_id="content_image", ) with gr.Row(): # Button to trigger the generation process submit = gr.Button("Generate", elem_id="generate", variant="primary") with gr.Row(variant="panel"): # Examples panel to provide sample images for users gr.Examples( examples=[ './demo_examples/bedroom_01.png', './demo_examples/kitti_02.png', './demo_examples/kitti_03.png', './demo_examples/re10k_04.jpg', './demo_examples/re10k_05.jpg', './demo_examples/re10k_06.jpg', ], inputs=[input_image], cache_examples=False, label="Examples", examples_per_page=20, ) with gr.Row(): # Display the preprocessed image (after resizing and padding) processed_image = gr.Image(label="Processed Image", interactive=False) with gr.Column(scale=2): with gr.Row(): with gr.Tab("Reconstruction"): # 3D model viewer to display the reconstructed model output_model = gr.Model3D( height=512, label="Output Model", interactive=False ) # Define the workflow for the Generate button submit.click(fn=check_input_image, inputs=[input_image]).success( fn=preprocess, inputs=[input_image], outputs=[processed_image], ).success( fn=reconstruct_and_export, inputs=[processed_image], outputs=[output_model], ) # Queue the requests to handle them sequentially (to avoid GPU resource conflicts) demo.queue(max_size=1) print("[INFO] Launching Gradio demo...") demo.launch(share=True) # Launch the Gradio interface and allow public sharing if __name__ == "__main__": print("[INFO] Running application...") main()