Add mAP_reid_detector.py

mAP_reid_detector.py

@@ -0,0 +1,107 @@
+#####
+# 保存query和前十个配准的detector图像，并计算mAP(mean average precision)
+######
+import shutil
+import matplotlib.pyplot as plt
+import torch
+from torchvision import transforms
+import numpy as np
+import cv2
+import os
+import param
+import param_v1
+from fastreid.my_extractor import MyExtractor
+
+
+def extractidfeature(id_path: str, extractor: MyExtractor):
+    img_id = cv2.imread(id_path)
+    tensor_id = extractor([img_id])
+    return tensor_id.detach().cpu().numpy()
+
+
+if __name__ == '__main__':
+    dislist = {}
+    my_option = param_v1.Parameters()  # load model in head
+    extractor = MyExtractor(my_option.weights_reid, use_cuda=True)
+    detector_index = 3
+    camera_index = 2 # 第几个摄像头
+    query_path = f"../ReIDMixLearning/mydataset/query/cam{camera_index}/"
+    detect_path = f"../yolov7-inferonly/detector/cam{detector_index}/"
+
+    detect_list = os.listdir(detect_path)
+    count = 0
+    number_query = 24
+#####计算mAP使用的变量
+    mAP = 0
+    precision = 0
+    sum_precision = 0
+    AP = 0
+    sum_AP = 0
+    number_true_detector = 0
+   #####
+
+    with torch.no_grad():
+        for i in range(1, 1+number_query):
+            count = 0
+            dislist.clear()
+            id_try = extractidfeature(query_path+f"/{camera_index}_{i}.jpg", extractor)
+            for det in detect_list:
+                count += 1
+                # if count % 5 == 0:
+                id_tar = extractidfeature(detect_path+'/'+det, extractor)
+                dist = np.sum(np.abs(id_try-id_tar))
+                dislist[det] = float(dist)
+                # if count > 10000:
+                #     break
+            matched = sorted(dislist.items(), key=lambda x : x[1])
+            print(matched)
+            try:
+                os.mkdir("match/" + f"c{camera_index}_q{i}_d{detector_index}")
+            except:
+                pass
+
+            img = plt.imread(query_path+f"/{camera_index}_{i}.jpg")
+            plt.subplot(1, 11, 1)
+            plt.title(f"q{i}c{camera_index}_d{detector_index}")
+            plt.imshow(img)
+            plt.xticks([])
+            plt.yticks([])
+            ranking = 1
+            ranking_true_detector = 1
+            for path in matched[0:10]:
+                s_path = detect_path+'/'+path[0]
+                t_path = "match/" + f"c{camera_index}_q{i}_d{detector_index}" + "/" + path[0]
+                # print(s_path, t_path)
+                shutil.copy(s_path, t_path)
+                ID = path[0].split('_')
+                # print(ID)
+                ID = int(ID[2].split(".")[0])
+                print(i, ID)
+                img = plt.imread(s_path)
+                plt.subplot(1, 11, ranking+1)
+                if ID == i:
+                    plt.title(ranking, color='black')
+                    precision = ranking_true_detector/ranking
+                    sum_precision += precision
+                    ranking_true_detector += 1
+
+                else:
+                    plt.title(ranking, color='red')
+                plt.imshow(img)
+                plt.xticks([])
+                plt.yticks([])
+                ranking += 1
+            # plt.show()
+            plt.draw()
+            plt.savefig(f"./mAP_detector/q{i}c{camera_index}_dc{detector_index}.jpg")
+            number_true_detector = ranking_true_detector-1
+            if 0 != number_true_detector:
+                AP = sum_precision / number_true_detector
+            else:
+                AP = 0
+            sum_precision = 0
+            sum_AP += AP
+            print("query", i, "AP==", AP)
+        mAP = sum_AP/number_query
+        print("mAP == ", mAP)
+