PyTorch Implementation of OIMNet++ (ECCV 2022)

This is an official PyTorch implementation of "OIMNet++: Prototypical Normalization and Localization-aware Learning for Person Search", ECCV 2022.

For more details, visit our project site or see our paper.
Our main contributions can be found in models/custom_modules.py and losses/oim.py.

Requirements

Python 3.8
PyTorch 1.7.1
GPU memory >= 22GB

Features

Re-implementation of vanilla OIMNet in pure PyTorch
Using automatic mixed precision to train with larger batch size under limited GPU memory and accelerate training

Getting Started

First, clone our git repository.

Docker

We highly recommend using our Dockerfile to set up the environment.

# build docker image
$ docker build -t oimnetplus:latest . 

# execute docker container
$ docker run --ipc=host -it -v <working_dir>:/workspace/work -v <dataset_dir>:/workspace/dataset -w /workspace/work oimnetplus:latest /bin/bash

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.

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'

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.
(PRs are always welcomed!)

Pretrained Models

We provide pretrained weights and the correponding configs below.

	OIMNet++	OIMNet+++
PRW	model config	model config
CUHK-SYSU	model config	model config

Citation

@inproceedings{lee2022oimnet++,
  title={OIMNet++: Prototypical Normalization and Localization-Aware Learning for Person Search},
  author={Lee, Sanghoon and Oh, Youngmin and Baek, Donghyeon and Lee, Junghyup and Ham, Bumsub},
  booktitle={European Conference on Computer Vision},
  pages={621--637},
  year={2022},
  organization={Springer}
}

Credits

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

cvlab-yonsei/OIMNetPlus