119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
#####
|
||
# 保存query和前十个配准的gallery图像，并计算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
|
||
|
||
from model import my_model
|
||
import param
|
||
|
||
|
||
def extractidfeature(id_path: str, extractor, p_device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
|
||
img_id = cv2.imread(id_path)
|
||
raw_transformer = transforms.Compose([
|
||
transforms.ToPILImage(),
|
||
transforms.Resize((128,64)),# hxw
|
||
transforms.ToTensor(),
|
||
])
|
||
tensor_id = extractor(raw_transformer(img_id).unsqueeze(dim=0).to(p_device))
|
||
|
||
return tensor_id.detach().cpu().numpy()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
my_option = param.Parameters()
|
||
dislist = {}
|
||
|
||
extractor = my_model.RestNet18() # load model in head
|
||
extractor.load_state_dict(torch.load(my_option.weights_reid))
|
||
extractor = extractor.to(my_option.device)
|
||
|
||
gallery_index = 4
|
||
camera_index = 1 # 第几个摄像头
|
||
query_path = f"./mydataset/query/cam{camera_index}/"
|
||
galla_path = f"./mydataset/gallery/{gallery_index}/"
|
||
|
||
gallary_list = os.listdir(galla_path)
|
||
count = 0
|
||
number_query = 12
|
||
|
||
#####计算mAP使用的变量
|
||
mAP = 0
|
||
precision = 0
|
||
sum_precision = 0
|
||
AP = 0
|
||
sum_AP = 0
|
||
number_true_gallery = 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 gal in gallary_list:
|
||
count += 1
|
||
if count % 5 == 0:
|
||
id_tar = extractidfeature(galla_path+'/'+gal, extractor)
|
||
dist = np.sum(np.abs(id_try-id_tar))
|
||
dislist[gal] = float(dist)
|
||
# if count > 10000:
|
||
# break
|
||
matched = sorted(dislist.items(), key=lambda x : x[1])
|
||
print(matched)
|
||
try:
|
||
os.mkdir("match/" + f"{camera_index}_{i}")
|
||
except:
|
||
pass
|
||
|
||
img = plt.imread(query_path+f"/{camera_index}_{i}.jpg")
|
||
plt.subplot(1, 11, 1)
|
||
plt.title(f"q_{i}_g_{gallery_index}")
|
||
plt.imshow(img)
|
||
plt.xticks([])
|
||
plt.yticks([])
|
||
ranking = 1
|
||
ranking_true_gallery = 1
|
||
for path in matched[0:10]:
|
||
s_path = galla_path+'/'+path[0]
|
||
t_path = "match/" + f"{camera_index}_{i}" + "/" + 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_gallery/ranking
|
||
sum_precision += precision
|
||
ranking_true_gallery += 1
|
||
|
||
else:
|
||
plt.title(ranking, color='red')
|
||
plt.imshow(img)
|
||
plt.xticks([])
|
||
plt.yticks([])
|
||
ranking += 1
|
||
# plt.show()
|
||
plt.draw()
|
||
plt.savefig(f"query{i}_cam{camera_index}_and_gallery_{gallery_index}.jpg")
|
||
number_true_gallery = ranking_true_gallery-1
|
||
if 0 != number_true_gallery:
|
||
AP = sum_precision / number_true_gallery
|
||
else:
|
||
AP = 0
|
||
sum_precision = 0
|
||
sum_AP += AP
|
||
print("query_", i, "AP==", AP)
|
||
mAP = sum_AP/number_query
|
||
print("mAP == ", mAP)
|
||
|