Skip to content

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

License

Notifications You must be signed in to change notification settings

cvlab-yonsei/OIMNetPlus

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

19 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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

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.

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.

About

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

Resources

License

Stars

Watchers

Forks

Packages

No packages published