├── .flake8 ├── .gitignore ├── .isort.cfg ├── .style.yapf ├── Interact, Embed, and EnlargE-Boosting Modality-specific Representations for Multi-Modal Person Re-identification.pdf ├── LICENSE ├── README.md ├── README.rst ├── configs └── RGBNT_ieee_part_margin.yaml ├── linter.sh ├── requirements.txt ├── scripts ├── __pycache__ │ ├── default_config.cpython-36.pyc │ ├── default_config.cpython-37.pyc │ └── default_config.cpython-38.pyc ├── default_config.py ├── main.py ├── mainMultiModal.py └── mainMultiModalCudnn.py ├── setup.py ├── tools ├── compute_mean_std.py ├── parse_test_res.py ├── visualize_actmap.py └── visualize_actmap_ori.py └── torchreid ├── __init__.py ├── data ├── __init__.py ├── datamanager.py ├── datasets │ ├── __init__.py │ ├── dataset.py │ ├── image │ │ ├── AllDay.py │ │ ├── RGBNT201.py │ │ ├── UAV.py │ │ ├── __init__.py │ │ ├── cuhk01.py │ │ ├── cuhk02.py │ │ ├── cuhk03.py │ │ ├── cuhksysu.py │ │ ├── dukemtmcreid.py │ │ ├── grid.py │ │ ├── ilids.py │ │ ├── market1501.py │ │ ├── market_to_RGBNT201.py │ │ ├── msmt17.py │ │ ├── prid.py │ │ ├── sensereid.py │ │ ├── university1652.py │ │ └── viper.py │ └── video │ │ ├── __init__.py │ │ ├── dukemtmcvidreid.py │ │ ├── ilidsvid.py │ │ ├── mars.py │ │ └── prid2011.py ├── sampler.py └── transforms.py ├── engine ├── __init__.py ├── engine.py ├── image │ ├── __init__.py │ ├── hcloss.py │ ├── margin.py │ ├── softmax.py │ └── triplet.py └── video │ ├── __init__.py │ ├── softmax.py │ └── triplet.py ├── losses ├── __init__.py ├── cross_entropy_loss.py ├── hard_mine_triplet_loss.py ├── hcloss.py ├── multi_modal_margin_loss_new.py └── time_loss.py ├── metrics ├── accuracy.py ├── distance.py ├── rank.py └── rank_cylib │ ├── Makefile │ ├── rank_cy.c │ ├── rank_cy.cpython-36m-x86_64-linux-gnu.so │ ├── rank_cy.pyx │ ├── setup.py │ └── test_cython.py ├── models ├── __init__.py ├── densenet.py ├── hacnn.py ├── ieee3modalPart.py ├── inceptionresnetv2.py ├── inceptionv4.py ├── layers.py ├── mlfn.py ├── mobilenetv2.py ├── mudeep.py ├── nasnet.py ├── osnet.py ├── osnet_ain.py ├── pcb.py ├── pfnet.py ├── resnet.py ├── resnet_ibn_a.py ├── resnet_ibn_b.py ├── resnetmid.py ├── senet.py ├── shufflenet.py ├── shufflenetv2.py ├── squeezenet.py ├── util.py └── xception.py ├── optim ├── __init__.py ├── lr_scheduler.py ├── optimizer.py └── radam.py └── utils ├── GPU-Re-Ranking ├── README.md ├── extension │ ├── adjacency_matrix │ │ ├── build_adjacency_matrix.cpp │ │ ├── build_adjacency_matrix_kernel.cu │ │ └── setup.py │ ├── make.sh │ └── propagation │ │ ├── gnn_propagate.cpp │ │ ├── gnn_propagate_kernel.cu │ │ └── setup.py ├── gnn_reranking.py ├── main.py └── utils.py ├── __init__.py ├── avgmeter.py ├── feature_extractor.py ├── loggers.py ├── model_complexity.py ├── reidtools.py ├── rerank.py ├── tools.py └── torchtools.py /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | ignore = 3 | # At least two spaces before inline comment 4 | E261, 5 | # Line lengths are recommended to be no greater than 79 characters 6 | E501, 7 | # Missing whitespace around arithmetic operator 8 | E226, 9 | # Blank line contains whitespace 10 | W293, 11 | # Do not use bare 'except' 12 | E722, 13 | # Line break after binary operator 14 | W504, 15 | # isort found an import in the wrong position 16 | I001 17 | max-line-length = 79 18 | exclude = __init__.py, build, torchreid/metrics/rank_cylib/ -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | 62 | # Flask stuff: 63 | instance/ 64 | .webassets-cache 65 | 66 | # Scrapy stuff: 67 | .scrapy 68 | 69 | # Sphinx documentation 70 | docs/_build/ 71 | 72 | # PyBuilder 73 | target/ 74 | 75 | # Jupyter Notebook 76 | .ipynb_checkpoints 77 | 78 | # IPython 79 | profile_default/ 80 | ipython_config.py 81 | 82 | # pyenv 83 | .python-version 84 | 85 | # pipenv 86 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 87 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 88 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 89 | # install all needed dependencies. 90 | #Pipfile.lock 91 | 92 | # celery beat schedule file 93 | celerybeat-schedule 94 | 95 | # SageMath parsed files 96 | *.sage.py 97 | 98 | # Environments 99 | .env 100 | .venv 101 | env/ 102 | venv/ 103 | ENV/ 104 | env.bak/ 105 | venv.bak/ 106 | 107 | # Spyder project settings 108 | .spyderproject 109 | .spyproject 110 | 111 | # Rope project settings 112 | .ropeproject 113 | 114 | # mkdocs documentation 115 | /site 116 | 117 | # mypy 118 | .mypy_cache/ 119 | .dmypy.json 120 | dmypy.json 121 | 122 | # Pyre type checker 123 | .pyre/ 124 | 125 | # Cython eval code 126 | *.c 127 | *.html 128 | 129 | # OS X 130 | .DS_Store 131 | .Spotlight-V100 132 | .Trashes 133 | ._* 134 | 135 | # ReID 136 | reid-data/ 137 | log/ 138 | saved-models/ 139 | model-zoo/ 140 | debug* 141 | -------------------------------------------------------------------------------- /.isort.cfg: -------------------------------------------------------------------------------- 1 | [isort] 2 | line_length=79 3 | multi_line_output=3 4 | length_sort=true 5 | known_standard_library=numpy,setuptools 6 | known_myself=torchreid 7 | known_third_party=matplotlib,cv2,torch,torchvision,PIL,yacs 8 | no_lines_before=STDLIB,THIRDPARTY 9 | sections=FUTURE,STDLIB,THIRDPARTY,myself,FIRSTPARTY,LOCALFOLDER 10 | default_section=FIRSTPARTY -------------------------------------------------------------------------------- /.style.yapf: -------------------------------------------------------------------------------- 1 | [style] 2 | BASED_ON_STYLE = pep8 3 | BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true 4 | SPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN = true 5 | DEDENT_CLOSING_BRACKETS = true 6 | SPACES_BEFORE_COMMENT = 1 7 | ARITHMETIC_PRECEDENCE_INDICATION = true -------------------------------------------------------------------------------- /Interact, Embed, and EnlargE-Boosting Modality-specific Representations for Multi-Modal Person Re-identification.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ziwang1121/IEEE/830a3fc444eb72736ff338b7c6fa7951d94c8d17/Interact, Embed, and EnlargE-Boosting Modality-specific Representations for Multi-Modal Person Re-identification.pdf -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Kaiyang Zhou 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Interact, Embed, and EnlargE (IEEE): Boosting Modality-specific Representations for Multi-Modal Person Re-identification 2 | 3 | We provide the codes for reproducing result of our paper **Interact, Embed, and EnlargE (IEEE): Boosting Modality-specific Representations for Multi-Modal Person Re-identification**. 4 | 5 | 6 | 7 | ## Installation 8 | 9 | 1. Basic environments: `python3.6`, `pytorch1.8.0`, `cuda11.1`. 10 | 11 | 2. Our codes structure is based on `Torchreid`. (More details can be found in link: https://github.com/KaiyangZhou/deep-person-reid , you can download the packages according to `Torchreid` requirements.) 12 | 13 | ```python 14 | # create environment 15 | cd AAAI2022_IEEE/ 16 | conda create --name ieeeReid python=3.6 17 | conda activate ieeeReid 18 | 19 | # install dependencies 20 | # make sure `which python` and `which pip` point to the correct path 21 | pip install -r requirements.txt 22 | 23 | # install torch and torchvision (select the proper cuda version to suit your machine) 24 | conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge 25 | 26 | # install torchreid (don't need to re-build it if you modify the source code) 27 | python setup.py develop 28 | ``` 29 | 30 | 31 | 32 | ## Get start 33 | 34 | 1. You can use the setting in `im_r50_softmax_256x128_amsgrad_RGBNT_ieee_part_margin.yaml` to get the results of full IEEE. 35 | 36 | ```python 37 | python ./scripts/mainMultiModal.py --config-file ./configs/RGBNT_ieee_part_margin.yaml --seed 40 38 | ``` 39 | 40 | ## Details 41 | 42 | 1. The details of our **Cross-modal Interacting Module (CIM)** and **Relation-based Embedding Module (REM)** can be found in `.\torchreid\models\ieee3modalPart.py`. The design of **Multi-modal Margin Loss(3M loss)** can be found in `.\torchreid\losses\multi_modal_margin_loss_new.py`. 43 | 44 | 2. Ablation study settings. 45 | 46 | You can control these two modules and the loss by change the corresponding codes. 47 | 48 | 1) **Cross-modal Interacting Module (CIM)** and **Relation-based Embedding Module (REM)** 49 | 50 | ```python 51 | # change the code in .\torchreid\models\ieee3modalPart.py 52 | 53 | class IEEE3modalPart(nn.Module): 54 | def __init__(··· 55 | ): 56 | modal_number = 3 57 | fc_dims = [128] 58 | pooling_dims = 768 59 | super(IEEE3modalPart, self).__init__() 60 | self.loss = loss 61 | self.parts = 6 62 | 63 | self.backbone = nn.ModuleList(··· 64 | ) 65 | 66 | # using Cross-modal Interacting Module (CIM) 67 | self.interaction = True 68 | # using channel attention in CIM 69 | self.attention = True 70 | 71 | # using Relation-based Embedding Module (REM) 72 | self.using_REM = True 73 | 74 | ··· 75 | ``` 76 | 77 | 2) **Multi-modal Margin Loss(3M loss)** 78 | 79 | ```python 80 | # change the code in .\configs\your_config_file.yaml 81 | 82 | # using Multi-modal Margin Loss(3M loss), you can change the margin by modify the parameter of "ieee_margin". 83 | ··· 84 | loss: 85 | name: 'margin' 86 | softmax: 87 | label_smooth: True 88 | ieee_margin: 1 89 | weight_m: 1.0 90 | weight_x: 1.0 91 | ··· 92 | 93 | # using only CE loss 94 | ··· 95 | loss: 96 | name: 'softmax' 97 | softmax: 98 | label_smooth: True 99 | weight_x: 1.0 100 | ··· 101 | ``` 102 | 103 | ## Dataset 104 | #### RGBNT201 & reconstructed cross-modal datasets & Market1501 multi-modal version: 105 | Google Drive Link: https://drive.google.com/drive/folders/1EscBadX-wMAT56_It5lXY-S3-b5nK1wH?usp=sharing 106 | Please contact with Zi Wang (email address: ziwang1121@foxmail.com). 107 | 108 | -------------------------------------------------------------------------------- /configs/RGBNT_ieee_part_margin.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: 'ieee3modalPart' 3 | pretrained: True 4 | # load_weights: r"E:\reidCode\AAAI 2022-master\log\ablation_new\p6_CIM_GateSumOri_3M_NonLocal\model\model.pth.tar-50" 5 | 6 | data: 7 | type: 'image' 8 | sources: ['RGBNT201'] 9 | targets: ['RGBNT201'] 10 | height: 256 11 | width: 128 12 | combineall: False 13 | workers: 0 14 | transforms: ['random_flip'] 15 | save_dir: 'log\ours171' 16 | 17 | loss: 18 | name: 'margin' 19 | softmax: 20 | label_smooth: True 21 | ieee_margin: 1 22 | weight_m: 1.0 23 | weight_x: 1.0 24 | 25 | sampler: 26 | train_sampler: 'RandomIdentitySampler' 27 | num_instances: 4 28 | 29 | train: 30 | optim: 'sgd' 31 | lr: 0.001 32 | max_epoch: 61 33 | batch_size: 8 34 | fixbase_epoch: 0 35 | open_layers: ['classifier_R', 'classifier_N', 'classifier_T'] 36 | lr_scheduler: 'multi_step' 37 | stepsize: [20, 40] 38 | print_freq: 90 39 | 40 | test: 41 | batch_size: 100 42 | dist_metric: 'euclidean' 43 | normalize_feature: False 44 | evaluate: False 45 | visrank: False 46 | eval_freq: 6 47 | ranks: [1, 5, 10] 48 | rerank: False 49 | 50 | -------------------------------------------------------------------------------- /linter.sh: -------------------------------------------------------------------------------- 1 | echo "Running isort" 2 | isort -y -sp . 3 | echo "Done" 4 | 5 | echo "Running yapf" 6 | yapf -i -r -vv -e build . 7 | echo "Done" 8 | 9 | echo "Running flake8" 10 | flake8 . 11 | echo "Done" -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | Cython 3 | h5py 4 | Pillow 5 | six 6 | scipy 7 | opencv-python 8 | matplotlib 9 | tb-nightly 10 | future 11 | yacs 12 | gdown 13 | flake8 14 | yapf 15 | isort==4.3.21 16 | imageio -------------------------------------------------------------------------------- /scripts/__pycache__/default_config.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ziwang1121/IEEE/830a3fc444eb72736ff338b7c6fa7951d94c8d17/scripts/__pycache__/default_config.cpython-36.pyc -------------------------------------------------------------------------------- /scripts/__pycache__/default_config.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ziwang1121/IEEE/830a3fc444eb72736ff338b7c6fa7951d94c8d17/scripts/__pycache__/default_config.cpython-37.pyc -------------------------------------------------------------------------------- /scripts/__pycache__/default_config.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ziwang1121/IEEE/830a3fc444eb72736ff338b7c6fa7951d94c8d17/scripts/__pycache__/default_config.cpython-38.pyc -------------------------------------------------------------------------------- /scripts/main.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import time 3 | import os.path as osp 4 | import argparse 5 | import torch 6 | import torch.nn as nn 7 | 8 | import torchreid 9 | from torchreid.utils import ( 10 | Logger, check_isfile, set_random_seed, collect_env_info, 11 | resume_from_checkpoint, load_pretrained_weights, compute_model_complexity 12 | ) 13 | 14 | from scripts.default_config import ( 15 | imagedata_kwargs, optimizer_kwargs, videodata_kwargs, engine_run_kwargs, 16 | get_default_config, lr_scheduler_kwargs 17 | ) 18 | 19 | def build_datamanager(cfg): 20 | if cfg.data.type == 'image': 21 | return torchreid.data.ImageDataManager(**imagedata_kwargs(cfg)) 22 | else: 23 | return torchreid.data.VideoDataManager(**videodata_kwargs(cfg)) 24 | 25 | 26 | def build_engine(cfg, datamanager, model, optimizer, scheduler): 27 | if cfg.data.type == 'image': 28 | if cfg.loss.name == 'softmax': 29 | engine = torchreid.engine.ImageSoftmaxEngine( 30 | datamanager, 31 | model, 32 | optimizer=optimizer, 33 | scheduler=scheduler, 34 | use_gpu=cfg.use_gpu, 35 | label_smooth=cfg.loss.softmax.label_smooth 36 | ) 37 | 38 | else: 39 | engine = torchreid.engine.ImageTripletEngine( 40 | datamanager, 41 | model, 42 | optimizer=optimizer, 43 | margin=cfg.loss.triplet.margin, 44 | weight_t=cfg.loss.triplet.weight_t, 45 | weight_x=cfg.loss.triplet.weight_x, 46 | scheduler=scheduler, 47 | use_gpu=cfg.use_gpu, 48 | label_smooth=cfg.loss.softmax.label_smooth 49 | ) 50 | 51 | else: 52 | if cfg.loss.name == 'softmax': 53 | engine = torchreid.engine.VideoSoftmaxEngine( 54 | datamanager, 55 | model, 56 | optimizer=optimizer, 57 | scheduler=scheduler, 58 | use_gpu=cfg.use_gpu, 59 | label_smooth=cfg.loss.softmax.label_smooth, 60 | pooling_method=cfg.video.pooling_method 61 | ) 62 | 63 | else: 64 | engine = torchreid.engine.VideoTripletEngine( 65 | datamanager, 66 | model, 67 | optimizer=optimizer, 68 | margin=cfg.loss.triplet.margin, 69 | weight_t=cfg.loss.triplet.weight_t, 70 | weight_x=cfg.loss.triplet.weight_x, 71 | scheduler=scheduler, 72 | use_gpu=cfg.use_gpu, 73 | label_smooth=cfg.loss.softmax.label_smooth 74 | ) 75 | 76 | return engine 77 | 78 | 79 | def reset_config(cfg, args): 80 | if args.root: 81 | cfg.data.root = args.root 82 | if args.sources: 83 | cfg.data.sources = args.sources 84 | if args.targets: 85 | cfg.data.targets = args.targets 86 | if args.transforms: 87 | cfg.data.transforms = args.transforms 88 | 89 | 90 | def check_cfg(cfg): 91 | if cfg.loss.name == 'triplet' and cfg.loss.triplet.weight_x == 0: 92 | assert cfg.train.fixbase_epoch == 0, \ 93 | 'The output of classifier is not included in the computational graph' 94 | 95 | 96 | def main(): 97 | parser = argparse.ArgumentParser( 98 | formatter_class=argparse.ArgumentDefaultsHelpFormatter 99 | ) 100 | parser.add_argument( 101 | '--config-file', type=str, default='', help='path to config file' 102 | ) 103 | parser.add_argument( 104 | '-s', 105 | '--sources', 106 | type=str, 107 | nargs='+', 108 | help='source datasets (delimited by space)' 109 | ) 110 | parser.add_argument( 111 | '-t', 112 | '--targets', 113 | type=str, 114 | nargs='+', 115 | help='target datasets (delimited by space)' 116 | ) 117 | parser.add_argument( 118 | '--transforms', type=str, nargs='+', help='data augmentation' 119 | ) 120 | parser.add_argument( 121 | '--root', type=str, default='', help='path to data root' 122 | ) 123 | parser.add_argument( 124 | 'opts', 125 | default=None, 126 | nargs=argparse.REMAINDER, 127 | help='Modify config options using the command-line' 128 | ) 129 | args = parser.parse_args() 130 | 131 | cfg = get_default_config() 132 | cfg.use_gpu = torch.cuda.is_available() 133 | if args.config_file: 134 | cfg.merge_from_file(args.config_file) 135 | reset_config(cfg, args) 136 | cfg.merge_from_list(args.opts) 137 | set_random_seed(cfg.train.seed) 138 | check_cfg(cfg) 139 | 140 | log_name = 'test.log' if cfg.test.evaluate else 'train.log' 141 | log_name += time.strftime('-%Y-%m-%d-%H-%M-%S') 142 | sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name)) 143 | 144 | print('Show configuration\n{}\n'.format(cfg)) 145 | print('Collecting env info ...') 146 | print('** System info **\n{}\n'.format(collect_env_info())) 147 | 148 | if cfg.use_gpu: 149 | torch.backends.cudnn.benchmark = True 150 | 151 | datamanager = build_datamanager(cfg) 152 | 153 | print('Building model: {}'.format(cfg.model.name)) 154 | model = torchreid.models.build_model( 155 | name=cfg.model.name, 156 | num_classes=datamanager.num_train_pids, 157 | loss=cfg.loss.name, 158 | pretrained=cfg.model.pretrained, 159 | use_gpu=cfg.use_gpu 160 | ) 161 | 162 | if cfg.model.load_weights and check_isfile(cfg.model.load_weights): 163 | load_pretrained_weights(model, cfg.model.load_weights) 164 | 165 | if cfg.use_gpu: 166 | model = nn.DataParallel(model).cuda() 167 | 168 | optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg)) 169 | scheduler = torchreid.optim.build_lr_scheduler( 170 | optimizer, **lr_scheduler_kwargs(cfg) 171 | ) 172 | 173 | if cfg.model.resume and check_isfile(cfg.model.resume): 174 | cfg.train.start_epoch = resume_from_checkpoint( 175 | cfg.model.resume, model, optimizer=optimizer, scheduler=scheduler 176 | ) 177 | 178 | print( 179 | 'Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type) 180 | ) 181 | engine = build_engine(cfg, datamanager, model, optimizer, scheduler) 182 | engine.run(**engine_run_kwargs(cfg)) 183 | 184 | 185 | if __name__ == '__main__': 186 | main() 187 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os.path as osp 3 | from setuptools import setup, find_packages 4 | from distutils.extension import Extension 5 | from Cython.Build import cythonize 6 | 7 | 8 | def readme(): 9 | with open('README.rst') as f: 10 | content = f.read() 11 | return content 12 | 13 | 14 | def find_version(): 15 | version_file = 'torchreid/__init__.py' 16 | with open(version_file, 'r') as f: 17 | exec(compile(f.read(), version_file, 'exec')) 18 | return locals()['__version__'] 19 | 20 | 21 | def numpy_include(): 22 | try: 23 | numpy_include = np.get_include() 24 | except AttributeError: 25 | numpy_include = np.get_numpy_include() 26 | return numpy_include 27 | 28 | 29 | ext_modules = [ 30 | Extension( 31 | 'torchreid.metrics.rank_cylib.rank_cy', 32 | ['torchreid/metrics/rank_cylib/rank_cy.pyx'], 33 | include_dirs=[numpy_include()], 34 | ) 35 | ] 36 | 37 | 38 | def get_requirements(filename='requirements.txt'): 39 | here = osp.dirname(osp.realpath(__file__)) 40 | with open(osp.join(here, filename), 'r') as f: 41 | requires = [line.replace('\n', '') for line in f.readlines()] 42 | return requires 43 | 44 | 45 | setup( 46 | name='torchreid', 47 | version=find_version(), 48 | description='A library for deep learning person re-ID in PyTorch', 49 | author='Kaiyang Zhou', 50 | license='MIT', 51 | long_description=readme(), 52 | url='https://github.com/KaiyangZhou/deep-person-reid', 53 | packages=find_packages(), 54 | install_requires=get_requirements(), 55 | keywords=['Person Re-Identification', 'Deep Learning', 'Computer Vision'], 56 | ext_modules=cythonize(ext_modules) 57 | ) 58 | -------------------------------------------------------------------------------- /tools/compute_mean_std.py: -------------------------------------------------------------------------------- 1 | """ 2 | Compute channel-wise mean and standard deviation of a dataset. 3 | 4 | Usage: 5 | $ python compute_mean_std.py DATASET_ROOT DATASET_KEY 6 | 7 | - The first argument points to the root path where you put the datasets. 8 | - The second argument means the specific dataset key. 9 | 10 | For instance, your datasets are put under $DATA and you wanna 11 | compute the statistics of Market1501, do 12 | $ python compute_mean_std.py $DATA market1501 13 | """ 14 | import argparse 15 | 16 | import torchreid 17 | 18 | 19 | def main(): 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument('root', type=str) 22 | parser.add_argument('sources', type=str) 23 | args = parser.parse_args() 24 | 25 | datamanager = torchreid.data.ImageDataManager( 26 | root=args.root, 27 | sources=args.sources, 28 | targets=None, 29 | height=256, 30 | width=128, 31 | batch_size_train=100, 32 | batch_size_test=100, 33 | transforms=None, 34 | norm_mean=[0., 0., 0.], 35 | norm_std=[1., 1., 1.], 36 | train_sampler='SequentialSampler' 37 | ) 38 | train_loader = datamanager.train_loader 39 | 40 | print('Computing mean and std ...') 41 | mean = 0. 42 | std = 0. 43 | n_samples = 0. 44 | for data in train_loader: 45 | data = data['img'] 46 | batch_size = data.size(0) 47 | data = data.view(batch_size, data.size(1), -1) 48 | mean += data.mean(2).sum(0) 49 | std += data.std(2).sum(0) 50 | n_samples += batch_size 51 | 52 | mean /= n_samples 53 | std /= n_samples 54 | print('Mean: {}'.format(mean)) 55 | print('Std: {}'.format(std)) 56 | 57 | 58 | if __name__ == '__main__': 59 | main() 60 | -------------------------------------------------------------------------------- /tools/parse_test_res.py: -------------------------------------------------------------------------------- 1 | """ 2 | This script aims to automate the process of calculating average results 3 | stored in the test.log files over multiple splits. 4 | 5 | How to use: 6 | For example, you have done evaluation over 20 splits on VIPeR, leading to 7 | the following file structure 8 | 9 | log/ 10 | eval_viper/ 11 | split_0/ 12 | test.log-xxxx 13 | split_1/ 14 | test.log-xxxx 15 | split_2/ 16 | test.log-xxxx 17 | ... 18 | 19 | You can run the following command in your terminal to get the average performance: 20 | $ python tools/parse_test_res.py log/eval_viper 21 | """ 22 | import os 23 | import re 24 | import glob 25 | import numpy as np 26 | import argparse 27 | from collections import defaultdict 28 | 29 | from torchreid.utils import check_isfile, listdir_nohidden 30 | 31 | 32 | def parse_file(filepath, regex_mAP, regex_r1, regex_r5, regex_r10, regex_r20): 33 | results = {} 34 | 35 | with open(filepath, 'r') as f: 36 | lines = f.readlines() 37 | 38 | for line in lines: 39 | line = line.strip() 40 | 41 | match_mAP = regex_mAP.search(line) 42 | if match_mAP: 43 | mAP = float(match_mAP.group(1)) 44 | results['mAP'] = mAP 45 | 46 | match_r1 = regex_r1.search(line) 47 | if match_r1: 48 | r1 = float(match_r1.group(1)) 49 | results['r1'] = r1 50 | 51 | match_r5 = regex_r5.search(line) 52 | if match_r5: 53 | r5 = float(match_r5.group(1)) 54 | results['r5'] = r5 55 | 56 | match_r10 = regex_r10.search(line) 57 | if match_r10: 58 | r10 = float(match_r10.group(1)) 59 | results['r10'] = r10 60 | 61 | match_r20 = regex_r20.search(line) 62 | if match_r20: 63 | r20 = float(match_r20.group(1)) 64 | results['r20'] = r20 65 | 66 | return results 67 | 68 | 69 | def main(args): 70 | regex_mAP = re.compile(r'mAP: ([\.\deE+-]+)%') 71 | regex_r1 = re.compile(r'Rank-1 : ([\.\deE+-]+)%') 72 | regex_r5 = re.compile(r'Rank-5 : ([\.\deE+-]+)%') 73 | regex_r10 = re.compile(r'Rank-10 : ([\.\deE+-]+)%') 74 | regex_r20 = re.compile(r'Rank-20 : ([\.\deE+-]+)%') 75 | 76 | final_res = defaultdict(list) 77 | 78 | directories = listdir_nohidden(args.directory, sort=True) 79 | num_dirs = len(directories) 80 | for directory in directories: 81 | fullpath = os.path.join(args.directory, directory) 82 | filepath = glob.glob(os.path.join(fullpath, 'test.log*'))[0] 83 | check_isfile(filepath) 84 | print(f'Parsing {filepath}') 85 | res = parse_file( 86 | filepath, regex_mAP, regex_r1, regex_r5, regex_r10, regex_r20 87 | ) 88 | for key, value in res.items(): 89 | final_res[key].append(value) 90 | 91 | print('Finished parsing') 92 | print(f'The average results over {num_dirs} splits are shown below') 93 | 94 | for key, values in final_res.items(): 95 | mean_val = np.mean(values) 96 | print(f'{key}: {mean_val:.1f}') 97 | 98 | 99 | if __name__ == '__main__': 100 | parser = argparse.ArgumentParser() 101 | parser.add_argument('directory', type=str, help='Path to directory') 102 | args = parser.parse_args() 103 | main(args) 104 | -------------------------------------------------------------------------------- /tools/visualize_actmap_ori.py: -------------------------------------------------------------------------------- 1 | """Visualizes CNN activation maps to see where the CNN focuses on to extract features. 2 | 3 | Reference: 4 | - Zagoruyko and Komodakis. Paying more attention to attention: Improving the 5 | performance of convolutional neural networks via attention transfer. ICLR, 2017 6 | - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019. 7 | """ 8 | import numpy as np 9 | import os.path as osp 10 | import argparse 11 | import cv2 12 | import torch 13 | from torch.nn import functional as F 14 | 15 | import torchreid 16 | from torchreid.utils import ( 17 | check_isfile, mkdir_if_missing, load_pretrained_weights 18 | ) 19 | 20 | IMAGENET_MEAN = [0.485, 0.456, 0.406] 21 | IMAGENET_STD = [0.229, 0.224, 0.225] 22 | GRID_SPACING = 10 23 | 24 | 25 | @torch.no_grad() 26 | def visactmap( 27 | model, 28 | test_loader, 29 | save_dir, 30 | width, 31 | height, 32 | use_gpu, 33 | img_mean=None, 34 | img_std=None 35 | ): 36 | if img_mean is None or img_std is None: 37 | # use imagenet mean and std 38 | img_mean = IMAGENET_MEAN 39 | img_std = IMAGENET_STD 40 | 41 | model.eval() 42 | 43 | for target in list(test_loader.keys()): 44 | data_loader = test_loader[target]['query'] # only process query images 45 | # original images and activation maps are saved individually 46 | actmap_dir = osp.join(save_dir, 'actmap_' + target) 47 | mkdir_if_missing(actmap_dir) 48 | print('Visualizing activation maps for {} ...'.format(target)) 49 | 50 | for batch_idx, data in enumerate(data_loader): 51 | imgs, paths = data['img'], data['impath'] 52 | if use_gpu: 53 | imgs = imgs.cuda() 54 | 55 | # forward to get convolutional feature maps 56 | try: 57 | outputs = model(imgs, return_featuremaps=True) 58 | except TypeError: 59 | raise TypeError( 60 | 'forward() got unexpected keyword argument "return_featuremaps". ' 61 | 'Please add return_featuremaps as an input argument to forward(). When ' 62 | 'return_featuremaps=True, return feature maps only.' 63 | ) 64 | 65 | if outputs.dim() != 4: 66 | raise ValueError( 67 | 'The model output is supposed to have ' 68 | 'shape of (b, c, h, w), i.e. 4 dimensions, but got {} dimensions. ' 69 | 'Please make sure you set the model output at eval mode ' 70 | 'to be the last convolutional feature maps'.format( 71 | outputs.dim() 72 | ) 73 | ) 74 | 75 | # compute activation maps 76 | outputs = (outputs**2).sum(1) 77 | b, h, w = outputs.size() 78 | outputs = outputs.view(b, h * w) 79 | outputs = F.normalize(outputs, p=2, dim=1) 80 | outputs = outputs.view(b, h, w) 81 | 82 | if use_gpu: 83 | imgs, outputs = imgs.cpu(), outputs.cpu() 84 | 85 | for j in range(outputs.size(0)): 86 | # get image name 87 | path = paths[j] 88 | imname = osp.basename(osp.splitext(path)[0]) 89 | 90 | # RGB image 91 | img = imgs[j, ...] 92 | for t, m, s in zip(img, img_mean, img_std): 93 | t.mul_(s).add_(m).clamp_(0, 1) 94 | img_np = np.uint8(np.floor(img.numpy() * 255)) 95 | img_np = img_np.transpose((1, 2, 0)) # (c, h, w) -> (h, w, c) 96 | 97 | # activation map 98 | am = outputs[j, ...].numpy() 99 | am = cv2.resize(am, (width, height)) 100 | am = 255 * (am - np.min(am)) / ( 101 | np.max(am) - np.min(am) + 1e-12 102 | ) 103 | am = np.uint8(np.floor(am)) 104 | am = cv2.applyColorMap(am, cv2.COLORMAP_JET) 105 | 106 | # overlapped 107 | overlapped = img_np*0.3 + am*0.7 108 | overlapped[overlapped > 255] = 255 109 | overlapped = overlapped.astype(np.uint8) 110 | 111 | # save images in a single figure (add white spacing between images) 112 | # from left to right: original image, activation map, overlapped image 113 | grid_img = 255 * np.ones( 114 | (height, 3*width + 2*GRID_SPACING, 3), dtype=np.uint8 115 | ) 116 | grid_img[:, :width, :] = img_np[:, :, ::-1] 117 | grid_img[:, 118 | width + GRID_SPACING:2*width + GRID_SPACING, :] = am 119 | grid_img[:, 2*width + 2*GRID_SPACING:, :] = overlapped 120 | cv2.imwrite(osp.join(actmap_dir, imname + '.jpg'), grid_img) 121 | 122 | if (batch_idx+1) % 10 == 0: 123 | print( 124 | '- done batch {}/{}'.format( 125 | batch_idx + 1, len(data_loader) 126 | ) 127 | ) 128 | 129 | 130 | def main(): 131 | parser = argparse.ArgumentParser() 132 | parser.add_argument('--root', type=str) 133 | parser.add_argument('-d', '--dataset', type=str, default='market1501') 134 | parser.add_argument('-m', '--model', type=str, default='osnet_x1_0') 135 | parser.add_argument('--weights', type=str) 136 | parser.add_argument('--save-dir', type=str, default='log') 137 | parser.add_argument('--height', type=int, default=256) 138 | parser.add_argument('--width', type=int, default=128) 139 | args = parser.parse_args() 140 | 141 | use_gpu = torch.cuda.is_available() 142 | 143 | datamanager = torchreid.data.ImageDataManager( 144 | root=args.root, 145 | sources=args.dataset, 146 | height=args.height, 147 | width=args.width, 148 | batch_size_train=100, 149 | batch_size_test=100, 150 | transforms=None, 151 | train_sampler='SequentialSampler' 152 | ) 153 | test_loader = datamanager.test_loader 154 | 155 | model = torchreid.models.build_model( 156 | name=args.model, 157 | num_classes=datamanager.num_train_pids, 158 | use_gpu=use_gpu 159 | ) 160 | 161 | if use_gpu: 162 | model = model.cuda() 163 | 164 | if args.weights and check_isfile(args.weights): 165 | load_pretrained_weights(model, args.weights) 166 | 167 | visactmap( 168 | model, test_loader, args.save_dir, args.width, args.height, use_gpu 169 | ) 170 | 171 | 172 | if __name__ == '__main__': 173 | main() 174 | -------------------------------------------------------------------------------- /torchreid/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | 3 | from torchreid import data, optim, utils, engine, losses, models, metrics 4 | 5 | __version__ = '1.4.0' 6 | __author__ = 'Kaiyang Zhou' 7 | __homepage__ = 'https://kaiyangzhou.github.io/' 8 | __description__ = 'Deep learning person re-identification in PyTorch' 9 | __url__ = 'https://github.com/KaiyangZhou/deep-person-reid' 10 | -------------------------------------------------------------------------------- /torchreid/data/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | 3 | from .datasets import ( 4 | Dataset, ImageDataset, VideoDataset, register_image_dataset, 5 | register_video_dataset 6 | ) 7 | from .datamanager import ImageDataManager, VideoDataManager 8 | -------------------------------------------------------------------------------- /torchreid/data/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | from torchreid.data.datasets.image import UAV 3 | from torchreid.data.datasets.image import RGBNT201 4 | 5 | from .image import ( 6 | GRID, PRID, CUHK01, CUHK02, CUHK03, MSMT17, CUHKSYSU, VIPeR, SenseReID, 7 | Market1501, DukeMTMCreID, University1652, iLIDS, UAV, RGBNT201, market_to_RGBNT201, AllDay 8 | ) 9 | 10 | from .video import PRID2011, Mars, DukeMTMCVidReID, iLIDSVID 11 | from .dataset import Dataset, ImageDataset, VideoDataset 12 | 13 | __image_datasets = { 14 | 'market1501': Market1501, 15 | 'cuhk03': CUHK03, 16 | 'dukemtmcreid': DukeMTMCreID, 17 | 'msmt17': MSMT17, 18 | 'viper': VIPeR, 19 | 'grid': GRID, 20 | 'cuhk01': CUHK01, 21 | 'ilids': iLIDS, 22 | 'sensereid': SenseReID, 23 | 'prid': PRID, 24 | 'cuhk02': CUHK02, 25 | 'university1652': University1652, 26 | 'cuhksysu': CUHKSYSU, 27 | 'UAVdata': UAV, 28 | 'RGBNT201': RGBNT201, 29 | 'market_to_RGBNT201': market_to_RGBNT201, 30 | 'AllDay': AllDay, 31 | } 32 | 33 | __video_datasets = { 34 | 'mars': Mars, 35 | 'ilidsvid': iLIDSVID, 36 | 'prid2011': PRID2011, 37 | 'dukemtmcvidreid': DukeMTMCVidReID 38 | } 39 | 40 | 41 | def init_image_dataset(name, **kwargs): 42 | """Initializes an image dataset.""" 43 | avai_datasets = list(__image_datasets.keys()) 44 | if name not in avai_datasets: 45 | raise ValueError( 46 | 'Invalid dataset name. Received "{}", ' 47 | 'but expected to be one of {}'.format(name, avai_datasets) 48 | ) 49 | return __image_datasets[name](**kwargs) 50 | 51 | 52 | def init_video_dataset(name, **kwargs): 53 | """Initializes a video dataset.""" 54 | avai_datasets = list(__video_datasets.keys()) 55 | if name not in avai_datasets: 56 | raise ValueError( 57 | 'Invalid dataset name. Received "{}", ' 58 | 'but expected to be one of {}'.format(name, avai_datasets) 59 | ) 60 | return __video_datasets[name](**kwargs) 61 | 62 | 63 | def register_image_dataset(name, dataset): 64 | """Registers a new image dataset. 65 | 66 | Args: 67 | name (str): key corresponding to the new dataset. 68 | dataset (Dataset): the new dataset class. 69 | 70 | Examples:: 71 | 72 | import torchreid 73 | import NewDataset 74 | torchreid.data.register_image_dataset('new_dataset', NewDataset) 75 | # single dataset case 76 | datamanager = torchreid.data.ImageDataManager( 77 | root='reid-data', 78 | sources='new_dataset' 79 | ) 80 | # multiple dataset case 81 | datamanager = torchreid.data.ImageDataManager( 82 | root='reid-data', 83 | sources=['new_dataset', 'dukemtmcreid'] 84 | ) 85 | """ 86 | global __image_datasets 87 | curr_datasets = list(__image_datasets.keys()) 88 | if name in curr_datasets: 89 | raise ValueError( 90 | 'The given name already exists, please choose ' 91 | 'another name excluding {}'.format(curr_datasets) 92 | ) 93 | __image_datasets[name] = dataset 94 | 95 | 96 | def register_video_dataset(name, dataset): 97 | """Registers a new video dataset. 98 | 99 | Args: 100 | name (str): key corresponding to the new dataset. 101 | dataset (Dataset): the new dataset class. 102 | 103 | Examples:: 104 | 105 | import torchreid 106 | import NewDataset 107 | torchreid.data.register_video_dataset('new_dataset', NewDataset) 108 | # single dataset case 109 | datamanager = torchreid.data.VideoDataManager( 110 | root='reid-data', 111 | sources='new_dataset' 112 | ) 113 | # multiple dataset case 114 | datamanager = torchreid.data.VideoDataManager( 115 | root='reid-data', 116 | sources=['new_dataset', 'ilidsvid'] 117 | ) 118 | """ 119 | global __video_datasets 120 | curr_datasets = list(__video_datasets.keys()) 121 | if name in curr_datasets: 122 | raise ValueError( 123 | 'The given name already exists, please choose ' 124 | 'another name excluding {}'.format(curr_datasets) 125 | ) 126 | __video_datasets[name] = dataset 127 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/AllDay.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | from os import pardir 3 | import re 4 | import glob 5 | import os.path as osp 6 | 7 | from numpy.lib.twodim_base import tri 8 | from torchreid.data.datasets import image 9 | import warnings 10 | 11 | from ..dataset import MultiModalImageDataset 12 | 13 | 14 | class AllDay(MultiModalImageDataset): 15 | dataset_dir = 'AllDay' 16 | 17 | def __init__(self, root='', **kwargs): 18 | self.root = osp.abspath(osp.expanduser(root)) 19 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 20 | 21 | # allow alternative directory structure 22 | self.data_dir = self.dataset_dir 23 | data_dir = osp.join(self.data_dir) 24 | if osp.isdir(data_dir): 25 | self.data_dir = data_dir 26 | else: 27 | warnings.warn( 28 | 'The current data structure is deprecated.' 29 | ) 30 | 31 | 32 | self.train_dir = osp.join(self.data_dir, 'train') 33 | self.query_dir = osp.join(self.data_dir, 'query_all') 34 | self.gallery_dir = osp.join(self.data_dir, 'gallery_all') 35 | 36 | required_files = [ 37 | self.data_dir, self.train_dir, self.query_dir, self.gallery_dir 38 | ] 39 | self.check_before_run(required_files) 40 | 41 | train = self._process_dir_train(self.train_dir, relabel=True) 42 | query = self._process_dir_test(self.query_dir, relabel=False) 43 | gallery = self._process_dir_test(self.gallery_dir, relabel=False) 44 | 45 | super(AllDay, self).__init__(train, query, gallery, **kwargs) 46 | 47 | 48 | def _process_dir_train(self, dir_path, relabel=False): 49 | img_paths_RGB = glob.glob(osp.join(dir_path, 'RGB', '*.jpg')) 50 | pid_container = set() 51 | for img_path_RGB in img_paths_RGB: 52 | jpg_name = img_path_RGB.split('/')[-1] 53 | pid = int(jpg_name.split('_')[0][0:6]) 54 | pid_container.add(pid) 55 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 56 | 57 | data = [] 58 | for img_path_RGB in img_paths_RGB: 59 | img = [] 60 | jpg_name = img_path_RGB.split('/')[-1] 61 | img_path_NI = osp.join(dir_path, 'NI', jpg_name) 62 | img_path_TI = osp.join(dir_path, 'TI', jpg_name) 63 | img.append(img_path_RGB) 64 | img.append(img_path_NI) 65 | img.append(img_path_TI) 66 | pid = int(jpg_name.split('_')[0][0:6]) 67 | camid = int(jpg_name.split('_')[1][3]) 68 | camid -= 1 # index starts from 0 69 | timeid = int(jpg_name.split('_')[2]) 70 | # print(img, pid, camid, timeid) 71 | if relabel: 72 | pid = pid2label[pid] 73 | data.append((img, pid, camid, timeid)) 74 | # print((img, pid, camid, timeid)) 75 | return data 76 | 77 | def _process_dir_test(self, dir_path, relabel=False): 78 | img_paths_RGB = glob.glob(osp.join(dir_path, 'RGB', '*.jpg')) 79 | pid_container = set() 80 | for img_path_RGB in img_paths_RGB: 81 | jpg_name = img_path_RGB.split('/')[-1] 82 | pid = int(jpg_name.split('_')[0][0:6]) 83 | pid_container.add(pid) 84 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 85 | 86 | data = [] 87 | for img_path_RGB in img_paths_RGB: 88 | img = [] 89 | jpg_name = img_path_RGB.split('/')[-1] 90 | img_path_NI = osp.join(dir_path, 'NI_change_ID_final', jpg_name) 91 | img_path_TI = osp.join(dir_path, 'TI_change_ID_final', jpg_name) 92 | img.append(img_path_RGB) 93 | img.append(img_path_NI) 94 | img.append(img_path_TI) 95 | pid = int(jpg_name.split('_')[0][0:6]) 96 | camid = int(jpg_name.split('_')[1][3]) 97 | camid -= 1 # index starts from 0 98 | timeid = int(jpg_name.split('_')[2]) 99 | # print(img, pid, camid, timeid) 100 | if relabel: 101 | pid = pid2label[pid] 102 | data.append((img, pid, camid, timeid)) 103 | # print(data[0:10]) 104 | return data 105 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/RGBNT201.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | from os import pardir 3 | import re 4 | import glob 5 | import os.path as osp 6 | 7 | from numpy.lib.twodim_base import tri 8 | from torchreid.data.datasets import image 9 | import warnings 10 | 11 | from ..dataset import MultiModalImageDataset 12 | 13 | 14 | class RGBNT201(MultiModalImageDataset): 15 | dataset_dir = 'RGBNT201' 16 | 17 | def __init__(self, root='', **kwargs): 18 | self.root = osp.abspath(osp.expanduser(root)) 19 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 20 | 21 | # allow alternative directory structure 22 | self.data_dir = self.dataset_dir 23 | data_dir = osp.join(self.data_dir) 24 | if osp.isdir(data_dir): 25 | self.data_dir = data_dir 26 | else: 27 | warnings.warn( 28 | 'The current data structure is deprecated.' 29 | ) 30 | 31 | 32 | self.train_dir = osp.join(self.data_dir, 'train_171') 33 | self.query_dir = osp.join(self.data_dir, 'test') 34 | self.gallery_dir = osp.join(self.data_dir, 'test') 35 | 36 | required_files = [ 37 | self.data_dir, self.train_dir, self.query_dir, self.gallery_dir 38 | ] 39 | self.check_before_run(required_files) 40 | 41 | train = self.process_dir(self.train_dir, relabel=True) 42 | query = self.process_dir(self.query_dir, relabel=False) 43 | gallery = self.process_dir(self.gallery_dir, relabel=False) 44 | 45 | super(RGBNT201, self).__init__(train, query, gallery, **kwargs) 46 | 47 | def process_dir(self, dir_path, relabel=False): 48 | img_paths_RGB = glob.glob(osp.join(dir_path, 'RGB', '*.jpg')) 49 | pid_container = set() 50 | for img_path_RGB in img_paths_RGB: 51 | jpg_name = img_path_RGB.split('\\')[-1] 52 | pid = int(jpg_name.split('_')[0][0:6]) 53 | pid_container.add(pid) 54 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 55 | 56 | data = [] 57 | for img_path_RGB in img_paths_RGB: 58 | img = [] 59 | jpg_name = img_path_RGB.split('\\')[-1] 60 | img_path_NI = osp.join(dir_path, 'NI', jpg_name) 61 | img_path_TI = osp.join(dir_path, 'TI', jpg_name) 62 | img.append(img_path_RGB) 63 | img.append(img_path_NI) 64 | img.append(img_path_TI) 65 | pid = int(jpg_name.split('_')[0][0:6]) 66 | camid = int(jpg_name.split('_')[1][3]) 67 | camid -= 1 # index starts from 0 68 | if relabel: 69 | pid = pid2label[pid] 70 | data.append((img, pid, camid)) 71 | # print("11111") 72 | return data 73 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/UAV.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | from os import pardir 3 | import re 4 | import glob 5 | import os.path as osp 6 | from torchreid import data 7 | 8 | from numpy.lib.twodim_base import tri 9 | from torchreid.data.datasets import image 10 | import warnings 11 | 12 | from ..dataset import MultiModalImageDataset 13 | 14 | 15 | class UAV(MultiModalImageDataset): 16 | dataset_dir = 'UAVdata' 17 | 18 | def __init__(self, root='', **kwargs): 19 | self.root = osp.abspath(osp.expanduser(root)) 20 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 21 | 22 | # allow alternative directory structure 23 | self.data_dir = self.dataset_dir 24 | data_dir = osp.join(self.data_dir) 25 | if osp.isdir(data_dir): 26 | self.data_dir = data_dir 27 | else: 28 | warnings.warn( 29 | 'The current data structure is deprecated.' 30 | ) 31 | 32 | self.train_dir = osp.join(self.data_dir, 'reid_bounding_box_train', 'train') 33 | self.query_dir = osp.join(self.data_dir, 'reid_bounding_box_train', 'query') 34 | self.gallery_dir = osp.join(self.data_dir, 'reid_bounding_box_train', 'gallery') 35 | 36 | required_files = [ 37 | self.data_dir, self.train_dir, self.query_dir, self.gallery_dir 38 | ] 39 | self.check_before_run(required_files) 40 | 41 | train = self.process_dir(self.train_dir, relabel=True) 42 | query = self.process_dir(self.query_dir, relabel=False) 43 | gallery = self.process_dir(self.gallery_dir, relabel=False) 44 | 45 | super(UAV, self).__init__(train, query, gallery, **kwargs) 46 | 47 | 48 | def process_dir(self, dir_path, relabel=False): 49 | img_paths = glob.glob(osp.join(dir_path, 'RGB', '*.jpg')) 50 | pid_container = set() 51 | for img_path in img_paths: 52 | jpg_name = img_path.split('\\')[-1] 53 | pid = int(jpg_name.split('.')[0][1:4]) 54 | if pid == -1: 55 | continue # junk images are just ignored 56 | pid_container.add(pid) 57 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 58 | 59 | data = [] 60 | i = 0 61 | for img_path in img_paths: 62 | jpg_name = img_path.split('\\')[-1] 63 | img = [] 64 | img.append(img_path) 65 | img.append(osp.join(dir_path, 'Gray', jpg_name)) 66 | pid = int(jpg_name.split('.')[0][1:4]) 67 | i = i + 1 68 | camid = i 69 | if pid == -1: 70 | continue # junk images are just ignored 71 | assert 0 <= pid <= 1501 # pid == 0 means background 72 | camid -= 1 # index starts from 0 73 | if relabel: 74 | pid = pid2label[pid] 75 | data.append((img, pid, camid)) 76 | return data 77 | 78 | 79 | 80 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | 3 | from .grid import GRID 4 | from .prid import PRID 5 | from .ilids import iLIDS 6 | from .viper import VIPeR 7 | from .cuhk01 import CUHK01 8 | from .cuhk02 import CUHK02 9 | from .cuhk03 import CUHK03 10 | from .msmt17 import MSMT17 11 | from .cuhksysu import CUHKSYSU 12 | from .sensereid import SenseReID 13 | from .market1501 import Market1501 14 | from .dukemtmcreid import DukeMTMCreID 15 | from .university1652 import University1652 16 | from .UAV import UAV 17 | from .RGBNT201 import RGBNT201 18 | from .market_to_RGBNT201 import market_to_RGBNT201 19 | from .AllDay import AllDay 20 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/cuhk01.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import glob 3 | import numpy as np 4 | import os.path as osp 5 | import zipfile 6 | 7 | from torchreid.utils import read_json, write_json 8 | 9 | from ..dataset import ImageDataset 10 | 11 | 12 | class CUHK01(ImageDataset): 13 | """CUHK01. 14 | 15 | Reference: 16 | Li et al. Human Reidentification with Transferred Metric Learning. ACCV 2012. 17 | 18 | URL: ``_ 19 | 20 | Dataset statistics: 21 | - identities: 971. 22 | - images: 3884. 23 | - cameras: 4. 24 | 25 | Note: CUHK01 and CUHK02 overlap. 26 | """ 27 | dataset_dir = 'cuhk01' 28 | dataset_url = None 29 | 30 | def __init__(self, root='', split_id=0, **kwargs): 31 | self.root = osp.abspath(osp.expanduser(root)) 32 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 33 | self.download_dataset(self.dataset_dir, self.dataset_url) 34 | 35 | self.zip_path = osp.join(self.dataset_dir, 'CUHK01.zip') 36 | self.campus_dir = osp.join(self.dataset_dir, 'campus') 37 | self.split_path = osp.join(self.dataset_dir, 'splits.json') 38 | 39 | self.extract_file() 40 | 41 | required_files = [self.dataset_dir, self.campus_dir] 42 | self.check_before_run(required_files) 43 | 44 | self.prepare_split() 45 | splits = read_json(self.split_path) 46 | if split_id >= len(splits): 47 | raise ValueError( 48 | 'split_id exceeds range, received {}, but expected between 0 and {}' 49 | .format(split_id, 50 | len(splits) - 1) 51 | ) 52 | split = splits[split_id] 53 | 54 | train = split['train'] 55 | query = split['query'] 56 | gallery = split['gallery'] 57 | 58 | train = [tuple(item) for item in train] 59 | query = [tuple(item) for item in query] 60 | gallery = [tuple(item) for item in gallery] 61 | 62 | super(CUHK01, self).__init__(train, query, gallery, **kwargs) 63 | 64 | def extract_file(self): 65 | if not osp.exists(self.campus_dir): 66 | print('Extracting files') 67 | zip_ref = zipfile.ZipFile(self.zip_path, 'r') 68 | zip_ref.extractall(self.dataset_dir) 69 | zip_ref.close() 70 | 71 | def prepare_split(self): 72 | """ 73 | Image name format: 0001001.png, where first four digits represent identity 74 | and last four digits represent cameras. Camera 1&2 are considered the same 75 | view and camera 3&4 are considered the same view. 76 | """ 77 | if not osp.exists(self.split_path): 78 | print('Creating 10 random splits of train ids and test ids') 79 | img_paths = sorted(glob.glob(osp.join(self.campus_dir, '*.png'))) 80 | img_list = [] 81 | pid_container = set() 82 | for img_path in img_paths: 83 | img_name = osp.basename(img_path) 84 | pid = int(img_name[:4]) - 1 85 | camid = (int(img_name[4:7]) - 1) // 2 # result is either 0 or 1 86 | img_list.append((img_path, pid, camid)) 87 | pid_container.add(pid) 88 | 89 | num_pids = len(pid_container) 90 | num_train_pids = num_pids // 2 91 | 92 | splits = [] 93 | for _ in range(10): 94 | order = np.arange(num_pids) 95 | np.random.shuffle(order) 96 | train_idxs = order[:num_train_pids] 97 | train_idxs = np.sort(train_idxs) 98 | idx2label = { 99 | idx: label 100 | for label, idx in enumerate(train_idxs) 101 | } 102 | 103 | train, test_a, test_b = [], [], [] 104 | for img_path, pid, camid in img_list: 105 | if pid in train_idxs: 106 | train.append((img_path, idx2label[pid], camid)) 107 | else: 108 | if camid == 0: 109 | test_a.append((img_path, pid, camid)) 110 | else: 111 | test_b.append((img_path, pid, camid)) 112 | 113 | # use cameraA as query and cameraB as gallery 114 | split = { 115 | 'train': train, 116 | 'query': test_a, 117 | 'gallery': test_b, 118 | 'num_train_pids': num_train_pids, 119 | 'num_query_pids': num_pids - num_train_pids, 120 | 'num_gallery_pids': num_pids - num_train_pids 121 | } 122 | splits.append(split) 123 | 124 | # use cameraB as query and cameraA as gallery 125 | split = { 126 | 'train': train, 127 | 'query': test_b, 128 | 'gallery': test_a, 129 | 'num_train_pids': num_train_pids, 130 | 'num_query_pids': num_pids - num_train_pids, 131 | 'num_gallery_pids': num_pids - num_train_pids 132 | } 133 | splits.append(split) 134 | 135 | print('Totally {} splits are created'.format(len(splits))) 136 | write_json(splits, self.split_path) 137 | print('Split file saved to {}'.format(self.split_path)) 138 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/cuhk02.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import glob 3 | import os.path as osp 4 | 5 | from ..dataset import ImageDataset 6 | 7 | 8 | class CUHK02(ImageDataset): 9 | """CUHK02. 10 | 11 | Reference: 12 | Li and Wang. Locally Aligned Feature Transforms across Views. CVPR 2013. 13 | 14 | URL: ``_ 15 | 16 | Dataset statistics: 17 | - 5 camera view pairs each with two cameras 18 | - 971, 306, 107, 193 and 239 identities from P1 - P5 19 | - totally 1,816 identities 20 | - image format is png 21 | 22 | Protocol: Use P1 - P4 for training and P5 for evaluation. 23 | 24 | Note: CUHK01 and CUHK02 overlap. 25 | """ 26 | dataset_dir = 'cuhk02' 27 | cam_pairs = ['P1', 'P2', 'P3', 'P4', 'P5'] 28 | test_cam_pair = 'P5' 29 | 30 | def __init__(self, root='', **kwargs): 31 | self.root = osp.abspath(osp.expanduser(root)) 32 | self.dataset_dir = osp.join(self.root, self.dataset_dir, 'Dataset') 33 | 34 | required_files = [self.dataset_dir] 35 | self.check_before_run(required_files) 36 | 37 | train, query, gallery = self.get_data_list() 38 | 39 | super(CUHK02, self).__init__(train, query, gallery, **kwargs) 40 | 41 | def get_data_list(self): 42 | num_train_pids, camid = 0, 0 43 | train, query, gallery = [], [], [] 44 | 45 | for cam_pair in self.cam_pairs: 46 | cam_pair_dir = osp.join(self.dataset_dir, cam_pair) 47 | 48 | cam1_dir = osp.join(cam_pair_dir, 'cam1') 49 | cam2_dir = osp.join(cam_pair_dir, 'cam2') 50 | 51 | impaths1 = glob.glob(osp.join(cam1_dir, '*.png')) 52 | impaths2 = glob.glob(osp.join(cam2_dir, '*.png')) 53 | 54 | if cam_pair == self.test_cam_pair: 55 | # add images to query 56 | for impath in impaths1: 57 | pid = osp.basename(impath).split('_')[0] 58 | pid = int(pid) 59 | query.append((impath, pid, camid)) 60 | camid += 1 61 | 62 | # add images to gallery 63 | for impath in impaths2: 64 | pid = osp.basename(impath).split('_')[0] 65 | pid = int(pid) 66 | gallery.append((impath, pid, camid)) 67 | camid += 1 68 | 69 | else: 70 | pids1 = [ 71 | osp.basename(impath).split('_')[0] for impath in impaths1 72 | ] 73 | pids2 = [ 74 | osp.basename(impath).split('_')[0] for impath in impaths2 75 | ] 76 | pids = set(pids1 + pids2) 77 | pid2label = { 78 | pid: label + num_train_pids 79 | for label, pid in enumerate(pids) 80 | } 81 | 82 | # add images to train from cam1 83 | for impath in impaths1: 84 | pid = osp.basename(impath).split('_')[0] 85 | pid = pid2label[pid] 86 | train.append((impath, pid, camid)) 87 | camid += 1 88 | 89 | # add images to train from cam2 90 | for impath in impaths2: 91 | pid = osp.basename(impath).split('_')[0] 92 | pid = pid2label[pid] 93 | train.append((impath, pid, camid)) 94 | camid += 1 95 | num_train_pids += len(pids) 96 | 97 | return train, query, gallery 98 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/cuhksysu.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import copy 3 | import glob 4 | import os.path as osp 5 | 6 | from ..dataset import ImageDataset 7 | 8 | 9 | class CUHKSYSU(ImageDataset): 10 | """CUHKSYSU. 11 | 12 | This dataset can only be used for model training. 13 | 14 | Reference: 15 | Xiao et al. End-to-end deep learning for person search. 16 | 17 | URL: ``_ 18 | 19 | Dataset statistics: 20 | - identities: 11,934 21 | - images: 34,574 22 | """ 23 | _train_only = True 24 | dataset_dir = 'cuhksysu' 25 | 26 | def __init__(self, root='', **kwargs): 27 | self.root = osp.abspath(osp.expanduser(root)) 28 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 29 | self.data_dir = osp.join(self.dataset_dir, 'cropped_images') 30 | 31 | # image name format: p11422_s16929_1.jpg 32 | train = self.process_dir(self.data_dir) 33 | query = [copy.deepcopy(train[0])] 34 | gallery = [copy.deepcopy(train[0])] 35 | 36 | super(CUHKSYSU, self).__init__(train, query, gallery, **kwargs) 37 | 38 | def process_dir(self, dirname): 39 | img_paths = glob.glob(osp.join(dirname, '*.jpg')) 40 | # num_imgs = len(img_paths) 41 | 42 | # get all identities: 43 | pid_container = set() 44 | for img_path in img_paths: 45 | img_name = osp.basename(img_path) 46 | pid = img_name.split('_')[0] 47 | pid_container.add(pid) 48 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 49 | 50 | # num_pids = len(pid_container) 51 | 52 | # extract data 53 | data = [] 54 | for img_path in img_paths: 55 | img_name = osp.basename(img_path) 56 | pid = img_name.split('_')[0] 57 | label = pid2label[pid] 58 | data.append((img_path, label, 0)) # dummy camera id 59 | 60 | return data 61 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/dukemtmcreid.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import re 3 | import glob 4 | import os.path as osp 5 | 6 | from ..dataset import ImageDataset 7 | 8 | 9 | class DukeMTMCreID(ImageDataset): 10 | """DukeMTMC-reID. 11 | 12 | Reference: 13 | - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016. 14 | - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017. 15 | 16 | URL: ``_ 17 | 18 | Dataset statistics: 19 | - identities: 1404 (train + query). 20 | - images:16522 (train) + 2228 (query) + 17661 (gallery). 21 | - cameras: 8. 22 | """ 23 | dataset_dir = 'dukemtmc-reid' 24 | dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip' 25 | 26 | def __init__(self, root='', **kwargs): 27 | self.root = osp.abspath(osp.expanduser(root)) 28 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 29 | self.download_dataset(self.dataset_dir, self.dataset_url) 30 | self.train_dir = osp.join( 31 | self.dataset_dir, 'DukeMTMC-reID/bounding_box_train' 32 | ) 33 | self.query_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/query') 34 | self.gallery_dir = osp.join( 35 | self.dataset_dir, 'DukeMTMC-reID/bounding_box_test' 36 | ) 37 | 38 | required_files = [ 39 | self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir 40 | ] 41 | self.check_before_run(required_files) 42 | 43 | train = self.process_dir(self.train_dir, relabel=True) 44 | query = self.process_dir(self.query_dir, relabel=False) 45 | gallery = self.process_dir(self.gallery_dir, relabel=False) 46 | 47 | super(DukeMTMCreID, self).__init__(train, query, gallery, **kwargs) 48 | 49 | def process_dir(self, dir_path, relabel=False): 50 | img_paths = glob.glob(osp.join(dir_path, '*.jpg')) 51 | pattern = re.compile(r'([-\d]+)_c(\d)') 52 | 53 | pid_container = set() 54 | for img_path in img_paths: 55 | pid, _ = map(int, pattern.search(img_path).groups()) 56 | pid_container.add(pid) 57 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 58 | 59 | data = [] 60 | for img_path in img_paths: 61 | pid, camid = map(int, pattern.search(img_path).groups()) 62 | assert 1 <= camid <= 8 63 | camid -= 1 # index starts from 0 64 | if relabel: 65 | pid = pid2label[pid] 66 | data.append((img_path, pid, camid)) 67 | 68 | return data 69 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/grid.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import glob 3 | import os.path as osp 4 | from scipy.io import loadmat 5 | 6 | from torchreid.utils import read_json, write_json 7 | 8 | from ..dataset import ImageDataset 9 | 10 | 11 | class GRID(ImageDataset): 12 | """GRID. 13 | 14 | Reference: 15 | Loy et al. Multi-camera activity correlation analysis. CVPR 2009. 16 | 17 | URL: ``_ 18 | 19 | Dataset statistics: 20 | - identities: 250. 21 | - images: 1275. 22 | - cameras: 8. 23 | """ 24 | dataset_dir = 'grid' 25 | dataset_url = 'http://personal.ie.cuhk.edu.hk/~ccloy/files/datasets/underground_reid.zip' 26 | _junk_pids = [0] 27 | 28 | def __init__(self, root='', split_id=0, **kwargs): 29 | self.root = osp.abspath(osp.expanduser(root)) 30 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 31 | self.download_dataset(self.dataset_dir, self.dataset_url) 32 | 33 | self.probe_path = osp.join( 34 | self.dataset_dir, 'underground_reid', 'probe' 35 | ) 36 | self.gallery_path = osp.join( 37 | self.dataset_dir, 'underground_reid', 'gallery' 38 | ) 39 | self.split_mat_path = osp.join( 40 | self.dataset_dir, 'underground_reid', 'features_and_partitions.mat' 41 | ) 42 | self.split_path = osp.join(self.dataset_dir, 'splits.json') 43 | 44 | required_files = [ 45 | self.dataset_dir, self.probe_path, self.gallery_path, 46 | self.split_mat_path 47 | ] 48 | self.check_before_run(required_files) 49 | 50 | self.prepare_split() 51 | splits = read_json(self.split_path) 52 | if split_id >= len(splits): 53 | raise ValueError( 54 | 'split_id exceeds range, received {}, ' 55 | 'but expected between 0 and {}'.format( 56 | split_id, 57 | len(splits) - 1 58 | ) 59 | ) 60 | split = splits[split_id] 61 | 62 | train = split['train'] 63 | query = split['query'] 64 | gallery = split['gallery'] 65 | 66 | train = [tuple(item) for item in train] 67 | query = [tuple(item) for item in query] 68 | gallery = [tuple(item) for item in gallery] 69 | 70 | super(GRID, self).__init__(train, query, gallery, **kwargs) 71 | 72 | def prepare_split(self): 73 | if not osp.exists(self.split_path): 74 | print('Creating 10 random splits') 75 | split_mat = loadmat(self.split_mat_path) 76 | trainIdxAll = split_mat['trainIdxAll'][0] # length = 10 77 | probe_img_paths = sorted( 78 | glob.glob(osp.join(self.probe_path, '*.jpeg')) 79 | ) 80 | gallery_img_paths = sorted( 81 | glob.glob(osp.join(self.gallery_path, '*.jpeg')) 82 | ) 83 | 84 | splits = [] 85 | for split_idx in range(10): 86 | train_idxs = trainIdxAll[split_idx][0][0][2][0].tolist() 87 | assert len(train_idxs) == 125 88 | idx2label = { 89 | idx: label 90 | for label, idx in enumerate(train_idxs) 91 | } 92 | 93 | train, query, gallery = [], [], [] 94 | 95 | # processing probe folder 96 | for img_path in probe_img_paths: 97 | img_name = osp.basename(img_path) 98 | img_idx = int(img_name.split('_')[0]) 99 | camid = int( 100 | img_name.split('_')[1] 101 | ) - 1 # index starts from 0 102 | if img_idx in train_idxs: 103 | train.append((img_path, idx2label[img_idx], camid)) 104 | else: 105 | query.append((img_path, img_idx, camid)) 106 | 107 | # process gallery folder 108 | for img_path in gallery_img_paths: 109 | img_name = osp.basename(img_path) 110 | img_idx = int(img_name.split('_')[0]) 111 | camid = int( 112 | img_name.split('_')[1] 113 | ) - 1 # index starts from 0 114 | if img_idx in train_idxs: 115 | train.append((img_path, idx2label[img_idx], camid)) 116 | else: 117 | gallery.append((img_path, img_idx, camid)) 118 | 119 | split = { 120 | 'train': train, 121 | 'query': query, 122 | 'gallery': gallery, 123 | 'num_train_pids': 125, 124 | 'num_query_pids': 125, 125 | 'num_gallery_pids': 900 126 | } 127 | splits.append(split) 128 | 129 | print('Totally {} splits are created'.format(len(splits))) 130 | write_json(splits, self.split_path) 131 | print('Split file saved to {}'.format(self.split_path)) 132 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/ilids.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import copy 3 | import glob 4 | import random 5 | import os.path as osp 6 | from collections import defaultdict 7 | 8 | from torchreid.utils import read_json, write_json 9 | 10 | from ..dataset import ImageDataset 11 | 12 | 13 | class iLIDS(ImageDataset): 14 | """QMUL-iLIDS. 15 | 16 | Reference: 17 | Zheng et al. Associating Groups of People. BMVC 2009. 18 | 19 | Dataset statistics: 20 | - identities: 119. 21 | - images: 476. 22 | - cameras: 8 (not explicitly provided). 23 | """ 24 | dataset_dir = 'ilids' 25 | dataset_url = 'http://www.eecs.qmul.ac.uk/~jason/data/i-LIDS_Pedestrian.tgz' 26 | 27 | def __init__(self, root='', split_id=0, **kwargs): 28 | self.root = osp.abspath(osp.expanduser(root)) 29 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 30 | self.download_dataset(self.dataset_dir, self.dataset_url) 31 | 32 | self.data_dir = osp.join(self.dataset_dir, 'i-LIDS_Pedestrian/Persons') 33 | self.split_path = osp.join(self.dataset_dir, 'splits.json') 34 | 35 | required_files = [self.dataset_dir, self.data_dir] 36 | self.check_before_run(required_files) 37 | 38 | self.prepare_split() 39 | splits = read_json(self.split_path) 40 | if split_id >= len(splits): 41 | raise ValueError( 42 | 'split_id exceeds range, received {}, but ' 43 | 'expected between 0 and {}'.format(split_id, 44 | len(splits) - 1) 45 | ) 46 | split = splits[split_id] 47 | 48 | train, query, gallery = self.process_split(split) 49 | 50 | super(iLIDS, self).__init__(train, query, gallery, **kwargs) 51 | 52 | def prepare_split(self): 53 | if not osp.exists(self.split_path): 54 | print('Creating splits ...') 55 | 56 | paths = glob.glob(osp.join(self.data_dir, '*.jpg')) 57 | img_names = [osp.basename(path) for path in paths] 58 | num_imgs = len(img_names) 59 | assert num_imgs == 476, 'There should be 476 images, but ' \ 60 | 'got {}, please check the data'.format(num_imgs) 61 | 62 | # store image names 63 | # image naming format: 64 | # the first four digits denote the person ID 65 | # the last four digits denote the sequence index 66 | pid_dict = defaultdict(list) 67 | for img_name in img_names: 68 | pid = int(img_name[:4]) 69 | pid_dict[pid].append(img_name) 70 | pids = list(pid_dict.keys()) 71 | num_pids = len(pids) 72 | assert num_pids == 119, 'There should be 119 identities, ' \ 73 | 'but got {}, please check the data'.format(num_pids) 74 | 75 | num_train_pids = int(num_pids * 0.5) 76 | 77 | splits = [] 78 | for _ in range(10): 79 | # randomly choose num_train_pids train IDs and the rest for test IDs 80 | pids_copy = copy.deepcopy(pids) 81 | random.shuffle(pids_copy) 82 | train_pids = pids_copy[:num_train_pids] 83 | test_pids = pids_copy[num_train_pids:] 84 | 85 | train = [] 86 | query = [] 87 | gallery = [] 88 | 89 | # for train IDs, all images are used in the train set. 90 | for pid in train_pids: 91 | img_names = pid_dict[pid] 92 | train.extend(img_names) 93 | 94 | # for each test ID, randomly choose two images, one for 95 | # query and the other one for gallery. 96 | for pid in test_pids: 97 | img_names = pid_dict[pid] 98 | samples = random.sample(img_names, 2) 99 | query.append(samples[0]) 100 | gallery.append(samples[1]) 101 | 102 | split = {'train': train, 'query': query, 'gallery': gallery} 103 | splits.append(split) 104 | 105 | print('Totally {} splits are created'.format(len(splits))) 106 | write_json(splits, self.split_path) 107 | print('Split file is saved to {}'.format(self.split_path)) 108 | 109 | def get_pid2label(self, img_names): 110 | pid_container = set() 111 | for img_name in img_names: 112 | pid = int(img_name[:4]) 113 | pid_container.add(pid) 114 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 115 | return pid2label 116 | 117 | def parse_img_names(self, img_names, pid2label=None): 118 | data = [] 119 | 120 | for img_name in img_names: 121 | pid = int(img_name[:4]) 122 | if pid2label is not None: 123 | pid = pid2label[pid] 124 | camid = int(img_name[4:7]) - 1 # 0-based 125 | img_path = osp.join(self.data_dir, img_name) 126 | data.append((img_path, pid, camid)) 127 | 128 | return data 129 | 130 | def process_split(self, split): 131 | train_pid2label = self.get_pid2label(split['train']) 132 | train = self.parse_img_names(split['train'], train_pid2label) 133 | query = self.parse_img_names(split['query']) 134 | gallery = self.parse_img_names(split['gallery']) 135 | return train, query, gallery 136 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/market1501.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | from os import pardir 3 | import re 4 | import glob 5 | import os.path as osp 6 | import warnings 7 | 8 | from ..dataset import ImageDataset 9 | 10 | 11 | class Market1501(ImageDataset): 12 | """Market1501. 13 | 14 | Reference: 15 | Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015. 16 | 17 | URL: ``_ 18 | 19 | Dataset statistics: 20 | - identities: 1501 (+1 for background). 21 | - images: 12936 (train) + 3368 (query) + 15913 (gallery). 22 | """ 23 | _junk_pids = [0, -1] 24 | dataset_dir = 'E:\datasetRelated\market1501' 25 | dataset_url = 'http://188.138.127.15:81/Datasets/Market-1501-v15.09.15.zip' 26 | 27 | def __init__(self, root='', market1501_500k=False, **kwargs): 28 | self.root = osp.abspath(osp.expanduser(root)) 29 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 30 | self.download_dataset(self.dataset_dir, self.dataset_url) 31 | 32 | # allow alternative directory structure 33 | self.data_dir = self.dataset_dir 34 | data_dir = osp.join(self.data_dir, 'Market-1501-v15.09.15') 35 | if osp.isdir(data_dir): 36 | self.data_dir = data_dir 37 | else: 38 | warnings.warn( 39 | 'The current data structure is deprecated. Please ' 40 | 'put data folders such as "bounding_box_train" under ' 41 | '"Market-1501-v15.09.15".' 42 | ) 43 | 44 | self.train_dir = osp.join(self.data_dir, 'bounding_box_train') 45 | self.query_dir = osp.join(self.data_dir, 'query') 46 | self.gallery_dir = osp.join(self.data_dir, 'bounding_box_test') 47 | self.extra_gallery_dir = osp.join(self.data_dir, 'images') 48 | self.market1501_500k = market1501_500k 49 | 50 | required_files = [ 51 | self.data_dir, self.train_dir, self.query_dir, self.gallery_dir 52 | ] 53 | if self.market1501_500k: 54 | required_files.append(self.extra_gallery_dir) 55 | self.check_before_run(required_files) 56 | 57 | train = self.process_dir(self.train_dir, relabel=True) 58 | query = self.process_dir(self.query_dir, relabel=False) 59 | gallery = self.process_dir(self.gallery_dir, relabel=False) 60 | if self.market1501_500k: 61 | gallery += self.process_dir(self.extra_gallery_dir, relabel=False) 62 | 63 | super(Market1501, self).__init__(train, query, gallery, **kwargs) 64 | 65 | def process_dir(self, dir_path, relabel=False): 66 | img_paths = glob.glob(osp.join(dir_path, '*.jpg')) 67 | pattern = re.compile(r'([-\d]+)_c(\d)') 68 | 69 | pid_container = set() 70 | for img_path in img_paths: 71 | pid, _ = map(int, pattern.search(img_path).groups()) 72 | if pid == -1: 73 | continue # junk images are just ignored 74 | pid_container.add(pid) 75 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 76 | 77 | data = [] 78 | for img_path in img_paths: 79 | pid, camid = map(int, pattern.search(img_path).groups()) 80 | if pid == -1: 81 | continue # junk images are just ignored 82 | assert 0 <= pid <= 1501 # pid == 0 means background 83 | assert 1 <= camid <= 6 84 | camid -= 1 # index starts from 0 85 | if relabel: 86 | pid = pid2label[pid] 87 | data.append((img_path, pid, camid)) 88 | return data 89 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/market_to_RGBNT201.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | from os import pardir 3 | import re 4 | import glob 5 | import os.path as osp 6 | 7 | from numpy.lib.twodim_base import tri 8 | from torchreid.data.datasets import image 9 | import warnings 10 | 11 | from ..dataset import MultiModalImageDataset 12 | 13 | 14 | class market_to_RGBNT201(MultiModalImageDataset): 15 | dataset_dir = 'E:\datasetRelated\market1501_to_RGBNT201_dark' 16 | 17 | def __init__(self, root='', **kwargs): 18 | self.root = osp.abspath(osp.expanduser(root)) 19 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 20 | 21 | # allow alternative directory structure 22 | self.data_dir = self.dataset_dir 23 | data_dir = osp.join(self.data_dir) 24 | if osp.isdir(data_dir): 25 | self.data_dir = data_dir 26 | else: 27 | warnings.warn( 28 | 'The current data structure is deprecated.' 29 | ) 30 | 31 | 32 | self.train_dir = osp.join(self.data_dir, 'train') 33 | self.query_dir = osp.join(self.data_dir, 'query') 34 | self.gallery_dir = osp.join(self.data_dir, 'gallery') 35 | 36 | required_files = [ 37 | self.data_dir, self.train_dir, self.query_dir, self.gallery_dir 38 | ] 39 | self.check_before_run(required_files) 40 | 41 | train = self.process_dir(self.train_dir, relabel=True) 42 | query = self.process_dir(self.query_dir, relabel=False) 43 | gallery = self.process_dir(self.gallery_dir, relabel=False) 44 | 45 | super(market_to_RGBNT201, self).__init__(train, query, gallery, **kwargs) 46 | 47 | def process_dir(self, dir_path, relabel=False): 48 | img_paths_RGB = glob.glob(osp.join(dir_path, 'RGB', '*.jpg')) 49 | pid_container = set() 50 | for img_path_RGB in img_paths_RGB: 51 | jpg_name = img_path_RGB.split('\\')[-1] 52 | pid = int(jpg_name.split('_')[0]) 53 | if pid == -1: 54 | continue # junk images are just ignored 55 | pid_container.add(pid) 56 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 57 | 58 | data = [] 59 | for img_path_RGB in img_paths_RGB: 60 | img = [] 61 | jpg_name = img_path_RGB.split('\\')[-1] 62 | img_path_NI = osp.join(dir_path, 'NI', jpg_name) 63 | img_path_TI = osp.join(dir_path, 'TI', jpg_name) 64 | img.append(img_path_RGB) 65 | img.append(img_path_NI) 66 | img.append(img_path_TI) 67 | pid = int(jpg_name.split('_')[0]) 68 | camid = int(jpg_name.split('_')[1][1]) 69 | if pid == -1: 70 | continue # junk images are just ignored 71 | assert 0 <= pid <= 1501 # pid == 0 means background 72 | assert 1 <= camid <= 6 73 | camid -= 1 # index starts from 0 74 | if relabel: 75 | pid = pid2label[pid] 76 | data.append((img, pid, camid)) 77 | # print(data[0], '\n', data[1000], '\n', data[2000]) 78 | return data 79 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/msmt17.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import os.path as osp 3 | 4 | from ..dataset import ImageDataset 5 | 6 | # Log 7 | # 22.01.2019 8 | # - add v2 9 | # - v1 and v2 differ in dir names 10 | # - note that faces in v2 are blurred 11 | TRAIN_DIR_KEY = 'train_dir' 12 | TEST_DIR_KEY = 'test_dir' 13 | VERSION_DICT = { 14 | 'MSMT17_V1': { 15 | TRAIN_DIR_KEY: 'train', 16 | TEST_DIR_KEY: 'test', 17 | }, 18 | 'MSMT17_V2': { 19 | TRAIN_DIR_KEY: 'mask_train_v2', 20 | TEST_DIR_KEY: 'mask_test_v2', 21 | } 22 | } 23 | 24 | 25 | class MSMT17(ImageDataset): 26 | """MSMT17. 27 | 28 | Reference: 29 | Wei et al. Person Transfer GAN to Bridge Domain Gap for Person Re-Identification. CVPR 2018. 30 | 31 | URL: ``_ 32 | 33 | Dataset statistics: 34 | - identities: 4101. 35 | - images: 32621 (train) + 11659 (query) + 82161 (gallery). 36 | - cameras: 15. 37 | """ 38 | dataset_dir = 'msmt17' 39 | dataset_url = None 40 | 41 | def __init__(self, root='', **kwargs): 42 | self.root = osp.abspath(osp.expanduser(root)) 43 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 44 | self.download_dataset(self.dataset_dir, self.dataset_url) 45 | 46 | has_main_dir = False 47 | for main_dir in VERSION_DICT: 48 | if osp.exists(osp.join(self.dataset_dir, main_dir)): 49 | train_dir = VERSION_DICT[main_dir][TRAIN_DIR_KEY] 50 | test_dir = VERSION_DICT[main_dir][TEST_DIR_KEY] 51 | has_main_dir = True 52 | break 53 | assert has_main_dir, 'Dataset folder not found' 54 | 55 | self.train_dir = osp.join(self.dataset_dir, main_dir, train_dir) 56 | self.test_dir = osp.join(self.dataset_dir, main_dir, test_dir) 57 | self.list_train_path = osp.join( 58 | self.dataset_dir, main_dir, 'list_train.txt' 59 | ) 60 | self.list_val_path = osp.join( 61 | self.dataset_dir, main_dir, 'list_val.txt' 62 | ) 63 | self.list_query_path = osp.join( 64 | self.dataset_dir, main_dir, 'list_query.txt' 65 | ) 66 | self.list_gallery_path = osp.join( 67 | self.dataset_dir, main_dir, 'list_gallery.txt' 68 | ) 69 | 70 | required_files = [self.dataset_dir, self.train_dir, self.test_dir] 71 | self.check_before_run(required_files) 72 | 73 | train = self.process_dir(self.train_dir, self.list_train_path) 74 | val = self.process_dir(self.train_dir, self.list_val_path) 75 | query = self.process_dir(self.test_dir, self.list_query_path) 76 | gallery = self.process_dir(self.test_dir, self.list_gallery_path) 77 | 78 | # Note: to fairly compare with published methods on the conventional ReID setting, 79 | # do not add val images to the training set. 80 | if 'combineall' in kwargs and kwargs['combineall']: 81 | train += val 82 | 83 | super(MSMT17, self).__init__(train, query, gallery, **kwargs) 84 | 85 | def process_dir(self, dir_path, list_path): 86 | with open(list_path, 'r') as txt: 87 | lines = txt.readlines() 88 | 89 | data = [] 90 | 91 | for img_idx, img_info in enumerate(lines): 92 | img_path, pid = img_info.split(' ') 93 | pid = int(pid) # no need to relabel 94 | camid = int(img_path.split('_')[2]) - 1 # index starts from 0 95 | img_path = osp.join(dir_path, img_path) 96 | data.append((img_path, pid, camid)) 97 | 98 | return data 99 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/prid.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import random 3 | import os.path as osp 4 | 5 | from torchreid.utils import read_json, write_json 6 | 7 | from ..dataset import ImageDataset 8 | 9 | 10 | class PRID(ImageDataset): 11 | """PRID (single-shot version of prid-2011) 12 | 13 | Reference: 14 | Hirzer et al. Person Re-Identification by Descriptive and Discriminative 15 | Classification. SCIA 2011. 16 | 17 | URL: ``_ 18 | 19 | Dataset statistics: 20 | - Two views. 21 | - View A captures 385 identities. 22 | - View B captures 749 identities. 23 | - 200 identities appear in both views (index starts from 1 to 200). 24 | """ 25 | dataset_dir = 'prid2011' 26 | dataset_url = None 27 | _junk_pids = list(range(201, 750)) 28 | 29 | def __init__(self, root='', split_id=0, **kwargs): 30 | self.root = osp.abspath(osp.expanduser(root)) 31 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 32 | self.download_dataset(self.dataset_dir, self.dataset_url) 33 | 34 | self.cam_a_dir = osp.join( 35 | self.dataset_dir, 'prid_2011', 'single_shot', 'cam_a' 36 | ) 37 | self.cam_b_dir = osp.join( 38 | self.dataset_dir, 'prid_2011', 'single_shot', 'cam_b' 39 | ) 40 | self.split_path = osp.join(self.dataset_dir, 'splits_single_shot.json') 41 | 42 | required_files = [self.dataset_dir, self.cam_a_dir, self.cam_b_dir] 43 | self.check_before_run(required_files) 44 | 45 | self.prepare_split() 46 | splits = read_json(self.split_path) 47 | if split_id >= len(splits): 48 | raise ValueError( 49 | 'split_id exceeds range, received {}, but expected between 0 and {}' 50 | .format(split_id, 51 | len(splits) - 1) 52 | ) 53 | split = splits[split_id] 54 | 55 | train, query, gallery = self.process_split(split) 56 | 57 | super(PRID, self).__init__(train, query, gallery, **kwargs) 58 | 59 | def prepare_split(self): 60 | if not osp.exists(self.split_path): 61 | print('Creating splits ...') 62 | 63 | splits = [] 64 | for _ in range(10): 65 | # randomly sample 100 IDs for train and use the rest 100 IDs for test 66 | # (note: there are only 200 IDs appearing in both views) 67 | pids = [i for i in range(1, 201)] 68 | train_pids = random.sample(pids, 100) 69 | train_pids.sort() 70 | test_pids = [i for i in pids if i not in train_pids] 71 | split = {'train': train_pids, 'test': test_pids} 72 | splits.append(split) 73 | 74 | print('Totally {} splits are created'.format(len(splits))) 75 | write_json(splits, self.split_path) 76 | print('Split file is saved to {}'.format(self.split_path)) 77 | 78 | def process_split(self, split): 79 | train_pids = split['train'] 80 | test_pids = split['test'] 81 | 82 | train_pid2label = {pid: label for label, pid in enumerate(train_pids)} 83 | 84 | # train 85 | train = [] 86 | for pid in train_pids: 87 | img_name = 'person_' + str(pid).zfill(4) + '.png' 88 | pid = train_pid2label[pid] 89 | img_a_path = osp.join(self.cam_a_dir, img_name) 90 | train.append((img_a_path, pid, 0)) 91 | img_b_path = osp.join(self.cam_b_dir, img_name) 92 | train.append((img_b_path, pid, 1)) 93 | 94 | # query and gallery 95 | query, gallery = [], [] 96 | for pid in test_pids: 97 | img_name = 'person_' + str(pid).zfill(4) + '.png' 98 | img_a_path = osp.join(self.cam_a_dir, img_name) 99 | query.append((img_a_path, pid, 0)) 100 | img_b_path = osp.join(self.cam_b_dir, img_name) 101 | gallery.append((img_b_path, pid, 1)) 102 | for pid in range(201, 750): 103 | img_name = 'person_' + str(pid).zfill(4) + '.png' 104 | img_b_path = osp.join(self.cam_b_dir, img_name) 105 | gallery.append((img_b_path, pid, 1)) 106 | 107 | return train, query, gallery 108 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/sensereid.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import copy 3 | import glob 4 | import os.path as osp 5 | 6 | from ..dataset import ImageDataset 7 | 8 | 9 | class SenseReID(ImageDataset): 10 | """SenseReID. 11 | 12 | This dataset is used for test purpose only. 13 | 14 | Reference: 15 | Zhao et al. Spindle Net: Person Re-identification with Human Body 16 | Region Guided Feature Decomposition and Fusion. CVPR 2017. 17 | 18 | URL: ``_ 19 | 20 | Dataset statistics: 21 | - query: 522 ids, 1040 images. 22 | - gallery: 1717 ids, 3388 images. 23 | """ 24 | dataset_dir = 'sensereid' 25 | dataset_url = None 26 | 27 | def __init__(self, root='', **kwargs): 28 | self.root = osp.abspath(osp.expanduser(root)) 29 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 30 | self.download_dataset(self.dataset_dir, self.dataset_url) 31 | 32 | self.query_dir = osp.join(self.dataset_dir, 'SenseReID', 'test_probe') 33 | self.gallery_dir = osp.join( 34 | self.dataset_dir, 'SenseReID', 'test_gallery' 35 | ) 36 | 37 | required_files = [self.dataset_dir, self.query_dir, self.gallery_dir] 38 | self.check_before_run(required_files) 39 | 40 | query = self.process_dir(self.query_dir) 41 | gallery = self.process_dir(self.gallery_dir) 42 | 43 | # relabel 44 | g_pids = set() 45 | for _, pid, _ in gallery: 46 | g_pids.add(pid) 47 | pid2label = {pid: i for i, pid in enumerate(g_pids)} 48 | 49 | query = [ 50 | (img_path, pid2label[pid], camid) for img_path, pid, camid in query 51 | ] 52 | gallery = [ 53 | (img_path, pid2label[pid], camid) 54 | for img_path, pid, camid in gallery 55 | ] 56 | train = copy.deepcopy(query) + copy.deepcopy(gallery) # dummy variable 57 | 58 | super(SenseReID, self).__init__(train, query, gallery, **kwargs) 59 | 60 | def process_dir(self, dir_path): 61 | img_paths = glob.glob(osp.join(dir_path, '*.jpg')) 62 | data = [] 63 | 64 | for img_path in img_paths: 65 | img_name = osp.splitext(osp.basename(img_path))[0] 66 | pid, camid = img_name.split('_') 67 | pid, camid = int(pid), int(camid) 68 | data.append((img_path, pid, camid)) 69 | 70 | return data 71 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/university1652.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import os 3 | import glob 4 | import os.path as osp 5 | import gdown 6 | 7 | from ..dataset import ImageDataset 8 | 9 | 10 | class University1652(ImageDataset): 11 | """University-1652. 12 | 13 | Reference: 14 | - Zheng et al. University-1652: A Multi-view Multi-source Benchmark for Drone-based Geo-localization. ACM MM 2020. 15 | 16 | URL: ``_ 17 | OneDrive: 18 | https://studentutsedu-my.sharepoint.com/:u:/g/personal/12639605_student_uts_edu_au/Ecrz6xK-PcdCjFdpNb0T0s8B_9J5ynaUy3q63_XumjJyrA?e=z4hpcz 19 | [Backup] GoogleDrive: 20 | https://drive.google.com/file/d/1iVnP4gjw-iHXa0KerZQ1IfIO0i1jADsR/view?usp=sharing 21 | [Backup] Baidu Yun: 22 | https://pan.baidu.com/s/1H_wBnWwikKbaBY1pMPjoqQ password: hrqp 23 | 24 | Dataset statistics: 25 | - buildings: 1652 (train + query). 26 | - The dataset split is as follows: 27 | | Split | #imgs | #buildings | #universities| 28 | | -------- | ----- | ----| ----| 29 | | Training | 50,218 | 701 | 33 | 30 | | Query_drone | 37,855 | 701 | 39 | 31 | | Query_satellite | 701 | 701 | 39| 32 | | Query_ground | 2,579 | 701 | 39| 33 | | Gallery_drone | 51,355 | 951 | 39| 34 | | Gallery_satellite | 951 | 951 | 39| 35 | | Gallery_ground | 2,921 | 793 | 39| 36 | - cameras: None. 37 | 38 | datamanager = torchreid.data.ImageDataManager( 39 | root='reid-data', 40 | sources='university1652', 41 | targets='university1652', 42 | height=256, 43 | width=256, 44 | batch_size_train=32, 45 | batch_size_test=100, 46 | transforms=['random_flip', 'random_crop'] 47 | ) 48 | """ 49 | dataset_dir = 'university1652' 50 | dataset_url = 'https://drive.google.com/uc?id=1iVnP4gjw-iHXa0KerZQ1IfIO0i1jADsR' 51 | 52 | def __init__(self, root='', **kwargs): 53 | self.root = osp.abspath(osp.expanduser(root)) 54 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 55 | print(self.dataset_dir) 56 | if not os.path.isdir(self.dataset_dir): 57 | os.mkdir(self.dataset_dir) 58 | gdown.download( 59 | self.dataset_url, self.dataset_dir + 'data.zip', quiet=False 60 | ) 61 | os.system('unzip %s' % (self.dataset_dir + 'data.zip')) 62 | self.train_dir = osp.join( 63 | self.dataset_dir, 'University-Release/train/' 64 | ) 65 | self.query_dir = osp.join( 66 | self.dataset_dir, 'University-Release/test/query_drone' 67 | ) 68 | self.gallery_dir = osp.join( 69 | self.dataset_dir, 'University-Release/test/gallery_satellite' 70 | ) 71 | 72 | required_files = [ 73 | self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir 74 | ] 75 | self.check_before_run(required_files) 76 | 77 | self.fake_camid = 0 78 | train = self.process_dir(self.train_dir, relabel=True, train=True) 79 | query = self.process_dir(self.query_dir, relabel=False) 80 | gallery = self.process_dir(self.gallery_dir, relabel=False) 81 | 82 | super(University1652, self).__init__(train, query, gallery, **kwargs) 83 | 84 | def process_dir(self, dir_path, relabel=False, train=False): 85 | IMG_EXTENSIONS = ( 86 | '.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', 87 | '.webp' 88 | ) 89 | if train: 90 | img_paths = glob.glob(osp.join(dir_path, '*/*/*')) 91 | else: 92 | img_paths = glob.glob(osp.join(dir_path, '*/*')) 93 | pid_container = set() 94 | for img_path in img_paths: 95 | if not img_path.lower().endswith(IMG_EXTENSIONS): 96 | continue 97 | pid = int(os.path.basename(os.path.dirname(img_path))) 98 | pid_container.add(pid) 99 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 100 | data = [] 101 | # no camera for university 102 | for img_path in img_paths: 103 | if not img_path.lower().endswith(IMG_EXTENSIONS): 104 | continue 105 | pid = int(os.path.basename(os.path.dirname(img_path))) 106 | if relabel: 107 | pid = pid2label[pid] 108 | data.append((img_path, pid, self.fake_camid)) 109 | self.fake_camid += 1 110 | return data 111 | -------------------------------------------------------------------------------- /torchreid/data/datasets/image/viper.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import glob 3 | import numpy as np 4 | import os.path as osp 5 | 6 | from torchreid.utils import read_json, write_json 7 | 8 | from ..dataset import ImageDataset 9 | 10 | 11 | class VIPeR(ImageDataset): 12 | """VIPeR. 13 | 14 | Reference: 15 | Gray et al. Evaluating appearance models for recognition, reacquisition, and tracking. PETS 2007. 16 | 17 | URL: ``_ 18 | 19 | Dataset statistics: 20 | - identities: 632. 21 | - images: 632 x 2 = 1264. 22 | - cameras: 2. 23 | """ 24 | dataset_dir = 'viper' 25 | dataset_url = 'http://users.soe.ucsc.edu/~manduchi/VIPeR.v1.0.zip' 26 | 27 | def __init__(self, root='', split_id=0, **kwargs): 28 | self.root = osp.abspath(osp.expanduser(root)) 29 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 30 | self.download_dataset(self.dataset_dir, self.dataset_url) 31 | 32 | self.cam_a_dir = osp.join(self.dataset_dir, 'VIPeR', 'cam_a') 33 | self.cam_b_dir = osp.join(self.dataset_dir, 'VIPeR', 'cam_b') 34 | self.split_path = osp.join(self.dataset_dir, 'splits.json') 35 | 36 | required_files = [self.dataset_dir, self.cam_a_dir, self.cam_b_dir] 37 | self.check_before_run(required_files) 38 | 39 | self.prepare_split() 40 | splits = read_json(self.split_path) 41 | if split_id >= len(splits): 42 | raise ValueError( 43 | 'split_id exceeds range, received {}, ' 44 | 'but expected between 0 and {}'.format( 45 | split_id, 46 | len(splits) - 1 47 | ) 48 | ) 49 | split = splits[split_id] 50 | 51 | train = split['train'] 52 | query = split['query'] # query and gallery share the same images 53 | gallery = split['gallery'] 54 | 55 | train = [tuple(item) for item in train] 56 | query = [tuple(item) for item in query] 57 | gallery = [tuple(item) for item in gallery] 58 | 59 | super(VIPeR, self).__init__(train, query, gallery, **kwargs) 60 | 61 | def prepare_split(self): 62 | if not osp.exists(self.split_path): 63 | print('Creating 10 random splits of train ids and test ids') 64 | 65 | cam_a_imgs = sorted(glob.glob(osp.join(self.cam_a_dir, '*.bmp'))) 66 | cam_b_imgs = sorted(glob.glob(osp.join(self.cam_b_dir, '*.bmp'))) 67 | assert len(cam_a_imgs) == len(cam_b_imgs) 68 | num_pids = len(cam_a_imgs) 69 | print('Number of identities: {}'.format(num_pids)) 70 | num_train_pids = num_pids // 2 71 | """ 72 | In total, there will be 20 splits because each random split creates two 73 | sub-splits, one using cameraA as query and cameraB as gallery 74 | while the other using cameraB as query and cameraA as gallery. 75 | Therefore, results should be averaged over 20 splits (split_id=0~19). 76 | 77 | In practice, a model trained on split_id=0 can be applied to split_id=0&1 78 | as split_id=0&1 share the same training data (so on and so forth). 79 | """ 80 | splits = [] 81 | for _ in range(10): 82 | order = np.arange(num_pids) 83 | np.random.shuffle(order) 84 | train_idxs = order[:num_train_pids] 85 | test_idxs = order[num_train_pids:] 86 | assert not bool(set(train_idxs) & set(test_idxs)), \ 87 | 'Error: train and test overlap' 88 | 89 | train = [] 90 | for pid, idx in enumerate(train_idxs): 91 | cam_a_img = cam_a_imgs[idx] 92 | cam_b_img = cam_b_imgs[idx] 93 | train.append((cam_a_img, pid, 0)) 94 | train.append((cam_b_img, pid, 1)) 95 | 96 | test_a = [] 97 | test_b = [] 98 | for pid, idx in enumerate(test_idxs): 99 | cam_a_img = cam_a_imgs[idx] 100 | cam_b_img = cam_b_imgs[idx] 101 | test_a.append((cam_a_img, pid, 0)) 102 | test_b.append((cam_b_img, pid, 1)) 103 | 104 | # use cameraA as query and cameraB as gallery 105 | split = { 106 | 'train': train, 107 | 'query': test_a, 108 | 'gallery': test_b, 109 | 'num_train_pids': num_train_pids, 110 | 'num_query_pids': num_pids - num_train_pids, 111 | 'num_gallery_pids': num_pids - num_train_pids 112 | } 113 | splits.append(split) 114 | 115 | # use cameraB as query and cameraA as gallery 116 | split = { 117 | 'train': train, 118 | 'query': test_b, 119 | 'gallery': test_a, 120 | 'num_train_pids': num_train_pids, 121 | 'num_query_pids': num_pids - num_train_pids, 122 | 'num_gallery_pids': num_pids - num_train_pids 123 | } 124 | splits.append(split) 125 | 126 | print('Totally {} splits are created'.format(len(splits))) 127 | write_json(splits, self.split_path) 128 | print('Split file saved to {}'.format(self.split_path)) 129 | -------------------------------------------------------------------------------- /torchreid/data/datasets/video/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | 3 | from .mars import Mars 4 | from .ilidsvid import iLIDSVID 5 | from .prid2011 import PRID2011 6 | from .dukemtmcvidreid import DukeMTMCVidReID 7 | -------------------------------------------------------------------------------- /torchreid/data/datasets/video/dukemtmcvidreid.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import glob 3 | import os.path as osp 4 | import warnings 5 | 6 | from torchreid.utils import read_json, write_json 7 | 8 | from ..dataset import VideoDataset 9 | 10 | 11 | class DukeMTMCVidReID(VideoDataset): 12 | """DukeMTMCVidReID. 13 | 14 | Reference: 15 | - Ristani et al. Performance Measures and a Data Set for Multi-Target, 16 | Multi-Camera Tracking. ECCVW 2016. 17 | - Wu et al. Exploit the Unknown Gradually: One-Shot Video-Based Person 18 | Re-Identification by Stepwise Learning. CVPR 2018. 19 | 20 | URL: ``_ 21 | 22 | Dataset statistics: 23 | - identities: 702 (train) + 702 (test). 24 | - tracklets: 2196 (train) + 2636 (test). 25 | """ 26 | dataset_dir = 'dukemtmc-vidreid' 27 | dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-VideoReID.zip' 28 | 29 | def __init__(self, root='', min_seq_len=0, **kwargs): 30 | self.root = osp.abspath(osp.expanduser(root)) 31 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 32 | self.download_dataset(self.dataset_dir, self.dataset_url) 33 | 34 | self.train_dir = osp.join(self.dataset_dir, 'DukeMTMC-VideoReID/train') 35 | self.query_dir = osp.join(self.dataset_dir, 'DukeMTMC-VideoReID/query') 36 | self.gallery_dir = osp.join( 37 | self.dataset_dir, 'DukeMTMC-VideoReID/gallery' 38 | ) 39 | self.split_train_json_path = osp.join( 40 | self.dataset_dir, 'split_train.json' 41 | ) 42 | self.split_query_json_path = osp.join( 43 | self.dataset_dir, 'split_query.json' 44 | ) 45 | self.split_gallery_json_path = osp.join( 46 | self.dataset_dir, 'split_gallery.json' 47 | ) 48 | self.min_seq_len = min_seq_len 49 | 50 | required_files = [ 51 | self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir 52 | ] 53 | self.check_before_run(required_files) 54 | 55 | train = self.process_dir( 56 | self.train_dir, self.split_train_json_path, relabel=True 57 | ) 58 | query = self.process_dir( 59 | self.query_dir, self.split_query_json_path, relabel=False 60 | ) 61 | gallery = self.process_dir( 62 | self.gallery_dir, self.split_gallery_json_path, relabel=False 63 | ) 64 | 65 | super(DukeMTMCVidReID, self).__init__(train, query, gallery, **kwargs) 66 | 67 | def process_dir(self, dir_path, json_path, relabel): 68 | if osp.exists(json_path): 69 | split = read_json(json_path) 70 | return split['tracklets'] 71 | 72 | print('=> Generating split json file (** this might take a while **)') 73 | pdirs = glob.glob(osp.join(dir_path, '*')) # avoid .DS_Store 74 | print( 75 | 'Processing "{}" with {} person identities'.format( 76 | dir_path, len(pdirs) 77 | ) 78 | ) 79 | 80 | pid_container = set() 81 | for pdir in pdirs: 82 | pid = int(osp.basename(pdir)) 83 | pid_container.add(pid) 84 | pid2label = {pid: label for label, pid in enumerate(pid_container)} 85 | 86 | tracklets = [] 87 | for pdir in pdirs: 88 | pid = int(osp.basename(pdir)) 89 | if relabel: 90 | pid = pid2label[pid] 91 | tdirs = glob.glob(osp.join(pdir, '*')) 92 | for tdir in tdirs: 93 | raw_img_paths = glob.glob(osp.join(tdir, '*.jpg')) 94 | num_imgs = len(raw_img_paths) 95 | 96 | if num_imgs < self.min_seq_len: 97 | continue 98 | 99 | img_paths = [] 100 | for img_idx in range(num_imgs): 101 | # some tracklet starts from 0002 instead of 0001 102 | img_idx_name = 'F' + str(img_idx + 1).zfill(4) 103 | res = glob.glob( 104 | osp.join(tdir, '*' + img_idx_name + '*.jpg') 105 | ) 106 | if len(res) == 0: 107 | warnings.warn( 108 | 'Index name {} in {} is missing, skip'.format( 109 | img_idx_name, tdir 110 | ) 111 | ) 112 | continue 113 | img_paths.append(res[0]) 114 | img_name = osp.basename(img_paths[0]) 115 | if img_name.find('_') == -1: 116 | # old naming format: 0001C6F0099X30823.jpg 117 | camid = int(img_name[5]) - 1 118 | else: 119 | # new naming format: 0001_C6_F0099_X30823.jpg 120 | camid = int(img_name[6]) - 1 121 | img_paths = tuple(img_paths) 122 | tracklets.append((img_paths, pid, camid)) 123 | 124 | print('Saving split to {}'.format(json_path)) 125 | split_dict = {'tracklets': tracklets} 126 | write_json(split_dict, json_path) 127 | 128 | return tracklets 129 | -------------------------------------------------------------------------------- /torchreid/data/datasets/video/ilidsvid.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import glob 3 | import os.path as osp 4 | from scipy.io import loadmat 5 | 6 | from torchreid.utils import read_json, write_json 7 | 8 | from ..dataset import VideoDataset 9 | 10 | 11 | class iLIDSVID(VideoDataset): 12 | """iLIDS-VID. 13 | 14 | Reference: 15 | Wang et al. Person Re-Identification by Video Ranking. ECCV 2014. 16 | 17 | URL: ``_ 18 | 19 | Dataset statistics: 20 | - identities: 300. 21 | - tracklets: 600. 22 | - cameras: 2. 23 | """ 24 | dataset_dir = 'ilids-vid' 25 | dataset_url = 'http://www.eecs.qmul.ac.uk/~xiatian/iLIDS-VID/iLIDS-VID.tar' 26 | 27 | def __init__(self, root='', split_id=0, **kwargs): 28 | self.root = osp.abspath(osp.expanduser(root)) 29 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 30 | self.download_dataset(self.dataset_dir, self.dataset_url) 31 | 32 | self.data_dir = osp.join(self.dataset_dir, 'i-LIDS-VID') 33 | self.split_dir = osp.join(self.dataset_dir, 'train-test people splits') 34 | self.split_mat_path = osp.join( 35 | self.split_dir, 'train_test_splits_ilidsvid.mat' 36 | ) 37 | self.split_path = osp.join(self.dataset_dir, 'splits.json') 38 | self.cam_1_path = osp.join( 39 | self.dataset_dir, 'i-LIDS-VID/sequences/cam1' 40 | ) 41 | self.cam_2_path = osp.join( 42 | self.dataset_dir, 'i-LIDS-VID/sequences/cam2' 43 | ) 44 | 45 | required_files = [self.dataset_dir, self.data_dir, self.split_dir] 46 | self.check_before_run(required_files) 47 | 48 | self.prepare_split() 49 | splits = read_json(self.split_path) 50 | if split_id >= len(splits): 51 | raise ValueError( 52 | 'split_id exceeds range, received {}, but expected between 0 and {}' 53 | .format(split_id, 54 | len(splits) - 1) 55 | ) 56 | split = splits[split_id] 57 | train_dirs, test_dirs = split['train'], split['test'] 58 | 59 | train = self.process_data(train_dirs, cam1=True, cam2=True) 60 | query = self.process_data(test_dirs, cam1=True, cam2=False) 61 | gallery = self.process_data(test_dirs, cam1=False, cam2=True) 62 | 63 | super(iLIDSVID, self).__init__(train, query, gallery, **kwargs) 64 | 65 | def prepare_split(self): 66 | if not osp.exists(self.split_path): 67 | print('Creating splits ...') 68 | mat_split_data = loadmat(self.split_mat_path)['ls_set'] 69 | 70 | num_splits = mat_split_data.shape[0] 71 | num_total_ids = mat_split_data.shape[1] 72 | assert num_splits == 10 73 | assert num_total_ids == 300 74 | num_ids_each = num_total_ids // 2 75 | 76 | # pids in mat_split_data are indices, so we need to transform them 77 | # to real pids 78 | person_cam1_dirs = sorted( 79 | glob.glob(osp.join(self.cam_1_path, '*')) 80 | ) 81 | person_cam2_dirs = sorted( 82 | glob.glob(osp.join(self.cam_2_path, '*')) 83 | ) 84 | 85 | person_cam1_dirs = [ 86 | osp.basename(item) for item in person_cam1_dirs 87 | ] 88 | person_cam2_dirs = [ 89 | osp.basename(item) for item in person_cam2_dirs 90 | ] 91 | 92 | # make sure persons in one camera view can be found in the other camera view 93 | assert set(person_cam1_dirs) == set(person_cam2_dirs) 94 | 95 | splits = [] 96 | for i_split in range(num_splits): 97 | # first 50% for testing and the remaining for training, following Wang et al. ECCV'14. 98 | train_idxs = sorted( 99 | list(mat_split_data[i_split, num_ids_each:]) 100 | ) 101 | test_idxs = sorted( 102 | list(mat_split_data[i_split, :num_ids_each]) 103 | ) 104 | 105 | train_idxs = [int(i) - 1 for i in train_idxs] 106 | test_idxs = [int(i) - 1 for i in test_idxs] 107 | 108 | # transform pids to person dir names 109 | train_dirs = [person_cam1_dirs[i] for i in train_idxs] 110 | test_dirs = [person_cam1_dirs[i] for i in test_idxs] 111 | 112 | split = {'train': train_dirs, 'test': test_dirs} 113 | splits.append(split) 114 | 115 | print( 116 | 'Totally {} splits are created, following Wang et al. ECCV\'14' 117 | .format(len(splits)) 118 | ) 119 | print('Split file is saved to {}'.format(self.split_path)) 120 | write_json(splits, self.split_path) 121 | 122 | def process_data(self, dirnames, cam1=True, cam2=True): 123 | tracklets = [] 124 | dirname2pid = {dirname: i for i, dirname in enumerate(dirnames)} 125 | 126 | for dirname in dirnames: 127 | if cam1: 128 | person_dir = osp.join(self.cam_1_path, dirname) 129 | img_names = glob.glob(osp.join(person_dir, '*.png')) 130 | assert len(img_names) > 0 131 | img_names = tuple(img_names) 132 | pid = dirname2pid[dirname] 133 | tracklets.append((img_names, pid, 0)) 134 | 135 | if cam2: 136 | person_dir = osp.join(self.cam_2_path, dirname) 137 | img_names = glob.glob(osp.join(person_dir, '*.png')) 138 | assert len(img_names) > 0 139 | img_names = tuple(img_names) 140 | pid = dirname2pid[dirname] 141 | tracklets.append((img_names, pid, 1)) 142 | 143 | return tracklets 144 | -------------------------------------------------------------------------------- /torchreid/data/datasets/video/mars.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import os.path as osp 3 | import warnings 4 | from scipy.io import loadmat 5 | 6 | from ..dataset import VideoDataset 7 | 8 | 9 | class Mars(VideoDataset): 10 | """MARS. 11 | 12 | Reference: 13 | Zheng et al. MARS: A Video Benchmark for Large-Scale Person Re-identification. ECCV 2016. 14 | 15 | URL: ``_ 16 | 17 | Dataset statistics: 18 | - identities: 1261. 19 | - tracklets: 8298 (train) + 1980 (query) + 9330 (gallery). 20 | - cameras: 6. 21 | """ 22 | dataset_dir = 'mars' 23 | dataset_url = None 24 | 25 | def __init__(self, root='', **kwargs): 26 | self.root = osp.abspath(osp.expanduser(root)) 27 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 28 | self.download_dataset(self.dataset_dir, self.dataset_url) 29 | 30 | self.train_name_path = osp.join( 31 | self.dataset_dir, 'info/train_name.txt' 32 | ) 33 | self.test_name_path = osp.join(self.dataset_dir, 'info/test_name.txt') 34 | self.track_train_info_path = osp.join( 35 | self.dataset_dir, 'info/tracks_train_info.mat' 36 | ) 37 | self.track_test_info_path = osp.join( 38 | self.dataset_dir, 'info/tracks_test_info.mat' 39 | ) 40 | self.query_IDX_path = osp.join(self.dataset_dir, 'info/query_IDX.mat') 41 | 42 | required_files = [ 43 | self.dataset_dir, self.train_name_path, self.test_name_path, 44 | self.track_train_info_path, self.track_test_info_path, 45 | self.query_IDX_path 46 | ] 47 | self.check_before_run(required_files) 48 | 49 | train_names = self.get_names(self.train_name_path) 50 | test_names = self.get_names(self.test_name_path) 51 | track_train = loadmat(self.track_train_info_path 52 | )['track_train_info'] # numpy.ndarray (8298, 4) 53 | track_test = loadmat(self.track_test_info_path 54 | )['track_test_info'] # numpy.ndarray (12180, 4) 55 | query_IDX = loadmat(self.query_IDX_path 56 | )['query_IDX'].squeeze() # numpy.ndarray (1980,) 57 | query_IDX -= 1 # index from 0 58 | track_query = track_test[query_IDX, :] 59 | gallery_IDX = [ 60 | i for i in range(track_test.shape[0]) if i not in query_IDX 61 | ] 62 | track_gallery = track_test[gallery_IDX, :] 63 | 64 | train = self.process_data( 65 | train_names, track_train, home_dir='bbox_train', relabel=True 66 | ) 67 | query = self.process_data( 68 | test_names, track_query, home_dir='bbox_test', relabel=False 69 | ) 70 | gallery = self.process_data( 71 | test_names, track_gallery, home_dir='bbox_test', relabel=False 72 | ) 73 | 74 | super(Mars, self).__init__(train, query, gallery, **kwargs) 75 | 76 | def get_names(self, fpath): 77 | names = [] 78 | with open(fpath, 'r') as f: 79 | for line in f: 80 | new_line = line.rstrip() 81 | names.append(new_line) 82 | return names 83 | 84 | def process_data( 85 | self, names, meta_data, home_dir=None, relabel=False, min_seq_len=0 86 | ): 87 | assert home_dir in ['bbox_train', 'bbox_test'] 88 | num_tracklets = meta_data.shape[0] 89 | pid_list = list(set(meta_data[:, 2].tolist())) 90 | 91 | if relabel: 92 | pid2label = {pid: label for label, pid in enumerate(pid_list)} 93 | tracklets = [] 94 | 95 | for tracklet_idx in range(num_tracklets): 96 | data = meta_data[tracklet_idx, ...] 97 | start_index, end_index, pid, camid = data 98 | if pid == -1: 99 | continue # junk images are just ignored 100 | assert 1 <= camid <= 6 101 | if relabel: 102 | pid = pid2label[pid] 103 | camid -= 1 # index starts from 0 104 | img_names = names[start_index - 1:end_index] 105 | 106 | # make sure image names correspond to the same person 107 | pnames = [img_name[:4] for img_name in img_names] 108 | assert len( 109 | set(pnames) 110 | ) == 1, 'Error: a single tracklet contains different person images' 111 | 112 | # make sure all images are captured under the same camera 113 | camnames = [img_name[5] for img_name in img_names] 114 | assert len( 115 | set(camnames) 116 | ) == 1, 'Error: images are captured under different cameras!' 117 | 118 | # append image names with directory information 119 | img_paths = [ 120 | osp.join(self.dataset_dir, home_dir, img_name[:4], img_name) 121 | for img_name in img_names 122 | ] 123 | if len(img_paths) >= min_seq_len: 124 | img_paths = tuple(img_paths) 125 | tracklets.append((img_paths, pid, camid)) 126 | 127 | return tracklets 128 | 129 | def combine_all(self): 130 | warnings.warn( 131 | 'Some query IDs do not appear in gallery. Therefore, combineall ' 132 | 'does not make any difference to Mars' 133 | ) 134 | -------------------------------------------------------------------------------- /torchreid/data/datasets/video/prid2011.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import glob 3 | import os.path as osp 4 | 5 | from torchreid.utils import read_json 6 | 7 | from ..dataset import VideoDataset 8 | 9 | 10 | class PRID2011(VideoDataset): 11 | """PRID2011. 12 | 13 | Reference: 14 | Hirzer et al. Person Re-Identification by Descriptive and 15 | Discriminative Classification. SCIA 2011. 16 | 17 | URL: ``_ 18 | 19 | Dataset statistics: 20 | - identities: 200. 21 | - tracklets: 400. 22 | - cameras: 2. 23 | """ 24 | dataset_dir = 'prid2011' 25 | dataset_url = None 26 | 27 | def __init__(self, root='', split_id=0, **kwargs): 28 | self.root = osp.abspath(osp.expanduser(root)) 29 | self.dataset_dir = osp.join(self.root, self.dataset_dir) 30 | self.download_dataset(self.dataset_dir, self.dataset_url) 31 | 32 | self.split_path = osp.join(self.dataset_dir, 'splits_prid2011.json') 33 | self.cam_a_dir = osp.join( 34 | self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_a' 35 | ) 36 | self.cam_b_dir = osp.join( 37 | self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_b' 38 | ) 39 | 40 | required_files = [self.dataset_dir, self.cam_a_dir, self.cam_b_dir] 41 | self.check_before_run(required_files) 42 | 43 | splits = read_json(self.split_path) 44 | if split_id >= len(splits): 45 | raise ValueError( 46 | 'split_id exceeds range, received {}, but expected between 0 and {}' 47 | .format(split_id, 48 | len(splits) - 1) 49 | ) 50 | split = splits[split_id] 51 | train_dirs, test_dirs = split['train'], split['test'] 52 | 53 | train = self.process_dir(train_dirs, cam1=True, cam2=True) 54 | query = self.process_dir(test_dirs, cam1=True, cam2=False) 55 | gallery = self.process_dir(test_dirs, cam1=False, cam2=True) 56 | 57 | super(PRID2011, self).__init__(train, query, gallery, **kwargs) 58 | 59 | def process_dir(self, dirnames, cam1=True, cam2=True): 60 | tracklets = [] 61 | dirname2pid = {dirname: i for i, dirname in enumerate(dirnames)} 62 | 63 | for dirname in dirnames: 64 | if cam1: 65 | person_dir = osp.join(self.cam_a_dir, dirname) 66 | img_names = glob.glob(osp.join(person_dir, '*.png')) 67 | assert len(img_names) > 0 68 | img_names = tuple(img_names) 69 | pid = dirname2pid[dirname] 70 | tracklets.append((img_names, pid, 0)) 71 | 72 | if cam2: 73 | person_dir = osp.join(self.cam_b_dir, dirname) 74 | img_names = glob.glob(osp.join(person_dir, '*.png')) 75 | assert len(img_names) > 0 76 | img_names = tuple(img_names) 77 | pid = dirname2pid[dirname] 78 | tracklets.append((img_names, pid, 1)) 79 | 80 | return tracklets 81 | -------------------------------------------------------------------------------- /torchreid/engine/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | 3 | from .image import ImageSoftmaxEngine, ImageTripletEngine, MultiModalImageSoftmaxEngine, Image3MEngine, ImageHCEngine 4 | 5 | from .video import VideoSoftmaxEngine, VideoTripletEngine 6 | from .engine import Engine 7 | -------------------------------------------------------------------------------- /torchreid/engine/image/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from .softmax import ImageSoftmaxEngine, MultiModalImageSoftmaxEngine 4 | from .margin import Image3MEngine 5 | from .triplet import ImageTripletEngine 6 | from .hcloss import ImageHCEngine 7 | 8 | -------------------------------------------------------------------------------- /torchreid/engine/image/hcloss.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | 3 | from torchreid.metrics.accuracy import accuracy 4 | from torchreid.losses import TripletLoss, CrossEntropyLoss, hetero_loss 5 | 6 | from ..engine import * 7 | 8 | 9 | class ImageHCEngine(Engine): 10 | r"""Triplet-loss engine for image-reid. 11 | 12 | Args: 13 | datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager`` 14 | or ``torchreid.data.VideoDataManager``. 15 | model (nn.Module): model instance. 16 | optimizer (Optimizer): an Optimizer. 17 | margin (float, optional): margin for triplet loss. Default is 0.3. 18 | weight_t (float, optional): weight for triplet loss. Default is 1. 19 | weight_x (float, optional): weight for softmax loss. Default is 1. 20 | scheduler (LRScheduler, optional): if None, no learning rate decay will be performed. 21 | use_gpu (bool, optional): use gpu. Default is True. 22 | label_smooth (bool, optional): use label smoothing regularizer. Default is True. 23 | 24 | Examples:: 25 | 26 | import torchreid 27 | datamanager = torchreid.data.ImageDataManager( 28 | root='path/to/reid-data', 29 | sources='market1501', 30 | height=256, 31 | width=128, 32 | combineall=False, 33 | batch_size=32, 34 | num_instances=4, 35 | train_sampler='RandomIdentitySampler' # this is important 36 | ) 37 | model = torchreid.models.build_model( 38 | name='resnet50', 39 | num_classes=datamanager.num_train_pids, 40 | loss='triplet' 41 | ) 42 | model = model.cuda() 43 | optimizer = torchreid.optim.build_optimizer( 44 | model, optim='adam', lr=0.0003 45 | ) 46 | scheduler = torchreid.optim.build_lr_scheduler( 47 | optimizer, 48 | lr_scheduler='single_step', 49 | stepsize=20 50 | ) 51 | engine = torchreid.engine.ImageTripletEngine( 52 | datamanager, model, optimizer, margin=0.3, 53 | weight_t=0.7, weight_x=1, scheduler=scheduler 54 | ) 55 | engine.run( 56 | max_epoch=60, 57 | save_dir='log/resnet50-triplet-market1501', 58 | print_freq=10 59 | ) 60 | """ 61 | 62 | def __init__( 63 | self, 64 | datamanager, 65 | model, 66 | optimizer, 67 | margin=3, 68 | weight_m=1, 69 | weight_x=1, 70 | scheduler=None, 71 | use_gpu=True, 72 | label_smooth=True 73 | ): 74 | super(ImageHCEngine, self).__init__(datamanager, use_gpu) 75 | 76 | self.model = model 77 | self.optimizer = optimizer 78 | self.scheduler = scheduler 79 | self.register_model('model', model, optimizer, scheduler) 80 | 81 | assert weight_m >= 0 and weight_x >= 0 82 | assert weight_m + weight_x > 0 83 | self.weight_m = weight_m 84 | self.weight_x = weight_x 85 | 86 | self.criterion_m = hetero_loss() 87 | self.criterion_x = CrossEntropyLoss( 88 | num_classes=self.datamanager.num_train_pids, 89 | use_gpu=self.use_gpu, 90 | label_smooth=label_smooth 91 | ) 92 | 93 | def forward_backward(self, data): 94 | 95 | imgs, pids = self.parse_data_for_train(data) 96 | 97 | if self.use_gpu: 98 | for i in range(len(imgs)): 99 | imgs[i] = imgs[i].cuda() 100 | pids = pids.cuda() 101 | 102 | outputs, features_RGB, features_NI, features_TI = self.model(imgs) 103 | 104 | loss = 0 105 | loss_summary = {} 106 | 107 | if self.weight_m > 0: 108 | loss_m = 0 109 | loss_m += self.criterion_m(features_RGB, features_NI, pids) 110 | loss_m += self.criterion_m(features_RGB, features_TI, pids) 111 | loss_m += self.criterion_m(features_NI, features_TI, pids) 112 | # loss_m /= 3 113 | loss += self.weight_m * loss_m 114 | loss_summary['loss_m'] = loss_m.item() 115 | 116 | if self.weight_x > 0: 117 | loss_x = self.compute_loss(self.criterion_x, outputs, pids) 118 | loss += self.weight_x * loss_x 119 | loss_summary['loss_x'] = loss_x.item() 120 | loss_summary['acc'] = accuracy(outputs, pids)[0].item() 121 | 122 | assert loss_summary 123 | 124 | self.optimizer.zero_grad() 125 | loss.backward() 126 | self.optimizer.step() 127 | return loss_summary 128 | -------------------------------------------------------------------------------- /torchreid/engine/image/margin.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | 3 | from torchreid.metrics.accuracy import accuracy 4 | from torchreid.losses import TripletLoss, CrossEntropyLoss, multiModalMarginLossNew 5 | 6 | from ..engine import * 7 | 8 | 9 | class Image3MEngine(Engine): 10 | r"""Triplet-loss engine for image-reid. 11 | 12 | Args: 13 | datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager`` 14 | or ``torchreid.data.VideoDataManager``. 15 | model (nn.Module): model instance. 16 | optimizer (Optimizer): an Optimizer. 17 | margin (float, optional): margin for triplet loss. Default is 0.3. 18 | weight_t (float, optional): weight for triplet loss. Default is 1. 19 | weight_x (float, optional): weight for softmax loss. Default is 1. 20 | scheduler (LRScheduler, optional): if None, no learning rate decay will be performed. 21 | use_gpu (bool, optional): use gpu. Default is True. 22 | label_smooth (bool, optional): use label smoothing regularizer. Default is True. 23 | 24 | Examples:: 25 | 26 | import torchreid 27 | datamanager = torchreid.data.ImageDataManager( 28 | root='path/to/reid-data', 29 | sources='market1501', 30 | height=256, 31 | width=128, 32 | combineall=False, 33 | batch_size=32, 34 | num_instances=4, 35 | train_sampler='RandomIdentitySampler' # this is important 36 | ) 37 | model = torchreid.models.build_model( 38 | name='resnet50', 39 | num_classes=datamanager.num_train_pids, 40 | loss='triplet' 41 | ) 42 | model = model.cuda() 43 | optimizer = torchreid.optim.build_optimizer( 44 | model, optim='adam', lr=0.0003 45 | ) 46 | scheduler = torchreid.optim.build_lr_scheduler( 47 | optimizer, 48 | lr_scheduler='single_step', 49 | stepsize=20 50 | ) 51 | engine = torchreid.engine.ImageTripletEngine( 52 | datamanager, model, optimizer, margin=0.3, 53 | weight_t=0.7, weight_x=1, scheduler=scheduler 54 | ) 55 | engine.run( 56 | max_epoch=60, 57 | save_dir='log/resnet50-triplet-market1501', 58 | print_freq=10 59 | ) 60 | """ 61 | 62 | def __init__( 63 | self, 64 | datamanager, 65 | model, 66 | optimizer, 67 | margin=3, 68 | weight_m=1, 69 | weight_x=1, 70 | scheduler=None, 71 | use_gpu=True, 72 | label_smooth=True 73 | ): 74 | super(Image3MEngine, self).__init__(datamanager, use_gpu) 75 | 76 | self.model = model 77 | self.optimizer = optimizer 78 | self.scheduler = scheduler 79 | self.register_model('model', model, optimizer, scheduler) 80 | 81 | assert weight_m >= 0 and weight_x >= 0 82 | assert weight_m + weight_x > 0 83 | self.weight_m = weight_m 84 | self.weight_x = weight_x 85 | 86 | self.criterion_m = multiModalMarginLossNew(margin=margin) 87 | self.criterion_x = CrossEntropyLoss( 88 | num_classes=self.datamanager.num_train_pids, 89 | use_gpu=self.use_gpu, 90 | label_smooth=label_smooth 91 | ) 92 | 93 | 94 | def forward_backward(self, data): 95 | imgs, pids, timeids = self.parse_data_for_train(data) 96 | 97 | if self.use_gpu: 98 | for i in range(len(imgs)): 99 | imgs[i] = imgs[i].cuda() 100 | pids = pids.cuda() 101 | 102 | outputs_R, outputs_N, outputs_T, features_RGB, features_NI, features_TI = self.model(imgs) 103 | 104 | loss = 0 105 | loss_summary = {} 106 | 107 | if self.weight_m > 0: 108 | loss_m = 0 109 | 110 | loss_m = self.criterion_m(features_RGB, features_NI, features_TI, pids) 111 | loss += self.weight_m * loss_m 112 | 113 | 114 | if self.weight_x > 0: 115 | loss_R = self.compute_loss(self.criterion_x, outputs_R, pids) 116 | loss_N = self.compute_loss(self.criterion_x, outputs_N, pids) 117 | loss_T = self.compute_loss(self.criterion_x, outputs_T, pids) 118 | 119 | loss_x = loss_R + loss_N + loss_T 120 | 121 | loss += self.weight_x * loss_x 122 | 123 | self.optimizer.zero_grad() 124 | loss.backward() 125 | self.optimizer.step() 126 | 127 | acc_R = 0; acc_N = 0; acc_T = 0 128 | if isinstance(outputs_R, (tuple, list)): 129 | for i in range(len(outputs_R)): 130 | acc_R += accuracy(outputs_R[i], pids)[0] 131 | acc_N += accuracy(outputs_N[i], pids)[0] 132 | acc_T += accuracy(outputs_T[i], pids)[0] 133 | acc_R /= len(outputs_R) 134 | acc_N /= len(outputs_N) 135 | acc_T /= len(outputs_T) 136 | else: 137 | acc_R = accuracy(outputs_R, pids)[0] 138 | acc_N = accuracy(outputs_N, pids)[0] 139 | acc_T = accuracy(outputs_T, pids)[0] 140 | 141 | 142 | loss_summary = { 143 | 'loss': loss.item(), 144 | 'LossX': loss_x.item(), 145 | 'LossM': loss_m, 146 | 'accR': acc_R.item(), 147 | 'lossR': loss_R.item(), 148 | 'accN': acc_N.item(), 149 | 'lossN': loss_N.item(), 150 | 'accT': acc_T.item(), 151 | 'lossT': loss_T.item(), 152 | } 153 | 154 | return loss_summary 155 | -------------------------------------------------------------------------------- /torchreid/engine/image/triplet.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | 3 | from cv2 import findTransformECC 4 | 5 | from torchreid.metrics.accuracy import accuracy 6 | from torchreid.losses import TripletLoss, CrossEntropyLoss 7 | 8 | from ..engine import * 9 | from torchreid.losses import time_loss 10 | 11 | class ImageTripletEngine(Engine): 12 | r"""Triplet-loss engine for image-reid. 13 | 14 | Args: 15 | datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager`` 16 | or ``torchreid.data.VideoDataManager``. 17 | model (nn.Module): model instance. 18 | optimizer (Optimizer): an Optimizer. 19 | margin (float, optional): margin for triplet loss. Default is 0.3. 20 | weight_t (float, optional): weight for triplet loss. Default is 1. 21 | weight_x (float, optional): weight for softmax loss. Default is 1. 22 | scheduler (LRScheduler, optional): if None, no learning rate decay will be performed. 23 | use_gpu (bool, optional): use gpu. Default is True. 24 | label_smooth (bool, optional): use label smoothing regularizer. Default is True. 25 | 26 | Examples:: 27 | 28 | import torchreid 29 | datamanager = torchreid.data.ImageDataManager( 30 | root='path/to/reid-data', 31 | sources='market1501', 32 | height=256, 33 | width=128, 34 | combineall=False, 35 | batch_size=32, 36 | num_instances=4, 37 | train_sampler='RandomIdentitySampler' # this is important 38 | ) 39 | model = torchreid.models.build_model( 40 | name='resnet50', 41 | num_classes=datamanager.num_train_pids, 42 | loss='triplet' 43 | ) 44 | model = model.cuda() 45 | optimizer = torchreid.optim.build_optimizer( 46 | model, optim='adam', lr=0.0003 47 | ) 48 | scheduler = torchreid.optim.build_lr_scheduler( 49 | optimizer, 50 | lr_scheduler='single_step', 51 | stepsize=20 52 | ) 53 | engine = torchreid.engine.ImageTripletEngine( 54 | datamanager, model, optimizer, margin=0.3, 55 | weight_t=0.7, weight_x=1, scheduler=scheduler 56 | ) 57 | engine.run( 58 | max_epoch=60, 59 | save_dir='log/resnet50-triplet-market1501', 60 | print_freq=10 61 | ) 62 | """ 63 | 64 | def __init__( 65 | self, 66 | datamanager, 67 | model, 68 | optimizer, 69 | margin=0.3, 70 | weight_t=1, 71 | weight_x=1, 72 | scheduler=None, 73 | use_gpu=True, 74 | label_smooth=True 75 | ): 76 | super(ImageTripletEngine, self).__init__(datamanager, use_gpu) 77 | 78 | self.model = model 79 | self.optimizer = optimizer 80 | self.scheduler = scheduler 81 | self.register_model('model', model, optimizer, scheduler) 82 | 83 | assert weight_t >= 0 and weight_x >= 0 84 | assert weight_t + weight_x > 0 85 | self.weight_t = weight_t 86 | self.weight_x = weight_x 87 | 88 | self.criterion_t = TripletLoss(margin=margin) 89 | # self.criterion_time = CenterLoss() 90 | self.criterion_x = CrossEntropyLoss( 91 | num_classes=self.datamanager.num_train_pids, 92 | use_gpu=self.use_gpu, 93 | label_smooth=label_smooth 94 | ) 95 | 96 | 97 | def forward_backward(self, data): 98 | imgs, pids, timeids = self.parse_data_for_train(data) 99 | 100 | if self.use_gpu: 101 | for i in range(len(imgs)): 102 | imgs[i] = imgs[i].cuda() 103 | pids = pids.cuda() 104 | 105 | outputs, features = self.model(imgs) 106 | if self.weight_t > 0: 107 | # loss_t_all = self.criterion_t(features_all, pids) 108 | loss_t = self.criterion_t(features, pids) 109 | 110 | 111 | # print(loss_t, loss_t_R, loss_t_N, loss_t_T) 112 | 113 | 114 | loss_R = self.compute_loss(self.criterion_x, outputs, pids) 115 | 116 | loss = loss_R + loss_t 117 | # loss = loss_R 118 | 119 | self.optimizer.zero_grad() 120 | loss.backward() 121 | self.optimizer.step() 122 | 123 | acc = 0 124 | if isinstance(outputs, (tuple, list)): 125 | for i in range(len(outputs)): 126 | acc += accuracy(outputs[i], pids)[0] 127 | # print(acc) 128 | acc /= len(outputs) 129 | else: 130 | acc = accuracy(outputs, pids)[0] 131 | 132 | 133 | 134 | loss_summary = { 135 | 'loss_all': loss.item(), 136 | 'loss_t': loss_t.item(), 137 | 'acc': acc.item(), 138 | } 139 | 140 | 141 | return loss_summary 142 | -------------------------------------------------------------------------------- /torchreid/engine/video/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from .softmax import VideoSoftmaxEngine 4 | from .triplet import VideoTripletEngine 5 | -------------------------------------------------------------------------------- /torchreid/engine/video/softmax.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import torch 3 | 4 | from torchreid.engine.image import ImageSoftmaxEngine 5 | 6 | 7 | class VideoSoftmaxEngine(ImageSoftmaxEngine): 8 | """Softmax-loss engine for video-reid. 9 | 10 | Args: 11 | datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager`` 12 | or ``torchreid.data.VideoDataManager``. 13 | model (nn.Module): model instance. 14 | optimizer (Optimizer): an Optimizer. 15 | scheduler (LRScheduler, optional): if None, no learning rate decay will be performed. 16 | use_gpu (bool, optional): use gpu. Default is True. 17 | label_smooth (bool, optional): use label smoothing regularizer. Default is True. 18 | pooling_method (str, optional): how to pool features for a tracklet. 19 | Default is "avg" (average). Choices are ["avg", "max"]. 20 | 21 | Examples:: 22 | 23 | import torch 24 | import torchreid 25 | # Each batch contains batch_size*seq_len images 26 | datamanager = torchreid.data.VideoDataManager( 27 | root='path/to/reid-data', 28 | sources='mars', 29 | height=256, 30 | width=128, 31 | combineall=False, 32 | batch_size=8, # number of tracklets 33 | seq_len=15 # number of images in each tracklet 34 | ) 35 | model = torchreid.models.build_model( 36 | name='resnet50', 37 | num_classes=datamanager.num_train_pids, 38 | loss='softmax' 39 | ) 40 | model = model.cuda() 41 | optimizer = torchreid.optim.build_optimizer( 42 | model, optim='adam', lr=0.0003 43 | ) 44 | scheduler = torchreid.optim.build_lr_scheduler( 45 | optimizer, 46 | lr_scheduler='single_step', 47 | stepsize=20 48 | ) 49 | engine = torchreid.engine.VideoSoftmaxEngine( 50 | datamanager, model, optimizer, scheduler=scheduler, 51 | pooling_method='avg' 52 | ) 53 | engine.run( 54 | max_epoch=60, 55 | save_dir='log/resnet50-softmax-mars', 56 | print_freq=10 57 | ) 58 | """ 59 | 60 | def __init__( 61 | self, 62 | datamanager, 63 | model, 64 | optimizer, 65 | scheduler=None, 66 | use_gpu=True, 67 | label_smooth=True, 68 | pooling_method='avg' 69 | ): 70 | super(VideoSoftmaxEngine, self).__init__( 71 | datamanager, 72 | model, 73 | optimizer, 74 | scheduler=scheduler, 75 | use_gpu=use_gpu, 76 | label_smooth=label_smooth 77 | ) 78 | self.pooling_method = pooling_method 79 | 80 | def parse_data_for_train(self, data): 81 | imgs = data['img'] 82 | pids = data['pid'] 83 | if imgs.dim() == 5: 84 | # b: batch size 85 | # s: sqeuence length 86 | # c: channel depth 87 | # h: height 88 | # w: width 89 | b, s, c, h, w = imgs.size() 90 | imgs = imgs.view(b * s, c, h, w) 91 | pids = pids.view(b, 1).expand(b, s) 92 | pids = pids.contiguous().view(b * s) 93 | return imgs, pids 94 | 95 | def extract_features(self, input): 96 | # b: batch size 97 | # s: sqeuence length 98 | # c: channel depth 99 | # h: height 100 | # w: width 101 | b, s, c, h, w = input.size() 102 | input = input.view(b * s, c, h, w) 103 | features = self.model(input) 104 | features = features.view(b, s, -1) 105 | if self.pooling_method == 'avg': 106 | features = torch.mean(features, 1) 107 | else: 108 | features = torch.max(features, 1)[0] 109 | return features 110 | -------------------------------------------------------------------------------- /torchreid/engine/video/triplet.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import torch 3 | 4 | from torchreid.engine.image import ImageTripletEngine 5 | 6 | 7 | class VideoTripletEngine(ImageTripletEngine): 8 | """Triplet-loss engine for video-reid. 9 | 10 | Args: 11 | datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager`` 12 | or ``torchreid.data.VideoDataManager``. 13 | model (nn.Module): model instance. 14 | optimizer (Optimizer): an Optimizer. 15 | margin (float, optional): margin for triplet loss. Default is 0.3. 16 | weight_t (float, optional): weight for triplet loss. Default is 1. 17 | weight_x (float, optional): weight for softmax loss. Default is 1. 18 | scheduler (LRScheduler, optional): if None, no learning rate decay will be performed. 19 | use_gpu (bool, optional): use gpu. Default is True. 20 | label_smooth (bool, optional): use label smoothing regularizer. Default is True. 21 | pooling_method (str, optional): how to pool features for a tracklet. 22 | Default is "avg" (average). Choices are ["avg", "max"]. 23 | 24 | Examples:: 25 | 26 | import torch 27 | import torchreid 28 | # Each batch contains batch_size*seq_len images 29 | # Each identity is sampled with num_instances tracklets 30 | datamanager = torchreid.data.VideoDataManager( 31 | root='path/to/reid-data', 32 | sources='mars', 33 | height=256, 34 | width=128, 35 | combineall=False, 36 | num_instances=4, 37 | train_sampler='RandomIdentitySampler' 38 | batch_size=8, # number of tracklets 39 | seq_len=15 # number of images in each tracklet 40 | ) 41 | model = torchreid.models.build_model( 42 | name='resnet50', 43 | num_classes=datamanager.num_train_pids, 44 | loss='triplet' 45 | ) 46 | model = model.cuda() 47 | optimizer = torchreid.optim.build_optimizer( 48 | model, optim='adam', lr=0.0003 49 | ) 50 | scheduler = torchreid.optim.build_lr_scheduler( 51 | optimizer, 52 | lr_scheduler='single_step', 53 | stepsize=20 54 | ) 55 | engine = torchreid.engine.VideoTripletEngine( 56 | datamanager, model, optimizer, margin=0.3, 57 | weight_t=0.7, weight_x=1, scheduler=scheduler, 58 | pooling_method='avg' 59 | ) 60 | engine.run( 61 | max_epoch=60, 62 | save_dir='log/resnet50-triplet-mars', 63 | print_freq=10 64 | ) 65 | """ 66 | 67 | def __init__( 68 | self, 69 | datamanager, 70 | model, 71 | optimizer, 72 | margin=0.3, 73 | weight_t=1, 74 | weight_x=1, 75 | scheduler=None, 76 | use_gpu=True, 77 | label_smooth=True, 78 | pooling_method='avg' 79 | ): 80 | super(VideoTripletEngine, self).__init__( 81 | datamanager, 82 | model, 83 | optimizer, 84 | margin=margin, 85 | weight_t=weight_t, 86 | weight_x=weight_x, 87 | scheduler=scheduler, 88 | use_gpu=use_gpu, 89 | label_smooth=label_smooth 90 | ) 91 | self.pooling_method = pooling_method 92 | 93 | def parse_data_for_train(self, data): 94 | imgs = data['img'] 95 | pids = data['pid'] 96 | if imgs.dim() == 5: 97 | # b: batch size 98 | # s: sqeuence length 99 | # c: channel depth 100 | # h: height 101 | # w: width 102 | b, s, c, h, w = imgs.size() 103 | imgs = imgs.view(b * s, c, h, w) 104 | pids = pids.view(b, 1).expand(b, s) 105 | pids = pids.contiguous().view(b * s) 106 | return imgs, pids 107 | 108 | def extract_features(self, input): 109 | # b: batch size 110 | # s: sqeuence length 111 | # c: channel depth 112 | # h: height 113 | # w: width 114 | b, s, c, h, w = input.size() 115 | input = input.view(b * s, c, h, w) 116 | features = self.model(input) 117 | features = features.view(b, s, -1) 118 | if self.pooling_method == 'avg': 119 | features = torch.mean(features, 1) 120 | else: 121 | features = torch.max(features, 1)[0] 122 | return features 123 | -------------------------------------------------------------------------------- /torchreid/losses/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | 3 | from .cross_entropy_loss import CrossEntropyLoss 4 | from .hard_mine_triplet_loss import TripletLoss 5 | from .hcloss import hetero_loss 6 | from .multi_modal_margin_loss_new import multiModalMarginLossNew 7 | from .time_loss import time_loss 8 | 9 | def DeepSupervision(criterion, xs, y): 10 | """DeepSupervision 11 | 12 | Applies criterion to each element in a list. 13 | 14 | Args: 15 | criterion: loss function 16 | xs: tuple of inputs 17 | y: ground truth 18 | """ 19 | 20 | loss = 0. 21 | for i in range(len(xs)): 22 | loss += criterion(xs[i], y) 23 | 24 | 25 | # loss = 0. 26 | # for i in range(len(xs)-1): 27 | # loss += criterion(xs[i], y) 28 | # loss += criterion(xs[i+1], y) * 3 29 | 30 | # loss = 0. 31 | # if len(xs) == 4: 32 | # for i in range(len(xs)-1): 33 | # loss += criterion(xs[i], y) 34 | # loss += criterion(xs[i+1], y) * 3 35 | # else: 36 | # for i in range(len(xs)): 37 | # loss += criterion(xs[i], y) 38 | 39 | # loss = 0. 40 | # for x in xs: 41 | # loss += criterion(x, y) 42 | # loss /= len(xs) 43 | 44 | return loss 45 | -------------------------------------------------------------------------------- /torchreid/losses/cross_entropy_loss.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, absolute_import 2 | import torch 3 | import torch.nn as nn 4 | 5 | 6 | class CrossEntropyLoss(nn.Module): 7 | r"""Cross entropy loss with label smoothing regularizer. 8 | 9 | Reference: 10 | Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016. 11 | 12 | With label smoothing, the label :math:`y` for a class is computed by 13 | 14 | .. math:: 15 | \begin{equation} 16 | (1 - \eps) \times y + \frac{\eps}{K}, 17 | \end{equation} 18 | 19 | where :math:`K` denotes the number of classes and :math:`\eps` is a weight. When 20 | :math:`\eps = 0`, the loss function reduces to the normal cross entropy. 21 | 22 | Args: 23 | num_classes (int): number of classes. 24 | eps (float, optional): weight. Default is 0.1. 25 | use_gpu (bool, optional): whether to use gpu devices. Default is True. 26 | label_smooth (bool, optional): whether to apply label smoothing. Default is True. 27 | """ 28 | 29 | def __init__(self, num_classes, eps=0.1, use_gpu=True, label_smooth=True): 30 | super(CrossEntropyLoss, self).__init__() 31 | self.num_classes = num_classes 32 | self.eps = eps if label_smooth else 0 33 | self.use_gpu = use_gpu 34 | self.logsoftmax = nn.LogSoftmax(dim=1) 35 | 36 | def forward(self, inputs, targets): 37 | """ 38 | Args: 39 | inputs (torch.Tensor): prediction matrix (before softmax) with 40 | shape (batch_size, num_classes). 41 | targets (torch.LongTensor): ground truth labels with shape (batch_size). 42 | Each position contains the label index. 43 | """ 44 | log_probs = self.logsoftmax(inputs) 45 | zeros = torch.zeros(log_probs.size()) 46 | targets = zeros.scatter_(1, targets.unsqueeze(1).data.cpu(), 1) 47 | if self.use_gpu: 48 | targets = targets.cuda() 49 | targets = (1 - self.eps) * targets + self.eps / self.num_classes 50 | return (-targets * log_probs).mean(0).sum() 51 | -------------------------------------------------------------------------------- /torchreid/losses/hard_mine_triplet_loss.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, absolute_import 2 | import torch 3 | import torch.nn as nn 4 | 5 | 6 | class TripletLoss(nn.Module): 7 | """Triplet loss with hard positive/negative mining. 8 | 9 | Reference: 10 | Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737. 11 | 12 | Imported from ``_. 13 | 14 | Args: 15 | margin (float, optional): margin for triplet. Default is 0.3. 16 | """ 17 | 18 | def __init__(self, margin=0.3): 19 | super(TripletLoss, self).__init__() 20 | self.margin = margin 21 | self.ranking_loss = nn.MarginRankingLoss(margin=margin) 22 | 23 | def forward(self, inputs, targets): 24 | """ 25 | Args: 26 | inputs (torch.Tensor): feature matrix with shape (batch_size, feat_dim). 27 | targets (torch.LongTensor): ground truth labels with shape (num_classes). 28 | """ 29 | n = inputs.size(0) 30 | 31 | # Compute pairwise distance, replace by the official when merged 32 | dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n) 33 | dist = dist + dist.t() 34 | dist.addmm_(inputs, inputs.t(), beta=1, alpha=-2) 35 | dist = dist.clamp(min=1e-12).sqrt() # for numerical stability 36 | 37 | # For each anchor, find the hardest positive and negative 38 | mask = targets.expand(n, n).eq(targets.expand(n, n).t()) 39 | dist_ap, dist_an = [], [] 40 | for i in range(n): 41 | dist_ap.append(dist[i][mask[i]].max().unsqueeze(0)) 42 | dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0)) 43 | dist_ap = torch.cat(dist_ap) 44 | dist_an = torch.cat(dist_an) 45 | 46 | # Compute ranking hinge loss 47 | y = torch.ones_like(dist_an) 48 | return self.ranking_loss(dist_an, dist_ap, y) 49 | -------------------------------------------------------------------------------- /torchreid/losses/hcloss.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch import nn, tensor 3 | import torch 4 | from torch.autograd import Variable 5 | 6 | class hetero_loss(nn.Module): 7 | def __init__(self, margin=0.1, dist_type = 'l2'): 8 | super(hetero_loss, self).__init__() 9 | self.margin = margin 10 | self.dist_type = dist_type 11 | if dist_type == 'l2': 12 | self.dist = nn.MSELoss(reduction='sum') 13 | if dist_type == 'cos': 14 | self.dist = nn.CosineSimilarity(dim=0) 15 | if dist_type == 'l1': 16 | self.dist = nn.L1Loss() 17 | 18 | def forward(self, feat1, feat2, label1): 19 | feat_size = feat1.size()[1] 20 | feat_num = feat1.size()[0] 21 | label_num = len(label1.unique()) 22 | feat1 = feat1.chunk(label_num, 0) 23 | feat2 = feat2.chunk(label_num, 0) 24 | #loss = Variable(.cuda()) 25 | for i in range(label_num): 26 | center1 = torch.mean(feat1[i], dim=0) 27 | center2 = torch.mean(feat2[i], dim=0) 28 | if self.dist_type == 'l2' or self.dist_type == 'l1': 29 | if i == 0: 30 | dist = max(0, abs(self.dist(center1, center2))) 31 | else: 32 | dist += max(0, abs(self.dist(center1, center2))) 33 | elif self.dist_type == 'cos': 34 | if i == 0: 35 | dist = max(0, 1-self.dist(center1, center2)) 36 | else: 37 | dist += max(0, 1-self.dist(center1, center2)) 38 | 39 | return dist 40 | -------------------------------------------------------------------------------- /torchreid/losses/multi_modal_margin_loss_new.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch import nn, tensor 3 | import torch 4 | from torch.autograd import Variable 5 | 6 | 7 | class multiModalMarginLossNew(nn.Module): 8 | def __init__(self, margin=3, dist_type='l2'): 9 | super(multiModalMarginLossNew, self).__init__() 10 | self.dist_type = dist_type 11 | self.margin = margin 12 | if dist_type == 'l2': 13 | self.dist = nn.MSELoss(reduction='sum') 14 | if dist_type == 'cos': 15 | self.dist = nn.CosineSimilarity(dim=0) 16 | if dist_type == 'l1': 17 | self.dist = nn.L1Loss() 18 | 19 | def forward(self, feat1, feat2, feat3, label1): 20 | # print("using 3MLoss") 21 | # print(feat1.shape, feat2.shape, label1.shape, label1) 22 | label_num = len(label1.unique()) 23 | feat1 = feat1.chunk(label_num, 0) 24 | feat2 = feat2.chunk(label_num, 0) 25 | feat3 = feat3.chunk(label_num, 0) 26 | # loss = Variable(.cuda()) 27 | for i in range(label_num): 28 | center1 = torch.mean(feat1[i], dim=0) 29 | center2 = torch.mean(feat2[i], dim=0) 30 | center3 = torch.mean(feat3[i], dim=0) 31 | # print(self.dist(center1, center2), self.dist(center1, center3), self.dist(center2, center3)) 32 | if self.dist_type == 'l2' or self.dist_type == 'l1': 33 | if i == 0: 34 | # print(self.dist(center1, center2), self.dist(center1, center3), self.dist(center2, center3)) 35 | dist = max(abs(self.margin - self.dist(center1, center2)), abs(self.margin - self.dist(center2, center3)), abs(self.margin - self.dist(center1, center3))) 36 | # dist = max(0, abs(self.margin - self.dist(center1, center2))) 37 | else: 38 | dist += max(abs(self.margin - self.dist(center1, center2)), abs(self.margin - self.dist(center2, center3)), abs(self.margin - self.dist(center1, center3))) 39 | # dist += max(0, abs(self.margin - self.dist(center1, center2))) 40 | return dist 41 | 42 | 43 | 44 | -------------------------------------------------------------------------------- /torchreid/losses/time_loss.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch import nn, tensor 3 | import torch 4 | from torch.autograd import Variable 5 | 6 | class time_loss(nn.Module): 7 | def __init__(self, margin=0.1, dist_type = 'l2'): 8 | super(time_loss, self).__init__() 9 | self.margin = margin 10 | self.dist_type = dist_type 11 | if dist_type == 'l2': 12 | self.dist = nn.MSELoss(reduction='sum') 13 | if dist_type == 'cos': 14 | self.dist = nn.CosineSimilarity(dim=0) 15 | if dist_type == 'l1': 16 | self.dist = nn.L1Loss() 17 | 18 | def forward(self, feat, label1): 19 | feat_size = feat.size()[1] 20 | feat_num = feat.size()[0] 21 | label_num = len(label1.unique()) 22 | feat = feat.chunk(label_num, 0) 23 | #loss = Variable(.cuda()) 24 | for i in range(label_num): 25 | center1 = torch.mean(feat[i], dim=0) 26 | if self.dist_type == 'l2' or self.dist_type == 'l1': 27 | if i == 0: 28 | dist = max(0, abs(self.dist(center1, center1))) 29 | else: 30 | dist += max(0, abs(self.dist(center1, center1))) 31 | elif self.dist_type == 'cos': 32 | if i == 0: 33 | dist = max(0, 1-self.dist(center1, center1)) 34 | else: 35 | dist += max(0, 1-self.dist(center1, center1)) 36 | return dist 37 | -------------------------------------------------------------------------------- /torchreid/metrics/accuracy.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import pdb 3 | 4 | def accuracy(output, target, topk=(1, )): 5 | """Computes the accuracy over the k top predictions for 6 | the specified values of k. 7 | 8 | Args: 9 | output (torch.Tensor): prediction matrix with shape (batch_size, num_classes). 10 | target (torch.LongTensor): ground truth labels with shape (batch_size). 11 | topk (tuple, optional): accuracy at top-k will be computed. For example, 12 | topk=(1, 5) means accuracy at top-1 and top-5 will be computed. 13 | 14 | Returns: 15 | list: accuracy at top-k. 16 | 17 | Examples:: 18 | >>> from torchreid import metrics 19 | >>> metrics.accuracy(output, target) 20 | """ 21 | # pdb.set_trace() 22 | maxk = max(topk) 23 | batch_size = target.size(0) 24 | 25 | if isinstance(output, (tuple, list)): 26 | output = output[0] 27 | 28 | _, pred = output.topk(maxk, 1, True, True) 29 | pred = pred.t() 30 | correct = pred.eq(target.view(1, -1).expand_as(pred)) 31 | 32 | res = [] 33 | for k in topk: 34 | correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) 35 | acc = correct_k.mul_(100.0 / batch_size) 36 | res.append(acc) 37 | 38 | return res 39 | -------------------------------------------------------------------------------- /torchreid/metrics/distance.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import torch 3 | from torch.nn import functional as F 4 | 5 | 6 | def compute_distance_matrix(input1, input2, metric='euclidean'): 7 | """A wrapper function for computing distance matrix. 8 | 9 | Args: 10 | input1 (torch.Tensor): 2-D feature matrix. 11 | input2 (torch.Tensor): 2-D feature matrix. 12 | metric (str, optional): "euclidean" or "cosine". 13 | Default is "euclidean". 14 | 15 | Returns: 16 | torch.Tensor: distance matrix. 17 | 18 | Examples:: 19 | >>> from torchreid import metrics 20 | >>> input1 = torch.rand(10, 2048) 21 | >>> input2 = torch.rand(100, 2048) 22 | >>> distmat = metrics.compute_distance_matrix(input1, input2) 23 | >>> distmat.size() # (10, 100) 24 | """ 25 | # check input 26 | assert isinstance(input1, torch.Tensor) 27 | assert isinstance(input2, torch.Tensor) 28 | assert input1.dim() == 2, 'Expected 2-D tensor, but got {}-D'.format( 29 | input1.dim() 30 | ) 31 | assert input2.dim() == 2, 'Expected 2-D tensor, but got {}-D'.format( 32 | input2.dim() 33 | ) 34 | assert input1.size(1) == input2.size(1) 35 | 36 | if metric == 'euclidean': 37 | distmat = euclidean_squared_distance(input1, input2) 38 | elif metric == 'cosine': 39 | distmat = cosine_distance(input1, input2) 40 | else: 41 | raise ValueError( 42 | 'Unknown distance metric: {}. ' 43 | 'Please choose either "euclidean" or "cosine"'.format(metric) 44 | ) 45 | 46 | return distmat 47 | 48 | 49 | def euclidean_squared_distance(input1, input2): 50 | """Computes euclidean squared distance. 51 | 52 | Args: 53 | input1 (torch.Tensor): 2-D feature matrix. 54 | input2 (torch.Tensor): 2-D feature matrix. 55 | 56 | Returns: 57 | torch.Tensor: distance matrix. 58 | """ 59 | m, n = input1.size(0), input2.size(0) 60 | mat1 = torch.pow(input1, 2).sum(dim=1, keepdim=True).expand(m, n) 61 | mat2 = torch.pow(input2, 2).sum(dim=1, keepdim=True).expand(n, m).t() 62 | distmat = mat1 + mat2 63 | distmat.addmm_(input1, input2.t(), beta=1, alpha=-2) 64 | return distmat 65 | 66 | 67 | def cosine_distance(input1, input2): 68 | """Computes cosine distance. 69 | 70 | Args: 71 | input1 (torch.Tensor): 2-D feature matrix. 72 | input2 (torch.Tensor): 2-D feature matrix. 73 | 74 | Returns: 75 | torch.Tensor: distance matrix. 76 | """ 77 | input1_normed = F.normalize(input1, p=2, dim=1) 78 | input2_normed = F.normalize(input2, p=2, dim=1) 79 | distmat = 1 - torch.mm(input1_normed, input2_normed.t()) 80 | return distmat 81 | -------------------------------------------------------------------------------- /torchreid/metrics/rank_cylib/Makefile: -------------------------------------------------------------------------------- 1 | all: 2 | $(PYTHON) setup.py build_ext --inplace 3 | rm -rf build 4 | clean: 5 | rm -rf build 6 | rm -f rank_cy.c *.so -------------------------------------------------------------------------------- /torchreid/metrics/rank_cylib/rank_cy.cpython-36m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ziwang1121/IEEE/830a3fc444eb72736ff338b7c6fa7951d94c8d17/torchreid/metrics/rank_cylib/rank_cy.cpython-36m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /torchreid/metrics/rank_cylib/setup.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from distutils.core import setup 3 | from distutils.extension import Extension 4 | from Cython.Build import cythonize 5 | 6 | 7 | def numpy_include(): 8 | try: 9 | numpy_include = np.get_include() 10 | except AttributeError: 11 | numpy_include = np.get_numpy_include() 12 | return numpy_include 13 | 14 | 15 | ext_modules = [ 16 | Extension( 17 | 'rank_cy', 18 | ['rank_cy.pyx'], 19 | include_dirs=[numpy_include()], 20 | ) 21 | ] 22 | 23 | setup( 24 | name='Cython-based reid evaluation code', 25 | ext_modules=cythonize(ext_modules) 26 | ) 27 | -------------------------------------------------------------------------------- /torchreid/metrics/rank_cylib/test_cython.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import sys 3 | import numpy as np 4 | import timeit 5 | import os.path as osp 6 | 7 | from torchreid import metrics 8 | 9 | sys.path.insert(0, osp.dirname(osp.abspath(__file__)) + '/../../..') 10 | """ 11 | Test the speed of cython-based evaluation code. The speed improvements 12 | can be much bigger when using the real reid data, which contains a larger 13 | amount of query and gallery images. 14 | 15 | Note: you might encounter the following error: 16 | 'AssertionError: Error: all query identities do not appear in gallery'. 17 | This is normal because the inputs are random numbers. Just try again. 18 | """ 19 | 20 | print('*** Compare running time ***') 21 | 22 | setup = ''' 23 | import sys 24 | import os.path as osp 25 | import numpy as np 26 | sys.path.insert(0, osp.dirname(osp.abspath(__file__)) + '/../../..') 27 | from torchreid import metrics 28 | num_q = 30 29 | num_g = 300 30 | max_rank = 5 31 | distmat = np.random.rand(num_q, num_g) * 20 32 | q_pids = np.random.randint(0, num_q, size=num_q) 33 | g_pids = np.random.randint(0, num_g, size=num_g) 34 | q_camids = np.random.randint(0, 5, size=num_q) 35 | g_camids = np.random.randint(0, 5, size=num_g) 36 | ''' 37 | 38 | print('=> Using market1501\'s metric') 39 | pytime = timeit.timeit( 40 | 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=False)', 41 | setup=setup, 42 | number=20 43 | ) 44 | cytime = timeit.timeit( 45 | 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=True)', 46 | setup=setup, 47 | number=20 48 | ) 49 | print('Python time: {} s'.format(pytime)) 50 | print('Cython time: {} s'.format(cytime)) 51 | print('Cython is {} times faster than python\n'.format(pytime / cytime)) 52 | 53 | print('=> Using cuhk03\'s metric') 54 | pytime = timeit.timeit( 55 | 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_metric_cuhk03=True, use_cython=False)', 56 | setup=setup, 57 | number=20 58 | ) 59 | cytime = timeit.timeit( 60 | 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_metric_cuhk03=True, use_cython=True)', 61 | setup=setup, 62 | number=20 63 | ) 64 | print('Python time: {} s'.format(pytime)) 65 | print('Cython time: {} s'.format(cytime)) 66 | print('Cython is {} times faster than python\n'.format(pytime / cytime)) 67 | """ 68 | print("=> Check precision") 69 | 70 | num_q = 30 71 | num_g = 300 72 | max_rank = 5 73 | distmat = np.random.rand(num_q, num_g) * 20 74 | q_pids = np.random.randint(0, num_q, size=num_q) 75 | g_pids = np.random.randint(0, num_g, size=num_g) 76 | q_camids = np.random.randint(0, 5, size=num_q) 77 | g_camids = np.random.randint(0, 5, size=num_g) 78 | 79 | cmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=False) 80 | print("Python:\nmAP = {} \ncmc = {}\n".format(mAP, cmc)) 81 | cmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=True) 82 | print("Cython:\nmAP = {} \ncmc = {}\n".format(mAP, cmc)) 83 | """ 84 | -------------------------------------------------------------------------------- /torchreid/models/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import torch 3 | 4 | from .pcb import * 5 | from .mlfn import * 6 | from .hacnn import * 7 | from .osnet import * 8 | from .senet import * 9 | from .mudeep import * 10 | from .nasnet import * 11 | from .resnet import * 12 | from .densenet import * 13 | from .xception import * 14 | from .osnet_ain import * 15 | from .resnetmid import * 16 | from .shufflenet import * 17 | from .squeezenet import * 18 | from .inceptionv4 import * 19 | from .mobilenetv2 import * 20 | from .resnet_ibn_a import * 21 | from .resnet_ibn_b import * 22 | from .shufflenetv2 import * 23 | from .inceptionresnetv2 import * 24 | from .ieee3modalPart import * 25 | from .pfnet import * 26 | 27 | 28 | __model_factory = { 29 | 'ieee3modalPart': ieee3modalPart, 30 | # image classification models 31 | 32 | 'resnet18': resnet18, 33 | 'resnet34': resnet34, 34 | 'resnet50': resnet50, 35 | 'resnet101': resnet101, 36 | 'resnet152': resnet152, 37 | 'resnext50_32x4d': resnext50_32x4d, 38 | 'resnext101_32x8d': resnext101_32x8d, 39 | 'resnet50_fc512': resnet50_fc512, 40 | 'se_resnet50': se_resnet50, 41 | 'se_resnet50_fc512': se_resnet50_fc512, 42 | 'se_resnet101': se_resnet101, 43 | 'se_resnext50_32x4d': se_resnext50_32x4d, 44 | 'se_resnext101_32x4d': se_resnext101_32x4d, 45 | 'densenet121': densenet121, 46 | 'densenet169': densenet169, 47 | 'densenet201': densenet201, 48 | 'densenet161': densenet161, 49 | 'densenet121_fc512': densenet121_fc512, 50 | 'inceptionresnetv2': inceptionresnetv2, 51 | 'inceptionv4': inceptionv4, 52 | 'xception': xception, 53 | 'resnet50_ibn_a': resnet50_ibn_a, 54 | 'resnet50_ibn_b': resnet50_ibn_b, 55 | # lightweight models 56 | 'nasnsetmobile': nasnetamobile, 57 | 'mobilenetv2_x1_0': mobilenetv2_x1_0, 58 | 'mobilenetv2_x1_4': mobilenetv2_x1_4, 59 | 'shufflenet': shufflenet, 60 | 'squeezenet1_0': squeezenet1_0, 61 | 'squeezenet1_0_fc512': squeezenet1_0_fc512, 62 | 'squeezenet1_1': squeezenet1_1, 63 | 'shufflenet_v2_x0_5': shufflenet_v2_x0_5, 64 | 'shufflenet_v2_x1_0': shufflenet_v2_x1_0, 65 | 'shufflenet_v2_x1_5': shufflenet_v2_x1_5, 66 | 'shufflenet_v2_x2_0': shufflenet_v2_x2_0, 67 | } 68 | 69 | 70 | def show_avai_models(): 71 | """Displays available models. 72 | 73 | Examples:: 74 | >>> from torchreid import models 75 | >>> models.show_avai_models() 76 | """ 77 | print(list(__model_factory.keys())) 78 | 79 | 80 | def build_model( 81 | name, num_classes, loss='softmax', pretrained=True, use_gpu=True 82 | ): 83 | """A function wrapper for building a model. 84 | 85 | Args: 86 | name (str): model name. 87 | num_classes (int): number of training identities. 88 | loss (str, optional): loss function to optimize the model. Currently 89 | supports "softmax" and "triplet". Default is "softmax". 90 | pretrained (bool, optional): whether to load ImageNet-pretrained weights. 91 | Default is True. 92 | use_gpu (bool, optional): whether to use gpu. Default is True. 93 | 94 | Returns: 95 | nn.Module 96 | 97 | Examples:: 98 | >>> from torchreid import models 99 | >>> model = models.build_model('resnet50', 751, loss='softmax') 100 | """ 101 | avai_models = list(__model_factory.keys()) 102 | if name not in avai_models: 103 | raise KeyError( 104 | 'Unknown model: {}. Must be one of {}'.format(name, avai_models) 105 | ) 106 | return __model_factory[name]( 107 | num_classes=num_classes, 108 | loss=loss, 109 | pretrained=pretrained, 110 | use_gpu=use_gpu 111 | ) 112 | 113 | 114 | -------------------------------------------------------------------------------- /torchreid/optim/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from .optimizer import build_optimizer 4 | from .lr_scheduler import build_lr_scheduler 5 | -------------------------------------------------------------------------------- /torchreid/optim/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | import torch 3 | 4 | AVAI_SCH = ['single_step', 'multi_step', 'cosine'] 5 | 6 | 7 | def build_lr_scheduler( 8 | optimizer, lr_scheduler='single_step', stepsize=1, gamma=0.1, max_epoch=1 9 | ): 10 | """A function wrapper for building a learning rate scheduler. 11 | 12 | Args: 13 | optimizer (Optimizer): an Optimizer. 14 | lr_scheduler (str, optional): learning rate scheduler method. Default is single_step. 15 | stepsize (int or list, optional): step size to decay learning rate. When ``lr_scheduler`` 16 | is "single_step", ``stepsize`` should be an integer. When ``lr_scheduler`` is 17 | "multi_step", ``stepsize`` is a list. Default is 1. 18 | gamma (float, optional): decay rate. Default is 0.1. 19 | max_epoch (int, optional): maximum epoch (for cosine annealing). Default is 1. 20 | 21 | Examples:: 22 | >>> # Decay learning rate by every 20 epochs. 23 | >>> scheduler = torchreid.optim.build_lr_scheduler( 24 | >>> optimizer, lr_scheduler='single_step', stepsize=20 25 | >>> ) 26 | >>> # Decay learning rate at 30, 50 and 55 epochs. 27 | >>> scheduler = torchreid.optim.build_lr_scheduler( 28 | >>> optimizer, lr_scheduler='multi_step', stepsize=[30, 50, 55] 29 | >>> ) 30 | """ 31 | if lr_scheduler not in AVAI_SCH: 32 | raise ValueError( 33 | 'Unsupported scheduler: {}. Must be one of {}'.format( 34 | lr_scheduler, AVAI_SCH 35 | ) 36 | ) 37 | 38 | if lr_scheduler == 'single_step': 39 | if isinstance(stepsize, list): 40 | stepsize = stepsize[-1] 41 | 42 | if not isinstance(stepsize, int): 43 | raise TypeError( 44 | 'For single_step lr_scheduler, stepsize must ' 45 | 'be an integer, but got {}'.format(type(stepsize)) 46 | ) 47 | 48 | scheduler = torch.optim.lr_scheduler.StepLR( 49 | optimizer, step_size=stepsize, gamma=gamma 50 | ) 51 | 52 | elif lr_scheduler == 'multi_step': 53 | if not isinstance(stepsize, list): 54 | raise TypeError( 55 | 'For multi_step lr_scheduler, stepsize must ' 56 | 'be a list, but got {}'.format(type(stepsize)) 57 | ) 58 | 59 | scheduler = torch.optim.lr_scheduler.MultiStepLR( 60 | optimizer, milestones=stepsize, gamma=gamma 61 | ) 62 | 63 | elif lr_scheduler == 'cosine': 64 | scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( 65 | optimizer, float(max_epoch) 66 | ) 67 | 68 | return scheduler 69 | -------------------------------------------------------------------------------- /torchreid/optim/optimizer.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | import warnings 3 | import torch 4 | import torch.nn as nn 5 | 6 | from .radam import RAdam 7 | 8 | AVAI_OPTIMS = ['adam', 'amsgrad', 'sgd', 'rmsprop', 'radam'] 9 | 10 | 11 | def build_optimizer( 12 | model, 13 | optim='adam', 14 | lr=0.0003, 15 | weight_decay=5e-04, 16 | momentum=0.9, 17 | sgd_dampening=0, 18 | sgd_nesterov=True, 19 | rmsprop_alpha=0.99, 20 | adam_beta1=0.9, 21 | adam_beta2=0.99, 22 | staged_lr=False, 23 | new_layers='', 24 | base_lr_mult=0.1 25 | ): 26 | """A function wrapper for building an optimizer. 27 | 28 | Args: 29 | model (nn.Module): model. 30 | optim (str, optional): optimizer. Default is "adam". 31 | lr (float, optional): learning rate. Default is 0.0003. 32 | weight_decay (float, optional): weight decay (L2 penalty). Default is 5e-04. 33 | momentum (float, optional): momentum factor in sgd. Default is 0.9, 34 | sgd_dampening (float, optional): dampening for momentum. Default is 0. 35 | sgd_nesterov (bool, optional): enables Nesterov momentum. Default is False. 36 | rmsprop_alpha (float, optional): smoothing constant for rmsprop. Default is 0.99. 37 | adam_beta1 (float, optional): beta-1 value in adam. Default is 0.9. 38 | adam_beta2 (float, optional): beta-2 value in adam. Default is 0.99, 39 | staged_lr (bool, optional): uses different learning rates for base and new layers. Base 40 | layers are pretrained layers while new layers are randomly initialized, e.g. the 41 | identity classification layer. Enabling ``staged_lr`` can allow the base layers to 42 | be trained with a smaller learning rate determined by ``base_lr_mult``, while the new 43 | layers will take the ``lr``. Default is False. 44 | new_layers (str or list): attribute names in ``model``. Default is empty. 45 | base_lr_mult (float, optional): learning rate multiplier for base layers. Default is 0.1. 46 | 47 | Examples:: 48 | >>> # A normal optimizer can be built by 49 | >>> optimizer = torchreid.optim.build_optimizer(model, optim='sgd', lr=0.01) 50 | >>> # If you want to use a smaller learning rate for pretrained layers 51 | >>> # and the attribute name for the randomly initialized layer is 'classifier', 52 | >>> # you can do 53 | >>> optimizer = torchreid.optim.build_optimizer( 54 | >>> model, optim='sgd', lr=0.01, staged_lr=True, 55 | >>> new_layers='classifier', base_lr_mult=0.1 56 | >>> ) 57 | >>> # Now the `classifier` has learning rate 0.01 but the base layers 58 | >>> # have learning rate 0.01 * 0.1. 59 | >>> # new_layers can also take multiple attribute names. Say the new layers 60 | >>> # are 'fc' and 'classifier', you can do 61 | >>> optimizer = torchreid.optim.build_optimizer( 62 | >>> model, optim='sgd', lr=0.01, staged_lr=True, 63 | >>> new_layers=['fc', 'classifier'], base_lr_mult=0.1 64 | >>> ) 65 | """ 66 | if optim not in AVAI_OPTIMS: 67 | raise ValueError( 68 | 'Unsupported optim: {}. Must be one of {}'.format( 69 | optim, AVAI_OPTIMS 70 | ) 71 | ) 72 | 73 | if not isinstance(model, nn.Module): 74 | raise TypeError( 75 | 'model given to build_optimizer must be an instance of nn.Module' 76 | ) 77 | 78 | if staged_lr: 79 | if isinstance(new_layers, str): 80 | if new_layers is None: 81 | warnings.warn( 82 | 'new_layers is empty, therefore, staged_lr is useless' 83 | ) 84 | new_layers = [new_layers] 85 | 86 | if isinstance(model, nn.DataParallel): 87 | model = model.module 88 | 89 | base_params = [] 90 | base_layers = [] 91 | new_params = [] 92 | 93 | for name, module in model.named_children(): 94 | if name in new_layers: 95 | new_params += [p for p in module.parameters()] 96 | else: 97 | base_params += [p for p in module.parameters()] 98 | base_layers.append(name) 99 | 100 | param_groups = [ 101 | { 102 | 'params': base_params, 103 | 'lr': lr * base_lr_mult 104 | }, 105 | { 106 | 'params': new_params 107 | }, 108 | ] 109 | 110 | else: 111 | param_groups = model.parameters() 112 | 113 | if optim == 'adam': 114 | optimizer = torch.optim.Adam( 115 | param_groups, 116 | lr=lr, 117 | weight_decay=weight_decay, 118 | betas=(adam_beta1, adam_beta2), 119 | ) 120 | 121 | elif optim == 'amsgrad': 122 | optimizer = torch.optim.Adam( 123 | param_groups, 124 | lr=lr, 125 | weight_decay=weight_decay, 126 | betas=(adam_beta1, adam_beta2), 127 | amsgrad=True, 128 | ) 129 | 130 | elif optim == 'sgd': 131 | optimizer = torch.optim.SGD( 132 | param_groups, 133 | lr=lr, 134 | momentum=momentum, 135 | weight_decay=weight_decay, 136 | dampening=sgd_dampening, 137 | nesterov=True, 138 | ) 139 | 140 | elif optim == 'rmsprop': 141 | optimizer = torch.optim.RMSprop( 142 | param_groups, 143 | lr=lr, 144 | momentum=momentum, 145 | weight_decay=weight_decay, 146 | alpha=rmsprop_alpha, 147 | ) 148 | 149 | elif optim == 'radam': 150 | optimizer = RAdam( 151 | param_groups, 152 | lr=lr, 153 | weight_decay=weight_decay, 154 | betas=(adam_beta1, adam_beta2) 155 | ) 156 | 157 | return optimizer 158 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/README.md: -------------------------------------------------------------------------------- 1 | # Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective 2 | 3 | [[Paper]](https://arxiv.org/abs/2012.07620v2) 4 | 5 | On the Market-1501 dataset, we accelerate the re-ranking processing from **89.2s** to **9.4ms** with one K40m GPU, facilitating the real-time post-processing. 6 | Similarly, we observe that our method achieves comparable or even better retrieval results on the other four image retrieval benchmarks, 7 | i.e., VeRi-776, Oxford-5k, Paris-6k and University-1652, with limited time cost. 8 | 9 | ## Prerequisites 10 | 11 | The code was mainly developed and tested with python 3.7, PyTorch 1.4.1, CUDA 10.2, and CentOS release 6.10. 12 | 13 | The code has been included in `/extension`. To compile it: 14 | 15 | ```shell 16 | cd extension 17 | sh make.sh 18 | ``` 19 | 20 | ## Demo 21 | 22 | The demo script `main.py` provides the gnn re-ranking method using the prepared feature. 23 | 24 | ```shell 25 | python main.py --data_path PATH_TO_DATA --k1 26 --k2 7 26 | ``` 27 | 28 | ## Citation 29 | ```bibtex 30 | @article{zhang2020understanding, 31 | title={Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective}, 32 | author={Xuanmeng Zhang, Minyue Jiang, Zhedong Zheng, Xiao Tan, Errui Ding, Yi Yang}, 33 | journal={arXiv preprint arXiv:2012.07620}, 34 | year={2020} 35 | } 36 | ``` 37 | 38 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/extension/adjacency_matrix/build_adjacency_matrix.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | at::Tensor build_adjacency_matrix_forward(torch::Tensor initial_rank); 6 | 7 | 8 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 9 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 10 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 11 | 12 | at::Tensor build_adjacency_matrix(at::Tensor initial_rank) { 13 | CHECK_INPUT(initial_rank); 14 | return build_adjacency_matrix_forward(initial_rank); 15 | } 16 | 17 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 18 | m.def("forward", &build_adjacency_matrix, "build_adjacency_matrix (CUDA)"); 19 | } 20 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/extension/adjacency_matrix/build_adjacency_matrix_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | #define CUDA_1D_KERNEL_LOOP(i, n) for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) 8 | 9 | 10 | __global__ void build_adjacency_matrix_kernel(float* initial_rank, float* A, const int total_num, const int topk, const int nthreads, const int all_num) { 11 | int index = blockIdx.x * blockDim.x + threadIdx.x; 12 | int stride = blockDim.x * gridDim.x; 13 | for (int i = index; i < all_num; i += stride) { 14 | int ii = i / topk; 15 | A[ii * total_num + int(initial_rank[i])] = float(1.0); 16 | } 17 | } 18 | 19 | at::Tensor build_adjacency_matrix_forward(at::Tensor initial_rank) { 20 | const auto total_num = initial_rank.size(0); 21 | const auto topk = initial_rank.size(1); 22 | const auto all_num = total_num * topk; 23 | auto A = torch::zeros({total_num, total_num}, at::device(initial_rank.device()).dtype(at::ScalarType::Float)); 24 | 25 | const int threads = 1024; 26 | const int blocks = (all_num + threads - 1) / threads; 27 | 28 | build_adjacency_matrix_kernel<<>>(initial_rank.data_ptr(), A.data_ptr(), total_num, topk, threads, all_num); 29 | return A; 30 | 31 | } 32 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/extension/adjacency_matrix/setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective 3 | 4 | Xuanmeng Zhang, Minyue Jiang, Zhedong Zheng, Xiao Tan, Errui Ding, Yi Yang 5 | 6 | Project Page : https://github.com/Xuanmeng-Zhang/gnn-re-ranking 7 | 8 | Paper: https://arxiv.org/abs/2012.07620v2 9 | 10 | ====================================================================== 11 | 12 | On the Market-1501 dataset, we accelerate the re-ranking processing from 89.2s to 9.4ms 13 | with one K40m GPU, facilitating the real-time post-processing. Similarly, we observe 14 | that our method achieves comparable or even better retrieval results on the other four 15 | image retrieval benchmarks, i.e., VeRi-776, Oxford-5k, Paris-6k and University-1652, 16 | with limited time cost. 17 | """ 18 | 19 | from setuptools import Extension, setup 20 | import torch 21 | import torch.nn as nn 22 | from torch.autograd import Function 23 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension 24 | 25 | setup( 26 | name='build_adjacency_matrix', 27 | ext_modules=[ 28 | CUDAExtension( 29 | 'build_adjacency_matrix', [ 30 | 'build_adjacency_matrix.cpp', 31 | 'build_adjacency_matrix_kernel.cu', 32 | ] 33 | ), 34 | ], 35 | cmdclass={'build_ext': BuildExtension} 36 | ) 37 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/extension/make.sh: -------------------------------------------------------------------------------- 1 | cd adjacency_matrix 2 | python setup.py install 3 | cd ../propagation 4 | python setup.py install -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/extension/propagation/gnn_propagate.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | at::Tensor gnn_propagate_forward(at::Tensor A, at::Tensor initial_rank, at::Tensor S); 6 | 7 | 8 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 9 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 10 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 11 | 12 | at::Tensor gnn_propagate(at::Tensor A ,at::Tensor initial_rank, at::Tensor S) { 13 | CHECK_INPUT(A); 14 | CHECK_INPUT(initial_rank); 15 | CHECK_INPUT(S); 16 | return gnn_propagate_forward(A, initial_rank, S); 17 | } 18 | 19 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 20 | m.def("forward", &gnn_propagate, "gnn propagate (CUDA)"); 21 | } -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/extension/propagation/gnn_propagate_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | __global__ void gnn_propagate_forward_kernel(float* initial_rank, float* A, float* A_qe, float* S, const int sample_num, const int topk, const int total_num) { 9 | int index = blockIdx.x * blockDim.x + threadIdx.x; 10 | int stride = blockDim.x * gridDim.x; 11 | for (int i = index; i < total_num; i += stride) { 12 | int fea = i % sample_num; 13 | int sample_index = i / sample_num; 14 | float sum = 0.0; 15 | for (int j = 0; j < topk ; j++) { 16 | int topk_fea_index = int(initial_rank[sample_index*topk+j]) * sample_num + fea; 17 | sum += A[ topk_fea_index] * S[sample_index*topk+j]; 18 | } 19 | A_qe[i] = sum; 20 | } 21 | } 22 | 23 | at::Tensor gnn_propagate_forward(at::Tensor A, at::Tensor initial_rank, at::Tensor S) { 24 | const auto sample_num = A.size(0); 25 | const auto topk = initial_rank.size(1); 26 | 27 | const auto total_num = sample_num * sample_num ; 28 | auto A_qe = torch::zeros({sample_num, sample_num}, at::device(initial_rank.device()).dtype(at::ScalarType::Float)); 29 | 30 | const int threads = 1024; 31 | const int blocks = (total_num + threads - 1) / threads; 32 | 33 | gnn_propagate_forward_kernel<<>>(initial_rank.data_ptr(), A.data_ptr(), A_qe.data_ptr(), S.data_ptr(), sample_num, topk, total_num); 34 | return A_qe; 35 | 36 | } -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/extension/propagation/setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective 3 | 4 | Xuanmeng Zhang, Minyue Jiang, Zhedong Zheng, Xiao Tan, Errui Ding, Yi Yang 5 | 6 | Project Page : https://github.com/Xuanmeng-Zhang/gnn-re-ranking 7 | 8 | Paper: https://arxiv.org/abs/2012.07620v2 9 | 10 | ====================================================================== 11 | 12 | On the Market-1501 dataset, we accelerate the re-ranking processing from 89.2s to 9.4ms 13 | with one K40m GPU, facilitating the real-time post-processing. Similarly, we observe 14 | that our method achieves comparable or even better retrieval results on the other four 15 | image retrieval benchmarks, i.e., VeRi-776, Oxford-5k, Paris-6k and University-1652, 16 | with limited time cost. 17 | """ 18 | 19 | from setuptools import Extension, setup 20 | import torch 21 | import torch.nn as nn 22 | from torch.autograd import Function 23 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension 24 | 25 | setup( 26 | name='gnn_propagate', 27 | ext_modules=[ 28 | CUDAExtension( 29 | 'gnn_propagate', [ 30 | 'gnn_propagate.cpp', 31 | 'gnn_propagate_kernel.cu', 32 | ] 33 | ), 34 | ], 35 | cmdclass={'build_ext': BuildExtension} 36 | ) 37 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/gnn_reranking.py: -------------------------------------------------------------------------------- 1 | """ 2 | Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective 3 | 4 | Xuanmeng Zhang, Minyue Jiang, Zhedong Zheng, Xiao Tan, Errui Ding, Yi Yang 5 | 6 | Project Page : https://github.com/Xuanmeng-Zhang/gnn-re-ranking 7 | 8 | Paper: https://arxiv.org/abs/2012.07620v2 9 | 10 | ====================================================================== 11 | 12 | On the Market-1501 dataset, we accelerate the re-ranking processing from 89.2s to 9.4ms 13 | with one K40m GPU, facilitating the real-time post-processing. Similarly, we observe 14 | that our method achieves comparable or even better retrieval results on the other four 15 | image retrieval benchmarks, i.e., VeRi-776, Oxford-5k, Paris-6k and University-1652, 16 | with limited time cost. 17 | """ 18 | 19 | import numpy as np 20 | import torch 21 | 22 | import gnn_propagate 23 | import build_adjacency_matrix 24 | from utils import * 25 | 26 | 27 | def gnn_reranking(X_q, X_g, k1, k2): 28 | query_num, gallery_num = X_q.shape[0], X_g.shape[0] 29 | 30 | X_u = torch.cat((X_q, X_g), axis=0) 31 | original_score = torch.mm(X_u, X_u.t()) 32 | del X_u, X_q, X_g 33 | 34 | # initial ranking list 35 | S, initial_rank = original_score.topk( 36 | k=k1, dim=-1, largest=True, sorted=True 37 | ) 38 | 39 | # stage 1 40 | A = build_adjacency_matrix.forward(initial_rank.float()) 41 | S = S * S 42 | 43 | # stage 2 44 | if k2 != 1: 45 | for i in range(2): 46 | A = A + A.T 47 | A = gnn_propagate.forward( 48 | A, initial_rank[:, :k2].contiguous().float(), 49 | S[:, :k2].contiguous().float() 50 | ) 51 | A_norm = torch.norm(A, p=2, dim=1, keepdim=True) 52 | A = A.div(A_norm.expand_as(A)) 53 | 54 | cosine_similarity = torch.mm(A[:query_num, ], A[query_num:, ].t()) 55 | del A, S 56 | 57 | L = torch.sort(-cosine_similarity, dim=1)[1] 58 | L = L.data.cpu().numpy() 59 | return L 60 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/main.py: -------------------------------------------------------------------------------- 1 | """ 2 | Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective 3 | 4 | Xuanmeng Zhang, Minyue Jiang, Zhedong Zheng, Xiao Tan, Errui Ding, Yi Yang 5 | 6 | Project Page : https://github.com/Xuanmeng-Zhang/gnn-re-ranking 7 | 8 | Paper: https://arxiv.org/abs/2012.07620v2 9 | 10 | ====================================================================== 11 | 12 | On the Market-1501 dataset, we accelerate the re-ranking processing from 89.2s to 9.4ms 13 | with one K40m GPU, facilitating the real-time post-processing. Similarly, we observe 14 | that our method achieves comparable or even better retrieval results on the other four 15 | image retrieval benchmarks, i.e., VeRi-776, Oxford-5k, Paris-6k and University-1652, 16 | with limited time cost. 17 | """ 18 | 19 | import os 20 | import numpy as np 21 | import argparse 22 | import torch 23 | 24 | from utils import * 25 | from gnn_reranking import * 26 | 27 | parser = argparse.ArgumentParser(description='Reranking_is_GNN') 28 | parser.add_argument( 29 | '--data_path', 30 | type=str, 31 | default='../xm_rerank_gpu_2/features/market_88_test.pkl', 32 | help='path to dataset' 33 | ) 34 | parser.add_argument( 35 | '--k1', 36 | type=int, 37 | default=26, # Market-1501 38 | # default=60, # Veri-776 39 | help='parameter k1' 40 | ) 41 | parser.add_argument( 42 | '--k2', 43 | type=int, 44 | default=7, # Market-1501 45 | # default=10, # Veri-776 46 | help='parameter k2' 47 | ) 48 | 49 | args = parser.parse_args() 50 | 51 | 52 | def main(): 53 | data = load_pickle(args.data_path) 54 | 55 | query_cam = data['query_cam'] 56 | query_label = data['query_label'] 57 | gallery_cam = data['gallery_cam'] 58 | gallery_label = data['gallery_label'] 59 | 60 | gallery_feature = torch.FloatTensor(data['gallery_f']) 61 | query_feature = torch.FloatTensor(data['query_f']) 62 | query_feature = query_feature.cuda() 63 | gallery_feature = gallery_feature.cuda() 64 | 65 | indices = gnn_reranking(query_feature, gallery_feature, args.k1, args.k2) 66 | evaluate_ranking_list( 67 | indices, query_label, query_cam, gallery_label, gallery_cam 68 | ) 69 | 70 | 71 | if __name__ == '__main__': 72 | main() 73 | -------------------------------------------------------------------------------- /torchreid/utils/GPU-Re-Ranking/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Understanding Image Retrieval Re-Ranking: A Graph Neural Network Perspective 3 | 4 | Xuanmeng Zhang, Minyue Jiang, Zhedong Zheng, Xiao Tan, Errui Ding, Yi Yang 5 | 6 | Project Page : https://github.com/Xuanmeng-Zhang/gnn-re-ranking 7 | 8 | Paper: https://arxiv.org/abs/2012.07620v2 9 | 10 | ====================================================================== 11 | 12 | On the Market-1501 dataset, we accelerate the re-ranking processing from 89.2s to 9.4ms 13 | with one K40m GPU, facilitating the real-time post-processing. Similarly, we observe 14 | that our method achieves comparable or even better retrieval results on the other four 15 | image retrieval benchmarks, i.e., VeRi-776, Oxford-5k, Paris-6k and University-1652, 16 | with limited time cost. 17 | """ 18 | 19 | import os 20 | import numpy as np 21 | import pickle 22 | import torch 23 | 24 | 25 | def load_pickle(pickle_path): 26 | with open(pickle_path, 'rb') as f: 27 | data = pickle.load(f) 28 | return data 29 | 30 | 31 | def save_pickle(pickle_path, data): 32 | with open(pickle_path, 'wb') as f: 33 | pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL) 34 | 35 | 36 | def pairwise_squared_distance(x): 37 | ''' 38 | x : (n_samples, n_points, dims) 39 | return : (n_samples, n_points, n_points) 40 | ''' 41 | x2s = (x * x).sum(-1, keepdim=True) 42 | return x2s + x2s.transpose(-1, -2) - 2 * x @ x.transpose(-1, -2) 43 | 44 | 45 | def pairwise_distance(x, y): 46 | m, n = x.size(0), y.size(0) 47 | 48 | x = x.view(m, -1) 49 | y = y.view(n, -1) 50 | 51 | dist = torch.pow(x, 2).sum( 52 | dim=1, keepdim=True 53 | ).expand(m, n) + torch.pow(y, 2).sum( 54 | dim=1, keepdim=True 55 | ).expand(n, m).t() 56 | dist.addmm_(1, -2, x, y.t()) 57 | 58 | return dist 59 | 60 | 61 | def cosine_similarity(x, y): 62 | m, n = x.size(0), y.size(0) 63 | 64 | x = x.view(m, -1) 65 | y = y.view(n, -1) 66 | 67 | y = y.t() 68 | score = torch.mm(x, y) 69 | 70 | return score 71 | 72 | 73 | def evaluate_ranking_list( 74 | indices, query_label, query_cam, gallery_label, gallery_cam 75 | ): 76 | CMC = np.zeros((len(gallery_label)), dtype=np.int) 77 | ap = 0.0 78 | 79 | for i in range(len(query_label)): 80 | ap_tmp, CMC_tmp = evaluate( 81 | indices[i], query_label[i], query_cam[i], gallery_label, 82 | gallery_cam 83 | ) 84 | if CMC_tmp[0] == -1: 85 | continue 86 | CMC = CMC + CMC_tmp 87 | ap += ap_tmp 88 | 89 | CMC = CMC.astype(np.float32) 90 | CMC = CMC / len(query_label) #average CMC 91 | print( 92 | 'Rank@1:%f Rank@5:%f Rank@10:%f mAP:%f' % 93 | (CMC[0], CMC[4], CMC[9], ap / len(query_label)) 94 | ) 95 | 96 | 97 | def evaluate(index, ql, qc, gl, gc): 98 | query_index = np.argwhere(gl == ql) 99 | camera_index = np.argwhere(gc == qc) 100 | 101 | good_index = np.setdiff1d(query_index, camera_index, assume_unique=True) 102 | junk_index1 = np.argwhere(gl == -1) 103 | junk_index2 = np.intersect1d(query_index, camera_index) 104 | junk_index = np.append(junk_index2, junk_index1) #.flatten()) 105 | 106 | CMC_tmp = compute_mAP(index, good_index, junk_index) 107 | return CMC_tmp 108 | 109 | 110 | def compute_mAP(index, good_index, junk_index): 111 | ap = 0 112 | cmc = np.zeros((len(index)), dtype=np.int) 113 | if good_index.size == 0: # if empty 114 | cmc[0] = -1 115 | return ap, cmc 116 | 117 | # remove junk_index 118 | mask = np.in1d(index, junk_index, invert=True) 119 | index = index[mask] 120 | 121 | # find good_index index 122 | ngood = len(good_index) 123 | mask = np.in1d(index, good_index) 124 | rows_good = np.argwhere(mask == True) 125 | rows_good = rows_good.flatten() 126 | 127 | cmc[rows_good[0]:] = 1 128 | for i in range(ngood): 129 | d_recall = 1.0 / ngood 130 | precision = (i+1) * 1.0 / (rows_good[i] + 1) 131 | if rows_good[i] != 0: 132 | old_precision = i * 1.0 / rows_good[i] 133 | else: 134 | old_precision = 1.0 135 | ap = ap + d_recall * (old_precision+precision) / 2 136 | 137 | return ap, cmc 138 | -------------------------------------------------------------------------------- /torchreid/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from .tools import * 4 | from .rerank import re_ranking 5 | from .loggers import * 6 | from .avgmeter import * 7 | from .reidtools import * 8 | from .torchtools import * 9 | from .model_complexity import compute_model_complexity 10 | from .feature_extractor import FeatureExtractor 11 | -------------------------------------------------------------------------------- /torchreid/utils/avgmeter.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, absolute_import 2 | from collections import defaultdict 3 | import torch 4 | 5 | __all__ = ['AverageMeter', 'MetricMeter'] 6 | 7 | 8 | class AverageMeter(object): 9 | """Computes and stores the average and current value. 10 | 11 | Examples:: 12 | >>> # Initialize a meter to record loss 13 | >>> losses = AverageMeter() 14 | >>> # Update meter after every minibatch update 15 | >>> losses.update(loss_value, batch_size) 16 | """ 17 | 18 | def __init__(self): 19 | self.reset() 20 | 21 | def reset(self): 22 | self.val = 0 23 | self.avg = 0 24 | self.sum = 0 25 | self.count = 0 26 | 27 | def update(self, val, n=1): 28 | self.val = val 29 | self.sum += val * n 30 | self.count += n 31 | self.avg = self.sum / self.count 32 | 33 | 34 | class MetricMeter(object): 35 | """A collection of metrics. 36 | 37 | Source: https://github.com/KaiyangZhou/Dassl.pytorch 38 | 39 | Examples:: 40 | >>> # 1. Create an instance of MetricMeter 41 | >>> metric = MetricMeter() 42 | >>> # 2. Update using a dictionary as input 43 | >>> input_dict = {'loss_1': value_1, 'loss_2': value_2} 44 | >>> metric.update(input_dict) 45 | >>> # 3. Convert to string and print 46 | >>> print(str(metric)) 47 | """ 48 | 49 | def __init__(self, delimiter='\t'): 50 | self.meters = defaultdict(AverageMeter) 51 | self.delimiter = delimiter 52 | 53 | def update(self, input_dict): 54 | if input_dict is None: 55 | return 56 | 57 | if not isinstance(input_dict, dict): 58 | raise TypeError( 59 | 'Input to MetricMeter.update() must be a dictionary' 60 | ) 61 | 62 | for k, v in input_dict.items(): 63 | if isinstance(v, torch.Tensor): 64 | v = v.item() 65 | self.meters[k].update(v) 66 | 67 | def __str__(self): 68 | output_str = [] 69 | for name, meter in self.meters.items(): 70 | output_str.append( 71 | '{} {:.4f} ({:.4f})'.format(name, meter.val, meter.avg) 72 | ) 73 | return self.delimiter.join(output_str) 74 | -------------------------------------------------------------------------------- /torchreid/utils/feature_extractor.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import numpy as np 3 | import torch 4 | import torchvision.transforms as T 5 | from PIL import Image 6 | 7 | from torchreid.utils import ( 8 | check_isfile, load_pretrained_weights, compute_model_complexity 9 | ) 10 | from torchreid.models import build_model 11 | 12 | 13 | class FeatureExtractor(object): 14 | """A simple API for feature extraction. 15 | 16 | FeatureExtractor can be used like a python function, which 17 | accepts input of the following types: 18 | - a list of strings (image paths) 19 | - a list of numpy.ndarray each with shape (H, W, C) 20 | - a single string (image path) 21 | - a single numpy.ndarray with shape (H, W, C) 22 | - a torch.Tensor with shape (B, C, H, W) or (C, H, W) 23 | 24 | Returned is a torch tensor with shape (B, D) where D is the 25 | feature dimension. 26 | 27 | Args: 28 | model_name (str): model name. 29 | model_path (str): path to model weights. 30 | image_size (sequence or int): image height and width. 31 | pixel_mean (list): pixel mean for normalization. 32 | pixel_std (list): pixel std for normalization. 33 | pixel_norm (bool): whether to normalize pixels. 34 | device (str): 'cpu' or 'cuda' (could be specific gpu devices). 35 | verbose (bool): show model details. 36 | 37 | Examples:: 38 | 39 | from torchreid.utils import FeatureExtractor 40 | 41 | extractor = FeatureExtractor( 42 | model_name='osnet_x1_0', 43 | model_path='a/b/c/model.pth.tar', 44 | device='cuda' 45 | ) 46 | 47 | image_list = [ 48 | 'a/b/c/image001.jpg', 49 | 'a/b/c/image002.jpg', 50 | 'a/b/c/image003.jpg', 51 | 'a/b/c/image004.jpg', 52 | 'a/b/c/image005.jpg' 53 | ] 54 | 55 | features = extractor(image_list) 56 | print(features.shape) # output (5, 512) 57 | """ 58 | 59 | def __init__( 60 | self, 61 | model_name='', 62 | model_path='', 63 | image_size=(256, 128), 64 | pixel_mean=[0.485, 0.456, 0.406], 65 | pixel_std=[0.229, 0.224, 0.225], 66 | pixel_norm=True, 67 | device='cuda', 68 | verbose=True 69 | ): 70 | # Build model 71 | model = build_model( 72 | model_name, 73 | num_classes=1, 74 | pretrained=True, 75 | use_gpu=device.startswith('cuda') 76 | ) 77 | model.eval() 78 | 79 | if verbose: 80 | num_params, flops = compute_model_complexity( 81 | model, (1, 3, image_size[0], image_size[1]) 82 | ) 83 | print('Model: {}'.format(model_name)) 84 | print('- params: {:,}'.format(num_params)) 85 | print('- flops: {:,}'.format(flops)) 86 | 87 | if model_path and check_isfile(model_path): 88 | load_pretrained_weights(model, model_path) 89 | 90 | # Build transform functions 91 | transforms = [] 92 | transforms += [T.Resize(image_size)] 93 | transforms += [T.ToTensor()] 94 | if pixel_norm: 95 | transforms += [T.Normalize(mean=pixel_mean, std=pixel_std)] 96 | preprocess = T.Compose(transforms) 97 | 98 | to_pil = T.ToPILImage() 99 | 100 | device = torch.device(device) 101 | model.to(device) 102 | 103 | # Class attributes 104 | self.model = model 105 | self.preprocess = preprocess 106 | self.to_pil = to_pil 107 | self.device = device 108 | 109 | def __call__(self, input): 110 | if isinstance(input, list): 111 | images = [] 112 | 113 | for element in input: 114 | if isinstance(element, str): 115 | image = Image.open(element).convert('RGB') 116 | 117 | elif isinstance(element, np.ndarray): 118 | image = self.to_pil(element) 119 | 120 | else: 121 | raise TypeError( 122 | 'Type of each element must belong to [str | numpy.ndarray]' 123 | ) 124 | 125 | image = self.preprocess(image) 126 | images.append(image) 127 | 128 | images = torch.stack(images, dim=0) 129 | images = images.to(self.device) 130 | 131 | elif isinstance(input, str): 132 | image = Image.open(input).convert('RGB') 133 | image = self.preprocess(image) 134 | images = image.unsqueeze(0).to(self.device) 135 | 136 | elif isinstance(input, np.ndarray): 137 | image = self.to_pil(input) 138 | image = self.preprocess(image) 139 | images = image.unsqueeze(0).to(self.device) 140 | 141 | elif isinstance(input, torch.Tensor): 142 | if input.dim() == 3: 143 | input = input.unsqueeze(0) 144 | images = input.to(self.device) 145 | 146 | else: 147 | raise NotImplementedError 148 | 149 | with torch.no_grad(): 150 | features = self.model(images) 151 | 152 | return features 153 | -------------------------------------------------------------------------------- /torchreid/utils/loggers.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import os 3 | import sys 4 | import os.path as osp 5 | 6 | from .tools import mkdir_if_missing 7 | 8 | __all__ = ['Logger', 'RankLogger'] 9 | 10 | 11 | class Logger(object): 12 | """Writes console output to external text file. 13 | 14 | Imported from ``_ 15 | 16 | Args: 17 | fpath (str): directory to save logging file. 18 | 19 | Examples:: 20 | >>> import sys 21 | >>> import os 22 | >>> import os.path as osp 23 | >>> from torchreid.utils import Logger 24 | >>> save_dir = 'log/resnet50-softmax-market1501' 25 | >>> log_name = 'train.log' 26 | >>> sys.stdout = Logger(osp.join(args.save_dir, log_name)) 27 | """ 28 | 29 | def __init__(self, fpath=None): 30 | self.console = sys.stdout 31 | self.file = None 32 | if fpath is not None: 33 | mkdir_if_missing(osp.dirname(fpath)) 34 | self.file = open(fpath, 'w') 35 | 36 | def __del__(self): 37 | self.close() 38 | 39 | def __enter__(self): 40 | pass 41 | 42 | def __exit__(self, *args): 43 | self.close() 44 | 45 | def write(self, msg): 46 | self.console.write(msg) 47 | if self.file is not None: 48 | self.file.write(msg) 49 | 50 | def flush(self): 51 | self.console.flush() 52 | if self.file is not None: 53 | self.file.flush() 54 | os.fsync(self.file.fileno()) 55 | 56 | def close(self): 57 | self.console.close() 58 | if self.file is not None: 59 | self.file.close() 60 | 61 | 62 | class RankLogger(object): 63 | """Records the rank1 matching accuracy obtained for each 64 | test dataset at specified evaluation steps and provides a function 65 | to show the summarized results, which are convenient for analysis. 66 | 67 | Args: 68 | sources (str or list): source dataset name(s). 69 | targets (str or list): target dataset name(s). 70 | 71 | Examples:: 72 | >>> from torchreid.utils import RankLogger 73 | >>> s = 'market1501' 74 | >>> t = 'market1501' 75 | >>> ranklogger = RankLogger(s, t) 76 | >>> ranklogger.write(t, 10, 0.5) 77 | >>> ranklogger.write(t, 20, 0.7) 78 | >>> ranklogger.write(t, 30, 0.9) 79 | >>> ranklogger.show_summary() 80 | >>> # You will see: 81 | >>> # => Show performance summary 82 | >>> # market1501 (source) 83 | >>> # - epoch 10 rank1 50.0% 84 | >>> # - epoch 20 rank1 70.0% 85 | >>> # - epoch 30 rank1 90.0% 86 | >>> # If there are multiple test datasets 87 | >>> t = ['market1501', 'dukemtmcreid'] 88 | >>> ranklogger = RankLogger(s, t) 89 | >>> ranklogger.write(t[0], 10, 0.5) 90 | >>> ranklogger.write(t[0], 20, 0.7) 91 | >>> ranklogger.write(t[0], 30, 0.9) 92 | >>> ranklogger.write(t[1], 10, 0.1) 93 | >>> ranklogger.write(t[1], 20, 0.2) 94 | >>> ranklogger.write(t[1], 30, 0.3) 95 | >>> ranklogger.show_summary() 96 | >>> # You can see: 97 | >>> # => Show performance summary 98 | >>> # market1501 (source) 99 | >>> # - epoch 10 rank1 50.0% 100 | >>> # - epoch 20 rank1 70.0% 101 | >>> # - epoch 30 rank1 90.0% 102 | >>> # dukemtmcreid (target) 103 | >>> # - epoch 10 rank1 10.0% 104 | >>> # - epoch 20 rank1 20.0% 105 | >>> # - epoch 30 rank1 30.0% 106 | """ 107 | 108 | def __init__(self, sources, targets): 109 | self.sources = sources 110 | self.targets = targets 111 | 112 | if isinstance(self.sources, str): 113 | self.sources = [self.sources] 114 | 115 | if isinstance(self.targets, str): 116 | self.targets = [self.targets] 117 | 118 | self.logger = { 119 | name: { 120 | 'epoch': [], 121 | 'rank1': [] 122 | } 123 | for name in self.targets 124 | } 125 | 126 | def write(self, name, epoch, rank1): 127 | """Writes result. 128 | 129 | Args: 130 | name (str): dataset name. 131 | epoch (int): current epoch. 132 | rank1 (float): rank1 result. 133 | """ 134 | self.logger[name]['epoch'].append(epoch) 135 | self.logger[name]['rank1'].append(rank1) 136 | 137 | def show_summary(self): 138 | """Shows saved results.""" 139 | print('=> Show performance summary') 140 | for name in self.targets: 141 | from_where = 'source' if name in self.sources else 'target' 142 | print('{} ({})'.format(name, from_where)) 143 | for epoch, rank1 in zip( 144 | self.logger[name]['epoch'], self.logger[name]['rank1'] 145 | ): 146 | print('- epoch {}\t rank1 {:.1%}'.format(epoch, rank1)) 147 | -------------------------------------------------------------------------------- /torchreid/utils/reidtools.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function, absolute_import 2 | import numpy as np 3 | import shutil 4 | import os.path as osp 5 | import cv2 6 | 7 | from .tools import mkdir_if_missing 8 | 9 | __all__ = ['visualize_ranked_results'] 10 | 11 | GRID_SPACING = 10 12 | QUERY_EXTRA_SPACING = 90 13 | BW = 5 # border width 14 | GREEN = (0, 255, 0) 15 | RED = (0, 0, 255) 16 | 17 | 18 | def visualize_ranked_results( 19 | distmat, dataset, data_type, width=128, height=256, save_dir='', topk=10 20 | ): 21 | """Visualizes ranked results. 22 | 23 | Supports both image-reid and video-reid. 24 | 25 | For image-reid, ranks will be plotted in a single figure. For video-reid, ranks will be 26 | saved in folders each containing a tracklet. 27 | 28 | Args: 29 | distmat (numpy.ndarray): distance matrix of shape (num_query, num_gallery). 30 | dataset (tuple): a 2-tuple containing (query, gallery), each of which contains 31 | tuples of (img_path(s), pid, camid, dsetid). 32 | data_type (str): "image" or "video". 33 | width (int, optional): resized image width. Default is 128. 34 | height (int, optional): resized image height. Default is 256. 35 | save_dir (str): directory to save output images. 36 | topk (int, optional): denoting top-k images in the rank list to be visualized. 37 | Default is 10. 38 | """ 39 | num_q, num_g = distmat.shape 40 | mkdir_if_missing(save_dir) 41 | 42 | print('# query: {}\n# gallery {}'.format(num_q, num_g)) 43 | print('Visualizing top-{} ranks ...'.format(topk)) 44 | 45 | query, gallery = dataset 46 | assert num_q == len(query) 47 | assert num_g == len(gallery) 48 | 49 | indices = np.argsort(distmat, axis=1) 50 | 51 | def _cp_img_to(src, dst, rank, prefix, matched=False): 52 | """ 53 | Args: 54 | src: image path or tuple (for vidreid) 55 | dst: target directory 56 | rank: int, denoting ranked position, starting from 1 57 | prefix: string 58 | matched: bool 59 | """ 60 | if isinstance(src, (tuple, list)): 61 | if prefix == 'gallery': 62 | suffix = 'TRUE' if matched else 'FALSE' 63 | dst = osp.join( 64 | dst, prefix + '_top' + str(rank).zfill(3) 65 | ) + '_' + suffix 66 | else: 67 | dst = osp.join(dst, prefix + '_top' + str(rank).zfill(3)) 68 | mkdir_if_missing(dst) 69 | for img_path in src: 70 | shutil.copy(img_path, dst) 71 | else: 72 | dst = osp.join( 73 | dst, prefix + '_top' + str(rank).zfill(3) + '_name_' + 74 | osp.basename(src) 75 | ) 76 | shutil.copy(src, dst) 77 | 78 | for q_idx in range(num_q): 79 | qimg_path, qpid, qcamid = query[q_idx][:3] 80 | qimg_path_name = qimg_path[0] if isinstance( 81 | qimg_path, (tuple, list) 82 | ) else qimg_path 83 | 84 | if data_type == 'image': 85 | qimg = cv2.imread(qimg_path[0]) 86 | qimg = cv2.resize(qimg, (width, height)) 87 | qimg = cv2.copyMakeBorder( 88 | qimg, BW, BW, BW, BW, cv2.BORDER_CONSTANT, value=(0, 0, 0) 89 | ) 90 | # resize twice to ensure that the border width is consistent across images 91 | qimg = cv2.resize(qimg, (width, height)) 92 | num_cols = topk + 1 93 | grid_img = 255 * np.ones( 94 | ( 95 | height, 96 | num_cols*width + topk*GRID_SPACING + QUERY_EXTRA_SPACING, 3 97 | ), 98 | dtype=np.uint8 99 | ) 100 | grid_img[:, :width, :] = qimg 101 | else: 102 | qdir = osp.join( 103 | save_dir, osp.basename(osp.splitext(qimg_path_name)[0]) 104 | ) 105 | mkdir_if_missing(qdir) 106 | _cp_img_to(qimg_path, qdir, rank=0, prefix='query') 107 | 108 | rank_idx = 1 109 | for g_idx in indices[q_idx, :]: 110 | gimg_path, gpid, gcamid = gallery[g_idx][:3] 111 | invalid = (qpid == gpid) & (qcamid == gcamid) 112 | 113 | if not invalid: 114 | matched = gpid == qpid 115 | if data_type == 'image': 116 | border_color = GREEN if matched else RED 117 | gimg = cv2.imread(gimg_path[0]) 118 | gimg = cv2.resize(gimg, (width, height)) 119 | gimg = cv2.copyMakeBorder( 120 | gimg, 121 | BW, 122 | BW, 123 | BW, 124 | BW, 125 | cv2.BORDER_CONSTANT, 126 | value=border_color 127 | ) 128 | gimg = cv2.resize(gimg, (width, height)) 129 | start = rank_idx*width + rank_idx*GRID_SPACING + QUERY_EXTRA_SPACING 130 | end = ( 131 | rank_idx+1 132 | ) * width + rank_idx*GRID_SPACING + QUERY_EXTRA_SPACING 133 | grid_img[:, start:end, :] = gimg 134 | else: 135 | _cp_img_to( 136 | gimg_path, 137 | qdir, 138 | rank=rank_idx, 139 | prefix='gallery', 140 | matched=matched 141 | ) 142 | 143 | rank_idx += 1 144 | if rank_idx > topk: 145 | break 146 | 147 | if data_type == 'image': 148 | imname = osp.basename(osp.splitext(qimg_path_name)[0]) 149 | cv2.imwrite(osp.join(save_dir, imname + '.jpg'), grid_img) 150 | 151 | if (q_idx+1) % 100 == 0: 152 | print('- done {}/{}'.format(q_idx + 1, num_q)) 153 | 154 | print('Done. Images have been saved to "{}" ...'.format(save_dir)) 155 | -------------------------------------------------------------------------------- /torchreid/utils/rerank.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2/python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Source: https://github.com/zhunzhong07/person-re-ranking 5 | 6 | Created on Mon Jun 26 14:46:56 2017 7 | @author: luohao 8 | Modified by Houjing Huang, 2017-12-22. 9 | - This version accepts distance matrix instead of raw features. 10 | - The difference of `/` division between python 2 and 3 is handled. 11 | - numpy.float16 is replaced by numpy.float32 for numerical precision. 12 | 13 | CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017. 14 | url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf 15 | Matlab version: https://github.com/zhunzhong07/person-re-ranking 16 | 17 | API 18 | q_g_dist: query-gallery distance matrix, numpy array, shape [num_query, num_gallery] 19 | q_q_dist: query-query distance matrix, numpy array, shape [num_query, num_query] 20 | g_g_dist: gallery-gallery distance matrix, numpy array, shape [num_gallery, num_gallery] 21 | k1, k2, lambda_value: parameters, the original paper is (k1=20, k2=6, lambda_value=0.3) 22 | Returns: 23 | final_dist: re-ranked distance, numpy array, shape [num_query, num_gallery] 24 | """ 25 | from __future__ import division, print_function, absolute_import 26 | import numpy as np 27 | 28 | __all__ = ['re_ranking'] 29 | 30 | 31 | def re_ranking(q_g_dist, q_q_dist, g_g_dist, k1=20, k2=6, lambda_value=0.3): 32 | 33 | # The following naming, e.g. gallery_num, is different from outer scope. 34 | # Don't care about it. 35 | 36 | original_dist = np.concatenate( 37 | [ 38 | np.concatenate([q_q_dist, q_g_dist], axis=1), 39 | np.concatenate([q_g_dist.T, g_g_dist], axis=1) 40 | ], 41 | axis=0 42 | ) 43 | original_dist = np.power(original_dist, 2).astype(np.float32) 44 | original_dist = np.transpose( 45 | 1. * original_dist / np.max(original_dist, axis=0) 46 | ) 47 | V = np.zeros_like(original_dist).astype(np.float32) 48 | initial_rank = np.argsort(original_dist).astype(np.int32) 49 | 50 | query_num = q_g_dist.shape[0] 51 | gallery_num = q_g_dist.shape[0] + q_g_dist.shape[1] 52 | all_num = gallery_num 53 | 54 | for i in range(all_num): 55 | # k-reciprocal neighbors 56 | forward_k_neigh_index = initial_rank[i, :k1 + 1] 57 | backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1] 58 | fi = np.where(backward_k_neigh_index == i)[0] 59 | k_reciprocal_index = forward_k_neigh_index[fi] 60 | k_reciprocal_expansion_index = k_reciprocal_index 61 | for j in range(len(k_reciprocal_index)): 62 | candidate = k_reciprocal_index[j] 63 | candidate_forward_k_neigh_index = initial_rank[ 64 | candidate, :int(np.around(k1 / 2.)) + 1] 65 | candidate_backward_k_neigh_index = initial_rank[ 66 | candidate_forward_k_neigh_index, :int(np.around(k1 / 2.)) + 1] 67 | fi_candidate = np.where( 68 | candidate_backward_k_neigh_index == candidate 69 | )[0] 70 | candidate_k_reciprocal_index = candidate_forward_k_neigh_index[ 71 | fi_candidate] 72 | if len( 73 | np. 74 | intersect1d(candidate_k_reciprocal_index, k_reciprocal_index) 75 | ) > 2. / 3 * len(candidate_k_reciprocal_index): 76 | k_reciprocal_expansion_index = np.append( 77 | k_reciprocal_expansion_index, candidate_k_reciprocal_index 78 | ) 79 | 80 | k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) 81 | weight = np.exp(-original_dist[i, k_reciprocal_expansion_index]) 82 | V[i, k_reciprocal_expansion_index] = 1. * weight / np.sum(weight) 83 | original_dist = original_dist[:query_num, ] 84 | if k2 != 1: 85 | V_qe = np.zeros_like(V, dtype=np.float32) 86 | for i in range(all_num): 87 | V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0) 88 | V = V_qe 89 | del V_qe 90 | del initial_rank 91 | invIndex = [] 92 | for i in range(gallery_num): 93 | invIndex.append(np.where(V[:, i] != 0)[0]) 94 | 95 | jaccard_dist = np.zeros_like(original_dist, dtype=np.float32) 96 | 97 | for i in range(query_num): 98 | temp_min = np.zeros(shape=[1, gallery_num], dtype=np.float32) 99 | indNonZero = np.where(V[i, :] != 0)[0] 100 | indImages = [] 101 | indImages = [invIndex[ind] for ind in indNonZero] 102 | for j in range(len(indNonZero)): 103 | temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum( 104 | V[i, indNonZero[j]], V[indImages[j], indNonZero[j]] 105 | ) 106 | jaccard_dist[i] = 1 - temp_min / (2.-temp_min) 107 | 108 | final_dist = jaccard_dist * (1-lambda_value) + original_dist*lambda_value 109 | del original_dist 110 | del V 111 | del jaccard_dist 112 | final_dist = final_dist[:query_num, query_num:] 113 | return final_dist 114 | -------------------------------------------------------------------------------- /torchreid/utils/tools.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, print_function, absolute_import 2 | import os 3 | import sys 4 | import json 5 | import time 6 | import errno 7 | import numpy as np 8 | import random 9 | import os.path as osp 10 | import warnings 11 | import PIL 12 | import torch 13 | from PIL import Image 14 | 15 | __all__ = [ 16 | 'mkdir_if_missing', 'check_isfile', 'read_json', 'write_json', 17 | 'set_random_seed', 'download_url', 'read_image', 'collect_env_info', 18 | 'listdir_nohidden' 19 | ] 20 | 21 | 22 | def mkdir_if_missing(dirname): 23 | """Creates dirname if it is missing.""" 24 | if not osp.exists(dirname): 25 | try: 26 | os.makedirs(dirname) 27 | except OSError as e: 28 | if e.errno != errno.EEXIST: 29 | raise 30 | 31 | 32 | def check_isfile(fpath): 33 | """Checks if the given path is a file. 34 | 35 | Args: 36 | fpath (str): file path. 37 | 38 | Returns: 39 | bool 40 | """ 41 | isfile = osp.isfile(fpath) 42 | if not isfile: 43 | warnings.warn('No file found at "{}"'.format(fpath)) 44 | return isfile 45 | 46 | 47 | def read_json(fpath): 48 | """Reads json file from a path.""" 49 | with open(fpath, 'r') as f: 50 | obj = json.load(f) 51 | return obj 52 | 53 | 54 | def write_json(obj, fpath): 55 | """Writes to a json file.""" 56 | mkdir_if_missing(osp.dirname(fpath)) 57 | with open(fpath, 'w') as f: 58 | json.dump(obj, f, indent=4, separators=(',', ': ')) 59 | 60 | 61 | def set_random_seed(seed): 62 | random.seed(seed) 63 | np.random.seed(seed) 64 | torch.manual_seed(seed) 65 | torch.cuda.manual_seed_all(seed) 66 | 67 | 68 | def download_url(url, dst): 69 | """Downloads file from a url to a destination. 70 | 71 | Args: 72 | url (str): url to download file. 73 | dst (str): destination path. 74 | """ 75 | from six.moves import urllib 76 | print('* url="{}"'.format(url)) 77 | print('* destination="{}"'.format(dst)) 78 | 79 | def _reporthook(count, block_size, total_size): 80 | global start_time 81 | if count == 0: 82 | start_time = time.time() 83 | return 84 | duration = time.time() - start_time 85 | progress_size = int(count * block_size) 86 | speed = int(progress_size / (1024*duration)) 87 | percent = int(count * block_size * 100 / total_size) 88 | sys.stdout.write( 89 | '\r...%d%%, %d MB, %d KB/s, %d seconds passed' % 90 | (percent, progress_size / (1024*1024), speed, duration) 91 | ) 92 | sys.stdout.flush() 93 | 94 | urllib.request.urlretrieve(url, dst, _reporthook) 95 | sys.stdout.write('\n') 96 | 97 | 98 | def read_image(path): 99 | """Reads image from path using ``PIL.Image``. 100 | 101 | Args: 102 | path (str): path to an image. 103 | 104 | Returns: 105 | PIL image 106 | """ 107 | got_img = False 108 | if not osp.exists(path): 109 | raise IOError('"{}" does not exist'.format(path)) 110 | while not got_img: 111 | try: 112 | img = Image.open(path).convert('RGB') 113 | got_img = True 114 | except IOError: 115 | print( 116 | 'IOError incurred when reading "{}". Will redo. Don\'t worry. Just chill.' 117 | .format(path) 118 | ) 119 | return img 120 | 121 | 122 | def collect_env_info(): 123 | """Returns env info as a string. 124 | 125 | Code source: github.com/facebookresearch/maskrcnn-benchmark 126 | """ 127 | from torch.utils.collect_env import get_pretty_env_info 128 | env_str = get_pretty_env_info() 129 | env_str += '\n Pillow ({})'.format(PIL.__version__) 130 | return env_str 131 | 132 | 133 | def listdir_nohidden(path, sort=False): 134 | """List non-hidden items in a directory. 135 | 136 | Args: 137 | path (str): directory path. 138 | sort (bool): sort the items. 139 | """ 140 | items = [f for f in os.listdir(path) if not f.startswith('.')] 141 | if sort: 142 | items.sort() 143 | return items 144 | --------------------------------------------------------------------------------