liuguoyou / tom-net

TOM-Net: Learning Transparent Object Matting from a Single Image, CVPR 2018

TOM-Net

TOM-Net: Learning Transparent Object Matting from a Single Image, CVPR 2018 (Spotlight),
Guanying Chen*, Kai Han*, Kwan-Yee K. Wong
(* equal contribution)

This paper addresses the problem of transparent object matting from a single image:

Dependencies

TOM-Net is implemented in Torch and tested with Ubuntu 14.04, please install Torch first following the official document.

• python 2.7
• numpy
• cv2
• CUDA-8.0
• CUDNN v5.1
• Torch STN (qassemoquab/stnbhwd)

  # Basic installation steps for stn
  git clone https://github.com/qassemoquab/stnbhwd.git
  cd stnbhwd
  luarocks make

Overview:

We provide:

• Pretrained model
• Datasets: Train (40GB), Validation (196MB), Test (179MB)
• Code to test model on new images
• Evaluation code on both the validation and testing data
• Instructions to train the model
• Example code for synthetic data rendering

Testing

Download Pretrained Model

sh scritps/download_pretrained_model.sh

Test on New Images

# Replace ${gpu} with the selected GPU ID (starting from 0)

# Test a single image without having the background image
CUDA_VISIBLE_DEVICES=${gpu} th eval/run_model.lua -input_img images/bull.jpg 

# You can find the results in data/TOM-Net_model/

Evaluation on Synthetic Validation Data

# Download synthetic validation dataset
sh scripts/download_validation_dataset.sh

# Quantitatively evaluate TOM-Net on different categories of synthetic object 
# Replace ${class} with one of the four object categories (glass, water, lens, cplx)
CUDA_VISIBLE_DEVICES=${gpu} th eval/run_synth_data.lua -img_list ${class}.txt

# Similarly, you can find the results in data/TOM-Net_model/

Evaluation on Real Testing Data

# Download real testing dataset, 
sh scripts/download_testing_dataset.sh

# Test on sample images used in the paper
CUDA_VISIBLE_DEVICES=${gpu} th eval/run_model.lua -img_list Sample_paper.txt

# Quantitatively evaluate TOM-Net on different categories of real-world object 
# Replace ${class} with one of the four object categories (Glass, Water, Lens, Cplx)
CUDA_VISIBLE_DEVICES=${gpu} th eval/run_model.lua -img_list ${class}.txt  

Training

To train a new TOM-Net model, you have to follow the following steps:

• Download the training data

# The size of the zipped training dataset is 40 GB and you need about 207 GB to unzip it.
sh scripts/download_training_dataset.sh

• Call main.lua to train CoarseNet on simple objects

CUDA_VISIBLE_DEVICES=$gpu th main.lua -train_list train_simple_98k.txt -nEpochs 13 -prefix 'simple'
# Please refer to opt.lua for more information about the training options

# You can find log file, checkpoints and visualization results in data/training/simple_*

• Call main.lua to train CoarseNet on both simple and complex objects

# Finetune CoarseNet with all of the data
CUDA_VISIBLE_DEVICES=$gpu th main.lua -train_list train_all_178k.txt -nEpochs 7 -prefix 'all' -retrain data/training/simple_*/checkpointdir/checkpoint13.t7

# You can find log file, checkpoints and visualization results in data/training/all_*

• Call main_refine.lua to train RefineNet on both simple and complex objects

CUDA_VISIBLE_DEVICES=$gpu th refine/main_refine.lua -train_list train_all_178k.txt -nEpochs 20 -coarse_net data/training/all_*/checkpointdir/checkpoint7.t7 
# Train RefineNet with all of the data
# Please refer to refine/opt_refine.lua for more information about the training options

# You can find log file, checkpoints and visualization results in data/training/all_*/refinement/

Synthetic Data Rendering

Please refer to TOM-Net_Rendering for sample rendering codes.

Citation

If you find this code or the provided data useful in your research, please consider cite:

@inproceedings{chen2018tomnet,
  title={TOM-Net: Learning Transparent Object Matting from a Single Image},
  author={Chen, Guanying and Han, Kai and Wong, Kwan-Yee K.},
  booktitle={CVPR},
  year={2018}
}