cvlab-yonsei/oimnetplus

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

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.

oimnetplus's Issues

mAP reference

Thank you for making the code public . it is such a great work .
is mAP used for person search is similar to the one used for person re-identification task. do you have any reference explaining it?

Demo code for visualizing the resulted images with bounding boxes

I'm very impressed by your outstanding work. Can you please upload a Demo code for visualizing the resulted images with bounding boxes.I really appreciate your effort doing this project.

demo

你的示例图是怎么得出来的