IEEE TPAMI 2021: MobileSal: Extremely Efficient RGB-D Salient Object Detection

This repository contains full training & testing code, and pretrained saliency maps. We have achieved competitive performance on the RGB-D salient object detection task with a speed of 450fps.

If you run into any problems or feel any difficulties to run this code, do not hesitate to leave issues in this repository.

My e-mail is: wuyuhuan @ mail.nankai (dot) edu.cn

[PDF]

This repository contains:

• Full code, data, pretrained models for training and testing
• MobileSal deployment, achieving 420FPS (fp32) and 800FPS (fp16) with batch size 1 on a single RTX 2080Ti.

• Requirements
  • PyTorch
  • Jittor
  • Deployment
• Data Preparing
• Train
• Test
  • Pretrained Models
  • Generate Saliency Maps
  • Deployment
  • Speed Test
  • Pretrained Saliency maps
• Others
  • TODO
  • Contact
  • License
  • Citation
  • Acknowlogdement

• Python 3.6+
• PyTorch >=0.4.1, OpenCV-Python
• Tested on PyTorch 1.7.1

• Python 3.7+
• Jittor, OpenCV-Python
• Tested on Jittor 1.3.1

• torch2trt

For Jittor users, we create a branch jittor. So please run the following command first:

git checkout jittor

To install MobileSal, please run:

pip install -r envs/requirements.txt

Before training/testing our network, please download the training data:

• Preprocessed data of 6 datasets: [Google Drive], [Baidu Pan, 9nxi]

Note: if you are blocked by Google and Baidu services, you can contact me via e-mail and I will send you a copy of data and model weights.

We have processed the data well so you can use them without any preprocessing steps. After completion of downloading, extract the data and put them to ./data/ folder. Then, the ./datasets/ folder should contain six folders: NJU2K/, NLPR/, STERE/, SSD/, SIP/, DUT-RGBD/, representing NJU2K, NLPR, STEREO, SSD, SIP, DUTLF-D datasets, respectively.

It is very simple to train our network. We have prepared a script to run the training step:

bash ./tools/train.sh

As in our paper, we train our model on the NJU2K_NLPR training set, and test our model on NJU2K_test, NLPR_test, STEREO, SIP, and SSD datasets. For DUTLF-D, we train our model on DUTLF-D training set and evaluate on its testing test.

(Default) Trained on NJU2K_NLPR training set:

• Single-scale Training: [Google Drive], [Baidu Pan, 9nxi]
• Multi-scale Training: [Google Drive], [Baidu Pan, 9nxi]

(Custom) Training on DUTLF-D training set:

• Multi-scale Training: [Google Drive], [Baidu Pan, 9nxi]

Download them and put them into the pretrained/ folder.

After preparing the pretrained models, it is also very simple to generate saliency maps via MobileSal:

bash ./tools/test.sh

The scripts will automatically generate saliency maps on the maps/ directory.

The deployment largely speeds up MobileSal with batch size of 1. An example script is located at: tools/test_trt.sh. Run:

bash ./tools/test_trt.sh

This script will automatically convert PyTorch MobileSal to TensorRT-based MobileSal. Then it will generate saliency maps via the TensorRT-based MobileSal. On deployment for real-world applications, you can load the converted TensorRT MobileSal for inference:

from torch2trt import torch2trt, TRTModule
model = TRTModule(); trt_model_path = "pretrained/mobilesal_trt.pth"
model.load_state_dict(torch.load(trt_model_path))
result = model(image, depth) # get result with [torch.Tensor] input

We provide a speed test script on MobileSal:

python speed_test.py

The speed result on a single RTX 2080Ti is as below:

For covenience, we provide the pretrained saliency maps on several datasets as below:

• Single-scale Training: [Google Drive], [Baidu Pan, 9nxi]

• Multi-scale Training: [Google Drive], [Baidu Pan, 9nxi]

1. Release the pretrained models and saliency maps on COME15K dataset.
2. Add results with the P2T transformer backbone.

• I encourage everyone to contact me via my e-mail. My e-mail is: wuyuhuan @ mail.nankai (dot) edu.cn

The code is released under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License for NonCommercial use only.

If you are using the code/model/data provided here in a publication, please consider citing our work:

@ARTICLE{wu2021mobilesal,
  author={Wu, Yu-Huan and Liu, Yun and Xu, Jun and Bian, Jia-Wang and Gu, Yu-Chao and Cheng, Ming-Ming},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 
  title={MobileSal: Extremely Efficient RGB-D Salient Object Detection}, 
  year={2022},
  volume={44},
  number={12},
  pages={10261--10269},
  doi={10.1109/TPAMI.2021.3134684}
}

This repository is built under the help of the following five projects for academic use only:

• PyTorch

• Jittor

• torch2trt