Material Palette: Extraction of Materials from a Single Image (CVPR 2024)

Ivan Lopes¹   Fabio Pizzati²   Raoul de Charette^1
1 Inria, ² Oxford Uni.

TL;DR, Material Palette extracts a palette of PBR materials -
albedo, normals, and roughness - from a single real-world image.

teaser.mp4

• Overview
• 1. Installation
• 2. Quick Start
  • Generation
  • Complete Pipeline
• 3. Project Structure
• 4. (optional) Retraining
• Acknowledgments
• Licence

🚨 Todo

• 3D rendering script.

Overview

This is the official repository of Material Palette. In a nutshell, the method works in three stages: first, concepts are extracted from an input image based on a user-provided mask; then, those concepts are used to generate texture images; finally, the generations are decomposed into SVBRDF maps (albedo, normals, and roughness). Visit our project page or consult our paper for more details!

Content: This repository allows the extraction of texture concepts from image and region mask sets. It also allows generation at different resolutions. Finally, it proposes a decomposition step thanks to our decomposition model, for which we share the training weights.

Tip

We propose a "Quick Start" section: before diving straight into the full pipeline, we share four pretrained concepts ⚡ so you can go ahead and experiment with the texture generation step of the method: see "§ Generation". Then you can try out the full method with your own image and masks = concept learning + generation + decomposition, see "§ Complete Pipeline".

1. Installation

1. Download the source code with git

   git clone https://github.com/astra-vision/MaterialPalette.git
   

   The repo can also be downloaded as a zip here.

2. Create a conda environment with the dependencies.

   conda env create --verbose -f deps.yml
   

   This repo was tested with Python 3.10.8, PyTorch 1.13, diffusers 0.19.3, peft 0.5, and PyTorch Lightning 1.8.3.

3. Load the conda environment:

   conda activate matpal
   

4. If you are looking to perform decomposition, download our pre-trained model and untar the archive:

   wget https://github.com/astra-vision/MaterialPalette/releases/download/weights/model.tar.gz
   

   ^{This is not required if you are only looking to perform texture extraction}

2. Quick start

Here are instructions to get you started using Material Palette. First, we provide some optimized concepts so you can experiment with the generation pipeline. We also show how to run the method on user-selected images and masks (concept learning + generation + decomposition)

§ Generation

Input image	1K	2K	4K	8K	⬇️ LoRA ~8Kb

^{All generations were downscaled for memory constraints.}

Go ahead and download one of the above LoRA concept checkpoints, example for "blue_tiles":

wget https://github.com/astra-vision/MaterialPalette/files/14601640/blue_tiles.zip;
unzip blue_tiles.zip

To generate from a checkpoint, use the concept module either via the command line interface or the functional interface in python:

• 

  python concept/infer.py path/to/LoRA/checkpoint
  

• 

  import concept
  concept.infer(path_to_LoRA_checkpoint)
  

Results will be placed relative to the checkpoint directory in a outputs folder.

You have control over the following parameters:

• stitch_mode: concatenation, average, or weighted average (default);
• resolution: the output resolution of the generated texture;
• prompt: one of the four prompt templates:
  • "p1": "top view realistic texture of S*",
  • "p2": "top view realistic S* texture",
  • "p3": "high resolution realistic S* texture in top view",
  • "p4": "realistic S* texture in top view";
• seed: inference seed when sampling noise;
• renorm: whether or not to renormalize the generated samples generations based on input image (this option can only be used when called from inside the pipeline, ie. when the input image is available);
• num_inference_steps: number of denoising steps.

^{A complete list of parameters can be viewed with python concept/infer.py --help}

§ Complete Pipeline

We provide an example (input image with user masks used for the pipeline figure). You can download it here: mansion.zip (credits photograph: Max Rahubovskiy).

To help you get started with your own images, you should follow this simple data structure: one folder per inverted image, inside should be the input image (.jpg, .jpeg, or .png) and a subdirectory named masks containing the different region masks as .png (these must all have the same aspect ratio as the RGB image). Here is an overview of our mansion example:

├── masks/
│ ├── wood.png
│ ├── grass.png
│ └── stone.png
└── mansion.jpg

region	mask	overlay	generation	albedo	normals	roughness

To invert and generate textures from a folder, use pipeline.py:

• 

  python pipeline.py path/to/folder
  

Under the hood, it uses two modules:

1. concept, to extract and generate the texture (concept.crop, concept.invert, and concept.infer);
2. capture, to perform the BRDF decomposition.

A minimal example is provided here:

• 

  ## Extract square crops from image for each of the binary masks located in <path>/masks
  regions = concept.crop(args.path)
  
  ## Iterate through regions to invert the concept and generate texture views
  for region in regions.iterdir():
      lora = concept.invert(region)
      concept.infer(lora, renorm=True)
  
  ## Construct a dataset with all generations and load pretrained decomposition model
  data = capture.get_data(predict_dir=args.path, predict_ds='sd')
  module = capture.get_inference_module(pt='model.ckpt')
  
  ## Proceed with inference on decomposition model
  decomp = Trainer(default_root_dir=args.path, accelerator='gpu', devices=1, precision=16)
  decomp.predict(module, data)
  

^{To view options available for the concept learning, use PYTHONPATH=. python concept/invert.py --help}

Important

By default, both train_text_encoder and gradient_checkpointing are set to True. Also, this implementation does not include the LPIPS filter/ranking of the generations. The code will only output a single sample per region. You may experiment with different prompts and parameters (see "Generation" section).

3. Project structure

The pipeline.py file is the entry point to run the whole pipeline on a folder containing the input image at its root and a masks/ sub-directory containing all user defined masks. The train.py file is used to train the decomposition model. The most important files are shown here:

.
├── capture/        % Module for decomposition
│ ├── callbacks/    % Lightning trainer callbacks
│ ├── data/         % Dataset, subsets, Lightning datamodules
│ ├── render/       % 2D physics based renderer
│ ├── utils/        % Utility functions
│ └── source/       % Network, loss, and LightningModule
│   └── routine.py  % Training loop
│
└── concept/        % Module for inversion and texture generation
  ├── crop.py       % Square crop extraction from image and masks
  ├── invert.py     % Optimization code to learn the concept S*
  └── infer.py      % Inference code to generate texture from S*

If you have any questions, post via the issues tracker or contact the corresponding author.

4. (optional) Training

We provide the pre-trained decomposition weights (see "Installation"). However, if you are looking to retrain the domain adaptive model for your own purposes, we provide the code to do so. Our method relies on the training of a multi-task network on labeled (real) and unlabeled (synthetic) images, jointly. In case you wish to retrain on the same datasets, you will have to download both the AmbientCG and TexSD datasets.

First download the PBR materials (source) dataset from AmbientCG:

python capture/data/download.py path/to/target/directory

To run the training script, use:

python train.py --config=path/to/yml/config

^{Additional options can be found with python train.py --help.}

Note

The decomposition model allows estimating the pixel-wise BRDF maps from a single texture image input.

Acknowledgments

This research project was mainly funded by the French Agence Nationale de la Recherche (ANR) as part of project SIGHT (ANR-20-CE23-0016). Fabio Pizzati was partially funded by KAUST (Grant DFR07910). Results were obtained using HPC resources from GENCI-IDRIS (Grant 2023-AD011014389).

The repository contains code taken from PEFT, SVBRDF-Estimation, DenseMTL. As for visualization, we used DeepBump and Blender. Credit to Runway for providing us all the stable-diffusion-v1-5 model weights. All images and 3D scenes used in this work have permissive licenses. Special credits to AmbientCG for the huge work.

The authors would also like to thank all members of the Astra-Vision team for their valuable feedback.

License

If you find this code useful, please cite our paper:

@inproceedings{lopes2024material,
    author = {Lopes, Ivan and Pizzati, Fabio and de Charette, Raoul},
    title = {Material Palette: Extraction of Materials from a Single Image},
    booktitle = {CVPR},
    year = {2024},
    project = {https://astra-vision.github.io/MaterialPalette/}
}

Material Palette is released under MIT License.

---

🢁 jump to top