update

util/FeatureExtractor.py

+import torch.nn as nn
+from IPython import embed
+
+class FeatureExtractor(nn.Module):
+    def __init__(self,submodule,extracted_layers):
+        super(FeatureExtractor,self).__init__()
+        self.submodule = submodule
+        self.extracted_layers = extracted_layers
+
+    def forward(self, x):
+        outputs = []
+        for name, module in self.submodule._modules.items():
+            if name is "classfier":
+                x = x.view(x.size(0),-1)
+            if name is "base":
+                for block_name, cnn_block in module._modules.items():
+                    x = cnn_block(x)
+                    if block_name in self.extracted_layers:
+                        outputs.append(x)
+        return outputs
\ No newline at end of file

util/data_manager.py

+from __future__ import print_function, absolute_import
+import os
+import glob
+import re
+import sys
+import urllib
+import tarfile
+import zipfile
+import os.path as osp
+from scipy.io import loadmat
+import numpy as np
+import h5py
+from scipy.misc import imsave
+
+from util.utils import mkdir_if_missing, write_json, read_json
+
+"""Image ReID"""
+
+class Market1501(object):
+    """
+    Market1501
+
+    Reference:
+    Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
+
+    URL: http://www.liangzheng.org/Project/project_reid.html
+    
+    Dataset statistics:
+    # identities: 1501 (+1 for background)
+    # images: 12936 (train) + 3368 (query) + 15913 (gallery)
+    """
+    dataset_dir = 'market1501'
+
+    def __init__(self, root='data', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
+        self.query_dir = osp.join(self.dataset_dir, 'query')
+        self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
+
+        self._check_before_run()
+
+        train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, relabel=True)
+        query, num_query_pids, num_query_imgs = self._process_dir(self.query_dir, relabel=False)
+        gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.gallery_dir, relabel=False)
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
+
+        print("=> Market1501 loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # images")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.train_dir):
+            raise RuntimeError("'{}' is not available".format(self.train_dir))
+        if not osp.exists(self.query_dir):
+            raise RuntimeError("'{}' is not available".format(self.query_dir))
+        if not osp.exists(self.gallery_dir):
+            raise RuntimeError("'{}' is not available".format(self.gallery_dir))
+
+    def _process_dir(self, dir_path, relabel=False):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'([-\d]+)_c(\d)')
+
+        pid_container = set()
+        for img_path in img_paths:
+            pid, _ = map(int, pattern.search(img_path).groups())
+            if pid == -1: continue  # junk images are just ignored
+            pid_container.add(pid)
+        pid2label = {pid:label for label, pid in enumerate(pid_container)}
+
+        dataset = []
+        for img_path in img_paths:
+            pid, camid = map(int, pattern.search(img_path).groups())
+            if pid == -1: continue  # junk images are just ignored
+            assert 0 <= pid <= 1501  # pid == 0 means background
+            assert 1 <= camid <= 6
+            camid -= 1 # index starts from 0
+            if relabel: pid = pid2label[pid]
+            dataset.append((img_path, pid, camid))
+
+        num_pids = len(pid_container)
+        num_imgs = len(dataset)
+        return dataset, num_pids, num_imgs
+
+
+class Market1501_Partial(object):
+    """
+    Market1501
+
+    Reference:
+    Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
+
+    URL: http://www.liangzheng.org/Project/project_reid.html
+
+    Dataset statistics:
+    # identities: 1501 (+1 for background)
+    # images: 12936 (train) + 3368 (query) + 15913 (gallery)
+    """
+    dataset_dir = 'market1501_partial'
+
+    def __init__(self, root='data', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
+        self.query_dir = osp.join(self.dataset_dir, 'query')
+        self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
+
+        self._check_before_run()
+
+        train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, relabel=True)
+        query, num_query_pids, num_query_imgs = self._process_dir(self.query_dir, relabel=False)
+        gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.gallery_dir, relabel=False)
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
+
+        print("=> Market1501 loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # images")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.train_dir):
+            raise RuntimeError("'{}' is not available".format(self.train_dir))
+        if not osp.exists(self.query_dir):
+            raise RuntimeError("'{}' is not available".format(self.query_dir))
+        if not osp.exists(self.gallery_dir):
+            raise RuntimeError("'{}' is not available".format(self.gallery_dir))
+
+    def _process_dir(self, dir_path, relabel=False):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'([-\d]+)_c(\d)')
+
+        pid_container = set()
+        for img_path in img_paths:
+            pid, _ = map(int, pattern.search(img_path).groups())
+            if pid == -1: continue  # junk images are just ignored
+            pid_container.add(pid)
+        pid2label = {pid: label for label, pid in enumerate(pid_container)}
+
+        dataset = []
+        for img_path in img_paths:
+            pid, camid = map(int, pattern.search(img_path).groups())
+            if pid == -1: continue  # junk images are just ignored
+            assert 0 <= pid <= 1501  # pid == 0 means background
+            assert 1 <= camid <= 6
+            camid -= 1  # index starts from 0
+            if relabel: pid = pid2label[pid]
+            dataset.append((img_path, pid, camid))
+
+        num_pids = len(pid_container)
+        num_imgs = len(dataset)
+        return dataset, num_pids, num_imgs
+
+class CUHK03(object):
+    """
+    CUHK03
+
+    Reference:
+    Li et al. DeepReID: Deep Filter Pairing Neural Network for Person Re-identification. CVPR 2014.
+
+    URL: http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html#!
+    
+    Dataset statistics:
+    # identities: 1360
+    # images: 13164
+    # cameras: 6
+    # splits: 20 (classic)
+
+    Args:
+        split_id (int): split index (default: 0)
+        cuhk03_labeled (bool): whether to load labeled images; if false, detected images are loaded (default: False)
+    """
+    dataset_dir = 'cuhk03'
+
+    def __init__(self, root='data', split_id=0, cuhk03_labeled=False, cuhk03_classic_split=False, **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.data_dir = osp.join(self.dataset_dir, 'cuhk03_release')
+        self.raw_mat_path = osp.join(self.data_dir, 'cuhk-03.mat')
+        
+        self.imgs_detected_dir = osp.join(self.dataset_dir, 'images_detected')
+        self.imgs_labeled_dir = osp.join(self.dataset_dir, 'images_labeled')
+        
+        self.split_classic_det_json_path = osp.join(self.dataset_dir, 'splits_classic_detected.json')
+        self.split_classic_lab_json_path = osp.join(self.dataset_dir, 'splits_classic_labeled.json')
+        
+        self.split_new_det_json_path = osp.join(self.dataset_dir, 'splits_new_detected.json')
+        self.split_new_lab_json_path = osp.join(self.dataset_dir, 'splits_new_labeled.json')
+        
+        self.split_new_det_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_detected.mat')
+        self.split_new_lab_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_labeled.mat')
+
+        self._check_before_run()
+        self._preprocess()
+
+        if cuhk03_labeled:
+            image_type = 'labeled'
+            split_path = self.split_classic_lab_json_path if cuhk03_classic_split else self.split_new_lab_json_path
+        else:
+            image_type = 'detected'
+            split_path = self.split_classic_det_json_path if cuhk03_classic_split else self.split_new_det_json_path
+
+        splits = read_json(split_path)
+        assert split_id < len(splits), "Condition split_id ({}) < len(splits) ({}) is false".format(split_id, len(splits))
+        split = splits[split_id]
+        print("Split index = {}".format(split_id))
+
+        train = split['train']
+        query = split['query']
+        gallery = split['gallery']
+
+        num_train_pids = split['num_train_pids']
+        num_query_pids = split['num_query_pids']
+        num_gallery_pids = split['num_gallery_pids']
+        num_total_pids = num_train_pids + num_query_pids
+
+        num_train_imgs = split['num_train_imgs']
+        num_query_imgs = split['num_query_imgs']
+        num_gallery_imgs = split['num_gallery_imgs']
+        num_total_imgs = num_train_imgs + num_query_imgs
+
+        print("=> CUHK03 ({}) loaded".format(image_type))
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # images")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.data_dir):
+            raise RuntimeError("'{}' is not available".format(self.data_dir))
+        if not osp.exists(self.raw_mat_path):
+            raise RuntimeError("'{}' is not available".format(self.raw_mat_path))
+        if not osp.exists(self.split_new_det_mat_path):
+            raise RuntimeError("'{}' is not available".format(self.split_new_det_mat_path))
+        if not osp.exists(self.split_new_lab_mat_path):
+            raise RuntimeError("'{}' is not available".format(self.split_new_lab_mat_path))
+
+    def _preprocess(self):
+        """
+        This function is a bit complex and ugly, what it does is
+        1. Extract data from cuhk-03.mat and save as png images.
+        2. Create 20 classic splits. (Li et al. CVPR'14)
+        3. Create new split. (Zhong et al. CVPR'17)
+        """
+        print("Note: if root path is changed, the previously generated json files need to be re-generated (delete them first)")
+        if osp.exists(self.imgs_labeled_dir) and \
+           osp.exists(self.imgs_detected_dir) and \
+           osp.exists(self.split_classic_det_json_path) and \
+           osp.exists(self.split_classic_lab_json_path) and \
+           osp.exists(self.split_new_det_json_path) and \
+           osp.exists(self.split_new_lab_json_path):
+            return
+
+        mkdir_if_missing(self.imgs_detected_dir)
+        mkdir_if_missing(self.imgs_labeled_dir)
+
+        print("Extract image data from {} and save as png".format(self.raw_mat_path))
+        mat = h5py.File(self.raw_mat_path, 'r')
+
+        def _deref(ref):
+            return mat[ref][:].T
+
+        def _process_images(img_refs, campid, pid, save_dir):
+            img_paths = [] # Note: some persons only have images for one view
+            for imgid, img_ref in enumerate(img_refs):
+                img = _deref(img_ref)
+                # skip empty cell
+                if img.size == 0 or img.ndim < 3: continue
+                # images are saved with the following format, index-1 (ensure uniqueness)
+                # campid: index of camera pair (1-5)
+                # pid: index of person in 'campid'-th camera pair
+                # viewid: index of view, {1, 2}
+                # imgid: index of image, (1-10)
+                viewid = 1 if imgid < 5 else 2
+                img_name = '{:01d}_{:03d}_{:01d}_{:02d}.png'.format(campid+1, pid+1, viewid, imgid+1)
+                img_path = osp.join(save_dir, img_name)
+                imsave(img_path, img)
+                img_paths.append(img_path)
+            return img_paths
+
+        def _extract_img(name):
+            print("Processing {} images (extract and save) ...".format(name))
+            meta_data = []
+            imgs_dir = self.imgs_detected_dir if name == 'detected' else self.imgs_labeled_dir
+            for campid, camp_ref in enumerate(mat[name][0]):
+                camp = _deref(camp_ref)
+                num_pids = camp.shape[0]
+                for pid in range(num_pids):
+                    img_paths = _process_images(camp[pid,:], campid, pid, imgs_dir)
+                    assert len(img_paths) > 0, "campid{}-pid{} has no images".format(campid, pid)
+                    meta_data.append((campid+1, pid+1, img_paths))
+                print("done camera pair {} with {} identities".format(campid+1, num_pids))
+            return meta_data
+
+        meta_detected = _extract_img('detected')
+        meta_labeled = _extract_img('labeled')
+
+        def _extract_classic_split(meta_data, test_split):
+            train, test = [], []
+            num_train_pids, num_test_pids = 0, 0
+            num_train_imgs, num_test_imgs = 0, 0
+            for i, (campid, pid, img_paths) in enumerate(meta_data):
+                
+                if [campid, pid] in test_split:
+                    for img_path in img_paths:
+                        camid = int(osp.basename(img_path).split('_')[2])
+                        test.append((img_path, num_test_pids, camid))
+                    num_test_pids += 1
+                    num_test_imgs += len(img_paths)
+                else:
+                    for img_path in img_paths:
+                        camid = int(osp.basename(img_path).split('_')[2])
+                        train.append((img_path, num_train_pids, camid))
+                    num_train_pids += 1
+                    num_train_imgs += len(img_paths)
+            return train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs
+
+        print("Creating classic splits (# = 20) ...")
+        splits_classic_det, splits_classic_lab = [], []
+        for split_ref in mat['testsets'][0]:
+            test_split = _deref(split_ref).tolist()
+
+            # create split for detected images
+            train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
+                _extract_classic_split(meta_detected, test_split)
+            splits_classic_det.append({
+                'train': train, 'query': test, 'gallery': test,
+                'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs,
+                'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs,
+                'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs,
+            })
+
+            # create split for labeled images
+            train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
+                _extract_classic_split(meta_labeled, test_split)
+            splits_classic_lab.append({
+                'train': train, 'query': test, 'gallery': test,
+                'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs,
+                'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs,
+                'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs,
+            })
+        
+        write_json(splits_classic_det, self.split_classic_det_json_path)
+        write_json(splits_classic_lab, self.split_classic_lab_json_path)
+
+        def _extract_set(filelist, pids, pid2label, idxs, img_dir, relabel):
+            tmp_set = []
+            unique_pids = set()
+            for idx in idxs:
+                img_name = filelist[idx][0]
+                camid = int(img_name.split('_')[2])
+                pid = pids[idx]
+                if relabel: pid = pid2label[pid]
+                img_path = osp.join(img_dir, img_name)
+                tmp_set.append((img_path, int(pid), camid))
+                unique_pids.add(pid)
+            return tmp_set, len(unique_pids), len(idxs)
+
+        def _extract_new_split(split_dict, img_dir):
+            train_idxs = split_dict['train_idx'].flatten() - 1 # index-0
+            pids = split_dict['labels'].flatten()
+            train_pids = set(pids[train_idxs])
+            pid2label = {pid: label for label, pid in enumerate(train_pids)}
+            query_idxs = split_dict['query_idx'].flatten() - 1
+            gallery_idxs = split_dict['gallery_idx'].flatten() - 1
+            filelist = split_dict['filelist'].flatten()
+            train_info = _extract_set(filelist, pids, pid2label, train_idxs, img_dir, relabel=True)
+            query_info = _extract_set(filelist, pids, pid2label, query_idxs, img_dir, relabel=False)
+            gallery_info = _extract_set(filelist, pids, pid2label, gallery_idxs, img_dir, relabel=False)
+            return train_info, query_info, gallery_info
+
+        print("Creating new splits for detected images (767/700) ...")
+        train_info, query_info, gallery_info = _extract_new_split(
+            loadmat(self.split_new_det_mat_path),
+            self.imgs_detected_dir,
+        )
+        splits = [{
+            'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0],
+            'num_train_pids': train_info[1], 'num_train_imgs': train_info[2],
+            'num_query_pids': query_info[1], 'num_query_imgs': query_info[2],
+            'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2],
+        }]
+        write_json(splits, self.split_new_det_json_path)
+
+        print("Creating new splits for labeled images (767/700) ...")
+        train_info, query_info, gallery_info = _extract_new_split(
+            loadmat(self.split_new_lab_mat_path),
+            self.imgs_labeled_dir,
+        )
+        splits = [{
+            'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0],
+            'num_train_pids': train_info[1], 'num_train_imgs': train_info[2],
+            'num_query_pids': query_info[1], 'num_query_imgs': query_info[2],
+            'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2],
+        }]
+        write_json(splits, self.split_new_lab_json_path)
+
+class DukeMTMCreID(object):
+    """
+    DukeMTMC-reID
+
+    Reference:
+    1. Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.
+    2. Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.
+
+    URL: https://github.com/layumi/DukeMTMC-reID_evaluation
+    
+    Dataset statistics:
+    # identities: 1404 (train + query)
+    # images:16522 (train) + 2228 (query) + 17661 (gallery)
+    # cameras: 8
+    """
+    dataset_dir = 'dukemtmc-reid'
+
+    def __init__(self, root='data', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.train_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/bounding_box_train')
+        self.query_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/query')
+        self.gallery_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/bounding_box_test')
+
+        self._check_before_run()
+
+        train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, relabel=True)
+        query, num_query_pids, num_query_imgs = self._process_dir(self.query_dir, relabel=False)
+        gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.gallery_dir, relabel=False)
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
+
+        print("=> DukeMTMC-reID loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # images")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.train_dir):
+            raise RuntimeError("'{}' is not available".format(self.train_dir))
+        if not osp.exists(self.query_dir):
+            raise RuntimeError("'{}' is not available".format(self.query_dir))
+        if not osp.exists(self.gallery_dir):
+            raise RuntimeError("'{}' is not available".format(self.gallery_dir))
+
+    def _process_dir(self, dir_path, relabel=False):
+        img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
+        pattern = re.compile(r'([-\d]+)_c(\d)')
+
+        pid_container = set()
+        for img_path in img_paths:
+            pid, _ = map(int, pattern.search(img_path).groups())
+            pid_container.add(pid)
+        pid2label = {pid:label for label, pid in enumerate(pid_container)}
+
+        dataset = []
+        for img_path in img_paths:
+            pid, camid = map(int, pattern.search(img_path).groups())
+            assert 1 <= camid <= 8
+            camid -= 1 # index starts from 0
+            if relabel: pid = pid2label[pid]
+            dataset.append((img_path, pid, camid))
+
+        num_pids = len(pid_container)
+        num_imgs = len(dataset)
+        return dataset, num_pids, num_imgs
+
+class MSMT17(object):
+    """
+    MSMT17
+
+    Reference:
+    Wei et al. Person Transfer GAN to Bridge Domain Gap for Person Re-Identification. CVPR 2018.
+
+    URL: http://www.pkuvmc.com/publications/msmt17.html
+    
+    Dataset statistics:
+    # identities: 4101
+    # images: 32621 (train) + 11659 (query) + 82161 (gallery)
+    # cameras: 15
+    """
+    dataset_dir = 'msmt17'
+
+    def __init__(self, root='data', **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.train_dir = osp.join(self.dataset_dir, 'MSMT17_V1/train')
+        self.test_dir = osp.join(self.dataset_dir, 'MSMT17_V1/test')
+        self.list_train_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_train.txt')
+        self.list_val_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_val.txt')
+        self.list_query_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_query.txt')
+        self.list_gallery_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_gallery.txt')
+
+        self._check_before_run()
+        train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, self.list_train_path)
+        #val, num_val_pids, num_val_imgs = self._process_dir(self.train_dir, self.list_val_path)
+        query, num_query_pids, num_query_imgs = self._process_dir(self.test_dir, self.list_query_path)
+        gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.test_dir, self.list_gallery_path)
+
+        #train += val
+        #num_train_imgs += num_val_imgs
+
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
+
+        print("=> MSMT17 loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # images")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.train_dir):
+            raise RuntimeError("'{}' is not available".format(self.train_dir))
+        if not osp.exists(self.test_dir):
+            raise RuntimeError("'{}' is not available".format(self.test_dir))
+
+    def _process_dir(self, dir_path, list_path):
+        with open(list_path, 'r') as txt:
+            lines = txt.readlines()
+        dataset = []
+        pid_container = set()
+        for img_idx, img_info in enumerate(lines):
+            img_path, pid = img_info.split(' ')
+            pid = int(pid) # no need to relabel
+            camid = int(img_path.split('_')[2])
+            img_path = osp.join(dir_path, img_path)
+            dataset.append((img_path, pid, camid))
+            pid_container.add(pid)
+        num_imgs = len(dataset)
+        num_pids = len(pid_container)
+        # check if pid starts from 0 and increments with 1
+        for idx, pid in enumerate(pid_container):
+            assert idx == pid, "See code comment for explanation"
+        return dataset, num_pids, num_imgs
+
+"""Video ReID"""
+
+class Mars(object):
+    """
+    MARS
+
+    Reference:
+    Zheng et al. MARS: A Video Benchmark for Large-Scale Person Re-identification. ECCV 2016.
+
+    URL: http://www.liangzheng.com.cn/Project/project_mars.html
+    
+    Dataset statistics:
+    # identities: 1261
+    # tracklets: 8298 (train) + 1980 (query) + 9330 (gallery)
+    # cameras: 6
+    """
+    dataset_dir = 'mars'
+
+    def __init__(self, root='data', min_seq_len=0, **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.train_name_path = osp.join(self.dataset_dir, 'info/train_name.txt')
+        self.test_name_path = osp.join(self.dataset_dir, 'info/test_name.txt')
+        self.track_train_info_path = osp.join(self.dataset_dir, 'info/tracks_train_info.mat')
+        self.track_test_info_path = osp.join(self.dataset_dir, 'info/tracks_test_info.mat')
+        self.query_IDX_path = osp.join(self.dataset_dir, 'info/query_IDX.mat')
+
+        self._check_before_run()
+
+        # prepare meta data
+        train_names = self._get_names(self.train_name_path)
+        test_names = self._get_names(self.test_name_path)
+        track_train = loadmat(self.track_train_info_path)['track_train_info'] # numpy.ndarray (8298, 4)
+        track_test = loadmat(self.track_test_info_path)['track_test_info'] # numpy.ndarray (12180, 4)
+        query_IDX = loadmat(self.query_IDX_path)['query_IDX'].squeeze() # numpy.ndarray (1980,)
+        query_IDX -= 1 # index from 0
+        track_query = track_test[query_IDX,:]
+        gallery_IDX = [i for i in range(track_test.shape[0]) if i not in query_IDX]
+        track_gallery = track_test[gallery_IDX,:]
+
+        train, num_train_tracklets, num_train_pids, num_train_imgs = \
+          self._process_data(train_names, track_train, home_dir='bbox_train', relabel=True, min_seq_len=min_seq_len)
+
+        query, num_query_tracklets, num_query_pids, num_query_imgs = \
+          self._process_data(test_names, track_query, home_dir='bbox_test', relabel=False, min_seq_len=min_seq_len)
+
+        gallery, num_gallery_tracklets, num_gallery_pids, num_gallery_imgs = \
+          self._process_data(test_names, track_gallery, home_dir='bbox_test', relabel=False, min_seq_len=min_seq_len)
+
+        num_imgs_per_tracklet = num_train_imgs + num_query_imgs + num_gallery_imgs
+        min_num = np.min(num_imgs_per_tracklet)
+        max_num = np.max(num_imgs_per_tracklet)
+        avg_num = np.mean(num_imgs_per_tracklet)
+
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
+
+        print("=> MARS loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # tracklets")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
+        print("  number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.train_name_path):
+            raise RuntimeError("'{}' is not available".format(self.train_name_path))
+        if not osp.exists(self.test_name_path):
+            raise RuntimeError("'{}' is not available".format(self.test_name_path))
+        if not osp.exists(self.track_train_info_path):
+            raise RuntimeError("'{}' is not available".format(self.track_train_info_path))
+        if not osp.exists(self.track_test_info_path):
+            raise RuntimeError("'{}' is not available".format(self.track_test_info_path))
+        if not osp.exists(self.query_IDX_path):
+            raise RuntimeError("'{}' is not available".format(self.query_IDX_path))
+
+    def _get_names(self, fpath):
+        names = []
+        with open(fpath, 'r') as f:
+            for line in f:
+                new_line = line.rstrip()
+                names.append(new_line)
+        return names
+
+    def _process_data(self, names, meta_data, home_dir=None, relabel=False, min_seq_len=0):
+        assert home_dir in ['bbox_train', 'bbox_test']
+        num_tracklets = meta_data.shape[0]
+        pid_list = list(set(meta_data[:,2].tolist()))
+        num_pids = len(pid_list)
+
+        if relabel: pid2label = {pid:label for label, pid in enumerate(pid_list)}
+        tracklets = []
+        num_imgs_per_tracklet = []
+
+        for tracklet_idx in range(num_tracklets):
+            data = meta_data[tracklet_idx,...]
+            start_index, end_index, pid, camid = data
+            if pid == -1: continue # junk images are just ignored
+            assert 1 <= camid <= 6
+            if relabel: pid = pid2label[pid]
+            camid -= 1 # index starts from 0
+            img_names = names[start_index-1:end_index]
+
+            # make sure image names correspond to the same person
+            pnames = [img_name[:4] for img_name in img_names]
+            assert len(set(pnames)) == 1, "Error: a single tracklet contains different person images"
+
+            # make sure all images are captured under the same camera
+            camnames = [img_name[5] for img_name in img_names]
+            assert len(set(camnames)) == 1, "Error: images are captured under different cameras!"
+
+            # append image names with directory information
+            img_paths = [osp.join(self.dataset_dir, home_dir, img_name[:4], img_name) for img_name in img_names]
+            if len(img_paths) >= min_seq_len:
+                img_paths = tuple(img_paths)
+                tracklets.append((img_paths, pid, camid))
+                num_imgs_per_tracklet.append(len(img_paths))
+
+        num_tracklets = len(tracklets)
+
+        return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
+
+class iLIDSVID(object):
+    """
+    iLIDS-VID
+
+    Reference:
+    Wang et al. Person Re-Identification by Video Ranking. ECCV 2014.
+
+    URL: http://www.eecs.qmul.ac.uk/~xiatian/downloads_qmul_iLIDS-VID_ReID_dataset.html
+    
+    Dataset statistics:
+    # identities: 300
+    # tracklets: 600
+    # cameras: 2
+    """
+    dataset_dir = 'ilids-vid'
+
+    def __init__(self, root='data', split_id=0, **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.dataset_url = 'http://www.eecs.qmul.ac.uk/~xiatian/iLIDS-VID/iLIDS-VID.tar'
+        self.data_dir = osp.join(self.dataset_dir, 'i-LIDS-VID')
+        self.split_dir = osp.join(self.dataset_dir, 'train-test people splits')
+        self.split_mat_path = osp.join(self.split_dir, 'train_test_splits_ilidsvid.mat')
+        self.split_path = osp.join(self.dataset_dir, 'splits.json')
+        self.cam_1_path = osp.join(self.dataset_dir, 'i-LIDS-VID/sequences/cam1')
+        self.cam_2_path = osp.join(self.dataset_dir, 'i-LIDS-VID/sequences/cam2')
+
+        self._download_data()
+        self._check_before_run()
+
+        self._prepare_split()
+        splits = read_json(self.split_path)
+        if split_id >= len(splits):
+            raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
+        split = splits[split_id]
+        train_dirs, test_dirs = split['train'], split['test']
+        print("# train identites: {}, # test identites {}".format(len(train_dirs), len(test_dirs)))
+
+        train, num_train_tracklets, num_train_pids, num_imgs_train = \
+          self._process_data(train_dirs, cam1=True, cam2=True)
+        query, num_query_tracklets, num_query_pids, num_imgs_query = \
+          self._process_data(test_dirs, cam1=True, cam2=False)
+        gallery, num_gallery_tracklets, num_gallery_pids, num_imgs_gallery = \
+          self._process_data(test_dirs, cam1=False, cam2=True)
+
+        num_imgs_per_tracklet = num_imgs_train + num_imgs_query + num_imgs_gallery
+        min_num = np.min(num_imgs_per_tracklet)
+        max_num = np.max(num_imgs_per_tracklet)
+        avg_num = np.mean(num_imgs_per_tracklet)
+
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
+
+        print("=> iLIDS-VID loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # tracklets")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
+        print("  number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _download_data(self):
+        if osp.exists(self.dataset_dir):
+            print("This dataset has been downloaded.")
+            return
+
+        mkdir_if_missing(self.dataset_dir)
+        fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url))
+
+        print("Downloading iLIDS-VID dataset")
+        url_opener = urllib.URLopener()
+        url_opener.retrieve(self.dataset_url, fpath)
+
+        print("Extracting files")
+        tar = tarfile.open(fpath)
+        tar.extractall(path=self.dataset_dir)
+        tar.close()
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.data_dir):
+            raise RuntimeError("'{}' is not available".format(self.data_dir))
+        if not osp.exists(self.split_dir):
+            raise RuntimeError("'{}' is not available".format(self.split_dir))
+
+    def _prepare_split(self):
+        if not osp.exists(self.split_path):
+            print("Creating splits")
+            mat_split_data = loadmat(self.split_mat_path)['ls_set']
+            
+            num_splits = mat_split_data.shape[0]
+            num_total_ids = mat_split_data.shape[1]
+            assert num_splits == 10
+            assert num_total_ids == 300
+            num_ids_each = num_total_ids/2
+
+            # pids in mat_split_data are indices, so we need to transform them
+            # to real pids
+            person_cam1_dirs = os.listdir(self.cam_1_path)
+            person_cam2_dirs = os.listdir(self.cam_2_path)
+
+            # make sure persons in one camera view can be found in the other camera view
+            assert set(person_cam1_dirs) == set(person_cam2_dirs)
+
+            splits = []
+            for i_split in range(num_splits):
+                # first 50% for testing and the remaining for training, following Wang et al. ECCV'14.
+                train_idxs = sorted(list(mat_split_data[i_split,num_ids_each:]))
+                test_idxs = sorted(list(mat_split_data[i_split,:num_ids_each]))
+                
+                train_idxs = [int(i)-1 for i in train_idxs]
+                test_idxs = [int(i)-1 for i in test_idxs]
+                
+                # transform pids to person dir names
+                train_dirs = [person_cam1_dirs[i] for i in train_idxs]
+                test_dirs = [person_cam1_dirs[i] for i in test_idxs]
+                
+                split = {'train': train_dirs, 'test': test_dirs}
+                splits.append(split)
+
+            print("Totally {} splits are created, following Wang et al. ECCV'14".format(len(splits)))
+            print("Split file is saved to {}".format(self.split_path))
+            write_json(splits, self.split_path)
+
+        print("Splits created")
+
+    def _process_data(self, dirnames, cam1=True, cam2=True):
+        tracklets = []
+        num_imgs_per_tracklet = []
+        dirname2pid = {dirname:i for i, dirname in enumerate(dirnames)}
+        
+        for dirname in dirnames:
+            if cam1:
+                person_dir = osp.join(self.cam_1_path, dirname)
+                img_names = glob.glob(osp.join(person_dir, '*.png'))
+                assert len(img_names) > 0
+                img_names = tuple(img_names)
+                pid = dirname2pid[dirname]
+                tracklets.append((img_names, pid, 0))
+                num_imgs_per_tracklet.append(len(img_names))
+
+            if cam2:
+                person_dir = osp.join(self.cam_2_path, dirname)
+                img_names = glob.glob(osp.join(person_dir, '*.png'))
+                assert len(img_names) > 0
+                img_names = tuple(img_names)
+                pid = dirname2pid[dirname]
+                tracklets.append((img_names, pid, 1))
+                num_imgs_per_tracklet.append(len(img_names))
+
+        num_tracklets = len(tracklets)
+        num_pids = len(dirnames)
+
+        return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
+
+class PRID(object):
+    """
+    PRID
+
+    Reference:
+    Hirzer et al. Person Re-Identification by Descriptive and Discriminative Classification. SCIA 2011.
+
+    URL: https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/PRID11/
+    
+    Dataset statistics:
+    # identities: 200
+    # tracklets: 400
+    # cameras: 2
+    """
+    dataset_dir = 'prid2011'
+
+    def __init__(self, root='data', split_id=0, min_seq_len=0, **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.dataset_url = 'https://files.icg.tugraz.at/f/6ab7e8ce8f/?raw=1'
+        self.split_path = osp.join(self.dataset_dir, 'splits_prid2011.json')
+        self.cam_a_path = osp.join(self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_a')
+        self.cam_b_path = osp.join(self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_b')
+
+        self._check_before_run()
+        splits = read_json(self.split_path)
+        if split_id >=  len(splits):
+            raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
+        split = splits[split_id]
+        train_dirs, test_dirs = split['train'], split['test']
+        print("# train identites: {}, # test identites {}".format(len(train_dirs), len(test_dirs)))
+
+        train, num_train_tracklets, num_train_pids, num_imgs_train = \
+          self._process_data(train_dirs, cam1=True, cam2=True)
+        query, num_query_tracklets, num_query_pids, num_imgs_query = \
+          self._process_data(test_dirs, cam1=True, cam2=False)
+        gallery, num_gallery_tracklets, num_gallery_pids, num_imgs_gallery = \
+          self._process_data(test_dirs, cam1=False, cam2=True)
+
+        num_imgs_per_tracklet = num_imgs_train + num_imgs_query + num_imgs_gallery
+        min_num = np.min(num_imgs_per_tracklet)
+        max_num = np.max(num_imgs_per_tracklet)
+        avg_num = np.mean(num_imgs_per_tracklet)
+
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
+
+        print("=> PRID-2011 loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # tracklets")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
+        print("  number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+
+    def _process_data(self, dirnames, cam1=True, cam2=True):
+        tracklets = []
+        num_imgs_per_tracklet = []
+        dirname2pid = {dirname:i for i, dirname in enumerate(dirnames)}
+        
+        for dirname in dirnames:
+            if cam1:
+                person_dir = osp.join(self.cam_a_path, dirname)
+                img_names = glob.glob(osp.join(person_dir, '*.png'))
+                assert len(img_names) > 0
+                img_names = tuple(img_names)
+                pid = dirname2pid[dirname]
+                tracklets.append((img_names, pid, 0))
+                num_imgs_per_tracklet.append(len(img_names))
+
+            if cam2:
+                person_dir = osp.join(self.cam_b_path, dirname)
+                img_names = glob.glob(osp.join(person_dir, '*.png'))
+                assert len(img_names) > 0
+                img_names = tuple(img_names)
+                pid = dirname2pid[dirname]
+                tracklets.append((img_names, pid, 1))
+                num_imgs_per_tracklet.append(len(img_names))
+
+        num_tracklets = len(tracklets)
+        num_pids = len(dirnames)
+
+        return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
+
+class DukeMTMCVidReID(object):
+    """
+    DukeMTMCVidReID
+
+    Reference:
+    Wu et al. Exploit the Unknown Gradually: One-Shot Video-Based Person
+    Re-Identification by Stepwise Learning. CVPR 2018.
+
+    URL: https://github.com/Yu-Wu/Exploit-Unknown-Gradually
+    
+    Dataset statistics:
+    # identities: 702 (train) + 702 (test)
+    # tracklets: 2196 (train) + 2636 (test)
+    """
+    dataset_dir = 'dukemtmc-vidreid'
+
+    def __init__(self, root='data', min_seq_len=0, **kwargs):
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        self.train_dir = osp.join(self.dataset_dir, 'dukemtmc_videoReID/train_split')
+        self.query_dir = osp.join(self.dataset_dir, 'dukemtmc_videoReID/query_split')
+        self.gallery_dir = osp.join(self.dataset_dir, 'dukemtmc_videoReID/gallery_split')
+        self.split_train_json_path = osp.join(self.dataset_dir, 'split_train.json')
+        self.split_query_json_path = osp.join(self.dataset_dir, 'split_query.json')
+        self.split_gallery_json_path = osp.join(self.dataset_dir, 'split_gallery.json')
+
+        self.min_seq_len = min_seq_len
+        self._check_before_run()
+        print("Note: if root path is changed, the previously generated json files need to be re-generated (so delete them first)")
+
+        train, num_train_tracklets, num_train_pids, num_imgs_train = \
+          self._process_dir(self.train_dir, self.split_train_json_path, relabel=True)
+        query, num_query_tracklets, num_query_pids, num_imgs_query = \
+          self._process_dir(self.query_dir, self.split_query_json_path, relabel=False)
+        gallery, num_gallery_tracklets, num_gallery_pids, num_imgs_gallery = \
+          self._process_dir(self.gallery_dir, self.split_gallery_json_path, relabel=False)
+
+        num_imgs_per_tracklet = num_imgs_train + num_imgs_query + num_imgs_gallery
+        min_num = np.min(num_imgs_per_tracklet)
+        max_num = np.max(num_imgs_per_tracklet)
+        avg_num = np.mean(num_imgs_per_tracklet)
+
+        num_total_pids = num_train_pids + num_query_pids
+        num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
+
+        print("=> DukeMTMC-VideoReID loaded")
+        print("Dataset statistics:")
+        print("  ------------------------------")
+        print("  subset   | # ids | # tracklets")
+        print("  ------------------------------")
+        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
+        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
+        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
+        print("  ------------------------------")
+        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
+        print("  number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
+        print("  ------------------------------")
+
+        self.train = train
+        self.query = query
+        self.gallery = gallery
+
+        self.num_train_pids = num_train_pids
+        self.num_query_pids = num_query_pids
+        self.num_gallery_pids = num_gallery_pids
+
+    def _check_before_run(self):
+        """Check if all files are available before going deeper"""
+        if not osp.exists(self.dataset_dir):
+            raise RuntimeError("'{}' is not available".format(self.dataset_dir))
+        if not osp.exists(self.train_dir):
+            raise RuntimeError("'{}' is not available".format(self.train_dir))
+        if not osp.exists(self.query_dir):
+            raise RuntimeError("'{}' is not available".format(self.query_dir))
+        if not osp.exists(self.gallery_dir):
+            raise RuntimeError("'{}' is not available".format(self.gallery_dir))
+
+    def _process_dir(self, dir_path, json_path, relabel):
+        if osp.exists(json_path):
+            print("=> {} generated before, awesome!".format(json_path))
+            split = read_json(json_path)
+            return split['tracklets'], split['num_tracklets'], split['num_pids'], split['num_imgs_per_tracklet']
+
+        print("=> Automatically generating split (might take a while for the first time, have a coffe)")
+        pdirs = glob.glob(osp.join(dir_path, '*')) # avoid .DS_Store
+        print("Processing {} with {} person identities".format(dir_path, len(pdirs)))
+
+        pid_container = set()
+        for pdir in pdirs:
+            pid = int(osp.basename(pdir))
+            pid_container.add(pid)
+        pid2label = {pid:label for label, pid in enumerate(pid_container)}
+
+        tracklets = []
+        num_imgs_per_tracklet = []
+        for pdir in pdirs:
+            pid = int(osp.basename(pdir))
+            if relabel: pid = pid2label[pid]
+            tdirs = glob.glob(osp.join(pdir, '*'))
+            for tdir in tdirs:
+                raw_img_paths = glob.glob(osp.join(tdir, '*.jpg'))
+                num_imgs = len(raw_img_paths)
+
+                if num_imgs < self.min_seq_len:
+                    continue
+
+                num_imgs_per_tracklet.append(num_imgs)
+                img_paths = []
+                for img_idx in range(num_imgs):
+                    # some tracklet starts from 0002 instead of 0001
+                    img_idx_name = 'F' + str(img_idx+1).zfill(4)
+                    res = glob.glob(osp.join(tdir, '*' + img_idx_name + '*.jpg'))
+                    if len(res) == 0:
+                        print("Warn: index name {} in {} is missing, jump to next".format(img_idx_name, tdir))
+                        continue
+                    img_paths.append(res[0])
+                img_name = osp.basename(img_paths[0])
+                camid = int(img_name[5]) - 1 # index-0
+                img_paths = tuple(img_paths)
+                tracklets.append((img_paths, pid, camid))
+
+        num_pids = len(pid_container)
+        num_tracklets = len(tracklets)
+
+        print("Saving split to {}".format(json_path))
+        split_dict = {
+            'tracklets': tracklets,
+            'num_tracklets': num_tracklets,
+            'num_pids': num_pids,
+            'num_imgs_per_tracklet': num_imgs_per_tracklet,
+        }
+        write_json(split_dict, json_path)
+
+        return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
+
+"""Create dataset"""
+
+__img_factory = {
+    'market1501': Market1501,
+    'market1501_partial': Market1501_Partial,
+    'cuhk03': CUHK03,
+    'dukemtmcreid': DukeMTMCreID,
+    'msmt17': MSMT17,
+}
+
+__vid_factory = {
+    'mars': Mars,
+    'ilidsvid': iLIDSVID,
+    'prid': PRID,
+    'dukemtmcvidreid': DukeMTMCVidReID,
+}
+
+def get_names():
+    return list(__img_factory.keys()) + list(__vid_factory.keys())
+
+def init_img_dataset(name, **kwargs):
+    if name not in __img_factory.keys():
+        raise KeyError("Invalid dataset, got '{}', but expected to be one of {}".format(name, __img_factory.keys()))
+    return __img_factory[name](**kwargs)
+
+def init_vid_dataset(name, **kwargs):
+    if name not in __vid_factory.keys():
+        raise KeyError("Invalid dataset, got '{}', but expected to be one of {}".format(name, __vid_factory.keys()))
+    return __vid_factory[name](**kwargs)

util/dataset_loader.py

+from __future__ import print_function, absolute_import
+import os
+from PIL import Image
+import numpy as np
+import os.path as osp
+
+import torch
+from torch.utils.data import Dataset
+
+def read_image(img_path):
+    """Keep reading image until succeed.
+    This can avoid IOError incurred by heavy IO process."""
+    got_img = False
+    if not osp.exists(img_path):
+        raise IOError("{} does not exist".format(img_path))
+    while not got_img:
+        try:
+            img = Image.open(img_path).convert('RGB')
+            got_img = True
+        except IOError:
+            print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
+            pass
+    return img
+
+class ImageDataset(Dataset):
+    """Image Person ReID Dataset"""
+    def __init__(self, dataset, transform=None):
+        self.dataset = dataset
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        img_path, pid, camid = self.dataset[index]
+        img = read_image(img_path)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, pid, camid
+
+class VideoDataset(Dataset):
+    """Video Person ReID Dataset.
+    Note batch data has shape (batch, seq_len, channel, height, width).
+    """
+    sample_methods = ['evenly', 'random', 'all']
+
+    def __init__(self, dataset, seq_len=15, sample='evenly', transform=None):
+        self.dataset = dataset
+        self.seq_len = seq_len
+        self.sample = sample
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        img_paths, pid, camid = self.dataset[index]
+        num = len(img_paths)
+
+        if self.sample == 'random':
+            """
+            Randomly sample seq_len items from num items,
+            if num is smaller than seq_len, then replicate items
+            """
+            indices = np.arange(num)
+            replace = False if num >= self.seq_len else True
+            indices = np.random.choice(indices, size=self.seq_len, replace=replace)
+            # sort indices to keep temporal order
+            # comment it to be order-agnostic
+            indices = np.sort(indices)
+        elif self.sample == 'evenly':
+            """Evenly sample seq_len items from num items."""
+            if num >= self.seq_len:
+                num -= num % self.seq_len
+                indices = np.arange(0, num, num/self.seq_len)
+            else:
+                # if num is smaller than seq_len, simply replicate the last image
+                # until the seq_len requirement is satisfied
+                indices = np.arange(0, num)
+                num_pads = self.seq_len - num
+                indices = np.concatenate([indices, np.ones(num_pads).astype(np.int32)*(num-1)])
+            assert len(indices) == self.seq_len
+        elif self.sample == 'all':
+            """
+            Sample all items, seq_len is useless now and batch_size needs
+            to be set to 1.
+            """
+            indices = np.arange(num)
+        else:
+            raise KeyError("Unknown sample method: {}. Expected one of {}".format(self.sample, self.sample_methods))
+
+        imgs = []
+        for index in indices:
+            img_path = img_paths[index]
+            img = read_image(img_path)
+            if self.transform is not None:
+                img = self.transform(img)
+            img = img.unsqueeze(0)
+            imgs.append(img)
+        imgs = torch.cat(imgs, dim=0)
+
+        return imgs, pid, camid
\ No newline at end of file

util/distance.py

+"""Numpy version of euclidean distance, shortest distance, etc.
+Notice the input/output shape of methods, so that you can better understand
+the meaning of these methods."""
+import numpy as np
+
+
+def normalize(nparray, order=2, axis=0):
+    """Normalize a N-D numpy array along the specified axis."""
+    norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
+    return nparray / (norm + np.finfo(np.float32).eps)
+
+
+def compute_dist(array1, array2, type='euclidean'):
+    """Compute the euclidean or cosine distance of all pairs.
+  Args:
+    array1: numpy array with shape [m1, n]
+    array2: numpy array with shape [m2, n]
+    type: one of ['cosine', 'euclidean']
+  Returns:
+    numpy array with shape [m1, m2]
+  """
+    assert type in ['cosine', 'euclidean']
+    if type == 'cosine':
+        array1 = normalize(array1, axis=1)
+        array2 = normalize(array2, axis=1)
+        dist = np.matmul(array1, array2.T)
+        return dist
+    else:
+        # shape [m1, 1]
+        square1 = np.sum(np.square(array1), axis=1)[..., np.newaxis]
+        # shape [1, m2]
+        square2 = np.sum(np.square(array2), axis=1)[np.newaxis, ...]
+        squared_dist = - 2 * np.matmul(array1, array2.T) + square1 + square2
+        squared_dist[squared_dist < 0] = 0
+        dist = np.sqrt(squared_dist)
+        return dist
+
+
+def shortest_dist(dist_mat):
+    """Parallel version.
+  Args:
+    dist_mat: numpy array, available shape
+      1) [m, n]
+      2) [m, n, N], N is batch size
+      3) [m, n, *], * can be arbitrary additional dimensions
+  Returns:
+    dist: three cases corresponding to `dist_mat`
+      1) scalar
+      2) numpy array, with shape [N]
+      3) numpy array with shape [*]
+  """
+    m, n = dist_mat.shape[:2]
+    dist = np.zeros_like(dist_mat)
+    for i in range(m):
+        for j in range(n):
+            if (i == 0) and (j == 0):
+                dist[i, j] = dist_mat[i, j]
+            elif (i == 0) and (j > 0):
+                dist[i, j] = dist[i, j - 1] + dist_mat[i, j]
+            elif (i > 0) and (j == 0):
+                dist[i, j] = dist[i - 1, j] + dist_mat[i, j]
+            else:
+                dist[i, j] = \
+                    np.min(np.stack([dist[i - 1, j], dist[i, j - 1]], axis=0), axis=0) \
+                    + dist_mat[i, j]
+    # I ran into memory disaster when returning this reference! I still don't
+    # know why.
+    # dist = dist[-1, -1]
+    dist = dist[-1, -1].copy()
+    return dist
+
+def unaligned_dist(dist_mat):
+    """Parallel version.
+    Args:
+      dist_mat: numpy array, available shape
+        1) [m, n]
+        2) [m, n, N], N is batch size
+        3) [m, n, *], * can be arbitrary additional dimensions
+    Returns:
+      dist: three cases corresponding to `dist_mat`
+        1) scalar
+        2) numpy array, with shape [N]
+        3) numpy array with shape [*]
+    """
+
+    m = dist_mat.shape[0]
+    dist = np.zeros_like(dist_mat[0])
+    for i in range(m):
+        dist[i] = dist_mat[i][i]
+    dist = np.sum(dist, axis=0).copy()
+    return dist
+
+
+def meta_local_dist(x, y, aligned):
+    """
+  Args:
+    x: numpy array, with shape [m, d]
+    y: numpy array, with shape [n, d]
+  Returns:
+    dist: scalar
+  """
+    eu_dist = compute_dist(x, y, 'euclidean')
+    dist_mat = (np.exp(eu_dist) - 1.) / (np.exp(eu_dist) + 1.)
+    if aligned:
+        dist = shortest_dist(dist_mat[np.newaxis])[0]
+    else:
+        dist = unaligned_dist(dist_mat[np.newaxis])[0]
+    return dist
+
+
+# Tooooooo slow!
+def serial_local_dist(x, y):
+    """
+  Args:
+    x: numpy array, with shape [M, m, d]
+    y: numpy array, with shape [N, n, d]
+  Returns:
+    dist: numpy array, with shape [M, N]
+  """
+    M, N = x.shape[0], y.shape[0]
+    dist_mat = np.zeros([M, N])
+    for i in range(M):
+        for j in range(N):
+            dist_mat[i, j] = meta_local_dist(x[i], y[j])
+    return dist_mat
+
+
+def parallel_local_dist(x, y, aligned):
+    """Parallel version.
+  Args:
+    x: numpy array, with shape [M, m, d]
+    y: numpy array, with shape [N, n, d]
+  Returns:
+    dist: numpy array, with shape [M, N]
+  """
+    M, m, d = x.shape
+    N, n, d = y.shape
+    x = x.reshape([M * m, d])
+    y = y.reshape([N * n, d])
+    # shape [M * m, N * n]
+    dist_mat = compute_dist(x, y, type='euclidean')
+    dist_mat = (np.exp(dist_mat) - 1.) / (np.exp(dist_mat) + 1.)
+    # shape [M * m, N * n] -> [M, m, N, n] -> [m, n, M, N]
+    dist_mat = dist_mat.reshape([M, m, N, n]).transpose([1, 3, 0, 2])
+    # shape [M, N]
+    if aligned:
+        dist_mat = shortest_dist(dist_mat)
+    else:
+        dist_mat = unaligned_dist(dist_mat)
+    return dist_mat
+
+
+def local_dist(x, y, aligned):
+    if (x.ndim == 2) and (y.ndim == 2):
+        return meta_local_dist(x, y, aligned)
+    elif (x.ndim == 3) and (y.ndim == 3):
+        return parallel_local_dist(x, y, aligned)
+    else:
+        raise NotImplementedError('Input shape not supported.')
+
+
+def low_memory_matrix_op(
+        func,
+        x, y,
+        x_split_axis, y_split_axis,
+        x_num_splits, y_num_splits,
+        verbose=False, aligned=True):
+    """
+  For matrix operation like multiplication, in order not to flood the memory
+  with huge data, split matrices into smaller parts (Divide and Conquer).
+
+  Note:
+    If still out of memory, increase `*_num_splits`.
+
+  Args:
+    func: a matrix function func(x, y) -> z with shape [M, N]
+    x: numpy array, the dimension to split has length M
+    y: numpy array, the dimension to split has length N
+    x_split_axis: The axis to split x into parts
+    y_split_axis: The axis to split y into parts
+    x_num_splits: number of splits. 1 <= x_num_splits <= M
+    y_num_splits: number of splits. 1 <= y_num_splits <= N
+    verbose: whether to print the progress
+
+  Returns:
+    mat: numpy array, shape [M, N]
+  """
+
+    if verbose:
+        import sys
+        import time
+        printed = False
+        st = time.time()
+        last_time = time.time()
+
+    mat = [[] for _ in range(x_num_splits)]
+    for i, part_x in enumerate(
+            np.array_split(x, x_num_splits, axis=x_split_axis)):
+        for j, part_y in enumerate(
+                np.array_split(y, y_num_splits, axis=y_split_axis)):
+            part_mat = func(part_x, part_y, aligned)
+            mat[i].append(part_mat)
+
+            if verbose:
+                if not printed:
+                    printed = True
+                else:
+                    # Clean the current line
+                    sys.stdout.write("\033[F\033[K")
+                print('Matrix part ({}, {}) / ({}, {}), +{:.2f}s, total {:.2f}s'
+                    .format(i + 1, j + 1, x_num_splits, y_num_splits,
+                            time.time() - last_time, time.time() - st))
+                last_time = time.time()
+        mat[i] = np.concatenate(mat[i], axis=1)
+    mat = np.concatenate(mat, axis=0)
+    return mat
+
+
+def low_memory_local_dist(x, y, aligned=True):
+    print('Computing local distance...')
+    x_num_splits = int(len(x) / 200) + 1
+    y_num_splits = int(len(y) / 200) + 1
+    z = low_memory_matrix_op(local_dist, x, y, 0, 0, x_num_splits, y_num_splits, verbose=True, aligned=aligned)
+    return z
\ No newline at end of file

util/eval_metrics.py

+from __future__ import print_function, absolute_import
+import numpy as np
+import copy
+from collections import defaultdict
+import sys
+
+def eval_cuhk03(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, N=100):
+    """Evaluation with cuhk03 metric
+    Key: one image for each gallery identity is randomly sampled for each query identity.
+    Random sampling is performed N times (default: N=100).
+    """
+    num_q, num_g = distmat.shape
+    if num_g < max_rank:
+        max_rank = num_g
+        print("Note: number of gallery samples is quite small, got {}".format(num_g))
+    indices = np.argsort(distmat, axis=1)
+    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)
+
+    # compute cmc curve for each query
+    all_cmc = []
+    all_AP = []
+    num_valid_q = 0. # number of valid query
+    for q_idx in range(num_q):
+        # get query pid and camid
+        q_pid = q_pids[q_idx]
+        q_camid = q_camids[q_idx]
+
+        # remove gallery samples that have the same pid and camid with query
+        order = indices[q_idx]
+        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
+        keep = np.invert(remove)
+
+        # compute cmc curve
+        orig_cmc = matches[q_idx][keep] # binary vector, positions with value 1 are correct matches
+        if not np.any(orig_cmc):
+            # this condition is true when query identity does not appear in gallery
+            continue
+
+        kept_g_pids = g_pids[order][keep]
+        g_pids_dict = defaultdict(list)
+        for idx, pid in enumerate(kept_g_pids):
+            g_pids_dict[pid].append(idx)
+
+        cmc, AP = 0., 0.
+        for repeat_idx in range(N):
+            mask = np.zeros(len(orig_cmc), dtype=np.bool)
+            for _, idxs in g_pids_dict.items():
+                # randomly sample one image for each gallery person
+                rnd_idx = np.random.choice(idxs)
+                mask[rnd_idx] = True
+            masked_orig_cmc = orig_cmc[mask]
+            _cmc = masked_orig_cmc.cumsum()
+            _cmc[_cmc > 1] = 1
+            cmc += _cmc[:max_rank].astype(np.float32)
+            # compute AP
+            num_rel = masked_orig_cmc.sum()
+            tmp_cmc = masked_orig_cmc.cumsum()
+            tmp_cmc = [x / (i+1.) for i, x in enumerate(tmp_cmc)]
+            tmp_cmc = np.asarray(tmp_cmc) * masked_orig_cmc
+            AP += tmp_cmc.sum() / num_rel
+        cmc /= N
+        AP /= N
+        all_cmc.append(cmc)
+        all_AP.append(AP)
+        num_valid_q += 1.
+
+    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"
+
+    all_cmc = np.asarray(all_cmc).astype(np.float32)
+    all_cmc = all_cmc.sum(0) / num_valid_q
+    mAP = np.mean(all_AP)
+
+    return all_cmc, mAP
+
+def eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank):
+    """Evaluation with market1501 metric
+    Key: for each query identity, its gallery images from the same camera view are discarded.
+    """
+    num_q, num_g = distmat.shape
+    if num_g < max_rank:
+        max_rank = num_g
+        print("Note: number of gallery samples is quite small, got {}".format(num_g))
+    indices = np.argsort(distmat, axis=1)
+    matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)
+
+    # compute cmc curve for each query
+    all_cmc = []
+    all_AP = []
+    num_valid_q = 0. # number of valid query
+    for q_idx in range(num_q):
+        # get query pid and camid
+        q_pid = q_pids[q_idx]
+        q_camid = q_camids[q_idx]
+
+        # remove gallery samples that have the same pid and camid with query
+        order = indices[q_idx]
+        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
+        keep = np.invert(remove)
+
+        # compute cmc curve
+        orig_cmc = matches[q_idx][keep] # binary vector, positions with value 1 are correct matches
+        if not np.any(orig_cmc):
+            # this condition is true when query identity does not appear in gallery
+            continue
+
+        cmc = orig_cmc.cumsum()
+        cmc[cmc > 1] = 1
+
+        all_cmc.append(cmc[:max_rank])
+        num_valid_q += 1.
+
+        # compute average precision
+        # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision
+        num_rel = orig_cmc.sum()
+        tmp_cmc = orig_cmc.cumsum()
+        tmp_cmc = [x / (i+1.) for i, x in enumerate(tmp_cmc)]
+        tmp_cmc = np.asarray(tmp_cmc) * orig_cmc
+        AP = tmp_cmc.sum() / num_rel
+        all_AP.append(AP)
+
+    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"
+
+    all_cmc = np.asarray(all_cmc).astype(np.float32)
+    all_cmc = all_cmc.sum(0) / num_valid_q
+    mAP = np.mean(all_AP)
+
+    return all_cmc, mAP
+
+def evaluate(distmat, q_pids, g_pids, q_camids, g_camids, max_rank=50, use_metric_cuhk03=False):
+    if use_metric_cuhk03:
+        return eval_cuhk03(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)
+    else:
+        return eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)
\ No newline at end of file

util/losses.py

+from __future__ import absolute_import
+from aligned.local_dist import *
+
+import torch
+from torch import nn
+
+"""
+Shorthands for loss:
+- CrossEntropyLabelSmooth: xent
+- TripletLoss: htri
+- CenterLoss: cent
+"""
+__all__ = ['DeepSupervision', 'CrossEntropyLoss','CrossEntropyLabelSmooth', 'TripletLoss', 'CenterLoss', 'RingLoss']
+
+def DeepSupervision(criterion, xs, y):
+    """
+    Args:
+        criterion: loss function
+        xs: tuple of inputs
+        y: ground truth
+    """
+    loss = 0.
+    for x in xs:
+        loss += criterion(x, y)
+    return loss
+
+class CrossEntropyLoss(nn.Module):
+    """Cross entropy loss.
+
+    """
+    def __init__(self, use_gpu=True):
+        super(CrossEntropyLoss, self).__init__()
+        self.use_gpu = use_gpu
+        self.crossentropy_loss = nn.CrossEntropyLoss()
+
+    def forward(self, inputs, targets):
+        """
+        Args:
+            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
+            targets: ground truth labels with shape (num_classes)
+        """
+        if self.use_gpu: targets = targets.cuda()
+        loss = self.crossentropy_loss(inputs, targets)
+        return loss
+
+class CrossEntropyLabelSmooth(nn.Module):
+    """Cross entropy loss with label smoothing regularizer.
+
+    Reference:
+    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
+    Equation: y = (1 - epsilon) * y + epsilon / K.
+
+    Args:
+        num_classes (int): number of classes.
+        epsilon (float): weight.
+    """
+    def __init__(self, num_classes, epsilon=0.1, use_gpu=True):
+        super(CrossEntropyLabelSmooth, self).__init__()
+        self.num_classes = num_classes
+        self.epsilon = epsilon
+        self.use_gpu = use_gpu
+        self.logsoftmax = nn.LogSoftmax(dim=1)
+
+    def forward(self, inputs, targets):
+        """
+        Args:
+            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
+            targets: ground truth labels with shape (num_classes)
+        """
+        log_probs = self.logsoftmax(inputs)
+        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1)
+        if self.use_gpu: targets = targets.cuda()
+        targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
+        loss = (- targets * log_probs).mean(0).sum()
+        return loss
+
+class TripletLoss(nn.Module):
+    """Triplet loss with hard positive/negative mining.
+
+    Reference:
+    Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.
+
+    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.
+
+    Args:
+        margin (float): margin for triplet.
+    """
+    def __init__(self, margin=0.3, mutual_flag = False):
+        super(TripletLoss, self).__init__()
+        self.margin = margin
+        self.ranking_loss = nn.MarginRankingLoss(margin=margin)
+        self.mutual = mutual_flag
+
+    def forward(self, inputs, targets):
+        """
+        Args:
+            inputs: feature matrix with shape (batch_size, feat_dim)
+            targets: ground truth labels with shape (num_classes)
+        """
+        n = inputs.size(0)
+        # inputs = 1. * inputs / (torch.norm(inputs, 2, dim=-1, keepdim=True).expand_as(inputs) + 1e-12)
+        # Compute pairwise distance, replace by the official when merged
+        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)
+        dist = dist + dist.t()
+        dist.addmm_(1, -2, inputs, inputs.t())
+        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
+        # For each anchor, find the hardest positive and negative
+        mask = targets.expand(n, n).eq(targets.expand(n, n).t())
+        dist_ap, dist_an = [], []
+        for i in range(n):
+            dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))
+            dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))
+        dist_ap = torch.cat(dist_ap)
+        dist_an = torch.cat(dist_an)
+        # Compute ranking hinge loss
+        y = torch.ones_like(dist_an)
+        loss = self.ranking_loss(dist_an, dist_ap, y)
+        if self.mutual:
+            return loss, dist
+        return loss
+
+class TripletLossAlignedReID(nn.Module):
+    """Triplet loss with hard positive/negative mining.
+
+    Reference:
+    Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.
+
+    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.
+
+    Args:
+        margin (float): margin for triplet.
+    """
+    def __init__(self, margin=0.3, mutual_flag = False):
+        super(TripletLossAlignedReID, self).__init__()
+        self.margin = margin
+        self.ranking_loss = nn.MarginRankingLoss(margin=margin)
+        self.ranking_loss_local = nn.MarginRankingLoss(margin=margin)
+        self.mutual = mutual_flag
+
+    def forward(self, inputs, targets, local_features):
+        """
+        Args:
+            inputs: feature matrix with shape (batch_size, feat_dim)
+            targets: ground truth labels with shape (num_classes)
+        """
+        n = inputs.size(0)
+        #inputs = 1. * inputs / (torch.norm(inputs, 2, dim=-1, keepdim=True).expand_as(inputs) + 1e-12)
+        # Compute pairwise distance, replace by the official when merged
+        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)
+        dist = dist + dist.t()
+        dist.addmm_(1, -2, inputs, inputs.t())
+        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
+        # For each anchor, find the hardest positive and negative
+        dist_ap,dist_an,p_inds,n_inds = hard_example_mining(dist,targets,return_inds=True)
+        local_features = local_features.permute(0,2,1)
+        p_local_features = local_features[p_inds]
+        n_local_features = local_features[n_inds]
+        local_dist_ap = batch_local_dist(local_features, p_local_features)
+        local_dist_an = batch_local_dist(local_features, n_local_features)
+
+        # Compute ranking hinge loss
+        y = torch.ones_like(dist_an)
+        global_loss = self.ranking_loss(dist_an, dist_ap, y)
+        local_loss = self.ranking_loss_local(local_dist_an,local_dist_ap, y)
+        if self.mutual:
+            return global_loss+local_loss,dist
+        return global_loss,local_loss
+
+class CenterLoss(nn.Module):
+    """Center loss.
+    
+    Reference:
+    Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016.
+    
+    Args:
+        num_classes (int): number of classes.
+        feat_dim (int): feature dimension.
+    """
+    def __init__(self, num_classes=10, feat_dim=2, use_gpu=True):
+        super(CenterLoss, self).__init__()
+        self.num_classes = num_classes
+        self.feat_dim = feat_dim
+        self.use_gpu = use_gpu
+
+        if self.use_gpu:
+            self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda())
+        else:
+            self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim))
+
+    def forward(self, x, labels):
+        """
+        Args:
+            x: feature matrix with shape (batch_size, feat_dim).
+            labels: ground truth labels with shape (num_classes).
+        """
+        batch_size = x.size(0)
+        distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \
+                  torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t()
+        distmat.addmm_(1, -2, x, self.centers.t())
+
+        classes = torch.arange(self.num_classes).long()
+        if self.use_gpu: classes = classes.cuda()
+        labels = labels.unsqueeze(1).expand(batch_size, self.num_classes)
+        mask = labels.eq(classes.expand(batch_size, self.num_classes))
+
+        dist = []
+        for i in range(batch_size):
+            value = distmat[i][mask[i]]
+            value = value.clamp(min=1e-12, max=1e+12) # for numerical stability
+            dist.append(value)
+        dist = torch.cat(dist)
+        loss = dist.mean()
+
+        return loss
+
+class RingLoss(nn.Module):
+    """Ring loss.
+    
+    Reference:
+    Zheng et al. Ring loss: Convex Feature Normalization for Face Recognition. CVPR 2018.
+    """
+    def __init__(self, weight_ring=1.):
+        super(RingLoss, self).__init__()
+        self.radius = nn.Parameter(torch.ones(1, dtype=torch.float))
+        self.weight_ring = weight_ring
+
+    def forward(self, x):
+        l = ((x.norm(p=2, dim=1) - self.radius)**2).mean()
+        return l * self.weight_ring
+
+class KLMutualLoss(nn.Module):
+    def __init__(self):
+        super(KLMutualLoss,self).__init__()
+        self.kl_loss = nn.KLDivLoss(size_average=False)
+        self.log_softmax = nn.functional.log_softmax
+        self.softmax = nn.functional.softmax
+    def forward(self, pred1, pred2):
+        pred1 = self.log_softmax(pred1, dim=1)
+        pred2 = self.softmax(pred2, dim=1)
+        #loss = self.kl_loss(pred1, torch.autograd.Variable(pred2.data))
+        loss = self.kl_loss(pred1, pred2.detach())
+        # from IPython import embed
+        # embed()
+        #print(loss)
+        return loss
+
+class MetricMutualLoss(nn.Module):
+    def __init__(self):
+        super(MetricMutualLoss, self).__init__()
+        self.l2_loss = nn.MSELoss()
+
+    def forward(self, dist1, dist2,pids):
+        loss = self.l2_loss(dist1, dist2)
+        # from IPython import embed
+        # embed()
+        print(loss)
+        return loss
+
+
+if __name__ == '__main__':
+    pass
\ No newline at end of file

util/optimizers.py

+import torch
+
+__all__ = ['init_optim']
+
+def init_optim(optim, params, lr, weight_decay):
+    if optim == 'adam':
+        return torch.optim.Adam(params, lr=lr, weight_decay=weight_decay)
+    elif optim == 'sgd':
+        return torch.optim.SGD(params, lr=lr, momentum=0.9, weight_decay=weight_decay)
+    elif optim == 'rmsprop':
+        return torch.optim.RMSprop(params, lr=lr, momentum=0.9, weight_decay=weight_decay)
+    else:
+        raise KeyError("Unsupported optim: {}".format(optim))
\ No newline at end of file

util/re_ranking.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri, 25 May 2018 20:29:09
+
+@author: luohao
+"""
+
+"""
+CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017.
+url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf
+Matlab version: https://github.com/zhunzhong07/person-re-ranking
+"""
+
+"""
+API
+
+probFea: all feature vectors of the query set (torch tensor)
+probFea: all feature vectors of the gallery set (torch tensor)
+k1,k2,lambda: parameters, the original paper is (k1=20,k2=6,lambda=0.3)
+MemorySave: set to 'True' when using MemorySave mode
+Minibatch: avaliable when 'MemorySave' is 'True'
+"""
+
+import numpy as np
+import torch
+
+def re_ranking(probFea, galFea, k1, k2, lambda_value, local_distmat = None, only_local = False):
+    # if feature vector is numpy, you should use 'torch.tensor' transform it to tensor
+    query_num = probFea.size(0)
+    all_num = query_num + galFea.size(0)
+    if only_local:
+        original_dist = local_distmat
+    else:
+        feat = torch.cat([probFea,galFea])
+        print('using GPU to compute original distance')
+        distmat = torch.pow(feat,2).sum(dim=1, keepdim=True).expand(all_num,all_num) + \
+                      torch.pow(feat, 2).sum(dim=1, keepdim=True).expand(all_num, all_num).t()
+        distmat.addmm_(1,-2,feat,feat.t())
+        original_dist = distmat.numpy()
+        del feat
+        if not local_distmat is None:
+            original_dist = original_dist + local_distmat
+    gallery_num = original_dist.shape[0]
+    original_dist = np.transpose(original_dist / np.max(original_dist, axis=0))
+    V = np.zeros_like(original_dist).astype(np.float16)
+    initial_rank = np.argsort(original_dist).astype(np.int32)
+
+    print('starting re_ranking')
+    for i in range(all_num):
+        # k-reciprocal neighbors
+        forward_k_neigh_index = initial_rank[i, :k1 + 1]
+        backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1]
+        fi = np.where(backward_k_neigh_index == i)[0]
+        k_reciprocal_index = forward_k_neigh_index[fi]
+        k_reciprocal_expansion_index = k_reciprocal_index
+        for j in range(len(k_reciprocal_index)):
+            candidate = k_reciprocal_index[j]
+            candidate_forward_k_neigh_index = initial_rank[candidate, :int(np.around(k1 / 2)) + 1]
+            candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index,
+                                               :int(np.around(k1 / 2)) + 1]
+            fi_candidate = np.where(candidate_backward_k_neigh_index == candidate)[0]
+            candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate]
+            if len(np.intersect1d(candidate_k_reciprocal_index, k_reciprocal_index)) > 2 / 3 * len(
+                    candidate_k_reciprocal_index):
+                k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index, candidate_k_reciprocal_index)
+
+        k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index)
+        weight = np.exp(-original_dist[i, k_reciprocal_expansion_index])
+        V[i, k_reciprocal_expansion_index] = weight / np.sum(weight)
+    original_dist = original_dist[:query_num, ]
+    if k2 != 1:
+        V_qe = np.zeros_like(V, dtype=np.float16)
+        for i in range(all_num):
+            V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0)
+        V = V_qe
+        del V_qe
+    del initial_rank
+    invIndex = []
+    for i in range(gallery_num):
+        invIndex.append(np.where(V[:, i] != 0)[0])
+
+    jaccard_dist = np.zeros_like(original_dist, dtype=np.float16)
+
+    for i in range(query_num):
+        temp_min = np.zeros(shape=[1, gallery_num], dtype=np.float16)
+        indNonZero = np.where(V[i, :] != 0)[0]
+        indImages = [invIndex[ind] for ind in indNonZero]
+        for j in range(len(indNonZero)):
+            temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum(V[i, indNonZero[j]],
+                                                                               V[indImages[j], indNonZero[j]])
+        jaccard_dist[i] = 1 - temp_min / (2 - temp_min)
+
+    final_dist = jaccard_dist * (1 - lambda_value) + original_dist * lambda_value
+    del original_dist
+    del V
+    del jaccard_dist
+    final_dist = final_dist[:query_num, query_num:]
+    return final_dist
+

util/samplers.py

+from __future__ import absolute_import
+from collections import defaultdict
+import numpy as np
+
+import torch
+from torch.utils.data.sampler import Sampler
+
+class RandomIdentitySampler(Sampler):
+    """
+    Randomly sample N identities, then for each identity,
+    randomly sample K instances, therefore batch size is N*K.
+
+    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/utils/data/sampler.py.
+
+    Args:
+        data_source (Dataset): dataset to sample from.
+        num_instances (int): number of instances per identity.
+    """
+    def __init__(self, data_source, num_instances=4):
+        self.data_source = data_source
+        self.num_instances = num_instances
+        self.index_dic = defaultdict(list)
+        for index, (_, pid, _) in enumerate(data_source):
+            self.index_dic[pid].append(index)
+        self.pids = list(self.index_dic.keys())
+        self.num_identities = len(self.pids)
+
+    def __iter__(self):
+        indices = torch.randperm(self.num_identities)
+        ret = []
+        for i in indices:
+            pid = self.pids[i]
+            t = self.index_dic[pid]
+            replace = False if len(t) >= self.num_instances else True
+            t = np.random.choice(t, size=self.num_instances, replace=replace)
+            ret.extend(t)
+        return iter(ret)
+
+    def __len__(self):
+        return self.num_identities * self.num_instances
\ No newline at end of file

util/transforms.py

+from __future__ import absolute_import
+
+from torchvision.transforms import *
+from PIL import Image
+import random
+import numpy as np
+
+class Random2DTranslation(object):
+    """
+    With a probability, first increase image size to (1 + 1/8), and then perform random crop.
+
+    Args:
+        height (int): target height.
+        width (int): target width.
+        p (float): probability of performing this transformation. Default: 0.5.
+    """
+    def __init__(self, height, width, p=0.5, interpolation=Image.BILINEAR):
+        self.height = height
+        self.width = width
+        self.p = p
+        self.interpolation = interpolation
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be cropped.
+
+        Returns:
+            PIL Image: Cropped image.
+        """
+        if random.random() < self.p:
+            return img.resize((self.width, self.height), self.interpolation)
+        new_width, new_height = int(round(self.width * 1.125)), int(round(self.height * 1.125))
+        resized_img = img.resize((new_width, new_height), self.interpolation)
+        x_maxrange = new_width - self.width
+        y_maxrange = new_height - self.height
+        x1 = int(round(random.uniform(0, x_maxrange)))
+        y1 = int(round(random.uniform(0, y_maxrange)))
+        croped_img = resized_img.crop((x1, y1, x1 + self.width, y1 + self.height))
+        return croped_img
+
+if __name__ == '__main__':
+    pass
\ No newline at end of file

util/utils.py

+from __future__ import absolute_import
+import os
+import sys
+import errno
+import shutil
+import json
+import os.path as osp
+from PIL import Image
+import matplotlib.pyplot as plt
+import cv2
+import numpy as np
+from numpy import array,argmin
+
+import torch
+
+def mkdir_if_missing(directory):
+    if not osp.exists(directory):
+        try:
+            os.makedirs(directory)
+        except OSError as e:
+            if e.errno != errno.EEXIST:
+                raise
+
+class AverageMeter(object):
+    """Computes and stores the average and current value.
+       
+       Code imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262
+    """
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+def save_checkpoint(state, is_best, fpath='checkpoint.pth.tar'):
+    mkdir_if_missing(osp.dirname(fpath))
+    torch.save(state, fpath)
+    if is_best:
+        shutil.copy(fpath, osp.join(osp.dirname(fpath), 'best_model.pth.tar'))
+
+class Logger(object):
+    """
+    Write console output to external text file.
+    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/utils/logging.py.
+    """
+    def __init__(self, fpath=None):
+        self.console = sys.stdout
+        self.file = None
+        if fpath is not None:
+            mkdir_if_missing(os.path.dirname(fpath))
+            self.file = open(fpath, 'w')
+
+    def __del__(self):
+        self.close()
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args):
+        self.close()
+
+    def write(self, msg):
+        self.console.write(msg)
+        if self.file is not None:
+            self.file.write(msg)
+
+    def flush(self):
+        self.console.flush()
+        if self.file is not None:
+            self.file.flush()
+            os.fsync(self.file.fileno())
+
+    def close(self):
+        self.console.close()
+        if self.file is not None:
+            self.file.close()
+
+def read_json(fpath):
+    with open(fpath, 'r') as f:
+        obj = json.load(f)
+    return obj
+
+def write_json(obj, fpath):
+    mkdir_if_missing(osp.dirname(fpath))
+    with open(fpath, 'w') as f:
+        json.dump(obj, f, indent=4, separators=(',', ': '))
+
+def _traceback(D):
+    i,j = array(D.shape)-1
+    p,q = [i],[j]
+    while (i>0) or (j>0):
+        tb = argmin((D[i,j-1], D[i-1,j]))
+        if tb == 0:
+            j -= 1
+        else: #(tb==1)
+            i -= 1
+        p.insert(0,i)
+        q.insert(0,j)
+    return array(p), array(q)
+
+def dtw(dist_mat):
+    m, n = dist_mat.shape[:2]
+    dist = np.zeros_like(dist_mat)
+    for i in range(m):
+        for j in range(n):
+            if (i == 0) and (j == 0):
+                dist[i, j] = dist_mat[i, j]
+            elif (i == 0) and (j > 0):
+                dist[i, j] = dist[i, j - 1] + dist_mat[i, j]
+            elif (i > 0) and (j == 0):
+                dist[i, j] = dist[i - 1, j] + dist_mat[i, j]
+            else:
+                dist[i, j] = \
+                    np.min(np.stack([dist[i - 1, j], dist[i, j - 1]], axis=0), axis=0) \
+                    + dist_mat[i, j]
+    path = _traceback(dist)
+    return dist[-1,-1]/sum(dist.shape), dist, path
+
+def read_image(img_path):
+    got_img = False
+    if not osp.exists(img_path):
+        raise IOError("{} does not exist".format(img_path))
+    while not got_img:
+        try:
+            img = Image.open(img_path).convert('RGB')
+            got_img = True
+        except IOError:
+            print("IOError incurred when reading '{}'. Will Redo. Don't worry. Just chill".format(img_path))
+            pass
+    return img
+
+def img_to_tensor(img,transform):
+    img = transform(img)
+    img = img.unsqueeze(0)
+    return img
+
+def show_feature(x):
+    for j in range(len(x)):
+        for i in range(len(64)):
+            ax = plt.subplot(4,16,i+1)
+            ax.set_title('No #{}'.format(i))
+            ax.axis('off')
+            plt.imshow(x[j].cpu().data.numpy()[0,i,:,:],cmap='jet')
+        plt.show()
+
+def feat_flatten(feat):
+    shp = feat.shape
+    feat = feat.reshape(shp[0] * shp[1], shp[2])
+    return feat
+
+def show_similar(local_img_path, img_path, similarity, bbox):
+    img1 = cv2.imread(local_img_path)
+    img2 = cv2.imread(img_path)
+    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
+    img1 = cv2.resize(img1, (64, 128))
+    img2 = cv2.resize(img2, (64, 128))
+    cv2.rectangle(img1, (bbox[0], bbox[1]), (bbox[0] + bbox[2], bbox[1] + bbox[3]), (0, 255, 0), 1)
+
+    p = np.where(similarity == np.max(similarity))
+    y, x = p[0][0], p[1][0]
+    cv2.rectangle(img2, (x - bbox[2] / 2, y - bbox[3] / 2), (x + bbox[2] / 2, y + bbox[3] / 2), (0, 255, 0), 1)
+    plt.subplot(1, 3, 1).set_title('patch')
+    plt.imshow(img1)
+    plt.subplot(1, 3, 2).set_title(('max similarity: ' + str(np.max(similarity))))
+    plt.imshow(img2)
+    plt.subplot(1, 3, 3).set_title('similarity')
+    plt.imshow(similarity)
+
+def show_alignedreid(local_img_path, img_path, dist):
+    def drow_line(img, similarity):
+        for i in range(1, len(similarity)):
+            cv2.line(img, (0, i*16), (63, i*16), color=(0,255,0))
+            cv2.line(img, (96, i*16), (160, i*16), color=(0,255,0))
+    def drow_path(img, path):
+        for i in range(len(path[0])):
+            cv2.line(img, (64, 8+16*path[0][i]), (96,8+16*path[1][i]), color=(255,255,0))
+    img1 = cv2.imread(local_img_path)
+    img2 = cv2.imread(img_path)
+    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
+    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
+    img1 = cv2.resize(img1, (64,128))
+    img2 = cv2.resize(img2, (64,128))
+    img = np.zeros((128,160,3)).astype(img1.dtype)
+    img[:,:64,:] = img1
+    img[:,-64:,:] = img2
+    drow_line(img, dist)
+    d,D,sp = dtw(dist)
+    origin_dist = np.mean(np.diag(dist))
+    drow_path(img, sp)
+    plt.subplot(1,2,1).set_title('Aligned distance: %.4f \n Original distance: %.4f' %(d,origin_dist))
+    plt.subplot(1,2,1).set_xlabel('Aligned Result')
+    plt.imshow(img)
+    plt.subplot(1,2,2).set_title('Distance Map')
+    plt.subplot(1,2,2).set_xlabel('Right Image')
+    plt.subplot(1,2,2).set_ylabel('Left Image')
+    plt.imshow(dist)
+    plt.subplots_adjust(bottom=0.1, left=0.075, right=0.85, top=0.9)
+    cax = plt.axes([0.9, 0.25, 0.025, 0.5])
+    plt.colorbar(cax = cax)
+    plt.show()
+
+def merge_feature(feature_list, shp, sample_rate = None):
+    def pre_process(torch_feature_map):
+        numpy_feature_map = torch_feature_map.cpu().data.numpy()[0]
+        numpy_feature_map = numpy_feature_map.transpose(1,2,0)
+        shp = numpy_feature_map.shape[:2]
+        return numpy_feature_map, shp
+    def resize_as(tfm, shp):
+        nfm, shp2 = pre_process(tfm)
+        scale = shp[0]/shp2[0]
+        nfm1 = nfm.repeat(scale, axis = 0).repeat(scale, axis=1)
+        return nfm1
+    final_nfm = resize_as(feature_list[0], shp)
+    for i in range(1, len(feature_list)):
+        temp_nfm = resize_as(feature_list[i],shp)
+        final_nfm = np.concatenate((final_nfm, temp_nfm),axis =-1)
+    if sample_rate > 0:
+        final_nfm = final_nfm[0:-1:sample_rate, 0:-1,sample_rate, :]
+    return final_nfm
\ No newline at end of file