Awesome
TSP-Transformer
Abstract
Holistic scene understanding includes semantic segmentation, surface normal estimation, object boundary detection, depth estimation, etc. The key aspect of this problem is to learn representation effectively, as each subtask builds upon not only correlated but also distinct attributes. Inspired by visual-prompt tuning, we propose a Task-Specific Prompts Transformer, dubbed TSP-Transformer, for holistic scene understanding. It features a vanilla transformer in the early stage and tasks-specific prompts transformer encoder in the lateral stage, where tasks-specific prompts are augmented. By doing so, the transformer layer learns the generic information from the shared parts and is endowed with task-specific capacity. First, the tasks-specific prompts serve as induced priors for each task effectively. Moreover, the task-specific prompts can be seen as switches to favor task-specific representation learning for different tasks. Extensive experiments on NYUD-v2 and PASCAL-Context show that our method achieves state-of-the-art performance, validating the effectiveness of our method for holistic scene understanding.
Paper
Setup
Tested with PyTorch 1.11 and CUDA 11.3:
git clone https://github.com/tb2-sy/TSP-Transformer.git
conda create -n tsp python=3.7
conda activate tsp
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
pip install tqdm Pillow easydict pyyaml imageio scikit-image tensorboard
pip install opencv-python==4.5.4.60 setuptools==59.5.0
pip install timm==0.5.4 einops==0.4.1
Dataset
We use the same dataset (PASCAL-Context and NYUD-v2) as InvPT. You can download the data from here. And then extract the datasets by:
tar xfvz NYUDv2.tar.gz
tar xfvz PASCALContext.tar.gz
You need to specify the dataset directory as db_root
variable in configs/mypath.py
.
Training
Set the config files in ./conifigs
, with PASCAL-Context and NYUD-v2 dataset.
# Train the NYUD-v2 dataset
./run_nyud.sh
# Train the PASCAL-Context dataset
./run_pascal.sh
Evaluation
# Evaluation the NYUD-v2 dataset
./infer_nyud.sh
# Evaluation the PASCAL-Context dataset
./infer_pascal.sh
Pre-trained models
Please download the weights of our SOTA results.
Version | Dataset | Download | Segmentation | Human parsing | Saliency | Normals | Boundary |
---|---|---|---|---|---|---|---|
TSP-Transformer (our paper) | PASCAL-Context | google drive | 81.48 | 70.64 | 84.86 | 13.69 | 74.80 |
TaskPrompter (ICLR 2023) | PASCAL-Context | - | 80.89 | 68.89 | 84.83 | 13.72 | 73.50 |
InvPT (ECCV 2022) | PASCAL-Context | - | 79.03 | 67.61 | 84.81 | 14.15 | 73.00 |
Version | Dataset | Download | Segmentation | Depth | Normals | Boundary |
---|---|---|---|---|---|---|
TSP-Transformer (our paper) | NYUD-v2 | google drive | 55.39 | 0.4961 | 18.44 | 77.50 |
TaskPrompter (ICLR 2023) | NYUD-v2 | - | 55.30 | 0.5152 | 18.47 | 78.20 |
InvPT (ECCV 2022) | NYUD-v2 | - | 53.56 | 0.5183 | 19.04 | 78.10 |
TODO/Future work
- Upload paper and init project
- Training and Inference code
- Reproducible checkpoints
- Speed training and inference
- Reduce cuda memory usage
Contact
For any questions related to our paper and implementation, please email wangshuo2022@shanghaitech.edu.cn.
Citation
If you find our code or paper helps, please consider citing:
@article{wang2023tsp,
title={TSP-Transformer: Task-Specific Prompts Boosted Transformer for Holistic Scene Understanding},
author={Wang, Shuo and Li, Jing and Zhao, Zibo and Lian, Dongze and Huang, Binbin and Wang, Xiaomei and Li, Zhengxin and Gao, Shenghua},
journal={arXiv preprint arXiv:2311.03427},
year={2023}
}
Acknowledgements
The code is available under the MIT license and draws from InvPT, ATRC, and MIT-Net, which are also licensed under the MIT license.