open source codes
parent
31ce67e61d
commit
061be381c1
config
engine
modeling
|
@ -0,0 +1,3 @@
|
|||
.idea
|
||||
__pycache__
|
||||
.DS_Store
|
|
@ -0,0 +1,12 @@
|
|||
# Experiment 6 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on (=raw all trick, softmax_triplet.yml)
|
||||
# Dataset 2: dukemtmc
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
# SOLVER.WARMUP_ITERS 10 MODEL.LAST_STRIDE 1 INPUT.RE_PROB 0.5 MODEL.NECK "('bnneck')"
|
||||
|
||||
python3 tools/train.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('1')" DATASETS.NAMES "('dukemtmc')" OUTPUT_DIR "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/dukemtmc/Experiment-all-tricks-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment parameters : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on (=raw all trick, softmax_triplet.yml)
|
||||
# Dataset 1: market1501
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/train.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('0')" DATASETS.NAMES "('market1501')" OUTPUT_DIR "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/market1501/Experiment-all-tricks-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 9-4 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005
|
||||
# Dataset 2: dukemtmc
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/train.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('3')" DATASETS.NAMES "('dukemtmc')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" OUTPUT_DIR "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/dukemtmc/Experiment-all-tricks-tri_center-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 9-4 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005
|
||||
# Dataset 1: market1501
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/train.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('2')" DATASETS.NAMES "('market1501')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" OUTPUT_DIR "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/market1501/Experiment-all-tricks-tri_center-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005')"
|
114
README.md
114
README.md
|
@ -1,2 +1,112 @@
|
|||
# reid-strong-baseline
|
||||
Bag of Tricks and A Strong Baseline for Person Re-identification
|
||||
# ReID Strong Baseline
|
||||
Paper:
|
||||
The codes are expanded on a [ReID-baseline](https://github.com/L1aoXingyu/reid_baseline) , which is open sourced by our co-author Xingyu Liao.
|
||||
|
||||
We support
|
||||
- [x] easy dataset preparation
|
||||
- [x] end-to-end training and evaluation
|
||||
- [x] high modular management
|
||||
|
||||
Bag of tricks
|
||||
- Warm up learning rate
|
||||
- Random erasing augmentation
|
||||
- Label smoothing
|
||||
- Last stride
|
||||
- BNNeck
|
||||
- Center loss
|
||||
|
||||
## Get Started
|
||||
The designed architecture follows this guide [PyTorch-Project-Template](https://github.com/L1aoXingyu/PyTorch-Project-Template), you can check each folder's purpose by yourself.
|
||||
|
||||
1. `cd` to folder where you want to download this repo
|
||||
|
||||
2. Run `git clone... `
|
||||
|
||||
3. Install dependencies:
|
||||
- [pytorch 1.0](https://pytorch.org/)
|
||||
- torchvision
|
||||
- [ignite](https://github.com/pytorch/ignite)
|
||||
- [yacs](https://github.com/rbgirshick/yacs)
|
||||
|
||||
4. Prepare dataset
|
||||
|
||||
Create a directory to store reid datasets under this repo via
|
||||
```bash
|
||||
cd reid_baseline
|
||||
mkdir data
|
||||
```
|
||||
(1)Market1501
|
||||
|
||||
* Download dataset to `data/` from http://www.liangzheng.org/Project/project_reid.html
|
||||
* Extract dataset and rename to `market1501`. The data structure would like:
|
||||
|
||||
```bash
|
||||
data
|
||||
market1501 # this folder contains 6 files.
|
||||
bounding_box_test/
|
||||
bounding_box_train/
|
||||
......
|
||||
```
|
||||
(2)DukeMTMC-reID
|
||||
|
||||
* Download dataset to `data/` from https://github.com/layumi/DukeMTMC-reID_evaluation#download-dataset
|
||||
* Extract dataset and rename to `dukemtmc-reid`. The data structure would like:
|
||||
|
||||
```bash
|
||||
data
|
||||
dukemtmc-reid
|
||||
DukeMTMC-reID # this folder contains 8 files.
|
||||
bounding_box_test/
|
||||
bounding_box_train/
|
||||
......
|
||||
```
|
||||
|
||||
5. Prepare pretrained model if you don't have
|
||||
```python
|
||||
from torchvision import models
|
||||
models.resnet50(pretrained=True)
|
||||
```
|
||||
Then it will automatically download model in `~/.torch/models/`, you should set this path in `config/defaults.py` for all training or set in every single training config file in `configs/`.
|
||||
|
||||
6. If you want to know the detained configurations and their meaning, please refer to `config/defaults.py`. If you want to set your own parameters, you can follow our method: create a new yml file, then set your own parameters. Add `--config_file='configs/your yml file'` int the commands described below, then our code will merge your configuration. automatically.
|
||||
|
||||
## Train
|
||||
You can run these commands in `.sh ` files for training different datasets of differernt loss. You can also directly run code `sh *.sh` to run our demo.
|
||||
|
||||
1. Market1501, cross entropy loss + triplet loss
|
||||
|
||||
```bash
|
||||
python3 tools/train.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('your device id')" DATASETS.NAMES "('market1501')" OUTPUT_DIR "('your path to save checkpoints and logs')"
|
||||
```
|
||||
|
||||
2. DukeMTMC-reID, cross entropy loss + triplet loss + center loss
|
||||
|
||||
|
||||
```bash
|
||||
python3 tools/train.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('your device id')" DATASETS.NAMES "('dukemtmc')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" OUTPUT_DIR "('your path to save checkpoints and logs')"
|
||||
```
|
||||
|
||||
## Test
|
||||
You can test your model's performance directly by running these commands in `.sh ` files. You can also change the configuration to determine which feature of BNNeck and whether the feature is normalized (equivalent to use Cosine distance or Euclidean distance) for testing.
|
||||
|
||||
1. Test with Euclidean distance using feature before BN without re-ranking,.
|
||||
|
||||
```bash
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('your device id')" DATASETS.NAMES "('market1501')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" TEST.NECK_FEAT "('before')" TEST.FEAT_NORM "('no')" TEST.WEIGHT "('your path to trained checkpoints')"
|
||||
```
|
||||
2. Test with Cosine distance using feature after BN without re-ranking,.
|
||||
|
||||
```bash
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('your device id')" DATASETS.NAMES "('market1501')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.WEIGHT "('your path to trained checkpoints')"
|
||||
```
|
||||
3. Test with Cosine distance using feature after BN with re-ranking
|
||||
|
||||
```bash
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('your device id')" DATASETS.NAMES "('dukemtmc')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.RE_RANKING "('yes')" TEST.WEIGHT "('your path to trained checkpoints')"
|
||||
```
|
||||
|
||||
## Results
|
||||
|
||||
**Network architecture**
|
||||
|
||||
|
||||
|
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 6 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on (=raw all trick, softmax_triplet.yml)
|
||||
# Dataset 2: dukemtmc
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('1')" DATASETS.NAMES "('dukemtmc')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.WEIGHT "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/dukemtmc/Experiment-all-tricks-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on/resnet50_model_120.pth')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 6 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on (=raw all trick, softmax_triplet.yml)
|
||||
# Dataset 1: market1501
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('0')" DATASETS.NAMES "('market1501')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.WEIGHT "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/market1501/Experiment-all-tricks-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on/resnet50_model_120.pth')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 9-4 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005
|
||||
# Dataset 2: dukemtmc
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('1')" DATASETS.NAMES "('dukemtmc')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.WEIGHT "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/dukemtmc/Experiment-all-tricks-tri_center-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005/resnet50_model_120.pth')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 9-4 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005
|
||||
# Dataset 1: market1501
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('0')" DATASETS.NAMES "('market1501')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.WEIGHT "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/market1501/Experiment-all-tricks-tri_center-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005/resnet50_model_120.pth')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 9-4 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005
|
||||
# Dataset 1: market1501
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('0')" DATASETS.NAMES "('market1501')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.RE_RANKING "('yes')" TEST.WEIGHT "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/market1501/Experiment-all-tricks-tri_center-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005/resnet50_model_120.pth')"
|
|
@ -0,0 +1,11 @@
|
|||
# Experiment 9-4 : 256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005
|
||||
# Dataset 2: dukemtmc
|
||||
# imagesize: 256x128
|
||||
# batchsize: 16x4
|
||||
# warmup_step 10
|
||||
# random erase prob 0.5
|
||||
# labelsmooth: on
|
||||
# last stride 1
|
||||
# bnneck on
|
||||
|
||||
python3 tools/test.py --config_file='configs/softmax_triplet.yml' MODEL.DEVICE_ID "('1')" DATASETS.NAMES "('dukemtmc')" MODEL.IF_WITH_CENTER "('yes')" MODEL.METRIC_LOSS_TYPE "('triplet_center')" TEST.NECK_FEAT "('after')" TEST.FEAT_NORM "('yes')" TEST.RE_RANKING "('yes')" TEST.WEIGHT "('/home/haoluo/log/gu/reid_baseline_review/Opensource_test/dukemtmc/Experiment-all-tricks-tri_center-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on-triplet_centerloss0_0005/resnet50_model_120.pth')"
|
|
@ -0,0 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
from .defaults import _C as cfg
|
|
@ -0,0 +1,156 @@
|
|||
from yacs.config import CfgNode as CN
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Convention about Training / Test specific parameters
|
||||
# -----------------------------------------------------------------------------
|
||||
# Whenever an argument can be either used for training or for testing, the
|
||||
# corresponding name will be post-fixed by a _TRAIN for a training parameter,
|
||||
# or _TEST for a test-specific parameter.
|
||||
# For example, the number of images during training will be
|
||||
# IMAGES_PER_BATCH_TRAIN, while the number of images for testing will be
|
||||
# IMAGES_PER_BATCH_TEST
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Config definition
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
_C = CN()
|
||||
|
||||
_C.MODEL = CN()
|
||||
# Using cuda or cpu for training
|
||||
_C.MODEL.DEVICE = "cuda"
|
||||
# ID number of GPU
|
||||
_C.MODEL.DEVICE_ID = '0'
|
||||
# Name of backbone
|
||||
_C.MODEL.NAME = 'resnet50'
|
||||
# Last stride of backbone
|
||||
_C.MODEL.LAST_STRIDE = 1
|
||||
# Path to pretrained model of backbone
|
||||
_C.MODEL.PRETRAIN_PATH = ''
|
||||
# If train with BNNeck, options: 'bnneck' or 'no'
|
||||
_C.MODEL.NECK = 'bnneck'
|
||||
# If train loss include center loss, options: 'yes' or 'no'. Loss with center loss has different optimizer configuration
|
||||
_C.MODEL.IF_WITH_CENTER = 'no'
|
||||
# The loss type of metric loss
|
||||
# options:'triplet','cluster','triplet_cluster','center','range_center','triplet_center','triplet_range_center'
|
||||
_C.MODEL.METRIC_LOSS_TYPE = 'triplet'
|
||||
# For example, if loss type is cross entropy loss + triplet loss + center loss
|
||||
# the setting should be: _C.MODEL.METRIC_LOSS_TYPE = 'triplet_center' and _C.MODEL.IF_WITH_CENTER = 'yes'
|
||||
|
||||
# If train with label smooth, options: 'on', 'off'
|
||||
_C.MODEL.IF_LABELSMOOTH = 'on'
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# INPUT
|
||||
# -----------------------------------------------------------------------------
|
||||
_C.INPUT = CN()
|
||||
# Size of the image during training
|
||||
_C.INPUT.SIZE_TRAIN = [384, 128]
|
||||
# Size of the image during test
|
||||
_C.INPUT.SIZE_TEST = [384, 128]
|
||||
# Random probability for image horizontal flip
|
||||
_C.INPUT.PROB = 0.5
|
||||
# Random probability for random erasing
|
||||
_C.INPUT.RE_PROB = 0.5
|
||||
# Values to be used for image normalization
|
||||
_C.INPUT.PIXEL_MEAN = [0.485, 0.456, 0.406]
|
||||
# Values to be used for image normalization
|
||||
_C.INPUT.PIXEL_STD = [0.229, 0.224, 0.225]
|
||||
# Value of padding size
|
||||
_C.INPUT.PADDING = 10
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# -----------------------------------------------------------------------------
|
||||
_C.DATASETS = CN()
|
||||
# List of the dataset names for training, as present in paths_catalog.py
|
||||
_C.DATASETS.NAMES = ('market1501')
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# DataLoader
|
||||
# -----------------------------------------------------------------------------
|
||||
_C.DATALOADER = CN()
|
||||
# Number of data loading threads
|
||||
_C.DATALOADER.NUM_WORKERS = 8
|
||||
# Sampler for data loading
|
||||
_C.DATALOADER.SAMPLER = 'softmax'
|
||||
# Number of instance for one batch
|
||||
_C.DATALOADER.NUM_INSTANCE = 16
|
||||
|
||||
# ---------------------------------------------------------------------------- #
|
||||
# Solver
|
||||
# ---------------------------------------------------------------------------- #
|
||||
_C.SOLVER = CN()
|
||||
# Name of optimizer
|
||||
_C.SOLVER.OPTIMIZER_NAME = "Adam"
|
||||
# Number of max epoches
|
||||
_C.SOLVER.MAX_EPOCHS = 50
|
||||
# Base learning rate
|
||||
_C.SOLVER.BASE_LR = 3e-4
|
||||
# Factor of learning bias
|
||||
_C.SOLVER.BIAS_LR_FACTOR = 2
|
||||
# Momentum
|
||||
_C.SOLVER.MOMENTUM = 0.9
|
||||
# Margin of triplet loss
|
||||
_C.SOLVER.MARGIN = 0.3
|
||||
# Margin of cluster ;pss
|
||||
_C.SOLVER.CLUSTER_MARGIN = 0.3
|
||||
# Learning rate of SGD to learn the centers of center loss
|
||||
_C.SOLVER.CENTER_LR = 0.5
|
||||
# Balanced weight of center loss
|
||||
_C.SOLVER.CENTER_LOSS_WEIGHT = 0.0005
|
||||
# Settings of range loss
|
||||
_C.SOLVER.RANGE_K = 2
|
||||
_C.SOLVER.RANGE_MARGIN = 0.3
|
||||
_C.SOLVER.RANGE_ALPHA = 0
|
||||
_C.SOLVER.RANGE_BETA = 1
|
||||
_C.SOLVER.RANGE_LOSS_WEIGHT = 1
|
||||
|
||||
# Settings of weight decay
|
||||
_C.SOLVER.WEIGHT_DECAY = 0.0005
|
||||
_C.SOLVER.WEIGHT_DECAY_BIAS = 0.
|
||||
|
||||
# decay rate of learning rate
|
||||
_C.SOLVER.GAMMA = 0.1
|
||||
# decay step of learning rate
|
||||
_C.SOLVER.STEPS = (30, 55)
|
||||
|
||||
# warm up factor
|
||||
_C.SOLVER.WARMUP_FACTOR = 1.0 / 3
|
||||
# iterations of warm up
|
||||
_C.SOLVER.WARMUP_ITERS = 500
|
||||
# method of warm up, option: 'constant','linear'
|
||||
_C.SOLVER.WARMUP_METHOD = "linear"
|
||||
|
||||
# epoch number of saving checkpoints
|
||||
_C.SOLVER.CHECKPOINT_PERIOD = 50
|
||||
# iteration of display training log
|
||||
_C.SOLVER.LOG_PERIOD = 100
|
||||
# epoch number of validation
|
||||
_C.SOLVER.EVAL_PERIOD = 50
|
||||
|
||||
# Number of images per batch
|
||||
# This is global, so if we have 8 GPUs and IMS_PER_BATCH = 16, each GPU will
|
||||
# see 2 images per batch
|
||||
_C.SOLVER.IMS_PER_BATCH = 64
|
||||
|
||||
# This is global, so if we have 8 GPUs and IMS_PER_BATCH = 16, each GPU will
|
||||
# see 2 images per batch
|
||||
_C.TEST = CN()
|
||||
# Number of images per batch during test
|
||||
_C.TEST.IMS_PER_BATCH = 128
|
||||
# If test with re-ranking, options: 'yes','no'
|
||||
_C.TEST.RE_RANKING = 'no'
|
||||
# Path to trained model
|
||||
_C.TEST.WEIGHT = ""
|
||||
# Which feature of BNNeck to be used for test, before or after BNNneck, options: 'before' or 'after'
|
||||
_C.TEST.NECK_FEAT = 'after'
|
||||
# Whether feature is nomalized before test, if yes, it is equivalent to cosine distance
|
||||
_C.TEST.FEAT_NORM = 'yes'
|
||||
|
||||
# ---------------------------------------------------------------------------- #
|
||||
# Misc options
|
||||
# ---------------------------------------------------------------------------- #
|
||||
# Path to checkpoint and saved log of trained model
|
||||
_C.OUTPUT_DIR = ""
|
|
@ -0,0 +1,66 @@
|
|||
MODEL:
|
||||
PRETRAIN_PATH: '/home/haoluo/.torch/models/resnet50-19c8e357.pth'
|
||||
LAST_STRIDE: 2
|
||||
NECK: 'no'
|
||||
METRIC_LOSS_TYPE: 'triplet'
|
||||
IF_LABELSMOOTH: 'off'
|
||||
IF_WITH_CENTER: 'no'
|
||||
|
||||
|
||||
INPUT:
|
||||
SIZE_TRAIN: [256, 128]
|
||||
SIZE_TEST: [256, 128]
|
||||
PROB: 0.5 # random horizontal flip
|
||||
RE_PROB: 0.0 # random erasing
|
||||
PADDING: 10
|
||||
|
||||
DATASETS:
|
||||
NAMES: ('market1501')
|
||||
|
||||
DATALOADER:
|
||||
SAMPLER: 'softmax_triplet'
|
||||
NUM_INSTANCE: 4
|
||||
NUM_WORKERS: 8
|
||||
|
||||
SOLVER:
|
||||
OPTIMIZER_NAME: 'Adam'
|
||||
MAX_EPOCHS: 120
|
||||
BASE_LR: 0.00035
|
||||
|
||||
CLUSTER_MARGIN: 0.3
|
||||
|
||||
CENTER_LR: 0.5
|
||||
CENTER_LOSS_WEIGHT: 0.0005
|
||||
|
||||
RANGE_K: 2
|
||||
RANGE_MARGIN: 0.3
|
||||
RANGE_ALPHA: 0
|
||||
RANGE_BETA: 1
|
||||
RANGE_LOSS_WEIGHT: 1
|
||||
|
||||
BIAS_LR_FACTOR: 1
|
||||
WEIGHT_DECAY: 0.0005
|
||||
WEIGHT_DECAY_BIAS: 0.0005
|
||||
IMS_PER_BATCH: 64
|
||||
|
||||
STEPS: [40, 70]
|
||||
GAMMA: 0.1
|
||||
|
||||
WARMUP_FACTOR: 0.01
|
||||
WARMUP_ITERS: 0
|
||||
WARMUP_METHOD: 'linear'
|
||||
|
||||
CHECKPOINT_PERIOD: 40
|
||||
LOG_PERIOD: 20
|
||||
EVAL_PERIOD: 40
|
||||
|
||||
TEST:
|
||||
IMS_PER_BATCH: 128
|
||||
RE_RANKING: 'no'
|
||||
WEIGHT: "path"
|
||||
NECK_FEAT: 'after'
|
||||
FEAT_NORM: 'yes'
|
||||
|
||||
OUTPUT_DIR: "/home/haoluo/log/gu/reid_baseline_review/Opensource_test/market1501/Experiment-all-tricks-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on"
|
||||
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
MODEL:
|
||||
PRETRAIN_PATH: '/home/haoluo/.torch/models/resnet50-19c8e357.pth'
|
||||
|
||||
|
||||
INPUT:
|
||||
SIZE_TRAIN: [256, 128]
|
||||
SIZE_TEST: [256, 128]
|
||||
PROB: 0.5 # random horizontal flip
|
||||
RE_PROB: 0.5 # random erasing
|
||||
PADDING: 10
|
||||
|
||||
DATASETS:
|
||||
NAMES: ('market1501')
|
||||
|
||||
DATALOADER:
|
||||
SAMPLER: 'softmax'
|
||||
NUM_WORKERS: 8
|
||||
|
||||
SOLVER:
|
||||
OPTIMIZER_NAME: 'Adam'
|
||||
MAX_EPOCHS: 120
|
||||
BASE_LR: 0.00035
|
||||
BIAS_LR_FACTOR: 1
|
||||
WEIGHT_DECAY: 0.0005
|
||||
WEIGHT_DECAY_BIAS: 0.0005
|
||||
IMS_PER_BATCH: 64
|
||||
|
||||
STEPS: [30, 55]
|
||||
GAMMA: 0.1
|
||||
|
||||
WARMUP_FACTOR: 0.01
|
||||
WARMUP_ITERS: 5
|
||||
WARMUP_METHOD: 'linear'
|
||||
|
||||
CHECKPOINT_PERIOD: 20
|
||||
LOG_PERIOD: 20
|
||||
EVAL_PERIOD: 20
|
||||
|
||||
TEST:
|
||||
IMS_PER_BATCH: 128
|
||||
|
||||
OUTPUT_DIR: "/home/haoluo/log/reid/market1501/softmax_bs64_256x128"
|
||||
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
MODEL:
|
||||
PRETRAIN_PATH: '/home/haoluo/.torch/models/resnet50-19c8e357.pth'
|
||||
METRIC_LOSS_TYPE: 'triplet'
|
||||
IF_LABELSMOOTH: 'on'
|
||||
IF_WITH_CENTER: 'no'
|
||||
|
||||
|
||||
|
||||
|
||||
INPUT:
|
||||
SIZE_TRAIN: [256, 128]
|
||||
SIZE_TEST: [256, 128]
|
||||
PROB: 0.5 # random horizontal flip
|
||||
RE_PROB: 0.5 # random erasing
|
||||
PADDING: 10
|
||||
|
||||
DATASETS:
|
||||
NAMES: ('market1501')
|
||||
|
||||
DATALOADER:
|
||||
SAMPLER: 'softmax_triplet'
|
||||
NUM_INSTANCE: 4
|
||||
NUM_WORKERS: 8
|
||||
|
||||
SOLVER:
|
||||
OPTIMIZER_NAME: 'Adam'
|
||||
MAX_EPOCHS: 120
|
||||
BASE_LR: 0.00035
|
||||
|
||||
CLUSTER_MARGIN: 0.3
|
||||
|
||||
CENTER_LR: 0.5
|
||||
CENTER_LOSS_WEIGHT: 0.0005
|
||||
|
||||
RANGE_K: 2
|
||||
RANGE_MARGIN: 0.3
|
||||
RANGE_ALPHA: 0
|
||||
RANGE_BETA: 1
|
||||
RANGE_LOSS_WEIGHT: 1
|
||||
|
||||
BIAS_LR_FACTOR: 1
|
||||
WEIGHT_DECAY: 0.0005
|
||||
WEIGHT_DECAY_BIAS: 0.0005
|
||||
IMS_PER_BATCH: 64
|
||||
|
||||
STEPS: [40, 70]
|
||||
GAMMA: 0.1
|
||||
|
||||
WARMUP_FACTOR: 0.01
|
||||
WARMUP_ITERS: 10
|
||||
WARMUP_METHOD: 'linear'
|
||||
|
||||
CHECKPOINT_PERIOD: 40
|
||||
LOG_PERIOD: 20
|
||||
EVAL_PERIOD: 40
|
||||
|
||||
TEST:
|
||||
IMS_PER_BATCH: 128
|
||||
RE_RANKING: 'no'
|
||||
WEIGHT: "path"
|
||||
NECK_FEAT: 'after'
|
||||
FEAT_NORM: 'yes'
|
||||
|
||||
OUTPUT_DIR: "/home/haoluo/log/gu/reid_baseline_review/Opensource_test/market1501/Experiment-all-tricks-256x128-bs16x4-warmup10-erase0_5-labelsmooth_on-laststride1-bnneck_on"
|
||||
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
from .build import make_data_loader
|
|
@ -0,0 +1,45 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from .collate_batch import train_collate_fn, val_collate_fn
|
||||
from .datasets import init_dataset, ImageDataset
|
||||
from .samplers import RandomIdentitySampler, RandomIdentitySampler_alignedreid # New add by gu
|
||||
from .transforms import build_transforms
|
||||
|
||||
|
||||
def make_data_loader(cfg):
|
||||
train_transforms = build_transforms(cfg, is_train=True)
|
||||
val_transforms = build_transforms(cfg, is_train=False)
|
||||
num_workers = cfg.DATALOADER.NUM_WORKERS
|
||||
if len(cfg.DATASETS.NAMES) == 1:
|
||||
dataset = init_dataset(cfg.DATASETS.NAMES)
|
||||
else:
|
||||
# TODO: add multi dataset to train
|
||||
dataset = init_dataset(cfg.DATASETS.NAMES)
|
||||
|
||||
num_classes = dataset.num_train_pids
|
||||
train_set = ImageDataset(dataset.train, train_transforms)
|
||||
if cfg.DATALOADER.SAMPLER == 'softmax':
|
||||
train_loader = DataLoader(
|
||||
train_set, batch_size=cfg.SOLVER.IMS_PER_BATCH, shuffle=True, num_workers=num_workers,
|
||||
collate_fn=train_collate_fn
|
||||
)
|
||||
else:
|
||||
train_loader = DataLoader(
|
||||
train_set, batch_size=cfg.SOLVER.IMS_PER_BATCH,
|
||||
sampler=RandomIdentitySampler(dataset.train, cfg.SOLVER.IMS_PER_BATCH, cfg.DATALOADER.NUM_INSTANCE),
|
||||
# sampler=RandomIdentitySampler_alignedreid(dataset.train, cfg.DATALOADER.NUM_INSTANCE), # new add by gu
|
||||
num_workers=num_workers, collate_fn=train_collate_fn
|
||||
)
|
||||
|
||||
val_set = ImageDataset(dataset.query + dataset.gallery, val_transforms)
|
||||
val_loader = DataLoader(
|
||||
val_set, batch_size=cfg.TEST.IMS_PER_BATCH, shuffle=False, num_workers=num_workers,
|
||||
collate_fn=val_collate_fn
|
||||
)
|
||||
return train_loader, val_loader, len(dataset.query), num_classes
|
|
@ -0,0 +1,18 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def train_collate_fn(batch):
|
||||
imgs, pids, _, _, = zip(*batch)
|
||||
pids = torch.tensor(pids, dtype=torch.int64)
|
||||
return torch.stack(imgs, dim=0), pids
|
||||
|
||||
|
||||
def val_collate_fn(batch):
|
||||
imgs, pids, camids, _ = zip(*batch)
|
||||
return torch.stack(imgs, dim=0), pids, camids
|
|
@ -0,0 +1,27 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
from .cuhk03 import CUHK03
|
||||
from .dukemtmcreid import DukeMTMCreID
|
||||
from .market1501 import Market1501
|
||||
from .msmt17 import MSMT17
|
||||
from .dataset_loader import ImageDataset
|
||||
|
||||
__factory = {
|
||||
'market1501': Market1501,
|
||||
'cuhk03': CUHK03,
|
||||
'dukemtmc': DukeMTMCreID,
|
||||
'msmt17': MSMT17,
|
||||
}
|
||||
|
||||
|
||||
def get_names():
|
||||
return __factory.keys()
|
||||
|
||||
|
||||
def init_dataset(name, *args, **kwargs):
|
||||
if name not in __factory.keys():
|
||||
raise KeyError("Unknown datasets: {}".format(name))
|
||||
return __factory[name](*args, **kwargs)
|
|
@ -0,0 +1,95 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class BaseDataset(object):
|
||||
"""
|
||||
Base class of reid dataset
|
||||
"""
|
||||
|
||||
def get_imagedata_info(self, data):
|
||||
pids, cams = [], []
|
||||
for _, pid, camid in data:
|
||||
pids += [pid]
|
||||
cams += [camid]
|
||||
pids = set(pids)
|
||||
cams = set(cams)
|
||||
num_pids = len(pids)
|
||||
num_cams = len(cams)
|
||||
num_imgs = len(data)
|
||||
return num_pids, num_imgs, num_cams
|
||||
|
||||
def get_videodata_info(self, data, return_tracklet_stats=False):
|
||||
pids, cams, tracklet_stats = [], [], []
|
||||
for img_paths, pid, camid in data:
|
||||
pids += [pid]
|
||||
cams += [camid]
|
||||
tracklet_stats += [len(img_paths)]
|
||||
pids = set(pids)
|
||||
cams = set(cams)
|
||||
num_pids = len(pids)
|
||||
num_cams = len(cams)
|
||||
num_tracklets = len(data)
|
||||
if return_tracklet_stats:
|
||||
return num_pids, num_tracklets, num_cams, tracklet_stats
|
||||
return num_pids, num_tracklets, num_cams
|
||||
|
||||
def print_dataset_statistics(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class BaseImageDataset(BaseDataset):
|
||||
"""
|
||||
Base class of image reid dataset
|
||||
"""
|
||||
|
||||
def print_dataset_statistics(self, train, query, gallery):
|
||||
num_train_pids, num_train_imgs, num_train_cams = self.get_imagedata_info(train)
|
||||
num_query_pids, num_query_imgs, num_query_cams = self.get_imagedata_info(query)
|
||||
num_gallery_pids, num_gallery_imgs, num_gallery_cams = self.get_imagedata_info(gallery)
|
||||
|
||||
print("Dataset statistics:")
|
||||
print(" ----------------------------------------")
|
||||
print(" subset | # ids | # images | # cameras")
|
||||
print(" ----------------------------------------")
|
||||
print(" train | {:5d} | {:8d} | {:9d}".format(num_train_pids, num_train_imgs, num_train_cams))
|
||||
print(" query | {:5d} | {:8d} | {:9d}".format(num_query_pids, num_query_imgs, num_query_cams))
|
||||
print(" gallery | {:5d} | {:8d} | {:9d}".format(num_gallery_pids, num_gallery_imgs, num_gallery_cams))
|
||||
print(" ----------------------------------------")
|
||||
|
||||
|
||||
class BaseVideoDataset(BaseDataset):
|
||||
"""
|
||||
Base class of video reid dataset
|
||||
"""
|
||||
|
||||
def print_dataset_statistics(self, train, query, gallery):
|
||||
num_train_pids, num_train_tracklets, num_train_cams, train_tracklet_stats = \
|
||||
self.get_videodata_info(train, return_tracklet_stats=True)
|
||||
|
||||
num_query_pids, num_query_tracklets, num_query_cams, query_tracklet_stats = \
|
||||
self.get_videodata_info(query, return_tracklet_stats=True)
|
||||
|
||||
num_gallery_pids, num_gallery_tracklets, num_gallery_cams, gallery_tracklet_stats = \
|
||||
self.get_videodata_info(gallery, return_tracklet_stats=True)
|
||||
|
||||
tracklet_stats = train_tracklet_stats + query_tracklet_stats + gallery_tracklet_stats
|
||||
min_num = np.min(tracklet_stats)
|
||||
max_num = np.max(tracklet_stats)
|
||||
avg_num = np.mean(tracklet_stats)
|
||||
|
||||
print("Dataset statistics:")
|
||||
print(" -------------------------------------------")
|
||||
print(" subset | # ids | # tracklets | # cameras")
|
||||
print(" -------------------------------------------")
|
||||
print(" train | {:5d} | {:11d} | {:9d}".format(num_train_pids, num_train_tracklets, num_train_cams))
|
||||
print(" query | {:5d} | {:11d} | {:9d}".format(num_query_pids, num_query_tracklets, num_query_cams))
|
||||
print(" gallery | {:5d} | {:11d} | {:9d}".format(num_gallery_pids, num_gallery_tracklets, num_gallery_cams))
|
||||
print(" -------------------------------------------")
|
||||
print(" number of images per tracklet: {} ~ {}, average {:.2f}".format(min_num, max_num, avg_num))
|
||||
print(" -------------------------------------------")
|
|
@ -0,0 +1,259 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: liaoxingyu2@jd.com
|
||||
"""
|
||||
|
||||
import h5py
|
||||
import os.path as osp
|
||||
from scipy.io import loadmat
|
||||
from scipy.misc import imsave
|
||||
|
||||
from utils.iotools import mkdir_if_missing, write_json, read_json
|
||||
from .bases import BaseImageDataset
|
||||
|
||||
|
||||
class CUHK03(BaseImageDataset):
|
||||
"""
|
||||
CUHK03
|
||||
Reference:
|
||||
Li et al. DeepReID: Deep Filter Pairing Neural Network for Person Re-identification. CVPR 2014.
|
||||
URL: http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html#!
|
||||
|
||||
Dataset statistics:
|
||||
# identities: 1360
|
||||
# images: 13164
|
||||
# cameras: 6
|
||||
# splits: 20 (classic)
|
||||
Args:
|
||||
split_id (int): split index (default: 0)
|
||||
cuhk03_labeled (bool): whether to load labeled images; if false, detected images are loaded (default: False)
|
||||
"""
|
||||
dataset_dir = 'cuhk03'
|
||||
|
||||
def __init__(self, root='/home/haoluo/data', split_id=0, cuhk03_labeled=False,
|
||||
cuhk03_classic_split=False, verbose=True,
|
||||
**kwargs):
|
||||
super(CUHK03, self).__init__()
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.data_dir = osp.join(self.dataset_dir, 'cuhk03_release')
|
||||
self.raw_mat_path = osp.join(self.data_dir, 'cuhk-03.mat')
|
||||
|
||||
self.imgs_detected_dir = osp.join(self.dataset_dir, 'images_detected')
|
||||
self.imgs_labeled_dir = osp.join(self.dataset_dir, 'images_labeled')
|
||||
|
||||
self.split_classic_det_json_path = osp.join(self.dataset_dir, 'splits_classic_detected.json')
|
||||
self.split_classic_lab_json_path = osp.join(self.dataset_dir, 'splits_classic_labeled.json')
|
||||
|
||||
self.split_new_det_json_path = osp.join(self.dataset_dir, 'splits_new_detected.json')
|
||||
self.split_new_lab_json_path = osp.join(self.dataset_dir, 'splits_new_labeled.json')
|
||||
|
||||
self.split_new_det_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_detected.mat')
|
||||
self.split_new_lab_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_labeled.mat')
|
||||
|
||||
self._check_before_run()
|
||||
self._preprocess()
|
||||
|
||||
if cuhk03_labeled:
|
||||
image_type = 'labeled'
|
||||
split_path = self.split_classic_lab_json_path if cuhk03_classic_split else self.split_new_lab_json_path
|
||||
else:
|
||||
image_type = 'detected'
|
||||
split_path = self.split_classic_det_json_path if cuhk03_classic_split else self.split_new_det_json_path
|
||||
|
||||
splits = read_json(split_path)
|
||||
assert split_id < len(splits), "Condition split_id ({}) < len(splits) ({}) is false".format(split_id,
|
||||
len(splits))
|
||||
split = splits[split_id]
|
||||
print("Split index = {}".format(split_id))
|
||||
|
||||
train = split['train']
|
||||
query = split['query']
|
||||
gallery = split['gallery']
|
||||
|
||||
if verbose:
|
||||
print("=> CUHK03 ({}) loaded".format(image_type))
|
||||
self.print_dataset_statistics(train, query, gallery)
|
||||
|
||||
self.train = train
|
||||
self.query = query
|
||||
self.gallery = gallery
|
||||
|
||||
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
|
||||
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
|
||||
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
|
||||
|
||||
def _check_before_run(self):
|
||||
"""Check if all files are available before going deeper"""
|
||||
if not osp.exists(self.dataset_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
||||
if not osp.exists(self.data_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.data_dir))
|
||||
if not osp.exists(self.raw_mat_path):
|
||||
raise RuntimeError("'{}' is not available".format(self.raw_mat_path))
|
||||
if not osp.exists(self.split_new_det_mat_path):
|
||||
raise RuntimeError("'{}' is not available".format(self.split_new_det_mat_path))
|
||||
if not osp.exists(self.split_new_lab_mat_path):
|
||||
raise RuntimeError("'{}' is not available".format(self.split_new_lab_mat_path))
|
||||
|
||||
def _preprocess(self):
|
||||
"""
|
||||
This function is a bit complex and ugly, what it does is
|
||||
1. Extract data from cuhk-03.mat and save as png images.
|
||||
2. Create 20 classic splits. (Li et al. CVPR'14)
|
||||
3. Create new split. (Zhong et al. CVPR'17)
|
||||
"""
|
||||
print(
|
||||
"Note: if root path is changed, the previously generated json files need to be re-generated (delete them first)")
|
||||
if osp.exists(self.imgs_labeled_dir) and \
|
||||
osp.exists(self.imgs_detected_dir) and \
|
||||
osp.exists(self.split_classic_det_json_path) and \
|
||||
osp.exists(self.split_classic_lab_json_path) and \
|
||||
osp.exists(self.split_new_det_json_path) and \
|
||||
osp.exists(self.split_new_lab_json_path):
|
||||
return
|
||||
|
||||
mkdir_if_missing(self.imgs_detected_dir)
|
||||
mkdir_if_missing(self.imgs_labeled_dir)
|
||||
|
||||
print("Extract image data from {} and save as png".format(self.raw_mat_path))
|
||||
mat = h5py.File(self.raw_mat_path, 'r')
|
||||
|
||||
def _deref(ref):
|
||||
return mat[ref][:].T
|
||||
|
||||
def _process_images(img_refs, campid, pid, save_dir):
|
||||
img_paths = [] # Note: some persons only have images for one view
|
||||
for imgid, img_ref in enumerate(img_refs):
|
||||
img = _deref(img_ref)
|
||||
# skip empty cell
|
||||
if img.size == 0 or img.ndim < 3: continue
|
||||
# images are saved with the following format, index-1 (ensure uniqueness)
|
||||
# campid: index of camera pair (1-5)
|
||||
# pid: index of person in 'campid'-th camera pair
|
||||
# viewid: index of view, {1, 2}
|
||||
# imgid: index of image, (1-10)
|
||||
viewid = 1 if imgid < 5 else 2
|
||||
img_name = '{:01d}_{:03d}_{:01d}_{:02d}.png'.format(campid + 1, pid + 1, viewid, imgid + 1)
|
||||
img_path = osp.join(save_dir, img_name)
|
||||
if not osp.isfile(img_path):
|
||||
imsave(img_path, img)
|
||||
img_paths.append(img_path)
|
||||
return img_paths
|
||||
|
||||
def _extract_img(name):
|
||||
print("Processing {} images (extract and save) ...".format(name))
|
||||
meta_data = []
|
||||
imgs_dir = self.imgs_detected_dir if name == 'detected' else self.imgs_labeled_dir
|
||||
for campid, camp_ref in enumerate(mat[name][0]):
|
||||
camp = _deref(camp_ref)
|
||||
num_pids = camp.shape[0]
|
||||
for pid in range(num_pids):
|
||||
img_paths = _process_images(camp[pid, :], campid, pid, imgs_dir)
|
||||
assert len(img_paths) > 0, "campid{}-pid{} has no images".format(campid, pid)
|
||||
meta_data.append((campid + 1, pid + 1, img_paths))
|
||||
print("- done camera pair {} with {} identities".format(campid + 1, num_pids))
|
||||
return meta_data
|
||||
|
||||
meta_detected = _extract_img('detected')
|
||||
meta_labeled = _extract_img('labeled')
|
||||
|
||||
def _extract_classic_split(meta_data, test_split):
|
||||
train, test = [], []
|
||||
num_train_pids, num_test_pids = 0, 0
|
||||
num_train_imgs, num_test_imgs = 0, 0
|
||||
for i, (campid, pid, img_paths) in enumerate(meta_data):
|
||||
|
||||
if [campid, pid] in test_split:
|
||||
for img_path in img_paths:
|
||||
camid = int(osp.basename(img_path).split('_')[2]) - 1 # make it 0-based
|
||||
test.append((img_path, num_test_pids, camid))
|
||||
num_test_pids += 1
|
||||
num_test_imgs += len(img_paths)
|
||||
else:
|
||||
for img_path in img_paths:
|
||||
camid = int(osp.basename(img_path).split('_')[2]) - 1 # make it 0-based
|
||||
train.append((img_path, num_train_pids, camid))
|
||||
num_train_pids += 1
|
||||
num_train_imgs += len(img_paths)
|
||||
return train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs
|
||||
|
||||
print("Creating classic splits (# = 20) ...")
|
||||
splits_classic_det, splits_classic_lab = [], []
|
||||
for split_ref in mat['testsets'][0]:
|
||||
test_split = _deref(split_ref).tolist()
|
||||
|
||||
# create split for detected images
|
||||
train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
|
||||
_extract_classic_split(meta_detected, test_split)
|
||||
splits_classic_det.append({
|
||||
'train': train, 'query': test, 'gallery': test,
|
||||
'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs,
|
||||
'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs,
|
||||
'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs,
|
||||
})
|
||||
|
||||
# create split for labeled images
|
||||
train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
|
||||
_extract_classic_split(meta_labeled, test_split)
|
||||
splits_classic_lab.append({
|
||||
'train': train, 'query': test, 'gallery': test,
|
||||
'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs,
|
||||
'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs,
|
||||
'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs,
|
||||
})
|
||||
|
||||
write_json(splits_classic_det, self.split_classic_det_json_path)
|
||||
write_json(splits_classic_lab, self.split_classic_lab_json_path)
|
||||
|
||||
def _extract_set(filelist, pids, pid2label, idxs, img_dir, relabel):
|
||||
tmp_set = []
|
||||
unique_pids = set()
|
||||
for idx in idxs:
|
||||
img_name = filelist[idx][0]
|
||||
camid = int(img_name.split('_')[2]) - 1 # make it 0-based
|
||||
pid = pids[idx]
|
||||
if relabel: pid = pid2label[pid]
|
||||
img_path = osp.join(img_dir, img_name)
|
||||
tmp_set.append((img_path, int(pid), camid))
|
||||
unique_pids.add(pid)
|
||||
return tmp_set, len(unique_pids), len(idxs)
|
||||
|
||||
def _extract_new_split(split_dict, img_dir):
|
||||
train_idxs = split_dict['train_idx'].flatten() - 1 # index-0
|
||||
pids = split_dict['labels'].flatten()
|
||||
train_pids = set(pids[train_idxs])
|
||||
pid2label = {pid: label for label, pid in enumerate(train_pids)}
|
||||
query_idxs = split_dict['query_idx'].flatten() - 1
|
||||
gallery_idxs = split_dict['gallery_idx'].flatten() - 1
|
||||
filelist = split_dict['filelist'].flatten()
|
||||
train_info = _extract_set(filelist, pids, pid2label, train_idxs, img_dir, relabel=True)
|
||||
query_info = _extract_set(filelist, pids, pid2label, query_idxs, img_dir, relabel=False)
|
||||
gallery_info = _extract_set(filelist, pids, pid2label, gallery_idxs, img_dir, relabel=False)
|
||||
return train_info, query_info, gallery_info
|
||||
|
||||
print("Creating new splits for detected images (767/700) ...")
|
||||
train_info, query_info, gallery_info = _extract_new_split(
|
||||
loadmat(self.split_new_det_mat_path),
|
||||
self.imgs_detected_dir,
|
||||
)
|
||||
splits = [{
|
||||
'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0],
|
||||
'num_train_pids': train_info[1], 'num_train_imgs': train_info[2],
|
||||
'num_query_pids': query_info[1], 'num_query_imgs': query_info[2],
|
||||
'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2],
|
||||
}]
|
||||
write_json(splits, self.split_new_det_json_path)
|
||||
|
||||
print("Creating new splits for labeled images (767/700) ...")
|
||||
train_info, query_info, gallery_info = _extract_new_split(
|
||||
loadmat(self.split_new_lab_mat_path),
|
||||
self.imgs_labeled_dir,
|
||||
)
|
||||
splits = [{
|
||||
'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0],
|
||||
'num_train_pids': train_info[1], 'num_train_imgs': train_info[2],
|
||||
'num_query_pids': query_info[1], 'num_query_imgs': query_info[2],
|
||||
'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2],
|
||||
}]
|
||||
write_json(splits, self.split_new_lab_json_path)
|
|
@ -0,0 +1,45 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
def read_image(img_path):
|
||||
"""Keep reading image until succeed.
|
||||
This can avoid IOError incurred by heavy IO process."""
|
||||
got_img = False
|
||||
if not osp.exists(img_path):
|
||||
raise IOError("{} does not exist".format(img_path))
|
||||
while not got_img:
|
||||
try:
|
||||
img = Image.open(img_path).convert('RGB')
|
||||
got_img = True
|
||||
except IOError:
|
||||
print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
|
||||
pass
|
||||
return img
|
||||
|
||||
|
||||
class ImageDataset(Dataset):
|
||||
"""Image Person ReID Dataset"""
|
||||
|
||||
def __init__(self, dataset, transform=None):
|
||||
self.dataset = dataset
|
||||
self.transform = transform
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
|
||||
def __getitem__(self, index):
|
||||
img_path, pid, camid = self.dataset[index]
|
||||
img = read_image(img_path)
|
||||
|
||||
if self.transform is not None:
|
||||
img = self.transform(img)
|
||||
|
||||
return img, pid, camid, img_path
|
|
@ -0,0 +1,106 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: liaoxingyu2@jd.com
|
||||
"""
|
||||
|
||||
import glob
|
||||
import re
|
||||
import urllib
|
||||
import zipfile
|
||||
|
||||
import os.path as osp
|
||||
|
||||
from utils.iotools import mkdir_if_missing
|
||||
from .bases import BaseImageDataset
|
||||
|
||||
|
||||
class DukeMTMCreID(BaseImageDataset):
|
||||
"""
|
||||
DukeMTMC-reID
|
||||
Reference:
|
||||
1. Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.
|
||||
2. Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.
|
||||
URL: https://github.com/layumi/DukeMTMC-reID_evaluation
|
||||
|
||||
Dataset statistics:
|
||||
# identities: 1404 (train + query)
|
||||
# images:16522 (train) + 2228 (query) + 17661 (gallery)
|
||||
# cameras: 8
|
||||
"""
|
||||
dataset_dir = 'dukemtmc-reid'
|
||||
|
||||
def __init__(self, root='/home/haoluo/data', verbose=True, **kwargs):
|
||||
super(DukeMTMCreID, self).__init__()
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip'
|
||||
self.train_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/bounding_box_train')
|
||||
self.query_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/query')
|
||||
self.gallery_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/bounding_box_test')
|
||||
|
||||
self._download_data()
|
||||
self._check_before_run()
|
||||
|
||||
train = self._process_dir(self.train_dir, relabel=True)
|
||||
query = self._process_dir(self.query_dir, relabel=False)
|
||||
gallery = self._process_dir(self.gallery_dir, relabel=False)
|
||||
|
||||
if verbose:
|
||||
print("=> DukeMTMC-reID loaded")
|
||||
self.print_dataset_statistics(train, query, gallery)
|
||||
|
||||
self.train = train
|
||||
self.query = query
|
||||
self.gallery = gallery
|
||||
|
||||
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
|
||||
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
|
||||
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
|
||||
|
||||
def _download_data(self):
|
||||
if osp.exists(self.dataset_dir):
|
||||
print("This dataset has been downloaded.")
|
||||
return
|
||||
|
||||
print("Creating directory {}".format(self.dataset_dir))
|
||||
mkdir_if_missing(self.dataset_dir)
|
||||
fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url))
|
||||
|
||||
print("Downloading DukeMTMC-reID dataset")
|
||||
urllib.urlretrieve(self.dataset_url, fpath)
|
||||
|
||||
print("Extracting files")
|
||||
zip_ref = zipfile.ZipFile(fpath, 'r')
|
||||
zip_ref.extractall(self.dataset_dir)
|
||||
zip_ref.close()
|
||||
|
||||
def _check_before_run(self):
|
||||
"""Check if all files are available before going deeper"""
|
||||
if not osp.exists(self.dataset_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
||||
if not osp.exists(self.train_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.train_dir))
|
||||
if not osp.exists(self.query_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.query_dir))
|
||||
if not osp.exists(self.gallery_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
|
||||
|
||||
def _process_dir(self, dir_path, relabel=False):
|
||||
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
|
||||
pattern = re.compile(r'([-\d]+)_c(\d)')
|
||||
|
||||
pid_container = set()
|
||||
for img_path in img_paths:
|
||||
pid, _ = map(int, pattern.search(img_path).groups())
|
||||
pid_container.add(pid)
|
||||
pid2label = {pid: label for label, pid in enumerate(pid_container)}
|
||||
|
||||
dataset = []
|
||||
for img_path in img_paths:
|
||||
pid, camid = map(int, pattern.search(img_path).groups())
|
||||
assert 1 <= camid <= 8
|
||||
camid -= 1 # index starts from 0
|
||||
if relabel: pid = pid2label[pid]
|
||||
dataset.append((img_path, pid, camid))
|
||||
|
||||
return dataset
|
|
@ -0,0 +1,63 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def eval_func(distmat, q_pids, g_pids, q_camids, g_camids, max_rank=50):
|
||||
"""Evaluation with market1501 metric
|
||||
Key: for each query identity, its gallery images from the same camera view are discarded.
|
||||
"""
|
||||
num_q, num_g = distmat.shape
|
||||
if num_g < max_rank:
|
||||
max_rank = num_g
|
||||
print("Note: number of gallery samples is quite small, got {}".format(num_g))
|
||||
indices = np.argsort(distmat, axis=1)
|
||||
matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)
|
||||
|
||||
# compute cmc curve for each query
|
||||
all_cmc = []
|
||||
all_AP = []
|
||||
num_valid_q = 0. # number of valid query
|
||||
for q_idx in range(num_q):
|
||||
# get query pid and camid
|
||||
q_pid = q_pids[q_idx]
|
||||
q_camid = q_camids[q_idx]
|
||||
|
||||
# remove gallery samples that have the same pid and camid with query
|
||||
order = indices[q_idx]
|
||||
remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
|
||||
keep = np.invert(remove)
|
||||
|
||||
# compute cmc curve
|
||||
# binary vector, positions with value 1 are correct matches
|
||||
orig_cmc = matches[q_idx][keep]
|
||||
if not np.any(orig_cmc):
|
||||
# this condition is true when query identity does not appear in gallery
|
||||
continue
|
||||
|
||||
cmc = orig_cmc.cumsum()
|
||||
cmc[cmc > 1] = 1
|
||||
|
||||
all_cmc.append(cmc[:max_rank])
|
||||
num_valid_q += 1.
|
||||
|
||||
# compute average precision
|
||||
# reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision
|
||||
num_rel = orig_cmc.sum()
|
||||
tmp_cmc = orig_cmc.cumsum()
|
||||
tmp_cmc = [x / (i + 1.) for i, x in enumerate(tmp_cmc)]
|
||||
tmp_cmc = np.asarray(tmp_cmc) * orig_cmc
|
||||
AP = tmp_cmc.sum() / num_rel
|
||||
all_AP.append(AP)
|
||||
|
||||
assert num_valid_q > 0, "Error: all query identities do not appear in gallery"
|
||||
|
||||
all_cmc = np.asarray(all_cmc).astype(np.float32)
|
||||
all_cmc = all_cmc.sum(0) / num_valid_q
|
||||
mAP = np.mean(all_AP)
|
||||
|
||||
return all_cmc, mAP
|
|
@ -0,0 +1,85 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import glob
|
||||
import re
|
||||
|
||||
import os.path as osp
|
||||
|
||||
from .bases import BaseImageDataset
|
||||
|
||||
|
||||
class Market1501(BaseImageDataset):
|
||||
"""
|
||||
Market1501
|
||||
Reference:
|
||||
Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
|
||||
URL: http://www.liangzheng.org/Project/project_reid.html
|
||||
|
||||
Dataset statistics:
|
||||
# identities: 1501 (+1 for background)
|
||||
# images: 12936 (train) + 3368 (query) + 15913 (gallery)
|
||||
"""
|
||||
dataset_dir = 'market1501'
|
||||
|
||||
def __init__(self, root='/home/haoluo/data', verbose=True, **kwargs):
|
||||
super(Market1501, self).__init__()
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
|
||||
self.query_dir = osp.join(self.dataset_dir, 'query')
|
||||
self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
|
||||
|
||||
self._check_before_run()
|
||||
|
||||
train = self._process_dir(self.train_dir, relabel=True)
|
||||
query = self._process_dir(self.query_dir, relabel=False)
|
||||
gallery = self._process_dir(self.gallery_dir, relabel=False)
|
||||
|
||||
if verbose:
|
||||
print("=> Market1501 loaded")
|
||||
self.print_dataset_statistics(train, query, gallery)
|
||||
|
||||
self.train = train
|
||||
self.query = query
|
||||
self.gallery = gallery
|
||||
|
||||
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
|
||||
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
|
||||
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
|
||||
|
||||
def _check_before_run(self):
|
||||
"""Check if all files are available before going deeper"""
|
||||
if not osp.exists(self.dataset_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
||||
if not osp.exists(self.train_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.train_dir))
|
||||
if not osp.exists(self.query_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.query_dir))
|
||||
if not osp.exists(self.gallery_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
|
||||
|
||||
def _process_dir(self, dir_path, relabel=False):
|
||||
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
|
||||
pattern = re.compile(r'([-\d]+)_c(\d)')
|
||||
|
||||
pid_container = set()
|
||||
for img_path in img_paths:
|
||||
pid, _ = map(int, pattern.search(img_path).groups())
|
||||
if pid == -1: continue # junk images are just ignored
|
||||
pid_container.add(pid)
|
||||
pid2label = {pid: label for label, pid in enumerate(pid_container)}
|
||||
|
||||
dataset = []
|
||||
for img_path in img_paths:
|
||||
pid, camid = map(int, pattern.search(img_path).groups())
|
||||
if pid == -1: continue # junk images are just ignored
|
||||
assert 0 <= pid <= 1501 # pid == 0 means background
|
||||
assert 1 <= camid <= 6
|
||||
camid -= 1 # index starts from 0
|
||||
if relabel: pid = pid2label[pid]
|
||||
dataset.append((img_path, pid, camid))
|
||||
|
||||
return dataset
|
|
@ -0,0 +1,83 @@
|
|||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/1/17 15:00
|
||||
# @Author : Hao Luo
|
||||
# @File : msmt17.py
|
||||
|
||||
import glob
|
||||
import re
|
||||
|
||||
import os.path as osp
|
||||
|
||||
from .bases import BaseImageDataset
|
||||
|
||||
|
||||
class MSMT17(BaseImageDataset):
|
||||
"""
|
||||
MSMT17
|
||||
|
||||
Reference:
|
||||
Wei et al. Person Transfer GAN to Bridge Domain Gap for Person Re-Identification. CVPR 2018.
|
||||
|
||||
URL: http://www.pkuvmc.com/publications/msmt17.html
|
||||
|
||||
Dataset statistics:
|
||||
# identities: 4101
|
||||
# images: 32621 (train) + 11659 (query) + 82161 (gallery)
|
||||
# cameras: 15
|
||||
"""
|
||||
dataset_dir = 'msmt17'
|
||||
|
||||
def __init__(self,root='/home/haoluo/data', verbose=True, **kwargs):
|
||||
super(MSMT17, self).__init__()
|
||||
self.dataset_dir = osp.join(root, self.dataset_dir)
|
||||
self.train_dir = osp.join(self.dataset_dir, 'MSMT17_V2/mask_train_v2')
|
||||
self.test_dir = osp.join(self.dataset_dir, 'MSMT17_V2/mask_test_v2')
|
||||
self.list_train_path = osp.join(self.dataset_dir, 'MSMT17_V2/list_train.txt')
|
||||
self.list_val_path = osp.join(self.dataset_dir, 'MSMT17_V2/list_val.txt')
|
||||
self.list_query_path = osp.join(self.dataset_dir, 'MSMT17_V2/list_query.txt')
|
||||
self.list_gallery_path = osp.join(self.dataset_dir, 'MSMT17_V2/list_gallery.txt')
|
||||
|
||||
self._check_before_run()
|
||||
train = self._process_dir(self.train_dir, self.list_train_path)
|
||||
#val, num_val_pids, num_val_imgs = self._process_dir(self.train_dir, self.list_val_path)
|
||||
query = self._process_dir(self.test_dir, self.list_query_path)
|
||||
gallery = self._process_dir(self.test_dir, self.list_gallery_path)
|
||||
if verbose:
|
||||
print("=> MSMT17 loaded")
|
||||
self.print_dataset_statistics(train, query, gallery)
|
||||
|
||||
self.train = train
|
||||
self.query = query
|
||||
self.gallery = gallery
|
||||
|
||||
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
|
||||
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
|
||||
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
|
||||
|
||||
def _check_before_run(self):
|
||||
"""Check if all files are available before going deeper"""
|
||||
if not osp.exists(self.dataset_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
||||
if not osp.exists(self.train_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.train_dir))
|
||||
if not osp.exists(self.test_dir):
|
||||
raise RuntimeError("'{}' is not available".format(self.test_dir))
|
||||
|
||||
def _process_dir(self, dir_path, list_path):
|
||||
with open(list_path, 'r') as txt:
|
||||
lines = txt.readlines()
|
||||
dataset = []
|
||||
pid_container = set()
|
||||
for img_idx, img_info in enumerate(lines):
|
||||
img_path, pid = img_info.split(' ')
|
||||
pid = int(pid) # no need to relabel
|
||||
camid = int(img_path.split('_')[2])
|
||||
img_path = osp.join(dir_path, img_path)
|
||||
dataset.append((img_path, pid, camid))
|
||||
pid_container.add(pid)
|
||||
|
||||
# check if pid starts from 0 and increments with 1
|
||||
for idx, pid in enumerate(pid_container):
|
||||
assert idx == pid, "See code comment for explanation"
|
||||
return dataset
|
|
@ -0,0 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
from .triplet_sampler import RandomIdentitySampler, RandomIdentitySampler_alignedreid # new add by gu
|
|
@ -0,0 +1,110 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: liaoxingyu2@jd.com
|
||||
"""
|
||||
|
||||
import copy
|
||||
import random
|
||||
import torch
|
||||
from collections import defaultdict
|
||||
|
||||
import numpy as np
|
||||
from torch.utils.data.sampler import Sampler
|
||||
|
||||
|
||||
class RandomIdentitySampler(Sampler):
|
||||
"""
|
||||
Randomly sample N identities, then for each identity,
|
||||
randomly sample K instances, therefore batch size is N*K.
|
||||
Args:
|
||||
- data_source (list): list of (img_path, pid, camid).
|
||||
- num_instances (int): number of instances per identity in a batch.
|
||||
- batch_size (int): number of examples in a batch.
|
||||
"""
|
||||
|
||||
def __init__(self, data_source, batch_size, num_instances):
|
||||
self.data_source = data_source
|
||||
self.batch_size = batch_size
|
||||
self.num_instances = num_instances
|
||||
self.num_pids_per_batch = self.batch_size // self.num_instances
|
||||
self.index_dic = defaultdict(list)
|
||||
for index, (_, pid, _) in enumerate(self.data_source):
|
||||
self.index_dic[pid].append(index)
|
||||
self.pids = list(self.index_dic.keys())
|
||||
|
||||
# estimate number of examples in an epoch
|
||||
self.length = 0
|
||||
for pid in self.pids:
|
||||
idxs = self.index_dic[pid]
|
||||
num = len(idxs)
|
||||
if num < self.num_instances:
|
||||
num = self.num_instances
|
||||
self.length += num - num % self.num_instances
|
||||
|
||||
def __iter__(self):
|
||||
batch_idxs_dict = defaultdict(list)
|
||||
|
||||
for pid in self.pids:
|
||||
idxs = copy.deepcopy(self.index_dic[pid])
|
||||
if len(idxs) < self.num_instances:
|
||||
idxs = np.random.choice(idxs, size=self.num_instances, replace=True)
|
||||
random.shuffle(idxs)
|
||||
batch_idxs = []
|
||||
for idx in idxs:
|
||||
batch_idxs.append(idx)
|
||||
if len(batch_idxs) == self.num_instances:
|
||||
batch_idxs_dict[pid].append(batch_idxs)
|
||||
batch_idxs = []
|
||||
|
||||
avai_pids = copy.deepcopy(self.pids)
|
||||
final_idxs = []
|
||||
|
||||
while len(avai_pids) >= self.num_pids_per_batch:
|
||||
selected_pids = random.sample(avai_pids, self.num_pids_per_batch)
|
||||
for pid in selected_pids:
|
||||
batch_idxs = batch_idxs_dict[pid].pop(0)
|
||||
final_idxs.extend(batch_idxs)
|
||||
if len(batch_idxs_dict[pid]) == 0:
|
||||
avai_pids.remove(pid)
|
||||
|
||||
return iter(final_idxs)
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
|
||||
|
||||
# New add by gu
|
||||
class RandomIdentitySampler_alignedreid(Sampler):
|
||||
"""
|
||||
Randomly sample N identities, then for each identity,
|
||||
randomly sample K instances, therefore batch size is N*K.
|
||||
|
||||
Code imported from https://github.com/Cysu/open-reid/blob/master/reid/utils/data/sampler.py.
|
||||
|
||||
Args:
|
||||
data_source (Dataset): dataset to sample from.
|
||||
num_instances (int): number of instances per identity.
|
||||
"""
|
||||
def __init__(self, data_source, num_instances):
|
||||
self.data_source = data_source
|
||||
self.num_instances = num_instances
|
||||
self.index_dic = defaultdict(list)
|
||||
for index, (_, pid, _) in enumerate(data_source):
|
||||
self.index_dic[pid].append(index)
|
||||
self.pids = list(self.index_dic.keys())
|
||||
self.num_identities = len(self.pids)
|
||||
|
||||
def __iter__(self):
|
||||
indices = torch.randperm(self.num_identities)
|
||||
ret = []
|
||||
for i in indices:
|
||||
pid = self.pids[i]
|
||||
t = self.index_dic[pid]
|
||||
replace = False if len(t) >= self.num_instances else True
|
||||
t = np.random.choice(t, size=self.num_instances, replace=replace)
|
||||
ret.extend(t)
|
||||
return iter(ret)
|
||||
|
||||
def __len__(self):
|
||||
return self.num_identities * self.num_instances
|
|
@ -0,0 +1,7 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
from .build import build_transforms
|
|
@ -0,0 +1,31 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: liaoxingyu2@jd.com
|
||||
"""
|
||||
|
||||
import torchvision.transforms as T
|
||||
|
||||
from .transforms import RandomErasing
|
||||
|
||||
|
||||
def build_transforms(cfg, is_train=True):
|
||||
normalize_transform = T.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD)
|
||||
if is_train:
|
||||
transform = T.Compose([
|
||||
T.Resize(cfg.INPUT.SIZE_TRAIN),
|
||||
T.RandomHorizontalFlip(p=cfg.INPUT.PROB),
|
||||
T.Pad(cfg.INPUT.PADDING),
|
||||
T.RandomCrop(cfg.INPUT.SIZE_TRAIN),
|
||||
T.ToTensor(),
|
||||
normalize_transform,
|
||||
RandomErasing(probability=cfg.INPUT.RE_PROB, mean=cfg.INPUT.PIXEL_MEAN)
|
||||
])
|
||||
else:
|
||||
transform = T.Compose([
|
||||
T.Resize(cfg.INPUT.SIZE_TEST),
|
||||
T.ToTensor(),
|
||||
normalize_transform
|
||||
])
|
||||
|
||||
return transform
|
|
@ -0,0 +1,55 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: liaoxingyu2@jd.com
|
||||
"""
|
||||
|
||||
import math
|
||||
import random
|
||||
|
||||
|
||||
class RandomErasing(object):
|
||||
""" Randomly selects a rectangle region in an image and erases its pixels.
|
||||
'Random Erasing Data Augmentation' by Zhong et al.
|
||||
See https://arxiv.org/pdf/1708.04896.pdf
|
||||
Args:
|
||||
probability: The probability that the Random Erasing operation will be performed.
|
||||
sl: Minimum proportion of erased area against input image.
|
||||
sh: Maximum proportion of erased area against input image.
|
||||
r1: Minimum aspect ratio of erased area.
|
||||
mean: Erasing value.
|
||||
"""
|
||||
|
||||
def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)):
|
||||
self.probability = probability
|
||||
self.mean = mean
|
||||
self.sl = sl
|
||||
self.sh = sh
|
||||
self.r1 = r1
|
||||
|
||||
def __call__(self, img):
|
||||
|
||||
if random.uniform(0, 1) >= self.probability:
|
||||
return img
|
||||
|
||||
for attempt in range(100):
|
||||
area = img.size()[1] * img.size()[2]
|
||||
|
||||
target_area = random.uniform(self.sl, self.sh) * area
|
||||
aspect_ratio = random.uniform(self.r1, 1 / self.r1)
|
||||
|
||||
h = int(round(math.sqrt(target_area * aspect_ratio)))
|
||||
w = int(round(math.sqrt(target_area / aspect_ratio)))
|
||||
|
||||
if w < img.size()[2] and h < img.size()[1]:
|
||||
x1 = random.randint(0, img.size()[1] - h)
|
||||
y1 = random.randint(0, img.size()[2] - w)
|
||||
if img.size()[0] == 3:
|
||||
img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
|
||||
img[1, x1:x1 + h, y1:y1 + w] = self.mean[1]
|
||||
img[2, x1:x1 + h, y1:y1 + w] = self.mean[2]
|
||||
else:
|
||||
img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
|
||||
return img
|
||||
|
||||
return img
|
|
@ -0,0 +1,72 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
import logging
|
||||
|
||||
import torch
|
||||
from ignite.engine import Engine
|
||||
|
||||
from utils.reid_metric import R1_mAP, R1_mAP_reranking
|
||||
|
||||
|
||||
def create_supervised_evaluator(model, metrics,
|
||||
device=None):
|
||||
"""
|
||||
Factory function for creating an evaluator for supervised models
|
||||
|
||||
Args:
|
||||
model (`torch.nn.Module`): the model to train
|
||||
metrics (dict of str - :class:`ignite.metrics.Metric`): a map of metric names to Metrics
|
||||
device (str, optional): device type specification (default: None).
|
||||
Applies to both model and batches.
|
||||
Returns:
|
||||
Engine: an evaluator engine with supervised inference function
|
||||
"""
|
||||
if device:
|
||||
model.to(device)
|
||||
|
||||
def _inference(engine, batch):
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
data, pids, camids = batch
|
||||
data = data.cuda()
|
||||
feat = model(data)
|
||||
return feat, pids, camids
|
||||
|
||||
engine = Engine(_inference)
|
||||
|
||||
for name, metric in metrics.items():
|
||||
metric.attach(engine, name)
|
||||
|
||||
return engine
|
||||
|
||||
|
||||
def inference(
|
||||
cfg,
|
||||
model,
|
||||
val_loader,
|
||||
num_query
|
||||
):
|
||||
device = cfg.MODEL.DEVICE
|
||||
|
||||
logger = logging.getLogger("reid_baseline.inference")
|
||||
logger.info("Enter inferencing")
|
||||
if cfg.TEST.RE_RANKING == 'no':
|
||||
print("Create evaluator")
|
||||
evaluator = create_supervised_evaluator(model, metrics={'r1_mAP': R1_mAP(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)},
|
||||
device=device)
|
||||
elif cfg.TEST.RE_RANKING == 'yes':
|
||||
print("Create evaluator for reranking")
|
||||
evaluator = create_supervised_evaluator(model, metrics={'r1_mAP': R1_mAP_reranking(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)},
|
||||
device=device)
|
||||
else:
|
||||
print("Unsupported re_ranking config. Only support for no or yes, but got {}.".format(cfg.TEST.RE_RANKING))
|
||||
|
||||
evaluator.run(val_loader)
|
||||
cmc, mAP = evaluator.state.metrics['r1_mAP']
|
||||
logger.info('Validation Results')
|
||||
logger.info("mAP: {:.1%}".format(mAP))
|
||||
for r in [1, 5, 10]:
|
||||
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1]))
|
|
@ -0,0 +1,262 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
from ignite.engine import Engine, Events
|
||||
from ignite.handlers import ModelCheckpoint, Timer
|
||||
from ignite.metrics import RunningAverage
|
||||
|
||||
from utils.reid_metric import R1_mAP
|
||||
|
||||
|
||||
def create_supervised_trainer(model, optimizer, loss_fn,
|
||||
device=None):
|
||||
"""
|
||||
Factory function for creating a trainer for supervised models
|
||||
|
||||
Args:
|
||||
model (`torch.nn.Module`): the model to train
|
||||
optimizer (`torch.optim.Optimizer`): the optimizer to use
|
||||
loss_fn (torch.nn loss function): the loss function to use
|
||||
device (str, optional): device type specification (default: None).
|
||||
Applies to both model and batches.
|
||||
|
||||
Returns:
|
||||
Engine: a trainer engine with supervised update function
|
||||
"""
|
||||
if device:
|
||||
model.to(device)
|
||||
|
||||
def _update(engine, batch):
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
img, target = batch
|
||||
img = img.cuda()
|
||||
target = target.cuda()
|
||||
score, feat = model(img)
|
||||
loss = loss_fn(score, feat, target)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
# compute acc
|
||||
acc = (score.max(1)[1] == target).float().mean()
|
||||
return loss.item(), acc.item()
|
||||
|
||||
return Engine(_update)
|
||||
|
||||
|
||||
def create_supervised_trainer_with_center(model, center_criterion, optimizer, optimizer_center, loss_fn, cetner_loss_weight,
|
||||
device=None):
|
||||
"""
|
||||
Factory function for creating a trainer for supervised models
|
||||
|
||||
Args:
|
||||
model (`torch.nn.Module`): the model to train
|
||||
optimizer (`torch.optim.Optimizer`): the optimizer to use
|
||||
loss_fn (torch.nn loss function): the loss function to use
|
||||
device (str, optional): device type specification (default: None).
|
||||
Applies to both model and batches.
|
||||
|
||||
Returns:
|
||||
Engine: a trainer engine with supervised update function
|
||||
"""
|
||||
if device:
|
||||
model.to(device)
|
||||
|
||||
def _update(engine, batch):
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
optimizer_center.zero_grad()
|
||||
img, target = batch
|
||||
img = img.cuda()
|
||||
target = target.cuda()
|
||||
score, feat = model(img)
|
||||
loss = loss_fn(score, feat, target)
|
||||
# print("Total loss is {}, center loss is {}".format(loss, center_criterion(feat, target)))
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
for param in center_criterion.parameters():
|
||||
param.grad.data *= (1. / cetner_loss_weight)
|
||||
optimizer_center.step()
|
||||
|
||||
# compute acc
|
||||
acc = (score.max(1)[1] == target).float().mean()
|
||||
return loss.item(), acc.item()
|
||||
|
||||
return Engine(_update)
|
||||
|
||||
|
||||
def create_supervised_evaluator(model, metrics,
|
||||
device=None):
|
||||
"""
|
||||
Factory function for creating an evaluator for supervised models
|
||||
|
||||
Args:
|
||||
model (`torch.nn.Module`): the model to train
|
||||
metrics (dict of str - :class:`ignite.metrics.Metric`): a map of metric names to Metrics
|
||||
device (str, optional): device type specification (default: None).
|
||||
Applies to both model and batches.
|
||||
Returns:
|
||||
Engine: an evaluator engine with supervised inference function
|
||||
"""
|
||||
if device:
|
||||
model.to(device)
|
||||
|
||||
def _inference(engine, batch):
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
data, pids, camids = batch
|
||||
data = data.cuda()
|
||||
feat = model(data)
|
||||
return feat, pids, camids
|
||||
|
||||
engine = Engine(_inference)
|
||||
|
||||
for name, metric in metrics.items():
|
||||
metric.attach(engine, name)
|
||||
|
||||
return engine
|
||||
|
||||
|
||||
def do_train(
|
||||
cfg,
|
||||
model,
|
||||
train_loader,
|
||||
val_loader,
|
||||
optimizer,
|
||||
scheduler,
|
||||
loss_fn,
|
||||
num_query
|
||||
):
|
||||
log_period = cfg.SOLVER.LOG_PERIOD
|
||||
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
|
||||
eval_period = cfg.SOLVER.EVAL_PERIOD
|
||||
output_dir = cfg.OUTPUT_DIR
|
||||
device = cfg.MODEL.DEVICE
|
||||
epochs = cfg.SOLVER.MAX_EPOCHS
|
||||
|
||||
logger = logging.getLogger("reid_baseline.train")
|
||||
logger.info("Start training")
|
||||
trainer = create_supervised_trainer(model, optimizer, loss_fn, device=device)
|
||||
evaluator = create_supervised_evaluator(model, metrics={'r1_mAP': R1_mAP(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)}, device=device)
|
||||
checkpointer = ModelCheckpoint(output_dir, cfg.MODEL.NAME, checkpoint_period, n_saved=10, require_empty=False)
|
||||
timer = Timer(average=True)
|
||||
|
||||
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpointer, {'model': model.state_dict(),
|
||||
'optimizer': optimizer.state_dict()})
|
||||
timer.attach(trainer, start=Events.EPOCH_STARTED, resume=Events.ITERATION_STARTED,
|
||||
pause=Events.ITERATION_COMPLETED, step=Events.ITERATION_COMPLETED)
|
||||
|
||||
# average metric to attach on trainer
|
||||
RunningAverage(output_transform=lambda x: x[0]).attach(trainer, 'avg_loss')
|
||||
RunningAverage(output_transform=lambda x: x[1]).attach(trainer, 'avg_acc')
|
||||
|
||||
@trainer.on(Events.EPOCH_STARTED)
|
||||
def adjust_learning_rate(engine):
|
||||
scheduler.step()
|
||||
|
||||
@trainer.on(Events.ITERATION_COMPLETED)
|
||||
def log_training_loss(engine):
|
||||
iter = (engine.state.iteration - 1) % len(train_loader) + 1
|
||||
|
||||
if iter % log_period == 0:
|
||||
logger.info("Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
|
||||
.format(engine.state.epoch, iter, len(train_loader),
|
||||
engine.state.metrics['avg_loss'], engine.state.metrics['avg_acc'],
|
||||
scheduler.get_lr()[0]))
|
||||
|
||||
# adding handlers using `trainer.on` decorator API
|
||||
@trainer.on(Events.EPOCH_COMPLETED)
|
||||
def print_times(engine):
|
||||
logger.info('Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]'
|
||||
.format(engine.state.epoch, timer.value() * timer.step_count,
|
||||
train_loader.batch_size / timer.value()))
|
||||
logger.info('-' * 10)
|
||||
timer.reset()
|
||||
|
||||
@trainer.on(Events.EPOCH_COMPLETED)
|
||||
def log_validation_results(engine):
|
||||
if engine.state.epoch % eval_period == 0:
|
||||
evaluator.run(val_loader)
|
||||
cmc, mAP = evaluator.state.metrics['r1_mAP']
|
||||
logger.info("Validation Results - Epoch: {}".format(engine.state.epoch))
|
||||
logger.info("mAP: {:.1%}".format(mAP))
|
||||
for r in [1, 5, 10]:
|
||||
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1]))
|
||||
|
||||
trainer.run(train_loader, max_epochs=epochs)
|
||||
|
||||
|
||||
def do_train_with_center(
|
||||
cfg,
|
||||
model,
|
||||
center_criterion,
|
||||
train_loader,
|
||||
val_loader,
|
||||
optimizer,
|
||||
optimizer_center,
|
||||
scheduler,
|
||||
loss_fn,
|
||||
num_query
|
||||
):
|
||||
log_period = cfg.SOLVER.LOG_PERIOD
|
||||
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
|
||||
eval_period = cfg.SOLVER.EVAL_PERIOD
|
||||
output_dir = cfg.OUTPUT_DIR
|
||||
device = cfg.MODEL.DEVICE
|
||||
epochs = cfg.SOLVER.MAX_EPOCHS
|
||||
|
||||
logger = logging.getLogger("reid_baseline.train")
|
||||
logger.info("Start training")
|
||||
trainer = create_supervised_trainer_with_center(model, center_criterion, optimizer, optimizer_center, loss_fn, cfg.SOLVER.CENTER_LOSS_WEIGHT, device=device)
|
||||
evaluator = create_supervised_evaluator(model, metrics={'r1_mAP': R1_mAP(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)}, device=device)
|
||||
checkpointer = ModelCheckpoint(output_dir, cfg.MODEL.NAME, checkpoint_period, n_saved=10, require_empty=False)
|
||||
timer = Timer(average=True)
|
||||
|
||||
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpointer, {'model': model.state_dict(),
|
||||
'optimizer': optimizer.state_dict()})
|
||||
timer.attach(trainer, start=Events.EPOCH_STARTED, resume=Events.ITERATION_STARTED,
|
||||
pause=Events.ITERATION_COMPLETED, step=Events.ITERATION_COMPLETED)
|
||||
|
||||
# average metric to attach on trainer
|
||||
RunningAverage(output_transform=lambda x: x[0]).attach(trainer, 'avg_loss')
|
||||
RunningAverage(output_transform=lambda x: x[1]).attach(trainer, 'avg_acc')
|
||||
|
||||
@trainer.on(Events.EPOCH_STARTED)
|
||||
def adjust_learning_rate(engine):
|
||||
scheduler.step()
|
||||
|
||||
@trainer.on(Events.ITERATION_COMPLETED)
|
||||
def log_training_loss(engine):
|
||||
iter = (engine.state.iteration - 1) % len(train_loader) + 1
|
||||
|
||||
if iter % log_period == 0:
|
||||
logger.info("Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
|
||||
.format(engine.state.epoch, iter, len(train_loader),
|
||||
engine.state.metrics['avg_loss'], engine.state.metrics['avg_acc'],
|
||||
scheduler.get_lr()[0]))
|
||||
|
||||
# adding handlers using `trainer.on` decorator API
|
||||
@trainer.on(Events.EPOCH_COMPLETED)
|
||||
def print_times(engine):
|
||||
logger.info('Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]'
|
||||
.format(engine.state.epoch, timer.value() * timer.step_count,
|
||||
train_loader.batch_size / timer.value()))
|
||||
logger.info('-' * 10)
|
||||
timer.reset()
|
||||
|
||||
@trainer.on(Events.EPOCH_COMPLETED)
|
||||
def log_validation_results(engine):
|
||||
if engine.state.epoch % eval_period == 0:
|
||||
evaluator.run(val_loader)
|
||||
cmc, mAP = evaluator.state.metrics['r1_mAP']
|
||||
logger.info("Validation Results - Epoch: {}".format(engine.state.epoch))
|
||||
logger.info("mAP: {:.1%}".format(mAP))
|
||||
for r in [1, 5, 10]:
|
||||
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(r, cmc[r - 1]))
|
||||
|
||||
trainer.run(train_loader, max_epochs=epochs)
|
|
@ -0,0 +1,142 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .triplet_loss import TripletLoss, CrossEntropyLabelSmooth
|
||||
from .cluster_loss import ClusterLoss
|
||||
from .center_loss import CenterLoss
|
||||
from .range_loss import RangeLoss
|
||||
|
||||
|
||||
def make_loss(cfg, num_classes): # modified by gu
|
||||
sampler = cfg.DATALOADER.SAMPLER
|
||||
if cfg.MODEL.METRIC_LOSS_TYPE == 'triplet':
|
||||
triplet = TripletLoss(cfg.SOLVER.MARGIN) # triplet loss
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'cluster':
|
||||
cluster = ClusterLoss(cfg.SOLVER.CLUSTER_MARGIN, True, True, cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE, cfg.DATALOADER.NUM_INSTANCE)
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_cluster':
|
||||
triplet = TripletLoss(cfg.SOLVER.MARGIN) # triplet loss
|
||||
cluster = ClusterLoss(cfg.SOLVER.CLUSTER_MARGIN, True, True, cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE, cfg.DATALOADER.NUM_INSTANCE)
|
||||
else:
|
||||
print('expected METRIC_LOSS_TYPE should be triplet, cluster, triplet_cluster'
|
||||
'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
|
||||
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
xent = CrossEntropyLabelSmooth(num_classes=num_classes) # new add by luo
|
||||
print("label smooth on, numclasses:", num_classes)
|
||||
|
||||
if sampler == 'softmax':
|
||||
def loss_func(score, feat, target):
|
||||
return F.cross_entropy(score, target)
|
||||
elif cfg.DATALOADER.SAMPLER == 'triplet':
|
||||
def loss_func(score, feat, target):
|
||||
return triplet(feat, target)[0]
|
||||
elif cfg.DATALOADER.SAMPLER == 'softmax_triplet':
|
||||
def loss_func(score, feat, target):
|
||||
if cfg.MODEL.METRIC_LOSS_TYPE == 'triplet':
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
return xent(score, target) + triplet(feat, target)[0] # new add by luo, open label smooth
|
||||
else:
|
||||
return F.cross_entropy(score, target) + triplet(feat, target)[0] # new add by luo, no label smooth
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'cluster':
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
return xent(score, target) + cluster(feat, target)[0] # new add by luo, open label smooth
|
||||
else:
|
||||
return F.cross_entropy(score, target) + cluster(feat, target)[0] # new add by luo, no label smooth
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_cluster':
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
return xent(score, target) + triplet(feat, target)[0] + cluster(feat, target)[0] # new add by luo, open label smooth
|
||||
else:
|
||||
return F.cross_entropy(score, target) + triplet(feat, target)[0] + cluster(feat, target)[0] # new add by luo, no label smooth
|
||||
else:
|
||||
print('expected METRIC_LOSS_TYPE should be triplet, cluster, triplet_cluster,'
|
||||
'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
|
||||
else:
|
||||
print('expected sampler should be softmax, triplet or softmax_triplet, '
|
||||
'but got {}'.format(cfg.DATALOADER.SAMPLER))
|
||||
return loss_func
|
||||
|
||||
|
||||
def make_loss_with_center(cfg, num_classes): # modified by gu
|
||||
if cfg.MODEL.METRIC_LOSS_TYPE == 'center':
|
||||
center_criterion = CenterLoss(num_classes=num_classes, feat_dim=2048, use_gpu=True) # center loss
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'range_center':
|
||||
center_criterion = CenterLoss(num_classes=num_classes, feat_dim=2048, use_gpu=True) # center_range loss
|
||||
range_criterion = RangeLoss(k=cfg.SOLVER.RANGE_K, margin=cfg.SOLVER.RANGE_MARGIN, alpha=cfg.SOLVER.RANGE_ALPHA,
|
||||
beta=cfg.SOLVER.RANGE_BETA, ordered=True, use_gpu=True,
|
||||
ids_per_batch=cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE,
|
||||
imgs_per_id=cfg.DATALOADER.NUM_INSTANCE)
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_center':
|
||||
triplet = TripletLoss(cfg.SOLVER.MARGIN) # triplet loss
|
||||
center_criterion = CenterLoss(num_classes=num_classes, feat_dim=2048, use_gpu=True) # center loss
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_range_center':
|
||||
triplet = TripletLoss(cfg.SOLVER.MARGIN) # triplet loss
|
||||
center_criterion = CenterLoss(num_classes=num_classes, feat_dim=2048, use_gpu=True) # center_range loss
|
||||
range_criterion = RangeLoss(k=cfg.SOLVER.RANGE_K, margin=cfg.SOLVER.RANGE_MARGIN, alpha=cfg.SOLVER.RANGE_ALPHA,
|
||||
beta=cfg.SOLVER.RANGE_BETA, ordered=True, use_gpu=True,
|
||||
ids_per_batch=cfg.SOLVER.IMS_PER_BATCH // cfg.DATALOADER.NUM_INSTANCE,
|
||||
imgs_per_id=cfg.DATALOADER.NUM_INSTANCE)
|
||||
else:
|
||||
print('expected METRIC_LOSS_TYPE with center should be center, '
|
||||
'range_center,triplet_center, triplet_range_center '
|
||||
'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
|
||||
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
xent = CrossEntropyLabelSmooth(num_classes=num_classes) # new add by luo
|
||||
print("label smooth on, numclasses:", num_classes)
|
||||
|
||||
def loss_func(score, feat, target):
|
||||
if cfg.MODEL.METRIC_LOSS_TYPE == 'center':
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
return xent(score, target) + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) # new add by luo, open label smooth
|
||||
else:
|
||||
return F.cross_entropy(score, target) + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) # new add by luo, no label smooth
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'range_center':
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
return xent(score, target) + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
|
||||
cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0] # new add by luo, open label smooth
|
||||
else:
|
||||
return F.cross_entropy(score, target) + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
|
||||
cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0] # new add by luo, no label smooth
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_center':
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
return xent(score, target) + \
|
||||
triplet(feat, target)[0] + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) # new add by luo, open label smooth
|
||||
else:
|
||||
return F.cross_entropy(score, target) + \
|
||||
triplet(feat, target)[0] + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) # new add by luo, no label smooth
|
||||
|
||||
elif cfg.MODEL.METRIC_LOSS_TYPE == 'triplet_range_center':
|
||||
if cfg.MODEL.IF_LABELSMOOTH == 'on':
|
||||
return xent(score, target) + \
|
||||
triplet(feat, target)[0] + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
|
||||
cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0] # new add by luo, open label smooth
|
||||
else:
|
||||
return F.cross_entropy(score, target) + \
|
||||
triplet(feat, target)[0] + \
|
||||
cfg.SOLVER.CENTER_LOSS_WEIGHT * center_criterion(feat, target) + \
|
||||
cfg.SOLVER.RANGE_LOSS_WEIGHT * range_criterion(feat, target)[0] # new add by luo, no label smooth
|
||||
|
||||
else:
|
||||
print('expected METRIC_LOSS_TYPE with center should be center,'
|
||||
' range_center, triplet_center, triplet_range_center '
|
||||
'but got {}'.format(cfg.MODEL.METRIC_LOSS_TYPE))
|
||||
return loss_func, center_criterion
|
|
@ -0,0 +1,67 @@
|
|||
from __future__ import absolute_import
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class CenterLoss(nn.Module):
|
||||
"""Center loss.
|
||||
|
||||
Reference:
|
||||
Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016.
|
||||
|
||||
Args:
|
||||
num_classes (int): number of classes.
|
||||
feat_dim (int): feature dimension.
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes=751, feat_dim=2048, use_gpu=True):
|
||||
super(CenterLoss, self).__init__()
|
||||
self.num_classes = num_classes
|
||||
self.feat_dim = feat_dim
|
||||
self.use_gpu = use_gpu
|
||||
|
||||
if self.use_gpu:
|
||||
self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda())
|
||||
else:
|
||||
self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim))
|
||||
|
||||
def forward(self, x, labels):
|
||||
"""
|
||||
Args:
|
||||
x: feature matrix with shape (batch_size, feat_dim).
|
||||
labels: ground truth labels with shape (num_classes).
|
||||
"""
|
||||
assert x.size(0) == labels.size(0), "features.size(0) is not equal to labels.size(0)"
|
||||
|
||||
batch_size = x.size(0)
|
||||
distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \
|
||||
torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t()
|
||||
distmat.addmm_(1, -2, x, self.centers.t())
|
||||
|
||||
classes = torch.arange(self.num_classes).long()
|
||||
if self.use_gpu: classes = classes.cuda()
|
||||
labels = labels.unsqueeze(1).expand(batch_size, self.num_classes)
|
||||
mask = labels.eq(classes.expand(batch_size, self.num_classes))
|
||||
|
||||
dist = []
|
||||
for i in range(batch_size):
|
||||
value = distmat[i][mask[i]]
|
||||
value = value.clamp(min=1e-12, max=1e+12) # for numerical stability
|
||||
dist.append(value)
|
||||
dist = torch.cat(dist)
|
||||
loss = dist.mean()
|
||||
return loss
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
use_gpu = False
|
||||
center_loss = CenterLoss(use_gpu=use_gpu)
|
||||
features = torch.rand(16, 2048)
|
||||
targets = torch.Tensor([0, 1, 2, 3, 2, 3, 1, 4, 5, 3, 2, 1, 0, 0, 5, 4]).long()
|
||||
if use_gpu:
|
||||
features = torch.rand(16, 2048).cuda()
|
||||
targets = torch.Tensor([0, 1, 2, 3, 2, 3, 1, 4, 5, 3, 2, 1, 0, 0, 5, 4]).cuda()
|
||||
|
||||
loss = center_loss(features, targets)
|
||||
print(loss)
|
|
@ -0,0 +1,269 @@
|
|||
from __future__ import absolute_import
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ClusterLoss(nn.Module):
|
||||
def __init__(self, margin=10, use_gpu=True, ordered=True, ids_per_batch=16, imgs_per_id=4):
|
||||
super(ClusterLoss, self).__init__()
|
||||
self.use_gpu = use_gpu
|
||||
self.margin = margin
|
||||
self.ordered = ordered
|
||||
self.ids_per_batch = ids_per_batch
|
||||
self.imgs_per_id = imgs_per_id
|
||||
|
||||
def _euclidean_dist(self, x, y):
|
||||
"""
|
||||
Args:
|
||||
x: pytorch Variable, with shape [m, d]
|
||||
y: pytorch Variable, with shape [n, d]
|
||||
Returns:
|
||||
dist: pytorch Variable, with shape [m, n]
|
||||
"""
|
||||
m, n = x.size(0), y.size(0)
|
||||
xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
|
||||
yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
|
||||
dist = xx + yy
|
||||
dist.addmm_(1, -2, x, y.t())
|
||||
dist = dist.clamp(min=1e-12).sqrt() # for numerical stability
|
||||
return dist
|
||||
|
||||
def _cluster_loss(self, features, targets, ordered=True, ids_per_batch=16, imgs_per_id=4):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
cluster_loss
|
||||
"""
|
||||
if self.use_gpu:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
|
||||
inter_min_distance = torch.zeros(unique_labels.size(0))
|
||||
intra_max_distance = torch.zeros(unique_labels.size(0))
|
||||
center_features = torch.zeros(unique_labels.size(0), features.size(1))
|
||||
|
||||
if self.use_gpu:
|
||||
inter_min_distance = inter_min_distance.cuda()
|
||||
intra_max_distance = intra_max_distance.cuda()
|
||||
center_features = center_features.cuda()
|
||||
|
||||
index = torch.range(0, unique_labels.size(0) - 1)
|
||||
for i in range(unique_labels.size(0)):
|
||||
label = unique_labels[i]
|
||||
same_class_features = features[targets == label]
|
||||
center_features[i] = same_class_features.mean(dim=0)
|
||||
intra_class_distance = self._euclidean_dist(center_features[index==i], same_class_features)
|
||||
# print('intra_class_distance', intra_class_distance)
|
||||
intra_max_distance[i] = intra_class_distance.max()
|
||||
# print('intra_max_distance:', intra_max_distance)
|
||||
|
||||
for i in range(unique_labels.size(0)):
|
||||
inter_class_distance = self._euclidean_dist(center_features[index==i], center_features[index != i])
|
||||
# print('inter_class_distance', inter_class_distance)
|
||||
inter_min_distance[i] = inter_class_distance.min()
|
||||
# print('inter_min_distance:', inter_min_distance)
|
||||
cluster_loss = torch.mean(torch.relu(intra_max_distance - inter_min_distance + self.margin))
|
||||
return cluster_loss, intra_max_distance, inter_min_distance
|
||||
|
||||
def forward(self, features, targets):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
cluster_loss
|
||||
"""
|
||||
assert features.size(0) == targets.size(0), "features.size(0) is not equal to targets.size(0)"
|
||||
cluster_loss, cluster_dist_ap, cluster_dist_an = self._cluster_loss(features, targets, self.ordered, self.ids_per_batch, self.imgs_per_id)
|
||||
return cluster_loss, cluster_dist_ap, cluster_dist_an
|
||||
|
||||
|
||||
class ClusterLoss_local(nn.Module):
|
||||
def __init__(self, margin=10, use_gpu=True, ordered=True, ids_per_batch=32, imgs_per_id=4):
|
||||
super(ClusterLoss_local, self).__init__()
|
||||
self.use_gpu = use_gpu
|
||||
self.margin = margin
|
||||
self.ordered = ordered
|
||||
self.ids_per_batch = ids_per_batch
|
||||
self.imgs_per_id = imgs_per_id
|
||||
|
||||
def _euclidean_dist(self, x, y):
|
||||
"""
|
||||
Args:
|
||||
x: pytorch Variable, with shape [m, d]
|
||||
y: pytorch Variable, with shape [n, d]
|
||||
Returns:
|
||||
dist: pytorch Variable, with shape [m, n]
|
||||
"""
|
||||
m, n = x.size(0), y.size(0)
|
||||
xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
|
||||
yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
|
||||
dist = xx + yy
|
||||
dist.addmm_(1, -2, x, y.t())
|
||||
dist = dist.clamp(min=1e-12).sqrt() # for numerical stability
|
||||
return dist
|
||||
|
||||
def _shortest_dist(self, dist_mat):
|
||||
"""Parallel version.
|
||||
Args:
|
||||
dist_mat: pytorch Variable, available shape:
|
||||
1) [m, n]
|
||||
2) [m, n, N], N is batch size
|
||||
3) [m, n, *], * can be arbitrary additional dimensions
|
||||
Returns:
|
||||
dist: three cases corresponding to `dist_mat`:
|
||||
1) scalar
|
||||
2) pytorch Variable, with shape [N]
|
||||
3) pytorch Variable, with shape [*]
|
||||
"""
|
||||
m, n = dist_mat.size()[:2]
|
||||
# Just offering some reference for accessing intermediate distance.
|
||||
dist = [[0 for _ in range(n)] for _ in range(m)]
|
||||
for i in range(m):
|
||||
for j in range(n):
|
||||
if (i == 0) and (j == 0):
|
||||
dist[i][j] = dist_mat[i, j]
|
||||
elif (i == 0) and (j > 0):
|
||||
dist[i][j] = dist[i][j - 1] + dist_mat[i, j]
|
||||
elif (i > 0) and (j == 0):
|
||||
dist[i][j] = dist[i - 1][j] + dist_mat[i, j]
|
||||
else:
|
||||
dist[i][j] = torch.min(dist[i - 1][j], dist[i][j - 1]) + dist_mat[i, j]
|
||||
dist = dist[-1][-1]
|
||||
return dist
|
||||
|
||||
def _local_dist(self, x, y):
|
||||
"""
|
||||
Args:
|
||||
x: pytorch Variable, with shape [M, m, d]
|
||||
y: pytorch Variable, with shape [N, n, d]
|
||||
Returns:
|
||||
dist: pytorch Variable, with shape [M, N]
|
||||
"""
|
||||
M, m, d = x.size()
|
||||
N, n, d = y.size()
|
||||
x = x.contiguous().view(M * m, d)
|
||||
y = y.contiguous().view(N * n, d)
|
||||
# shape [M * m, N * n]
|
||||
dist_mat = self._euclidean_dist(x, y)
|
||||
dist_mat = (torch.exp(dist_mat) - 1.) / (torch.exp(dist_mat) + 1.)
|
||||
# shape [M * m, N * n] -> [M, m, N, n] -> [m, n, M, N]
|
||||
dist_mat = dist_mat.contiguous().view(M, m, N, n).permute(1, 3, 0, 2)
|
||||
# shape [M, N]
|
||||
dist_mat = self._shortest_dist(dist_mat)
|
||||
return dist_mat
|
||||
|
||||
def _cluster_loss(self, features, targets,ordered=True, ids_per_batch=32, imgs_per_id=4):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, H, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
cluster_loss
|
||||
"""
|
||||
if self.use_gpu:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
|
||||
inter_min_distance = torch.zeros(unique_labels.size(0))
|
||||
intra_max_distance = torch.zeros(unique_labels.size(0))
|
||||
center_features = torch.zeros(unique_labels.size(0), features.size(1), features.size(2))
|
||||
|
||||
if self.use_gpu:
|
||||
inter_min_distance = inter_min_distance.cuda()
|
||||
intra_max_distance = intra_max_distance.cuda()
|
||||
center_features = center_features.cuda()
|
||||
|
||||
index = torch.range(0, unique_labels.size(0) - 1)
|
||||
for i in range(unique_labels.size(0)):
|
||||
label = unique_labels[i]
|
||||
same_class_features = features[targets == label]
|
||||
center_features[i] = same_class_features.mean(dim=0)
|
||||
intra_class_distance = self._local_dist(center_features[index==i], same_class_features)
|
||||
# print('intra_class_distance', intra_class_distance)
|
||||
intra_max_distance[i] = intra_class_distance.max()
|
||||
# print('intra_max_distance:', intra_max_distance)
|
||||
|
||||
for i in range(unique_labels.size(0)):
|
||||
inter_class_distance = self._local_dist(center_features[index==i], center_features[index != i])
|
||||
# print('inter_class_distance', inter_class_distance)
|
||||
inter_min_distance[i] = inter_class_distance.min()
|
||||
# print('inter_min_distance:', inter_min_distance)
|
||||
|
||||
cluster_loss = torch.mean(torch.relu(intra_max_distance - inter_min_distance + self.margin))
|
||||
return cluster_loss, intra_max_distance, inter_min_distance
|
||||
|
||||
def forward(self, features, targets):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, H, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
cluster_loss
|
||||
"""
|
||||
assert features.size(0) == targets.size(0), "features.size(0) is not equal to targets.size(0)"
|
||||
cluster_loss, cluster_dist_ap, cluster_dist_an = self._cluster_loss(features, targets, self.ordered, self.ids_per_batch, self.imgs_per_id)
|
||||
return cluster_loss, cluster_dist_ap, cluster_dist_an
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
use_gpu = True
|
||||
cluster_loss = ClusterLoss(use_gpu=use_gpu, ids_per_batch=4, imgs_per_id=4)
|
||||
features = torch.rand(16, 2048)
|
||||
targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3])
|
||||
if use_gpu:
|
||||
features = torch.rand(16, 2048).cuda()
|
||||
targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]).cuda()
|
||||
loss = cluster_loss(features, targets)
|
||||
print(loss)
|
||||
|
||||
cluster_loss_local = ClusterLoss_local(use_gpu=use_gpu, ids_per_batch=4, imgs_per_id=4)
|
||||
features = torch.rand(16, 8, 2048)
|
||||
targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3])
|
||||
if use_gpu:
|
||||
features = torch.rand(16, 8, 2048).cuda()
|
||||
targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]).cuda()
|
||||
loss = cluster_loss_local(features, targets)
|
||||
print(loss)
|
|
@ -0,0 +1,232 @@
|
|||
from __future__ import absolute_import
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class RangeLoss(nn.Module):
|
||||
"""
|
||||
Range_loss = alpha * intra_class_loss + beta * inter_class_loss
|
||||
intra_class_loss is the harmonic mean value of the top_k largest distances beturn intra_class_pairs
|
||||
inter_class_loss is the shortest distance between different class centers
|
||||
"""
|
||||
def __init__(self, k=2, margin=0.1, alpha=0.5, beta=0.5, use_gpu=True, ordered=True, ids_per_batch=32, imgs_per_id=4):
|
||||
super(RangeLoss, self).__init__()
|
||||
self.use_gpu = use_gpu
|
||||
self.margin = margin
|
||||
self.k = k
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.ordered = ordered
|
||||
self.ids_per_batch = ids_per_batch
|
||||
self.imgs_per_id = imgs_per_id
|
||||
|
||||
def _pairwise_distance(self, features):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
Return:
|
||||
pairwise distance matrix with shape(batch_size, batch_size)
|
||||
"""
|
||||
n = features.size(0)
|
||||
dist = torch.pow(features, 2).sum(dim=1, keepdim=True).expand(n, n)
|
||||
dist = dist + dist.t()
|
||||
dist.addmm_(1, -2, features, features.t())
|
||||
dist = dist.clamp(min=1e-12).sqrt() # for numerical stability
|
||||
return dist
|
||||
|
||||
def _compute_top_k(self, features):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
Return:
|
||||
top_k largest distances
|
||||
"""
|
||||
# reading the codes below can help understand better
|
||||
'''
|
||||
dist_array_2 = self._pairwise_distance(features)
|
||||
n = features.size(0)
|
||||
mask = torch.zeros(n, n)
|
||||
if self.use_gpu: mask=mask.cuda()
|
||||
for i in range(0, n):
|
||||
for j in range(i+1, n):
|
||||
mask[i, j] += 1
|
||||
dist_array_2 = dist_array_2 * mask
|
||||
dist_array_2 = dist_array_2.view(1, -1)
|
||||
dist_array_2 = dist_array_2[torch.gt(dist_array_2, 0)]
|
||||
top_k_2 = dist_array_2.sort()[0][-self.k:]
|
||||
print(top_k_2)
|
||||
'''
|
||||
dist_array = self._pairwise_distance(features)
|
||||
dist_array = dist_array.view(1, -1)
|
||||
top_k = dist_array.sort()[0][0, -self.k * 2::2] # Because there are 2 same value of same feature pair in the dist_array
|
||||
# print('top k intra class dist:', top_k)
|
||||
return top_k
|
||||
|
||||
def _compute_min_dist(self, center_features):
|
||||
"""
|
||||
Args:
|
||||
center_features: center matrix (before softmax) with shape (center_number, center_dim)
|
||||
Return:
|
||||
minimum center distance
|
||||
"""
|
||||
'''
|
||||
# reading codes below can help understand better
|
||||
dist_array = self._pairwise_distance(center_features)
|
||||
n = center_features.size(0)
|
||||
mask = torch.zeros(n, n)
|
||||
if self.use_gpu: mask=mask.cuda()
|
||||
for i in range(0, n):
|
||||
for j in range(i + 1, n):
|
||||
mask[i, j] += 1
|
||||
dist_array *= mask
|
||||
dist_array = dist_array.view(1, -1)
|
||||
dist_array = dist_array[torch.gt(dist_array, 0)]
|
||||
min_inter_class_dist = dist_array.min()
|
||||
print(min_inter_class_dist)
|
||||
'''
|
||||
n = center_features.size(0)
|
||||
dist_array2 = self._pairwise_distance(center_features)
|
||||
min_inter_class_dist2 = dist_array2.view(1, -1).sort()[0][0][n] # exclude self compare, the first one is the min_inter_class_dist
|
||||
return min_inter_class_dist2
|
||||
|
||||
def _calculate_centers(self, features, targets, ordered, ids_per_batch, imgs_per_id):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
center_features: center matrix (before softmax) with shape (center_number, center_dim)
|
||||
"""
|
||||
if self.use_gpu:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
|
||||
center_features = torch.zeros(unique_labels.size(0), features.size(1))
|
||||
if self.use_gpu:
|
||||
center_features = center_features.cuda()
|
||||
|
||||
for i in range(unique_labels.size(0)):
|
||||
label = unique_labels[i]
|
||||
same_class_features = features[targets == label]
|
||||
center_features[i] = same_class_features.mean(dim=0)
|
||||
return center_features
|
||||
|
||||
def _inter_class_loss(self, features, targets, ordered, ids_per_batch, imgs_per_id):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
margin: inter class ringe loss margin
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
inter_class_loss
|
||||
"""
|
||||
center_features = self._calculate_centers(features, targets, ordered, ids_per_batch, imgs_per_id)
|
||||
min_inter_class_center_distance = self._compute_min_dist(center_features)
|
||||
# print('min_inter_class_center_dist:', min_inter_class_center_distance)
|
||||
return torch.relu(self.margin - min_inter_class_center_distance)
|
||||
|
||||
def _intra_class_loss(self, features, targets, ordered, ids_per_batch, imgs_per_id):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
intra_class_loss
|
||||
"""
|
||||
if self.use_gpu:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
unique_labels = targets.cpu().unique().cuda()
|
||||
else:
|
||||
if ordered:
|
||||
if targets.size(0) == ids_per_batch * imgs_per_id:
|
||||
unique_labels = targets[0:targets.size(0):imgs_per_id]
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
else:
|
||||
unique_labels = targets.unique()
|
||||
|
||||
intra_distance = torch.zeros(unique_labels.size(0))
|
||||
if self.use_gpu:
|
||||
intra_distance = intra_distance.cuda()
|
||||
|
||||
for i in range(unique_labels.size(0)):
|
||||
label = unique_labels[i]
|
||||
same_class_distances = 1.0 / self._compute_top_k(features[targets == label])
|
||||
intra_distance[i] = self.k / torch.sum(same_class_distances)
|
||||
# print('intra_distace:', intra_distance)
|
||||
return torch.sum(intra_distance)
|
||||
|
||||
def _range_loss(self, features, targets, ordered, ids_per_batch, imgs_per_id):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
range_loss
|
||||
"""
|
||||
inter_class_loss = self._inter_class_loss(features, targets, ordered, ids_per_batch, imgs_per_id)
|
||||
intra_class_loss = self._intra_class_loss(features, targets, ordered, ids_per_batch, imgs_per_id)
|
||||
range_loss = self.alpha * intra_class_loss + self.beta * inter_class_loss
|
||||
return range_loss, intra_class_loss, inter_class_loss
|
||||
|
||||
def forward(self, features, targets):
|
||||
"""
|
||||
Args:
|
||||
features: prediction matrix (before softmax) with shape (batch_size, feature_dim)
|
||||
targets: ground truth labels with shape (batch_size)
|
||||
ordered: bool type. If the train data per batch are formed as p*k, where p is the num of ids per batch and k is the num of images per id.
|
||||
ids_per_batch: num of different ids per batch
|
||||
imgs_per_id: num of images per id
|
||||
Return:
|
||||
range_loss
|
||||
"""
|
||||
assert features.size(0) == targets.size(0), "features.size(0) is not equal to targets.size(0)"
|
||||
if self.use_gpu:
|
||||
features = features.cuda()
|
||||
targets = targets.cuda()
|
||||
|
||||
range_loss, intra_class_loss, inter_class_loss = self._range_loss(features, targets, self.ordered, self.ids_per_batch, self.imgs_per_id)
|
||||
return range_loss, intra_class_loss, inter_class_loss
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
use_gpu = False
|
||||
range_loss = RangeLoss(use_gpu=use_gpu, ids_per_batch=4, imgs_per_id=4)
|
||||
features = torch.rand(16, 2048)
|
||||
targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3])
|
||||
if use_gpu:
|
||||
features = torch.rand(16, 2048).cuda()
|
||||
targets = torch.Tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]).cuda()
|
||||
loss = range_loss(features, targets)
|
||||
print(loss)
|
|
@ -0,0 +1,147 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def normalize(x, axis=-1):
|
||||
"""Normalizing to unit length along the specified dimension.
|
||||
Args:
|
||||
x: pytorch Variable
|
||||
Returns:
|
||||
x: pytorch Variable, same shape as input
|
||||
"""
|
||||
x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
|
||||
return x
|
||||
|
||||
|
||||
def euclidean_dist(x, y):
|
||||
"""
|
||||
Args:
|
||||
x: pytorch Variable, with shape [m, d]
|
||||
y: pytorch Variable, with shape [n, d]
|
||||
Returns:
|
||||
dist: pytorch Variable, with shape [m, n]
|
||||
"""
|
||||
m, n = x.size(0), y.size(0)
|
||||
xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
|
||||
yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
|
||||
dist = xx + yy
|
||||
dist.addmm_(1, -2, x, y.t())
|
||||
dist = dist.clamp(min=1e-12).sqrt() # for numerical stability
|
||||
return dist
|
||||
|
||||
|
||||
def hard_example_mining(dist_mat, labels, return_inds=False):
|
||||
"""For each anchor, find the hardest positive and negative sample.
|
||||
Args:
|
||||
dist_mat: pytorch Variable, pair wise distance between samples, shape [N, N]
|
||||
labels: pytorch LongTensor, with shape [N]
|
||||
return_inds: whether to return the indices. Save time if `False`(?)
|
||||
Returns:
|
||||
dist_ap: pytorch Variable, distance(anchor, positive); shape [N]
|
||||
dist_an: pytorch Variable, distance(anchor, negative); shape [N]
|
||||
p_inds: pytorch LongTensor, with shape [N];
|
||||
indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1
|
||||
n_inds: pytorch LongTensor, with shape [N];
|
||||
indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1
|
||||
NOTE: Only consider the case in which all labels have same num of samples,
|
||||
thus we can cope with all anchors in parallel.
|
||||
"""
|
||||
|
||||
assert len(dist_mat.size()) == 2
|
||||
assert dist_mat.size(0) == dist_mat.size(1)
|
||||
N = dist_mat.size(0)
|
||||
|
||||
# shape [N, N]
|
||||
is_pos = labels.expand(N, N).eq(labels.expand(N, N).t())
|
||||
is_neg = labels.expand(N, N).ne(labels.expand(N, N).t())
|
||||
|
||||
# `dist_ap` means distance(anchor, positive)
|
||||
# both `dist_ap` and `relative_p_inds` with shape [N, 1]
|
||||
dist_ap, relative_p_inds = torch.max(
|
||||
dist_mat[is_pos].contiguous().view(N, -1), 1, keepdim=True)
|
||||
# `dist_an` means distance(anchor, negative)
|
||||
# both `dist_an` and `relative_n_inds` with shape [N, 1]
|
||||
dist_an, relative_n_inds = torch.min(
|
||||
dist_mat[is_neg].contiguous().view(N, -1), 1, keepdim=True)
|
||||
# shape [N]
|
||||
dist_ap = dist_ap.squeeze(1)
|
||||
dist_an = dist_an.squeeze(1)
|
||||
|
||||
if return_inds:
|
||||
# shape [N, N]
|
||||
ind = (labels.new().resize_as_(labels)
|
||||
.copy_(torch.arange(0, N).long())
|
||||
.unsqueeze(0).expand(N, N))
|
||||
# shape [N, 1]
|
||||
p_inds = torch.gather(
|
||||
ind[is_pos].contiguous().view(N, -1), 1, relative_p_inds.data)
|
||||
n_inds = torch.gather(
|
||||
ind[is_neg].contiguous().view(N, -1), 1, relative_n_inds.data)
|
||||
# shape [N]
|
||||
p_inds = p_inds.squeeze(1)
|
||||
n_inds = n_inds.squeeze(1)
|
||||
return dist_ap, dist_an, p_inds, n_inds
|
||||
|
||||
return dist_ap, dist_an
|
||||
|
||||
|
||||
class TripletLoss(object):
|
||||
"""Modified from Tong Xiao's open-reid (https://github.com/Cysu/open-reid).
|
||||
Related Triplet Loss theory can be found in paper 'In Defense of the Triplet
|
||||
Loss for Person Re-Identification'."""
|
||||
|
||||
def __init__(self, margin=None):
|
||||
self.margin = margin
|
||||
if margin is not None:
|
||||
self.ranking_loss = nn.MarginRankingLoss(margin=margin)
|
||||
else:
|
||||
self.ranking_loss = nn.SoftMarginLoss()
|
||||
|
||||
def __call__(self, global_feat, labels, normalize_feature=False):
|
||||
if normalize_feature:
|
||||
global_feat = normalize(global_feat, axis=-1)
|
||||
dist_mat = euclidean_dist(global_feat, global_feat)
|
||||
dist_ap, dist_an = hard_example_mining(
|
||||
dist_mat, labels)
|
||||
y = dist_an.new().resize_as_(dist_an).fill_(1)
|
||||
if self.margin is not None:
|
||||
loss = self.ranking_loss(dist_an, dist_ap, y)
|
||||
else:
|
||||
loss = self.ranking_loss(dist_an - dist_ap, y)
|
||||
return loss, dist_ap, dist_an
|
||||
|
||||
class CrossEntropyLabelSmooth(nn.Module):
|
||||
"""Cross entropy loss with label smoothing regularizer.
|
||||
|
||||
Reference:
|
||||
Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
|
||||
Equation: y = (1 - epsilon) * y + epsilon / K.
|
||||
|
||||
Args:
|
||||
num_classes (int): number of classes.
|
||||
epsilon (float): weight.
|
||||
"""
|
||||
def __init__(self, num_classes, epsilon=0.1, use_gpu=True):
|
||||
super(CrossEntropyLabelSmooth, self).__init__()
|
||||
self.num_classes = num_classes
|
||||
self.epsilon = epsilon
|
||||
self.use_gpu = use_gpu
|
||||
self.logsoftmax = nn.LogSoftmax(dim=1)
|
||||
|
||||
def forward(self, inputs, targets):
|
||||
"""
|
||||
Args:
|
||||
inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
|
||||
targets: ground truth labels with shape (num_classes)
|
||||
"""
|
||||
log_probs = self.logsoftmax(inputs)
|
||||
targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1)
|
||||
if self.use_gpu: targets = targets.cuda()
|
||||
targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
|
||||
loss = (- targets * log_probs).mean(0).sum()
|
||||
return loss
|
|
@ -0,0 +1,13 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
from .baseline import Baseline
|
||||
|
||||
|
||||
def build_model(cfg, num_classes):
|
||||
if cfg.MODEL.NAME == 'resnet50':
|
||||
model = Baseline(num_classes, cfg.MODEL.LAST_STRIDE, cfg.MODEL.PRETRAIN_PATH, cfg.MODEL.NECK, cfg.TEST.NECK_FEAT)
|
||||
return model
|
|
@ -0,0 +1,6 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
|
@ -0,0 +1,109 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * 4)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Module):
|
||||
def __init__(self, last_stride=2, block=Bottleneck, layers=[3, 4, 6, 3]):
|
||||
self.inplanes = 64
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(
|
||||
block, 512, layers[3], stride=last_stride)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
return x
|
||||
|
||||
def load_param(self, model_path):
|
||||
param_dict = torch.load(model_path)
|
||||
for i in param_dict:
|
||||
if 'fc' in i:
|
||||
continue
|
||||
self.state_dict()[i].copy_(param_dict[i])
|
||||
|
||||
def random_init(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||
m.weight.data.normal_(0, math.sqrt(2. / n))
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
|
@ -0,0 +1,87 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .backbones.resnet import ResNet
|
||||
|
||||
|
||||
def weights_init_kaiming(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find('Linear') != -1:
|
||||
nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
|
||||
nn.init.constant_(m.bias, 0.0)
|
||||
elif classname.find('Conv') != -1:
|
||||
nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
|
||||
if m.bias is not None:
|
||||
nn.init.constant_(m.bias, 0.0)
|
||||
elif classname.find('BatchNorm') != -1:
|
||||
if m.affine:
|
||||
nn.init.constant_(m.weight, 1.0)
|
||||
nn.init.constant_(m.bias, 0.0)
|
||||
|
||||
|
||||
def weights_init_classifier(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find('Linear') != -1:
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
if m.bias:
|
||||
nn.init.constant_(m.bias, 0.0)
|
||||
|
||||
|
||||
class Baseline(nn.Module):
|
||||
in_planes = 2048
|
||||
|
||||
def __init__(self, num_classes, last_stride, model_path, neck, neck_feat):
|
||||
super(Baseline, self).__init__()
|
||||
self.base = ResNet(last_stride)
|
||||
self.base.load_param(model_path)
|
||||
self.gap = nn.AdaptiveAvgPool2d(1)
|
||||
# self.gap = nn.AdaptiveMaxPool2d(1)
|
||||
self.num_classes = num_classes
|
||||
self.neck = neck
|
||||
self.neck_feat = neck_feat
|
||||
|
||||
if self.neck == 'no':
|
||||
self.classifier = nn.Linear(self.in_planes, self.num_classes)
|
||||
# self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False) # new add by luo
|
||||
# self.classifier.apply(weights_init_classifier) # new add by luo
|
||||
elif self.neck == 'bnneck':
|
||||
self.bottleneck = nn.BatchNorm1d(self.in_planes)
|
||||
self.bottleneck.bias.requires_grad_(False) # no shift
|
||||
self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False)
|
||||
|
||||
self.bottleneck.apply(weights_init_kaiming)
|
||||
self.classifier.apply(weights_init_classifier)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
global_feat = self.gap(self.base(x)) # (b, 2048, 1, 1)
|
||||
global_feat = global_feat.view(global_feat.shape[0], -1) # flatten to (bs, 2048)
|
||||
|
||||
if self.neck == 'no':
|
||||
feat = global_feat
|
||||
elif self.neck == 'bnneck':
|
||||
feat = self.bottleneck(global_feat) # normalize for angular softmax
|
||||
|
||||
if self.training:
|
||||
cls_score = self.classifier(feat)
|
||||
return cls_score, global_feat # global feature for triplet loss
|
||||
else:
|
||||
if self.neck_feat == 'after':
|
||||
# print("Test with feature after BN")
|
||||
return feat
|
||||
else:
|
||||
# print("Test with feature before BN")
|
||||
return global_feat
|
||||
|
||||
def load_param(self, trained_path):
|
||||
param_dict = torch.load(trained_path)
|
||||
for i in param_dict:
|
||||
if 'classifier' in i:
|
||||
continue
|
||||
self.state_dict()[i].copy_(param_dict[i])
|
|
@ -0,0 +1,8 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
from .build import make_optimizer, make_optimizer_with_center
|
||||
from .lr_scheduler import WarmupMultiStepLR
|
|
@ -0,0 +1,44 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def make_optimizer(cfg, model):
|
||||
params = []
|
||||
for key, value in model.named_parameters():
|
||||
if not value.requires_grad:
|
||||
continue
|
||||
lr = cfg.SOLVER.BASE_LR
|
||||
weight_decay = cfg.SOLVER.WEIGHT_DECAY
|
||||
if "bias" in key:
|
||||
lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BIAS_LR_FACTOR
|
||||
weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS
|
||||
params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
|
||||
if cfg.SOLVER.OPTIMIZER_NAME == 'SGD':
|
||||
optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params, momentum=cfg.SOLVER.MOMENTUM)
|
||||
else:
|
||||
optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params)
|
||||
return optimizer
|
||||
|
||||
|
||||
def make_optimizer_with_center(cfg, model, center_criterion):
|
||||
params = []
|
||||
for key, value in model.named_parameters():
|
||||
if not value.requires_grad:
|
||||
continue
|
||||
lr = cfg.SOLVER.BASE_LR
|
||||
weight_decay = cfg.SOLVER.WEIGHT_DECAY
|
||||
if "bias" in key:
|
||||
lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BIAS_LR_FACTOR
|
||||
weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS
|
||||
params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
|
||||
if cfg.SOLVER.OPTIMIZER_NAME == 'SGD':
|
||||
optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params, momentum=cfg.SOLVER.MOMENTUM)
|
||||
else:
|
||||
optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params)
|
||||
optimizer_center = torch.optim.SGD(center_criterion.parameters(), lr=cfg.SOLVER.CENTER_LR)
|
||||
return optimizer, optimizer_center
|
|
@ -0,0 +1,56 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
from bisect import bisect_right
|
||||
import torch
|
||||
|
||||
|
||||
# FIXME ideally this would be achieved with a CombinedLRScheduler,
|
||||
# separating MultiStepLR with WarmupLR
|
||||
# but the current LRScheduler design doesn't allow it
|
||||
|
||||
class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler):
|
||||
def __init__(
|
||||
self,
|
||||
optimizer,
|
||||
milestones,
|
||||
gamma=0.1,
|
||||
warmup_factor=1.0 / 3,
|
||||
warmup_iters=500,
|
||||
warmup_method="linear",
|
||||
last_epoch=-1,
|
||||
):
|
||||
if not list(milestones) == sorted(milestones):
|
||||
raise ValueError(
|
||||
"Milestones should be a list of" " increasing integers. Got {}",
|
||||
milestones,
|
||||
)
|
||||
|
||||
if warmup_method not in ("constant", "linear"):
|
||||
raise ValueError(
|
||||
"Only 'constant' or 'linear' warmup_method accepted"
|
||||
"got {}".format(warmup_method)
|
||||
)
|
||||
self.milestones = milestones
|
||||
self.gamma = gamma
|
||||
self.warmup_factor = warmup_factor
|
||||
self.warmup_iters = warmup_iters
|
||||
self.warmup_method = warmup_method
|
||||
super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self):
|
||||
warmup_factor = 1
|
||||
if self.last_epoch < self.warmup_iters:
|
||||
if self.warmup_method == "constant":
|
||||
warmup_factor = self.warmup_factor
|
||||
elif self.warmup_method == "linear":
|
||||
alpha = self.last_epoch / self.warmup_iters
|
||||
warmup_factor = self.warmup_factor * (1 - alpha) + alpha
|
||||
return [
|
||||
base_lr
|
||||
* warmup_factor
|
||||
* self.gamma ** bisect_right(self.milestones, self.last_epoch)
|
||||
for base_lr in self.base_lrs
|
||||
]
|
|
@ -0,0 +1,5 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
|
@ -0,0 +1,26 @@
|
|||
import sys
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
sys.path.append('.')
|
||||
from solver.lr_scheduler import WarmupMultiStepLR
|
||||
from solver.build import make_optimizer
|
||||
from config import cfg
|
||||
|
||||
|
||||
class MyTestCase(unittest.TestCase):
|
||||
def test_something(self):
|
||||
net = nn.Linear(10, 10)
|
||||
optimizer = make_optimizer(cfg, net)
|
||||
lr_scheduler = WarmupMultiStepLR(optimizer, [20, 40], warmup_iters=10)
|
||||
for i in range(50):
|
||||
lr_scheduler.step()
|
||||
for j in range(3):
|
||||
print(i, lr_scheduler.get_lr()[0])
|
||||
optimizer.step()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
|
@ -0,0 +1,5 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
|
@ -0,0 +1,67 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from os import mkdir
|
||||
|
||||
import torch
|
||||
from torch.backends import cudnn
|
||||
|
||||
sys.path.append('.')
|
||||
from config import cfg
|
||||
from data import make_data_loader
|
||||
from engine.inference import inference
|
||||
from modeling import build_model
|
||||
from utils.logger import setup_logger
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="ReID Baseline Inference")
|
||||
parser.add_argument(
|
||||
"--config_file", default="", help="path to config file", type=str
|
||||
)
|
||||
parser.add_argument("opts", help="Modify config options using the command-line", default=None,
|
||||
nargs=argparse.REMAINDER)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
|
||||
|
||||
if args.config_file != "":
|
||||
cfg.merge_from_file(args.config_file)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
|
||||
output_dir = cfg.OUTPUT_DIR
|
||||
if output_dir and not os.path.exists(output_dir):
|
||||
mkdir(output_dir)
|
||||
|
||||
logger = setup_logger("reid_baseline", output_dir, 0)
|
||||
logger.info("Using {} GPUS".format(num_gpus))
|
||||
logger.info(args)
|
||||
|
||||
if args.config_file != "":
|
||||
logger.info("Loaded configuration file {}".format(args.config_file))
|
||||
with open(args.config_file, 'r') as cf:
|
||||
config_str = "\n" + cf.read()
|
||||
logger.info(config_str)
|
||||
logger.info("Running with config:\n{}".format(cfg))
|
||||
|
||||
if cfg.MODEL.DEVICE == "cuda":
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
|
||||
cudnn.benchmark = True
|
||||
|
||||
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
|
||||
model = build_model(cfg, num_classes)
|
||||
model.load_param(cfg.TEST.WEIGHT)
|
||||
|
||||
inference(cfg, model, val_loader, num_query)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
|
@ -0,0 +1,115 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
from torch.backends import cudnn
|
||||
|
||||
sys.path.append('.')
|
||||
from config import cfg
|
||||
from data import make_data_loader
|
||||
from engine.trainer import do_train, do_train_with_center
|
||||
from modeling import build_model
|
||||
from layers import make_loss, make_loss_with_center
|
||||
from solver import make_optimizer, make_optimizer_with_center, WarmupMultiStepLR
|
||||
|
||||
from utils.logger import setup_logger
|
||||
|
||||
|
||||
def train(cfg):
|
||||
# prepare dataset
|
||||
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
|
||||
|
||||
# prepare model
|
||||
model = build_model(cfg, num_classes)
|
||||
|
||||
if cfg.MODEL.IF_WITH_CENTER == 'no':
|
||||
print('Train without center loss, the loss type is', cfg.MODEL.METRIC_LOSS_TYPE)
|
||||
optimizer = make_optimizer(cfg, model)
|
||||
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
|
||||
cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
|
||||
|
||||
loss_func = make_loss(cfg, num_classes) # modified by gu
|
||||
|
||||
arguments = {}
|
||||
|
||||
do_train(
|
||||
cfg,
|
||||
model,
|
||||
train_loader,
|
||||
val_loader,
|
||||
optimizer,
|
||||
scheduler,
|
||||
loss_func,
|
||||
num_query
|
||||
)
|
||||
elif cfg.MODEL.IF_WITH_CENTER == 'yes':
|
||||
print('Train with center loss, the loss type is', cfg.MODEL.METRIC_LOSS_TYPE)
|
||||
loss_func, center_criterion = make_loss_with_center(cfg, num_classes) # modified by gu
|
||||
optimizer, optimizer_center = make_optimizer_with_center(cfg, model, center_criterion)
|
||||
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
|
||||
cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
|
||||
|
||||
arguments = {}
|
||||
|
||||
do_train_with_center(
|
||||
cfg,
|
||||
model,
|
||||
center_criterion,
|
||||
train_loader,
|
||||
val_loader,
|
||||
optimizer,
|
||||
optimizer_center,
|
||||
scheduler,
|
||||
loss_func,
|
||||
num_query
|
||||
)
|
||||
else:
|
||||
print("Unsupported value for cfg.MODEL.IF_WITH_CENTER {}, only support yes or no!\n".format(cfg.MODEL.IF_WITH_CENTER))
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="ReID Baseline Training")
|
||||
parser.add_argument(
|
||||
"--config_file", default="", help="path to config file", type=str
|
||||
)
|
||||
parser.add_argument("opts", help="Modify config options using the command-line", default=None,
|
||||
nargs=argparse.REMAINDER)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
|
||||
|
||||
if args.config_file != "":
|
||||
cfg.merge_from_file(args.config_file)
|
||||
cfg.merge_from_list(args.opts)
|
||||
cfg.freeze()
|
||||
|
||||
output_dir = cfg.OUTPUT_DIR
|
||||
if output_dir and not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
logger = setup_logger("reid_baseline", output_dir, 0)
|
||||
logger.info("Using {} GPUS".format(num_gpus))
|
||||
logger.info(args)
|
||||
|
||||
if args.config_file != "":
|
||||
logger.info("Loaded configuration file {}".format(args.config_file))
|
||||
with open(args.config_file, 'r') as cf:
|
||||
config_str = "\n" + cf.read()
|
||||
logger.info(config_str)
|
||||
logger.info("Running with config:\n{}".format(cfg))
|
||||
|
||||
if cfg.MODEL.DEVICE == "cuda":
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID # new add by gu
|
||||
cudnn.benchmark = True
|
||||
train(cfg)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
|
@ -0,0 +1,6 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import errno
|
||||
import json
|
||||
import os
|
||||
|
||||
import os.path as osp
|
||||
|
||||
|
||||
def mkdir_if_missing(directory):
|
||||
if not osp.exists(directory):
|
||||
try:
|
||||
os.makedirs(directory)
|
||||
except OSError as e:
|
||||
if e.errno != errno.EEXIST:
|
||||
raise
|
||||
|
||||
|
||||
def check_isfile(path):
|
||||
isfile = osp.isfile(path)
|
||||
if not isfile:
|
||||
print("=> Warning: no file found at '{}' (ignored)".format(path))
|
||||
return isfile
|
||||
|
||||
|
||||
def read_json(fpath):
|
||||
with open(fpath, 'r') as f:
|
||||
obj = json.load(f)
|
||||
return obj
|
||||
|
||||
|
||||
def write_json(obj, fpath):
|
||||
mkdir_if_missing(osp.dirname(fpath))
|
||||
with open(fpath, 'w') as f:
|
||||
json.dump(obj, f, indent=4, separators=(',', ': '))
|
|
@ -0,0 +1,30 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: sherlock
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def setup_logger(name, save_dir, distributed_rank):
|
||||
logger = logging.getLogger(name)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
# don't log results for the non-master process
|
||||
if distributed_rank > 0:
|
||||
return logger
|
||||
ch = logging.StreamHandler(stream=sys.stdout)
|
||||
ch.setLevel(logging.DEBUG)
|
||||
formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s")
|
||||
ch.setFormatter(formatter)
|
||||
logger.addHandler(ch)
|
||||
|
||||
if save_dir:
|
||||
fh = logging.FileHandler(os.path.join(save_dir, "log.txt"), mode='w')
|
||||
fh.setLevel(logging.DEBUG)
|
||||
fh.setFormatter(formatter)
|
||||
logger.addHandler(fh)
|
||||
|
||||
return logger
|
|
@ -0,0 +1,101 @@
|
|||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Created on Fri, 25 May 2018 20:29:09
|
||||
|
||||
@author: luohao
|
||||
"""
|
||||
|
||||
"""
|
||||
CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017.
|
||||
url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf
|
||||
Matlab version: https://github.com/zhunzhong07/person-re-ranking
|
||||
"""
|
||||
|
||||
"""
|
||||
API
|
||||
|
||||
probFea: all feature vectors of the query set (torch tensor)
|
||||
probFea: all feature vectors of the gallery set (torch tensor)
|
||||
k1,k2,lambda: parameters, the original paper is (k1=20,k2=6,lambda=0.3)
|
||||
MemorySave: set to 'True' when using MemorySave mode
|
||||
Minibatch: avaliable when 'MemorySave' is 'True'
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def re_ranking(probFea, galFea, k1, k2, lambda_value, local_distmat=None, only_local=False):
|
||||
# if feature vector is numpy, you should use 'torch.tensor' transform it to tensor
|
||||
query_num = probFea.size(0)
|
||||
all_num = query_num + galFea.size(0)
|
||||
if only_local:
|
||||
original_dist = local_distmat
|
||||
else:
|
||||
feat = torch.cat([probFea,galFea])
|
||||
print('using GPU to compute original distance')
|
||||
distmat = torch.pow(feat,2).sum(dim=1, keepdim=True).expand(all_num,all_num) + \
|
||||
torch.pow(feat, 2).sum(dim=1, keepdim=True).expand(all_num, all_num).t()
|
||||
distmat.addmm_(1,-2,feat,feat.t())
|
||||
original_dist = distmat.cpu().numpy()
|
||||
del feat
|
||||
if not local_distmat is None:
|
||||
original_dist = original_dist + local_distmat
|
||||
gallery_num = original_dist.shape[0]
|
||||
original_dist = np.transpose(original_dist / np.max(original_dist, axis=0))
|
||||
V = np.zeros_like(original_dist).astype(np.float16)
|
||||
initial_rank = np.argsort(original_dist).astype(np.int32)
|
||||
|
||||
print('starting re_ranking')
|
||||
for i in range(all_num):
|
||||
# k-reciprocal neighbors
|
||||
forward_k_neigh_index = initial_rank[i, :k1 + 1]
|
||||
backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1]
|
||||
fi = np.where(backward_k_neigh_index == i)[0]
|
||||
k_reciprocal_index = forward_k_neigh_index[fi]
|
||||
k_reciprocal_expansion_index = k_reciprocal_index
|
||||
for j in range(len(k_reciprocal_index)):
|
||||
candidate = k_reciprocal_index[j]
|
||||
candidate_forward_k_neigh_index = initial_rank[candidate, :int(np.around(k1 / 2)) + 1]
|
||||
candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index,
|
||||
:int(np.around(k1 / 2)) + 1]
|
||||
fi_candidate = np.where(candidate_backward_k_neigh_index == candidate)[0]
|
||||
candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate]
|
||||
if len(np.intersect1d(candidate_k_reciprocal_index, k_reciprocal_index)) > 2 / 3 * len(
|
||||
candidate_k_reciprocal_index):
|
||||
k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index, candidate_k_reciprocal_index)
|
||||
|
||||
k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index)
|
||||
weight = np.exp(-original_dist[i, k_reciprocal_expansion_index])
|
||||
V[i, k_reciprocal_expansion_index] = weight / np.sum(weight)
|
||||
original_dist = original_dist[:query_num, ]
|
||||
if k2 != 1:
|
||||
V_qe = np.zeros_like(V, dtype=np.float16)
|
||||
for i in range(all_num):
|
||||
V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0)
|
||||
V = V_qe
|
||||
del V_qe
|
||||
del initial_rank
|
||||
invIndex = []
|
||||
for i in range(gallery_num):
|
||||
invIndex.append(np.where(V[:, i] != 0)[0])
|
||||
|
||||
jaccard_dist = np.zeros_like(original_dist, dtype=np.float16)
|
||||
|
||||
for i in range(query_num):
|
||||
temp_min = np.zeros(shape=[1, gallery_num], dtype=np.float16)
|
||||
indNonZero = np.where(V[i, :] != 0)[0]
|
||||
indImages = [invIndex[ind] for ind in indNonZero]
|
||||
for j in range(len(indNonZero)):
|
||||
temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum(V[i, indNonZero[j]],
|
||||
V[indImages[j], indNonZero[j]])
|
||||
jaccard_dist[i] = 1 - temp_min / (2 - temp_min)
|
||||
|
||||
final_dist = jaccard_dist * (1 - lambda_value) + original_dist * lambda_value
|
||||
del original_dist
|
||||
del V
|
||||
del jaccard_dist
|
||||
final_dist = final_dist[:query_num, query_num:]
|
||||
return final_dist
|
||||
|
|
@ -0,0 +1,97 @@
|
|||
# encoding: utf-8
|
||||
"""
|
||||
@author: liaoxingyu
|
||||
@contact: sherlockliao01@gmail.com
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from ignite.metrics import Metric
|
||||
|
||||
from data.datasets.eval_reid import eval_func
|
||||
from .re_ranking import re_ranking
|
||||
|
||||
|
||||
class R1_mAP(Metric):
|
||||
def __init__(self, num_query, max_rank=50, feat_norm='yes'):
|
||||
super(R1_mAP, self).__init__()
|
||||
self.num_query = num_query
|
||||
self.max_rank = max_rank
|
||||
self.feat_norm = feat_norm
|
||||
|
||||
def reset(self):
|
||||
self.feats = []
|
||||
self.pids = []
|
||||
self.camids = []
|
||||
|
||||
def update(self, output):
|
||||
feat, pid, camid = output
|
||||
self.feats.append(feat)
|
||||
self.pids.extend(np.asarray(pid))
|
||||
self.camids.extend(np.asarray(camid))
|
||||
|
||||
def compute(self):
|
||||
feats = torch.cat(self.feats, dim=0)
|
||||
if self.feat_norm == 'yes':
|
||||
print("The test feature is normalized")
|
||||
feats = torch.nn.functional.normalize(feats, dim=1, p=2)
|
||||
# query
|
||||
qf = feats[:self.num_query]
|
||||
q_pids = np.asarray(self.pids[:self.num_query])
|
||||
q_camids = np.asarray(self.camids[:self.num_query])
|
||||
# gallery
|
||||
gf = feats[self.num_query:]
|
||||
g_pids = np.asarray(self.pids[self.num_query:])
|
||||
g_camids = np.asarray(self.camids[self.num_query:])
|
||||
m, n = qf.shape[0], gf.shape[0]
|
||||
distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
|
||||
torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
|
||||
distmat.addmm_(1, -2, qf, gf.t())
|
||||
distmat = distmat.cpu().numpy()
|
||||
cmc, mAP = eval_func(distmat, q_pids, g_pids, q_camids, g_camids)
|
||||
|
||||
return cmc, mAP
|
||||
|
||||
|
||||
class R1_mAP_reranking(Metric):
|
||||
def __init__(self, num_query, max_rank=50, feat_norm='yes'):
|
||||
super(R1_mAP_reranking, self).__init__()
|
||||
self.num_query = num_query
|
||||
self.max_rank = max_rank
|
||||
self.feat_norm = feat_norm
|
||||
|
||||
def reset(self):
|
||||
self.feats = []
|
||||
self.pids = []
|
||||
self.camids = []
|
||||
|
||||
def update(self, output):
|
||||
feat, pid, camid = output
|
||||
self.feats.append(feat)
|
||||
self.pids.extend(np.asarray(pid))
|
||||
self.camids.extend(np.asarray(camid))
|
||||
|
||||
def compute(self):
|
||||
feats = torch.cat(self.feats, dim=0)
|
||||
if self.feat_norm == 'yes':
|
||||
print("The test feature is normalized")
|
||||
feats = torch.nn.functional.normalize(feats, dim=1, p=2)
|
||||
|
||||
# query
|
||||
qf = feats[:self.num_query]
|
||||
q_pids = np.asarray(self.pids[:self.num_query])
|
||||
q_camids = np.asarray(self.camids[:self.num_query])
|
||||
# gallery
|
||||
gf = feats[self.num_query:]
|
||||
g_pids = np.asarray(self.pids[self.num_query:])
|
||||
g_camids = np.asarray(self.camids[self.num_query:])
|
||||
# m, n = qf.shape[0], gf.shape[0]
|
||||
# distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
|
||||
# torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
|
||||
# distmat.addmm_(1, -2, qf, gf.t())
|
||||
# distmat = distmat.cpu().numpy()
|
||||
print("Enter reranking")
|
||||
distmat = re_ranking(qf, gf, k1=20, k2=6, lambda_value=0.3)
|
||||
cmc, mAP = eval_func(distmat, q_pids, g_pids, q_camids, g_camids)
|
||||
|
||||
return cmc, mAP
|
Loading…
Reference in New Issue