DOCSEARCHbyAlgolia
Submit Docs

Wonder3D

repository·main·Indexed 26 days ago

https://github.com/xxlong0/wonder3d

A CVPR 2024 Highlight project that reconstructs highly-detailed textured 3D meshes from a single-view image in 2-3 minutes using cross-domain diffusion. The repository includes tools for generating multi-view normal maps and color images, as well as mesh extraction via Instant-NSR or NeuS. It provides support for Docker, Conda, and BlenderProc for rendering.

Tokens
5.1K
Snippets
16
Records
26
Agent score
90%

What's inside Wonder3D

  1. Quickstart: Run the Wonder3D Pipeline via Python

    main

    You can use Wonder3D as a Python library to reconstruct multi-view normal maps and color images from a single image.

    Requirements:

    • diffusers[torch]==0.19.3 (other versions may cause conflicts).
    • A GPU with CUDA support.
    • An input image where the object is centered and resized to approximately 80% of the image height.

    Usage Example:

    import torch
import requests
from PIL import Image
import numpy as np
from torchvision.utils import make_grid, save_image
from diffusers import DiffusionPipeline  # only tested on diffusers[torch]==0.19.3

def load_wonder3d_pipeline():
    pipeline = DiffusionPipeline.from_pretrained(
    'flamehaze1115/wonder3d-v1.0', # or use local checkpoint './ckpts'
    custom_pipeline='flamehaze1115/wonder3d-pipeline',
    torch_dtype=torch.float16
    )

    # enable xformers
    pipeline.unet.enable_xformers_memory_efficient_attention()

    if torch.cuda.is_available():
        pipeline.to('cuda:0')
    return pipeline

pipeline = load_wonder3d_pipeline()

# Download an example image.
cond = Image.open(requests.get("https://d.skis.ltd/nrp/sample-data/lysol.png", stream=True).raw)

# The object should be located in the center and resized to 80% of image height.
cond = Image.fromarray(np.array(cond)[:, :, :3])

# Run the pipeline!
images = pipeline(cond, num_inference_steps=20, output_type='pt', guidance_scale=1.0).images

result = make_grid(images, nrow=6, ncol=2, padding=0, value_range=(0, 1))

save_image(result, 'result.png')

  2. Install tiny-cuda-nn inside the Docker container

    main

    The tiny-cuda-nn package cannot be installed during the docker build process. You must run this installation command manually after you have started the Docker container and entered its bash shell.

    pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch

  3. Set up the rendering environment

    main

    The rendering code is based on BlenderProc. Note that using the Blender Cycles engine may cause long-time hanging issues on certain GPUs (e.g., A800).

    To set up the environment, navigate to the render_codes directory and install the required dependencies.

    cd ./render_codes
pip install -r requirements.txt

  4. Install NVIDIA Container Toolkit for Docker GPU support

    main

    If you cannot enable GPUs within Docker, you must install the NVIDIA Container Toolkit on your local machine. Follow these steps to install it via apt:

    1. Configure the production repository.
    2. Enable experimental features in the toolkit list.
    3. Install the packages.
    4. Restart the Docker service.

  5. Extract 3D Mesh using Instant-NSR

    main

    After generating multi-view images, you can extract a mesh using instant-nsr-pl. Note that the generated views are in orthographic camera space.

    Example Command:

    cd ./instant-nsr-pl
python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../outputs/cropsize-192-cfg1.0/ dataset.scene=owl

    Tip: To improve texture quality, increase the optimization steps by modifying trainer.max_steps in instant-nsr-pl/configs/neuralangelo-ortho-wmask.yaml (e.g., to 10000).

  6. Build and run Wonder3D using Docker

    main

    To set up the Wonder3D environment using Docker on Ubuntu 20.04, build the image from the root directory of the repository and then run the container with GPU support enabled.

    Note: Ensure you are in the yourworkspace/Wonder3D/ directory before executing these commands.

  7. Prepare and train on custom COLMAP data

    main

    To use custom images, you must first process them with COLMAP.

    1. Place your images in an images/ folder.
    2. Run the scripts/imgs2poses.py script pointing to the directory containing the images/ folder.
    3. Use a -colmap.yaml configuration file for training.

    Configuration Tips:

    • Adjust root_dir and img_wh (or img_downscale) in the config to match your data.
    • Set model.radius=1.0 for normalized scenes; adjust if the object is not properly wrapped.
    • Choose a scene center method via dataset.center_est_method: camera (center of camera positions), lookat (approximate look-at point), or point (center of COLMAP points).

  8. Extract mesh using Instant-NSR

    main

    After generating views, you can perform mesh extraction using instant-nsr-pl. Note that the generated views are in orthographic camera space; when viewing in MeshLab, use 'Toggle Orthographic Camera'.

    To improve texture quality, you can increase the optimization steps by modifying trainer.max_steps in instant-nsr-pl/configs/neuralangelo-ortho-wmask.yaml (e.g., from 3000 to 10000).

    cd ./instant-nsr-pl
python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../{your_save_path}/cropsize-{crop_size}-cfg{guidance_scale:.1f}/ dataset.scene={scene}

  9. Prepare for Inference (Manual Model Download)

    main

    If you cannot connect to Hugging Face, follow these steps to prepare the models manually:

    1. Download Checkpoints: Download the required checkpoints and place them in a ckpts folder in the root directory:

      Wonder3D
|-- ckpts
    |-- unet
    |-- scheduler
    |-- vae
    ...

      Then, update ./configs/mvdiffusion-joint-ortho-6views.yaml by setting pretrained_model_name_or_path="./ckpts".

    2. Download SAM: Download the sam_vit_h_4b8939.pth model and place it in the sam_pt folder:

      Wonder3D
|-- sam_pt
    |-- sam_vit_h_4b8939.pth

    3. Segment Foreground: Use rembg or Clipdrop to remove the background from your input image to ensure high-quality mesh reconstruction.

  10. Prepare models and configuration for inference

    main

    Before running inference, ensure you have the following model files in place:

    1. Checkpoints: Download the checkpoints and place them in a ckpts folder in the root directory.
      • Structure: Wonder3D/ckpts/{unet, scheduler, vae, ...}
      • Configuration: Update ./configs/mvdiffusion-joint-ortho-6views.yaml and set pretrained_model_name_or_path="./ckpts".
    2. SAM: Download the SAM model (sam_vit_h_4b8939.pth) and place it in the sam_pt folder.
      • Structure: Wonder3D/sam_pt/sam_vit_h_4b8939.pth.
    3. Foreground Mask: It is highly recommended to remove the background from your input image using rembg or Clipdrop to improve mesh quality.

  12. Train Wonder3D Models

    main

    Training is performed in two stages using accelerate. You must modify the root_dir in the corresponding config files to point to your data.

    Stage 1: Train multi-view attentions

    accelerate launch --config_file 8gpu.yaml train_mvdiffusion_image.py --config configs/train/stage1-mix-6views-lvis.yaml

    Stage 2: Add cross-domain attention modules

    accelerate launch --config_file 8gpu.yaml train_mvdiffusion_joint.py --config configs/train/stage2-joint-6views-lvis.yaml
    # stage 1:
accelerate launch --config_file 8gpu.yaml train_mvdiffusion_image.py --config configs/train/stage1-mix-6views-lvis.yaml

# stage 2
accelerate launch --config_file 8gpu.yaml train_mvdiffusion_joint.py --config configs/train/stage2-joint-6views-lvis.yaml