def train_n_eval(feat_type='SIFT'):
q_valid = QueryExtractor(dataset=hp.mode,
image_dir=hp.image_dir,
label_dir=hp.label_dir,
subset='valid')
data_valid = [q_name for q_name, _ in q_valid.get_queries().items()]
data_train = [
fname for fname in os.listdir(hp.image_dir) if fname not in data_valid
]
if feat_type == 'SIFT':
# train and creat database
if len(os.listdir('./database/BoW/SIFT/')) == 0:
create_db_SIFT(data_train)
# evaluate on queries
eval(feat_type, q_valid)
elif feat_type == 'SURF':
# train and creat database
if len(os.listdir('./database/BoW/SURF/')) == 0:
create_db_SURF(data_train)
# evaluate on queries
eval(feat_type, q_valid)
else:
pass
map_location = 'cpu'
ckpt = torch.load(hp.logdir + 'model.pkl', map_location=map_location)
model.load_state_dict(ckpt['state_dict'])
print('Restore model')
model.eval()
cs_func = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
db_embeddings, maps = create_DB(model)
mAP, time_running = [], []
q_valid = QueryExtractor(dataset=hp.mode,
image_dir=hp.image_dir,
label_dir=hp.label_dir,
subset='valid')
for q_name, attribute in q_valid.get_queries().items():
start = time.time()
bbox, class_idx = attribute[0], attribute[1]
# create image tensor
query_img = image_preprocessing(hp.image_dir + q_name)
query_tensor = torch.FloatTensor(
np.transpose(np.expand_dims(query_img, axis=0), axes=[0, 3, 1, 2]))
# get embedding vector
if use_gpu:
query_tensor = query_tensor.cuda()
cs_func = cs_func.cuda()
query_embedding = inference(model, query_tensor)
similarity = cs_func(query_embedding, db_embeddings).topk(
len(q_valid.get_groundtruth()[class_idx]))
prediction = [maps[idx] for idx in similarity[1].cpu().numpy()]
end = time.time()