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
- 3. Project Structure
- 4. (optional) Retraining
- 5. Acknowledgments
- 6. Licence
- 🚨 Release of the TexSD dataset!
- 🚨 3D rendering script.
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 of 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".
-
Download the source code with git
git clone https://github.com/astra-vision/MaterialPalette.git
The repo can also be download as a zip here.
-
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.
-
Load the conda environment:
conda activate matpal
-
If you are looking to perform decomposition, download our pretrained 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
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)
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
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
:
Under the hood, it uses two modules:
concept
, to extract and generate the texture (concept.crop
,concept.invert
, andconcept.infer
);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).
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.
We provide the pretrained 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
In order 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.
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 where 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 licences. Special credits to AmbientCG for the huge work.
Authors would also like to thank all members of the Astra-Vision team for their valuable feedback.
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.