121 lines
6.0 KiB
Python
121 lines
6.0 KiB
Python
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()
|
||
|
||
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 = 3
|
||
gallery_path = f"./mydataset/gallery/"
|
||
classes = os.listdir(gallery_path)
|
||
cmc_sum = []
|
||
cmc_final = []
|
||
for query_index in range(1, 6+1): #query是从1开始的
|
||
|
||
query_path = f"./mydataset/query/cam{query_index}/" # Query path
|
||
#Gallary path，这里直接对gallary文件夹下所有图片进行分析
|
||
print("--------------query cam is ", query_index)
|
||
for cls in classes:
|
||
if cls != f"{query_index}":
|
||
print("——————————gallery cam is", cls)
|
||
gallery_path = f"./mydataset/gallery/{cls}/"
|
||
query_feature_list = [] # 计算每个query的特征向量保存在列表中
|
||
acck_l_list = [] # 二维列表（列表的列表），其元素（也是列表）存储单个query的"acck"("cmc")
|
||
cmc_k_list = [] # 存储最终的CMC数字（所有query取平均）
|
||
query_list = os.listdir(query_path) # 为每张图片生成路径并保存在列表变量中
|
||
query_list.sort()
|
||
gallery_list = os.listdir(gallery_path) # 为每张图片生成路径并保存在列表变量中
|
||
with torch.no_grad():
|
||
# 对于每个query，（这句话加在本循环内每个注释）
|
||
for path in query_list:
|
||
distance_dict = {}
|
||
parsed_query_id = path.split("_")[1]
|
||
parsed_query_id = parsed_query_id.split(".")[0]
|
||
parsed_query_id = int(parsed_query_id)
|
||
if parsed_query_id == 4 and cls == "5":
|
||
continue
|
||
print(parsed_query_id)
|
||
id_feature = extractidfeature(query_path+path, extractor)
|
||
query_feature_list.append(id_feature)
|
||
counter = 0
|
||
# 计算每个gallary图片到目标query的特征距离
|
||
for p in gallery_list:
|
||
if counter%10 == 9: # 抽样，相当于设置帧间隔
|
||
img_path = gallery_path+p
|
||
gallery_feature = extractidfeature(img_path, extractor)
|
||
distance_dict[p] = np.sum(np.abs(id_feature-gallery_feature))
|
||
counter += 1
|
||
matched = sorted(distance_dict.items(), key=lambda x : x[1]) # 按特征距离排序
|
||
# 储存k从前1到前10名的是否（存在命中的ID）
|
||
temp_acck = [int(matched[0][0].split("_")[2].split(".")[0])==parsed_query_id] # ID相同为True，否则False
|
||
# 获取前k个acck情况
|
||
for i in range(1,11):
|
||
parsed_gallery_id = matched[i][0].split("_")[2].split(".")[0]
|
||
parsed_gallery_id = int(parsed_gallery_id) # 解析出gallary的ID
|
||
print(parsed_gallery_id) # 解析出gallary的ID
|
||
temp_acck.append(parsed_gallery_id==parsed_query_id or temp_acck[i-1]) # or用意是表明如果Acc(k-1)已经为True，那么Acck必然为True（递推）
|
||
print(temp_acck)
|
||
acck_l_list.append(temp_acck)
|
||
# 计算不同k下cmc的具体数值
|
||
for k in range(1, 11):
|
||
acc_counter = 0.0
|
||
length = len(acck_l_list)
|
||
for results in acck_l_list:
|
||
acc_counter += results[k]
|
||
cmc_k_list.append(acc_counter/length)
|
||
#打印以供检查
|
||
print(cmc_k_list)
|
||
cmc_sum.append(cmc_k_list)
|
||
k_index = [k for k in range(1, 11)]
|
||
#plot figure并保存（不会显示）
|
||
plt.clf()
|
||
plt.plot(k_index, cmc_k_list)
|
||
plt.xlabel("k")
|
||
plt.ylabel("ACCK")
|
||
plt.ylim(0.0,1.05)
|
||
plt.title(f"CMC:query_cam{query_index} and gallery {cls}")
|
||
for x, y in zip(k_index, cmc_k_list):
|
||
plt.text(x, y + 0.02, str(round(y, 3)), ha='center', va='bottom', fontsize=10.5)
|
||
plt.draw()
|
||
plt.savefig(f"./CMC/query_cam{query_index}_and_gallery_{cls}.jpg")
|
||
for k in range(10):
|
||
acc_counter = 0.0
|
||
length = len(cmc_sum)
|
||
for results in cmc_sum:
|
||
acc_counter += results[k]
|
||
cmc_final.append(acc_counter / length)
|
||
print(cmc_final)
|
||
k_index = [k for k in range(1, 11)]
|
||
# plot figure并保存（不会显示）
|
||
plt.clf()
|
||
plt.plot(k_index, cmc_final)
|
||
plt.xlabel("k")
|
||
plt.ylabel("ACCK")
|
||
plt.ylim(0.0, 1.05)
|
||
plt.title(f"CMC:query and gallery")
|
||
for x, y in zip(k_index, cmc_final):
|
||
plt.text(x, y + 0.02, str(round(y,3)), ha='center', va='bottom', fontsize=10.5)
|
||
plt.draw()
|
||
plt.savefig(f"./CMC/query_gallery.jpg") |