PSTR (CVPR2022)

• This code is an official implementation of "PSTR: End-to-End One-Step Person Search With Transformers (CVPR2022)".
• End-to-end one-step person search with Transformers, which does not requre NMS post-processing.
• Pre-trained models with ResNet50, ResNet50-DCN, and PVTv2b2.
• Curves of different methods on CUHK under different gallery sizes (plot_cuhk.py). If you want to add new results, please feel free to contact us.

Installation

• We install this project using cuda11.1 and PyTorch1.8.0 (or PyTorch1.9.0) as follows.

# Download this project
git clone https://github.com/JialeCao001/PSTR.git

# Create a new conda enviroment for PSTR
conda create -n pstr python=3.7 -y
conda activate pstr
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
#conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge

# Comiple mmcv, which has been included in this project
cd PSTR/mmcv
MMCV_WITH_OPS=1 pip install -e .

# Comiple this project 
cd PSTR
pip install -r requirements/build.txt
pip install -v -e .  # or "python setup.py develop"
pip install sklearn

• If you have the problem local variable 'beta1' referenced before assignment with PyTorch1.8, add one table space in L110 of optim/adamw.py

Train and Inference

Datasets and Annotations

• Download PRW and CUHK-SYSU datasets.
• Download the json annotations provided by AlignPS.

Train with a single GPU

python tools/train.py ${CONFIG_FILE} --no-validate

• CONFIG_FILE about PSTR is in configs/PSTR

Test with a single GPU

PRW: sh run_test_prw.sh 
CUHK: sh run_test_cuhk.sh  

• If you want to output the results of different models, please change CONFIGPATH, MODELPATH, OUTPATH for diffferent models

Results

We provide some models with different backbones and results on PRW and CUHK-SYSU datasets, which have a little difference to CVPR version due to jitter.

• All the models are based on multi-scale training and all the results are based on single-scale inference.

• + indicates adding a re-scoring module during evaluation, where we modify the final matching score as the weighted score of CBGM score and originial matching scores.

Citation

If the project helps your research, please cite this paper.

@article{Cao_PSTR_CVPR_2022,
  author =       {Jiale Cao and Yanwei Pang and Rao Muhammad Anwer and Hisham Cholakkal and Jin Xie and Mubarak Shah and Fahad Shahbaz Khan},
  title =        {PSTR: End-to-End One-Step Person Search With Transformers},
  journal =      {Proc. IEEE Conference on Computer Vision and Pattern Recognition},
  year =         {2022}
}

Acknowledgement

Many thanks to the open source codes: mmdetection, AlignPS, and SeqNet.