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/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