| 27 | |
| 28 | |
| 29 | def re_ranking(probFea, galFea, k1, k2, lambda_value, local_distmat=None, only_local=False): |
| 30 | # if feature vector is numpy, you should use 'torch.tensor' transform it to tensor |
| 31 | query_num = probFea.size(0) |
| 32 | all_num = query_num + galFea.size(0) |
| 33 | if only_local: |
| 34 | original_dist = local_distmat |
| 35 | else: |
| 36 | feat = torch.cat([probFea, galFea]) |
| 37 | # print('using GPU to compute original distance') |
| 38 | distmat = torch.pow(feat, 2).sum(dim=1, keepdim=True).expand(all_num, all_num) + \ |
| 39 | torch.pow(feat, 2).sum(dim=1, keepdim=True).expand(all_num, all_num).t() |
| 40 | distmat.addmm_(1, -2, feat, feat.t()) |
| 41 | original_dist = distmat.cpu().numpy() |
| 42 | del feat |
| 43 | if not local_distmat is None: |
| 44 | original_dist = original_dist + local_distmat |
| 45 | gallery_num = original_dist.shape[0] |
| 46 | original_dist = np.transpose(original_dist / np.max(original_dist, axis=0)) |
| 47 | V = np.zeros_like(original_dist).astype(np.float16) |
| 48 | initial_rank = np.argsort(original_dist).astype(np.int32) |
| 49 | |
| 50 | # print('starting re_ranking') |
| 51 | for i in range(all_num): |
| 52 | # k-reciprocal neighbors |
| 53 | forward_k_neigh_index = initial_rank[i, :k1 + 1] |
| 54 | backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1] |
| 55 | fi = np.where(backward_k_neigh_index == i)[0] |
| 56 | k_reciprocal_index = forward_k_neigh_index[fi] |
| 57 | k_reciprocal_expansion_index = k_reciprocal_index |
| 58 | for j in range(len(k_reciprocal_index)): |
| 59 | candidate = k_reciprocal_index[j] |
| 60 | candidate_forward_k_neigh_index = initial_rank[candidate, :int(np.around(k1 / 2)) + 1] |
| 61 | candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index, |
| 62 | :int(np.around(k1 / 2)) + 1] |
| 63 | fi_candidate = np.where(candidate_backward_k_neigh_index == candidate)[0] |
| 64 | candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate] |
| 65 | if len(np.intersect1d(candidate_k_reciprocal_index, k_reciprocal_index)) > 2 / 3 * len( |
| 66 | candidate_k_reciprocal_index): |
| 67 | k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index, candidate_k_reciprocal_index) |
| 68 | |
| 69 | k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) |
| 70 | weight = np.exp(-original_dist[i, k_reciprocal_expansion_index]) |
| 71 | V[i, k_reciprocal_expansion_index] = weight / np.sum(weight) |
| 72 | original_dist = original_dist[:query_num, ] |
| 73 | if k2 != 1: |
| 74 | V_qe = np.zeros_like(V, dtype=np.float16) |
| 75 | for i in range(all_num): |
| 76 | V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0) |
| 77 | V = V_qe |
| 78 | del V_qe |
| 79 | del initial_rank |
| 80 | invIndex = [] |
| 81 | for i in range(gallery_num): |
| 82 | invIndex.append(np.where(V[:, i] != 0)[0]) |
| 83 | |
| 84 | jaccard_dist = np.zeros_like(original_dist, dtype=np.float16) |
| 85 | |
| 86 | for i in range(query_num): |