PyTorch Implementation of Person Search

This includes OIMNet, OIMNet++, NAE4PS

Requirements

• Python 3.8
• PyTorch 1.7.1
• GPU memory >= 22GB

Getting Started

First, clone our git repository.

Prepare datasets

Download PRW and CUHK-SYSU datasets.
Modify the dataset directories below if necessary.

• PRW: L4 of configs/prw.yaml
  
• CUHK-SYSU: L3 of configs/ssm.yaml
  

Your directories should look like:

    <working_dir>
    OIMNetPlus
    ├── configs/
    ├── datasets/
    ├── engines/
    ├── losses/
    ├── models/
    ├── utils/
    ├── defaults.py
    ├── Dockerfile
    └── train.py
    
    <dataset_dir>
    ├── CUHK-SYSU/
    │   ├── annotation/
    │   ├── Image/
    │   └── ...
    └── PRW-v16.04.20/
        ├── annotations/
        ├── frames/
        ├── query_box/
        └── ...

Training and Evaluation

By running the commands below, evaluation results and training losses will be logged into a .txt file in the output directory.
Or open run.ipynb for commands.

• OIMNet++
  $ python train.py --cfg configs/prw.yaml
$ python train.py --cfg configs/ssm.yaml

• OIMNet+++
  $ python train.py --cfg configs/prw.yaml MODEL.ROI_HEAD.AUGMENT True
$ python train.py --cfg configs/ssm.yaml MODEL.ROI_HEAD.AUGMENT True

• OIMNet
  $ python train.py --cfg configs/prw.yaml MODEL.ROI_HEAD.NORM_TYPE 'none' MODEL.LOSS.TYPE 'OIM'
$ python train.py --cfg configs/ssm.yaml MODEL.ROI_HEAD.NORM_TYPE 'none' MODEL.LOSS.TYPE 'OIM'

• NAE4PS
  $ python train.py --cfg configs/prw.yaml MODEL.ROI_HEAD.NORM_TYPE 'none' MODEL.LOSS.TYPE 'OIM' MODEL.ROI_HEAD.EMBD 'NAE'
$ python train.py --cfg configs/ssm.yaml MODEL.ROI_HEAD.NORM_TYPE 'none' MODEL.LOSS.TYPE 'OIM' MODEL.ROI_HEAD.EMBD 'NAE'

We support training/evaluation using single GPU only.
This is due to unsynchronized items across multiple GPUs in OIM loss (i.e., LUT and CQ) and ProtoNorm.
OIMNet and NAE4PS results are in exps folder.

Credits

Our person search implementation is heavily based on Di Chen's NAE and Zhengjia Li's SeqNet and Sanghoon Lee's OIMNet++.
ProtoNorm implementation is based on ptrblck's manual BatchNorm implementation here.