Home

Awesome

Introduction

The code for our paper EANet: Enhancing Alignment for Cross-Domain Person Re-identification. A Chinese version of introduction for this paper can be found here.

This repository provides

Installation

The other packages and versions are listed in requirements.txt. You can install them by pip install -r requirements.txt.

Dataset Structure

Prepare datasets to have following structure:

${project_dir}/dataset
    market1501
        Market-1501-v15.09.15                   # Extracted from Market-1501-v15.09.15.zip, http://www.liangzheng.org/Project/project_reid.html
        Market-1501-v15.09.15_ps_label
        bounding_box_train_duke_style
        im_path_to_kpt.pkl
    cuhk03_np_detected_jpg
        cuhk03-np                               # Extracted from cuhk03-np.zip, https://pan.baidu.com/s/1RNvebTccjmmj1ig-LVjw7A
        cuhk03-np-jpg_ps_label
        im_path_to_kpt.pkl
    duke
        DukeMTMC-reID                           # Extracted from DukeMTMC-reID.zip, https://github.com/layumi/DukeMTMC-reID_evaluation
        DukeMTMC-reID_ps_label
        bounding_box_train_market1501_style
        im_path_to_kpt.pkl
    msmt17
        MSMT17_V1                               # Extracted from MSMT17_V1.tar.gz, https://www.pkuvmc.com/publications/msmt17.html
        MSMT17_V1_ps_label
        im_path_to_kpt.pkl
    partial_reid
        Partial-REID_Dataset                    # Extracted from Partial-REID_Dataset.rar, http://isee.sysu.edu.cn/files/resource/Partial-REID_Dataset.rar
    partial_ilids
        Partial_iLIDS                           # Provided by https://github.com/lingxiao-he/Partial-Person-ReID
    coco
        images
        masks_7_parts
        im_name_to_kpt.pkl
        im_name_to_h_w.pkl

Keypoints and part segmentation labels: Baidu Cloud or Google Drive. Our keypoint model is currently not public, while this repo can be an alternative. For part segmentation, we release our code at this repo.

Keypoint Format

The following example shows the keypoint format.

from __future__ import print_function
import cPickle
res = cPickle.load(open('dataset/market1501/im_path_to_kpt.pkl'))
# 'Market-1501-v15.09.15/bounding_box_train/0742_c1s4_014906_01.jpg' ...
print(res.keys()[:5])
# ['kpt', 'im_h_w']
print(res.values()[0].keys())
# A numpy array with shape [17, 3], for 17 keypoints. Each row is (x, y, confidence); we treat those with confidence > 0.1 as visible.
print(res.values()[0]['kpt'])
# [h, w] of the image
print(res.values()[0]['im_h_w'])

For COCO, there is a bit difference.

from __future__ import print_function
import cPickle
res = cPickle.load(open('dataset/coco/im_name_to_kpt.pkl'))
# Currently only contain train set.
# 'COCO_train2014_000000426663_185693.jpg' ...
print(res.keys()[:5])
# A numpy array with shape [17, 3], each row is (x, y, visibility), visibility is one of [0, 1, 2], refer to COCO dataset for detail
print(res.values()[0])

# image size
res = cPickle.load(open('dataset/coco/im_name_to_h_w.pkl'))
print(res.keys()[0])
print(res.values()[0])

Part Segmentation Label Format

Part segmentation label for each image is a single-channel PNG file, with same resolution as the image. Label mapping is as follows

{
    'background': 0,
    'head': 1,
    'torso': 2,
    'upper_arm': 3,
    'lower_arm': 4,
    'upper_leg': 5,
    'lower_leg': 6,
    'foot': 7,
}

Train/Test/Inference Example

Code Explanation

Training

The training logic is defined in package/optim/eanet_trainer.py and package/optim/cft_trainer.py, the latter for Clustering and Finetuning. A training command looks like

cd ${project_dir}
CUDA_VISIBLE_DEVICES=0 python -m package.optim.${trainer} --exp_dir ${exp_dir} --cfg_file ${cfg_file} [--ow_file ${ow_file}] [--ow_str ${ow_str}]

Testing

Test sets and testing interval can be set in config file, and the training script will test the model during training. If you want to test a trained model, create a exp_dir and place the ckpt.pth inside it, then set cfg.only_test = True in ${config_file} and run package/optim/eanet_trainer.py. In this case, the code only performs testing.

Inference

To use a trained model for inference (extracting feature), the overall logic is

  1. Create a exp_dir and place the ckpt.pth inside it
  2. Set cfg.only_infer = True
  3. trainer = EANetTrainer(), then try trainer.infer_one_im or trainer.infer_im_list or trainer.infer_dataloader, which is flexible.

NOTE: Our model inference (feature extractor) can be used as an API, not just run in command line. For example, if you want to use PCB feature extractor trained on Market1501 in another project, you can add our package path to $PYTHONPATH. Then in your project, do something like this

from package.optim.eanet_trainer import EANetTrainer
from package.eval.np_distance import compute_dist
from easydict import EasyDict

args = EasyDict()
args.exp_dir = 'exp/try_pcb_trained_on_market1501_for_reid_feature'  # There should be the corresponding `ckpt.pth` in it
args.cfg_file = '${EANet_PROJECT_DIR}/package/config/default.py'  # Set this `${EANet_PROJECT_DIR}`
args.ow_file = '${EANet_PROJECT_DIR}/paper_configs/PCB.txt'  # Set this `${EANet_PROJECT_DIR}`
args.ow_str = "cfg.only_infer = True"

eanet_trainer = EANetTrainer(args=args)
feature_dict1 = eanet_trainer.infer_one_im(im_path='YOUR/IMAGE/PATH/1.jpg', squeeze=False)  # You can also provide PIL Image instead of image path
feature_dict2 = eanet_trainer.infer_one_im(im_path='YOUR/IMAGE/PATH/2.jpg', squeeze=False)  # You can also provide PIL Image instead of image path
cosine_distance = compute_dist(feature_dict1['feat'], feature_dict2['feat'])[0][0]

NOTE: For your own images, if you want to perform Part Aligned Pooling for inference, you have to provide keypoint generated pap_mask. Though, as alternatives you can use GlobalPool or PCB, which also achieve satisfactory performance in our implementation.

For details of inference, refer to package/optim/reid_trainer.py.

Extension: Dataset

Current datasets

You can create new dataset class in package/data/datasets/ and then register it in package/data/create_dataset.py.

Extension: Model

Current backbone is ResNet. You can implement new backbones and then register them in package/model/backbone.py. Or you can re-define the whole model package/model/model.py, as long as you implement the methods declared in package/model/base_model.py.

Note: If you re-implement the whole model, you may have to modify the model.forward() logic in following files as well

This model calling can be abstracted away in the future.

Extension: Training

package/optim/reid_trainer.py covers the common logic for ReID training, with some abstraction to be implemented by sub classes. package/optim/eanet_trainer.py and package/optim/cft_trainer.py are concrete implementations, which also demonstrate the usage of hybrid batches and multi losses, etc.

Design Logic

We heavily use dict for passing data / configuration arguments when calling functions. This simplifies function headers and reduces the amount of code to modify when changing experiment settings.

TODO

Misc

Citation

If you find our work useful, please kindly cite our paper:

@article{huang2018eanet,
  title={EANet: Enhancing Alignment for Cross-Domain Person Re-identification},
  author={Huang, Houjing and Yang, Wenjie and Chen, Xiaotang and Zhao, Xin and Huang, Kaiqi and Lin, Jinbin and Huang, Guan and Du, Dalong},
  journal={arXiv preprint arXiv:1812.11369},
  year={2018}
}