From d83d58f2b7d2385b1e25470634ab80748b7ad40a Mon Sep 17 00:00:00 2001
From: Dechao Meng <dechao.meng@vipl.ict.ac.cn>
Date: Wed, 20 Jan 2021 12:56:00 +0800
Subject: [PATCH] fix bug and add pretrained model

---
 README.md                           |  5 ++-
 examples/parsing_reid/main.py       | 10 ++++++
 examples/parsing_reid/math_tools.py | 50 +++++++++++++++--------------
 3 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/README.md b/README.md
index ffc1679..b4a016a 100644
--- a/README.md
+++ b/README.md
@@ -55,7 +55,6 @@ python generate_pkl.py veriwild --input-path <VeRi_PATH> --output-path ../output
 python generate_pkl.py vehicleid --input-path <VeRi_PATH> --output-path ../outputs/vehicleid.pkl
 ```
 
-
 ## Training the parsing model
 <!-- We provide the pre-trained segmentation model on `examples/parsing/best_model_trainval.pth` which you can use to generate parsing masks for different datasets. 
 If you want to use the model directly, just skip this section.
@@ -101,6 +100,10 @@ python main.py train -c configs/vehicleid_b256_pven.yml
 python main.py train -c configs/veriwild_b256_224_pven.yml 
 ```
 
+## Pretrained Models
+We provide the pretrained parsing model, VeRi776 ReID model and VERIWild ReID model ( the classification layer has been removed ) for your convinient.
+
+链接: https://pan.baidu.com/s/1Q2NMVfGZPCskh-E6vmy9Cw  密码: iiw1
 ## Evaluate PVEN
 ```shell
 cd examples/parsing_reid
diff --git a/examples/parsing_reid/main.py b/examples/parsing_reid/main.py
index 218650b..ce22073 100644
--- a/examples/parsing_reid/main.py
+++ b/examples/parsing_reid/main.py
@@ -513,5 +513,15 @@ def eval_vehicle_id_(model, valid_loader, query_length, cfg):
         logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r - 1]:.2%}")
     return cmc, mAP
 
+@clk.command()
+@click.option('-i', '--model-path')
+@click.option('-o', '--output-path')
+def drop_linear(model_path, output_path):
+    model = torch.load(model_path)
+    for key in model.keys():
+        if 'classifier' in key:
+            model[key] = None
+    torch.save(model, output_path)
+
 if __name__ == '__main__':
     clk()
diff --git a/examples/parsing_reid/math_tools.py b/examples/parsing_reid/math_tools.py
index 8023463..bb2f090 100644
--- a/examples/parsing_reid/math_tools.py
+++ b/examples/parsing_reid/math_tools.py
@@ -1,4 +1,5 @@
 import os
+import ipdb
 
 import torch
 from torch.nn import functional as F
@@ -12,31 +13,32 @@
 from vehicle_reid_pytorch.metrics.rerank import re_ranking
 
 
-def calc_dist_split(qf, gf, split=0):
-    qf = qf
-    m = qf.shape[0]
-    n = gf.shape[0]
-    distmat = gf.new(m, n)
+# def calc_dist_split(qf, gf, split=0):
+#     qf = qf
+#     m = qf.shape[0]
+#     n = gf.shape[0]
+#     distmat = gf.new(m, n)
 
-    if split == 0:
-        distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
-                torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
+#     if split == 0:
+#         distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
+#                 torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
+#         distmat.addmm_(x, y.t(), beta=1, alpha=-2)
 
-    # 用于测试时控制显存
-    else:
-        start = 0
-        while start < n:
-            end = start + split if (start + split) < n else n
-            num = end - start
+#     # 用于测试时控制显存
+#     else:
+#         start = 0
+#         while start < n:
+#             end = start + split if (start + split) < n else n
+#             num = end - start
 
-            sub_distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, num) + \
-                    torch.pow(gf[start:end], 2).sum(dim=1, keepdim=True).expand(num, m).t()
-            # sub_distmat.addmm_(1, -2, qf, gf[start:end].t())
-            sub_distmat.addmm_(qf, gf[start:end].t(), beta=1, alpha=-2)
-            distmat[:, start:end] = sub_distmat.cpu()
-            start += num
+#             sub_distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, num) + \
+#                     torch.pow(gf[start:end], 2).sum(dim=1, keepdim=True).expand(num, m).t()
+#             # sub_distmat.addmm_(1, -2, qf, gf[start:end].t())
+#             sub_distmat.addmm_(qf, gf[start:end].t(), beta=1, alpha=-2)
+#             distmat[:, start:end] = sub_distmat.cpu()
+#             start += num
 
-    return distmat
+#     return distmat
 
 
 def clck_dist(feat1, feat2, vis_score1, vis_score2, split=0):
@@ -58,7 +60,7 @@ def clck_dist(feat1, feat2, vis_score1, vis_score2, split=0):
         parse_feat2 = feat2[:, :, i]
         ckcl_ = torch.mm(vis_score1[:, i].view(-1, 1), vis_score2[:, i].view(1, -1))  # [N, N]
         ckcl += ckcl_
-        dist_mat += calc_dist_split(parse_feat1, parse_feat2, split=split).sqrt() * ckcl_
+        dist_mat += euclidean_dist(parse_feat1, parse_feat2, split=split).sqrt() * ckcl_
 
     return dist_mat / ckcl
 
@@ -165,7 +167,7 @@ def compute(self, split=0):
             if split == 0:
                 distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
                         torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
-                distmat.addmm_(1, -2, qf, gf.t())
+                distmat.addmm_(qf, gf.t(), beta=1, alpha=-2)
             else:
                 distmat = gf.new(m, n)
                 start = 0
@@ -207,7 +209,7 @@ def compute(self, split=0):
             torch.save(outputs, os.path.join(self.output_path, 'test_output.pkl'), pickle_protocol=4)
 
         print('Eval...')
-        cmc, mAP, all_AP = eval_func_mp(distmat + self.lambda_ * local_distmat, q_pids, g_pids, q_camids, g_camids,
+        cmc, mAP, all_AP = eval_func_mp(distmat + self.lambda_ * (local_distmat ** 2), q_pids, g_pids, q_camids, g_camids,
                              remove_junk=self.remove_junk)
 
         return {