Deep Cosine Metric Learning for Person Re-identification
Switch branches/tags
Clone or download
Latest commit eda0daa Oct 28, 2018
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
datasets Initial commit Feb 11, 2018
nets Initial commit Feb 11, 2018
.gitignore Initial commit Jan 26, 2018
LICENSE Initial commit Jan 26, 2018
README.md Updated download link Oct 28, 2018
losses.py Fixed label array length in loss documentation Apr 24, 2018
metrics.py Initial commit Feb 11, 2018
queued_trainer.py Initial commit Feb 11, 2018
train_app.py Initial commit Feb 11, 2018
train_market1501.py Initial commit Feb 11, 2018
train_mars.py Initial commit Feb 11, 2018
vis_tools.py Initial commit Feb 11, 2018

README.md

cosine_metric_learning

Introduction

This repository contains code for training a metric feature representation to be used with the deep_sort tracker. The approach is described in

@inproceedings{Wojke2018deep,
  title={Deep Cosine Metric Learning for Person Re-identification},
  author={Wojke, Nicolai and Bewley, Alex},
  booktitle={2018 IEEE Winter Conference on Applications of Computer Vision (WACV)},
  year={2018},
  pages={748--756},
  organization={IEEE},
  doi={10.1109/WACV.2018.00087}
}

Pre-trained models used in the paper can be found here. A preprint of the paper is available here. The repository comes with code to train a model on the Market1501 and MARS datasets.

Training on Market1501

The following description assumes you have downloaded the Market1501 dataset to ./Market-1501-v15.09.15. The following command starts training using the cosine-softmax classifier described in the above paper:

python train_market1501.py \
    --dataset_dir=./Market-1501-v15.09.15/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/market1501/ \
    --run_id=cosine-softmax

This will create a directory ./output/market1501/cosine-softmax where TensorFlow checkpoints are stored and which can be monitored using tensorboard:

tensorboard --logdir ./output/market1501/cosine-softmax --port 6006

The code splits off 10% of the training data for validation. Concurrently to training, run the following command to run CMC evaluation metrics on the validation set:

CUDA_VISIBLE_DEVICES="" python train_market1501.py \
    --mode=eval \
    --dataset_dir=./Market-1501-v15.09.15/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/market1501/ \
    --run_id=cosine-softmax \
    --eval_log_dir=./eval_output/market1501

The command will block indefinitely to monitor the training directory for saved checkpoints and each stored checkpoint in the training directory is evaluated on the validation set. The results of this evaluation are stored in ./eval_output/market1501/cosine-softmax to be monitored using tensorboard:

tensorboard --logdir ./eval_output/market1501/cosine-softmax --port 6007

Training on MARS

To train on MARS, download the evaluation software and extract bbox_train.zip and bbox_test.zip from the dataset website into the evaluation software directory. The following description assumes they are stored in ./MARS-evaluation-master/bbox_train and ./MARS-evaluation-master/bbox_test. Training can be started with the following command:

python train_mars.py \
    --dataset_dir=./MARS-evaluation-master \
    --loss_mode=cosine-softmax \
    --log_dir=./output/mars/ \
    --run_id=cosine-softmax

Again, this will create a directory ./output/mars/cosine-softmax where TensorFlow checkpoints are stored and which can be monitored using tensorboard:

tensorboard --logdir ./output/mars/cosine-softmax --port 7006

As for Market1501, 10% of the training data are split off for validation. Concurrently to training, run the following command to run CMC evaluation metrics on the validation set:

CUDA_VISIBLE_DEVICES="" python train_mars.py \
    --mode=eval \
    --dataset_dir=./MARS-evaluation-master/ \
    --loss_mode=cosine-softmax \
    --log_dir=./output/mars/ \
    --run_id=cosine-softmax \
    --eval_log_dir=./eval_output/mars

Evaluation metrics on the validation set can be monitored with tensorboard

tensorboard --logdir ./eval_output/mars/cosine-softmax

Testing

Final model testing has been carried out using evaluation software provided by the dataset authors. The training scripts can be used to write features of the test split. The following command exports MARS test features to ./MARS-evaluation-master/feat_test.mat

python train_mars.py \
    --mode=export \
    --dataset_dir=./MARS-evaluation-master \
    --loss_mode=cosine-softmax .\
    --restore_path=PATH_TO_CHECKPOINT

where PATH_TO_CHECKPOINT the checkpoint file to evaluate. Note that the evaluation script needs minor adjustments to apply the cosine similarity metric. More precisely, change the feature computation in utils/process_box_features.m to average pooling (line 8) and apply a re-normalization at the end of the file. The modified file should look like this:

function video_feat = process_box_feat(box_feat, video_info)

nVideo = size(video_info, 1);
video_feat = zeros(size(box_feat, 1), nVideo);
for n = 1:nVideo
    feature_set = box_feat(:, video_info(n, 1):video_info(n, 2));
%    video_feat(:, n) = max(feature_set, [], 2); % max pooling 
     video_feat(:, n) = mean(feature_set, 2); % avg pooling
end

%%% normalize train and test features
sum_val = sqrt(sum(video_feat.^2));
for n = 1:size(video_feat, 1)
    video_feat(n, :) = video_feat(n, :)./sum_val;
end

The Market1501 script contains a similar export functionality which can be applied in the same way as described for MARS:

python train_market1501.py \
    --mode=export \
    --dataset_dir=./Market-1501-v15.09.15/
    --sdk_dir=./Market-1501_baseline-v16.01.14/
    --loss_mode=cosine-softmax \
    --restore_path=PATH_TO_CHECKPOINT

This command creates ./Market-1501_baseline-v16.01.14/feat_query.mat and ./Market-1501_baseline-v16.01.14/feat_test.mat to be used with the Market1501 evaluation code.

Model export

To export your trained model for use with the deep_sort tracker, run the following command:

python train_mars.py --mode=freeze --restore_path=PATH_TO_CHECKPOINT

This will create a mars.pb file which can be supplied to Deep SORT. Again, the Market1501 script contains a similar function.