Update simple_reid.py

This commit is contained in:
wffjwbbf 2022-11-04 19:24:01 +08:00
parent ceaa0a5324
commit be7ba90ed5
1 changed files with 17 additions and 10 deletions

View File

@ -20,10 +20,10 @@ def extractidfeature(id_path: str, extractor: FeatureExtractionDemo):
if __name__ == '__main__':
dislist = {}
extractor_model_path = "./deep_sort/deep/checkpoint/market_bot_R50-ibn.pth"
cfg_path = "./fastreid/cfgs/Market1501/bagtricks_R50-ibn.yml"
query_path = "D:\\study\\python\\1zhouxue\\UESTC_ReID_Dataset_V3\\query\\cam3"
cfg_path = "./fastreid/cfgs/Market1501/bagtricks_R50-ibn.yml"
query_path = "D:\\study\\python\\1zhouxue\\UESTC_ReID_Dataset_V3\\query\\cam6"
galla_path = "D:\\study\\python\\1zhouxue\\UESTC_ReID_Dataset_V3\\gallery"
C=3
C = 6 // 第几个摄像头
gallary_list = os.listdir(galla_path)
count = 0
cfg = get_cfg()
@ -32,20 +32,20 @@ if __name__ == '__main__':
cfg.merge_from_file(cfg_path)
cfg.freeze()
my_option = param.Parameters() # load model in head
my_option = param.Parameters() # load model in head
extractor = FeatureExtractionDemo(cfg)
with torch.no_grad():
for i in range(1,6):
for i in range(1, 7):
count = 0
dislist.clear()
id_try = extractidfeature(query_path+f"/{C}_{i}.jpg",extractor)
id_try = extractidfeature(query_path+f"/{C}_{i}.jpg", extractor)
for gal in gallary_list:
count += 1
if count % 25 == 0:
id_tar = extractidfeature(galla_path+'/'+gal,extractor)
id_tar = extractidfeature(galla_path+'/'+gal, extractor)
dist = np.sum(np.abs(id_try-id_tar))
dislist[gal] = float(dist)
if count > 8000:
if count > 10000:
break
matched = sorted(dislist.items(), key=lambda x : x[1])
print(matched)
@ -53,6 +53,13 @@ if __name__ == '__main__':
os.mkdir("match/" + f"{C}_{i}")
except:
pass
img = plt.imread(query_path+f"/{C}_{i}.jpg")
plt.subplot(1, 11, 1)
plt.title("query")
plt.imshow(img)
plt.xticks([])
plt.yticks([])
x = 1
for path in matched[0:10]:
s_path = galla_path+'/'+path[0]
@ -64,8 +71,8 @@ if __name__ == '__main__':
# print(ID)
ID = int(ID[0])
img = plt.imread(s_path)
plt.subplot(1, 10, x)
if id == i:
plt.subplot(1, 11, x+1)
if ID == i:
plt.title(x)
else:
plt.title(f"{x}_F")