SeekingForJoy
/
ComDesignProject
forked from wffjwbbf/ComDesignProject
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update simple_reid.py

This commit is contained in:
wffjwbbf 2022-11-04 18:54:52 +08:00
parent fe39425990
commit ceaa0a5324
1 changed files with 32 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
simple_reid.py
View File

@ -1,19 +1,17 @@
import shutil
import matplotlib.pyplot as plt
import torch
from torchvision import transforms
import numpy as np
import cv2
import os
import param
from deep_sort.deep.feature_extractor import FastReIDExtractor
from fastreid.demo import FeatureExtractionDemo
from fastreid.config import get_cfg
def extractidfeature(id_path:str,extractor:FeatureExtractionDemo):
def extractidfeature(id_path: str, extractor: FeatureExtractionDemo):
img_id = cv2.imread(id_path)
tensor_id = extractor.run_on_image(img_id)
return tensor_id.detach().cpu().numpy()
@ -21,46 +19,60 @@ def extractidfeature(id_path:str,extractor:FeatureExtractionDemo):
# example
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 = "./query/cam1"
galla_path = "./gallary"
query_path = "D:\\study\\python\\1zhouxue\\UESTC_ReID_Dataset_V3\\query\\cam3"
galla_path = "D:\\study\\python\\1zhouxue\\UESTC_ReID_Dataset_V3\\gallery"
C=3
gallary_list = os.listdir(galla_path)
count = 0
cfg = get_cfg()
cfg.MODEL.DEVICE = 'cpu'
cfg.MODEL.WEIGHTS = extractor_model_path
cfg.merge_from_file(cfg_path)
cfg.freeze()
my_option = param.Parameters() # load model in head
extractor = FeatureExtractionDemo(cfg)
with torch.no_grad():
for i in range(3,26):
for i in range(1,6):
count = 0
dislist.clear()
id_try = extractidfeature(query_path+f"/1_{i}.jpg",extractor)
id_try = extractidfeature(query_path+f"/{C}_{i}.jpg",extractor)
for gal in gallary_list:
count += 1
if(count%10 == 0):
if count % 25 == 0:
id_tar = extractidfeature(galla_path+'/'+gal,extractor)
dist = np.sum(np.abs(id_try-id_tar))
dist = np.sum(np.abs(id_try-id_tar))
dislist[gal] = float(dist)
if count>5000:
if count > 8000:
break
matched = sorted(dislist.items(),key=lambda x:x[1])
matched = sorted(dislist.items(), key=lambda x : x[1])
print(matched)
try:
os.mkdir("match/"+ f"1_{i}")
os.mkdir("match/" + f"{C}_{i}")
except:
pass
x = 1
for path in matched[0:10]:
s_path = galla_path+'/'+path[0]
t_path = "match/"+ f"1_{i}" +"/"+path[0]
print(s_path,t_path)
shutil.copy(s_path,t_path)
t_path = "match/" + f"{C}_{i}" + "/" + path[0]
print(s_path, t_path)
shutil.copy(s_path, t_path)
ID = path[0].split('_')
# print(ID)
ID = int(ID[0])
img = plt.imread(s_path)
plt.subplot(1, 10, x)
if id == i:
plt.title(x)
else:
plt.title(f"{x}_F")
plt.imshow(img)
plt.xticks([])
plt.yticks([])
x += 1
plt.show()