Project Page | Paper

Neural Gaffer is an end-to-end 2D relighting diffusion model that accurately relights any object in a single image under various lighting conditions. Moreover, by combining with other generative methods, our model enables many downstream 2D tasks, such as text-based relighting and object insertion. Our model can also operate as a strong relighting prior for 3D tasks, such as relighting a radiance field.

teaser1.mp4

0. TODO List

I'll be updating the following list. If you have any urgent requirements, such as needing to compare this method for an upcoming submission, please contact me via my email.

Release the checkpoint and the inference script for in-the-wild single image input
Release the training dataset for the diffusion model
Release the training code for the diffusion model
Release the 3D relighting code

1. Preparation

1.1 Installation

conda create -n neural-gaffer python=3.9
conda activate neural-gaffer
pip install -r requirements.txt
pip3 install -U xformers==0.0.28 --index-url https://download.pytorch.org/whl/cu118

1.2 Downloading the checkpoint

The checkpoint file will be saved in the ./log/neural_gaffer_res256 folder.

cd logs
wget https://huggingface.co/coast01/Neural_Gaffer/resolve/main/neural_gaffer_res256_ckpt.zip
unzip neural_gaffer_res256_ckpt.zip
cd ..

1.3 Downloading the training and validation dataset

Here we provide a subset of our training dataset and full set of the validation dataset. (But we didn't use all of the validation dataset when training the model and computing the metrics because it's a little too big).

The training dataset subset here only has 1000 objects, which is a subset of the full training dataset used to train our model (~ 90,000 objects). We provide the subset here to help you test if the code can be run correctly. The full training dataset object list and our validatoin dataset list in the ./filtered_object_list folder.

The original training dataset is too large to be uploaded to Hugging Face. I have provided the rendering code and preprocessing code here.

2. Training

Before running, please change the dataset directories (training and validation) in configs/neural_gaffer_training.txt. The following command trains the diffusion model for 2D relighting with 8 GPUs.

export NCCL_P2P_DISABLE=1 && export NCCL_IB_DISABLE=1 &&  accelerate launch --main_process_port 25525 --config_file configs/8_16fp.yaml  neural_gaffer_training.py  --dataloader_num_workers 32  --use_ema --gradient_checkpointing   --config configs/neural_gaffer_training.txt

3. Inference commands for 2D relighting

3.1 Relighting in-the-wild single image input

3.1.1 Image preprocessing: segment, rescale, and recenter

Put the input images under the --img_dir folder and run the following command to segment the foreground. The preprocessed data will be saved in --out_dir. Here, we borrow code from One-2-3-45.

Note: If your input images have masks and you don't want to do rescale and recenter, you can skip this step by manually saving the three-channel foreground and mask of each input image in the {$out_dir}/img and {$out_dir}/mask folders, respectively.

# download the pre-trained SAM to segment the foreground
# only need to run once
cd models/checkpoints
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
cd ../..
#################################################################
# Segment the foreground
python scripts/segment_foreground.py --img_dir ./demo --sam_ckpt ./models/checkpoints/sam_vit_h_4b8939.pth --out_dir ./preprocessed_data  --gpu_idx 0
# The preprocessed data will be saved in ./'preprocessed_data'

3.1.2 Preprocessing the target environment maps

Place the target environment maps in the --lighting_dir folder, then run the following command to preprocess them. The preprocessed data will be saved in the --output_dir folder. Use --frame_num to specify the number of frames for rotating the environment maps 360 degrees along the azimuthal direction.

python scripts/generate_bg_and_rotate_envir_map.py --lighting_dir 'demo/environment_map_sample' --output_dir './preprocessed_lighting_data' --frame_num 120

3.1.3 Relighting

The following command relights the input images stored in the --val_img_dir folder using preprocessed target lighting data from the --val_lighting_dir. The relighted images will be saved in the --save_dir folder. The checkpoint file for the diffusion model is located in the --output_dir folder.

In total, this command will generate 2,400 relighted images (