Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
380 lines
12 KiB
Python
380 lines
12 KiB
Python
import time
|
||
import torch
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
from torch.cuda.amp import autocast
|
||
import torch.nn.functional as F
|
||
import faiss
|
||
import faiss.contrib.torch_utils
|
||
import h5py
|
||
import os
|
||
|
||
def predict(train_config, model, dataloader):
|
||
|
||
model.eval()
|
||
|
||
# wait before starting progress bar
|
||
# time.sleep(0.1)
|
||
bar = tqdm(dataloader, total=len(dataloader))
|
||
# if train_config.verbose:
|
||
# bar = tqdm(dataloader, total=len(dataloader))
|
||
# else:
|
||
# bar = dataloader
|
||
|
||
img_features_list = []
|
||
# import time
|
||
# torch.cuda.synchronize()
|
||
# st = time.time()
|
||
with torch.no_grad():
|
||
|
||
for img, pt in bar:
|
||
|
||
with autocast():
|
||
|
||
img_feature, _ = model(img, pt)
|
||
# print(f"Initial memory allocated: {torch.cuda.memory_allocated()} bytes")
|
||
|
||
# save features in fp32 for sim calculation
|
||
img_features_list.append(img_feature.to(torch.float32))
|
||
# torch.cuda.synchronize()
|
||
# et = time.time()
|
||
print('---------------------------------time---------------------------------')
|
||
# print('time cost: ', (et - st)/len(dataloader))
|
||
# keep Features on GPU
|
||
img_features = torch.cat(img_features_list, dim=0)
|
||
|
||
# if train_config.verbose:
|
||
bar.close()
|
||
|
||
return img_features
|
||
|
||
def predict_rerank(train_config, model, dataloader, name, mode):
|
||
|
||
model.eval()
|
||
|
||
# wait before starting progress bar
|
||
time.sleep(0.1)
|
||
|
||
if train_config.verbose:
|
||
bar = tqdm(dataloader, total=len(dataloader))
|
||
else:
|
||
bar = dataloader
|
||
|
||
img_features_list = []
|
||
h5_name = str(name)+'_'+mode + '.h5'
|
||
|
||
|
||
with torch.no_grad():
|
||
|
||
for img, pt, img_path in bar:
|
||
|
||
with autocast():
|
||
|
||
img_feature, geat_list = model(img, pt)
|
||
# save features in fp32 for sim calculation
|
||
img_features_list.append(img_feature.to(torch.float32))
|
||
# average_geats = torch.mean(geat_list, dim=2)
|
||
# average_geats = average_geats.reshape(geat_list.shape[1], geat_list.shape[3], geat_list.shape[4]).cpu()
|
||
feature_geats = geat_list.squeeze(0).cpu()
|
||
feature_geats = feature_geats[::60, :, :, :].reshape(-1, 24)
|
||
|
||
# if os.path.exists(h5_name):
|
||
# pass
|
||
with h5py.File(h5_name, 'a', libver='latest') as fd:
|
||
if img_path[0] in fd:
|
||
continue
|
||
grp = fd.create_group(img_path[0])
|
||
grp.create_dataset('global_feature', data=feature_geats.cpu())
|
||
|
||
# keep Features on GPU
|
||
img_features = torch.cat(img_features_list, dim=0)
|
||
print('---------------------------------save h5 file---------------------------------')
|
||
|
||
if train_config.verbose:
|
||
bar.close()
|
||
|
||
return img_features
|
||
|
||
def predict_backbone(train_config, model, dataloader, LPN):
|
||
|
||
model.eval()
|
||
|
||
# wait before starting progress bar
|
||
# time.sleep(0.1)
|
||
bar = tqdm(dataloader, total=len(dataloader))
|
||
|
||
# if train_config.verbose:
|
||
# bar = tqdm(dataloader, total=len(dataloader))
|
||
# else:
|
||
# bar = dataloader
|
||
|
||
img_features_list = []
|
||
# import time
|
||
# torch.cuda.synchronize()
|
||
# st = time.time()
|
||
with torch.no_grad():
|
||
|
||
for img in bar:
|
||
|
||
with autocast():
|
||
|
||
# img_feature = model(img)
|
||
# img_feature = model(img.to(train_config.device).half())
|
||
img_feature = model(img.to(train_config["device"]).half())
|
||
# img_feature = model(img.to(train_config["device"]))
|
||
|
||
# save features in fp32 for sim calculation
|
||
if LPN:
|
||
img_feature_tensor = torch.stack(img_feature, dim=2).reshape(img_feature[0].shape[0], -1)
|
||
img_features_list.append(img_feature_tensor.to(torch.float32))
|
||
else:
|
||
|
||
img_features_list.append(img_feature.to(torch.float32))
|
||
|
||
|
||
# torch.cuda.synchronize()
|
||
# et = time.time()
|
||
# print('---------------------------------time---------------------------------')
|
||
# print('time cost: ', (et - st)/len(dataloader))
|
||
# keep Features on GPU
|
||
img_features = torch.cat(img_features_list, dim=0)
|
||
|
||
# if train_config.verbose:
|
||
bar.close()
|
||
|
||
return img_features
|
||
|
||
def evaluate_reank(config,
|
||
model,
|
||
query_loader,
|
||
gallery_loader,
|
||
pos_gt,
|
||
ranks=[1, 5, 10],
|
||
name = None,
|
||
cleanup=True):
|
||
# 需要保存下来group中的特征，故重新书写此代码
|
||
|
||
|
||
print("Extract Features:")
|
||
img_features_query = predict_rerank(config, model, query_loader, name, 'query')
|
||
img_features_gallery = predict_rerank(config, model, gallery_loader, name, 'gallery')
|
||
|
||
|
||
gl = img_features_gallery.cpu()
|
||
ql = img_features_query.cpu()
|
||
|
||
# -------------------------init------------------------------------------
|
||
faiss_index = faiss.IndexFlatL2(gl.shape[1])
|
||
# add references
|
||
faiss_index.add(gl)
|
||
|
||
# search for queries in the index
|
||
_, predictions = faiss_index.search(ql, max(ranks))
|
||
|
||
|
||
|
||
correct_at_rank = np.zeros(len(ranks))
|
||
|
||
multi_num = ql.shape[0] / len(pos_gt)
|
||
really_pos_gt = pos_gt * int(multi_num)
|
||
|
||
for q_idx, pred in enumerate(predictions):
|
||
for i, n in enumerate(ranks):
|
||
if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1])):
|
||
correct_at_rank[i] += 1
|
||
|
||
correct_at_rank = correct_at_rank / len(predictions)
|
||
|
||
|
||
return correct_at_rank, predictions,really_pos_gt
|
||
|
||
|
||
def evaluate(config,
|
||
model,
|
||
query_loader,
|
||
gallery_loader,
|
||
pos_gt,
|
||
mode,
|
||
LPN,
|
||
ranks=[1, 5, 10],
|
||
name = None,
|
||
cleanup=True):
|
||
|
||
|
||
print("Extract Features:")
|
||
|
||
if mode == 'group':
|
||
|
||
img_features_query = predict(config, model, query_loader)
|
||
img_features_gallery = predict(config, model, gallery_loader)
|
||
elif mode == 'vanilia':
|
||
|
||
img_features_query = predict_backbone(config, model, query_loader, LPN)
|
||
img_features_gallery = predict_backbone(config, model, gallery_loader, LPN)
|
||
|
||
gl = img_features_gallery.cpu()
|
||
ql = img_features_query.cpu()
|
||
# t-sne
|
||
# import numpy as np
|
||
# from sklearn.manifold import TSNE
|
||
# from sklearn.preprocessing import StandardScaler
|
||
# import matplotlib.pyplot as plt
|
||
|
||
# ql_stand = StandardScaler().fit_transform(ql)
|
||
# num = int(ql_stand.shape[0] / 76)
|
||
# t_sne_save = config.dataset_root_dir + '/' + name + '/'
|
||
# y = list(range(0,10))
|
||
# reap_y = np.array([item for item in y for _ in range(num)])
|
||
|
||
# f_1 = ql_stand[::76, :]
|
||
# f_2 = ql_stand[5::76, :]
|
||
# f_3 = ql_stand[10::76, :]
|
||
# f_4 = ql_stand[15::76, :]
|
||
# f_5 = ql_stand[20::76, :]
|
||
# f_6 = ql_stand[25::76, :]
|
||
# f_7 = ql_stand[30::76, :]
|
||
# f_8 = ql_stand[35::76, :]
|
||
# f_9 = ql_stand[40::76, :]
|
||
# f_10 = ql_stand[45::76, :]
|
||
|
||
# x_stand = np.concatenate((f_1, f_2, f_3, f_4, f_5, f_6, f_7, f_8, f_9,f_10), axis=0)
|
||
|
||
# tsne = TSNE(n_components=2, perplexity=num-1, n_iter=5000, n_jobs=-1)
|
||
# X_tsne = tsne.fit_transform(x_stand)
|
||
# plt.figure(figsize=(8, 8))
|
||
# # 归一化颜色值
|
||
# norm = plt.Normalize(reap_y.min(), reap_y.max())
|
||
# # 选择不同的颜色映射
|
||
# cmap = plt.get_cmap('plasma')
|
||
|
||
# # 转换颜色值到[0, 1]区间内
|
||
# colors = cmap(norm(reap_y))
|
||
# scatter = plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=colors, alpha=0.7)
|
||
# plt.colorbar(scatter)
|
||
|
||
# plt.savefig(t_sne_save + 't_sne_' + 'dinov2'+ '.png')
|
||
|
||
|
||
# -------------------------init------------------------------------------
|
||
faiss_index = faiss.IndexFlatL2(gl.shape[1])
|
||
# add references
|
||
faiss_index.add(gl)
|
||
|
||
# search for queries in the index
|
||
_, predictions = faiss_index.search(ql, max(ranks))
|
||
|
||
|
||
|
||
correct_at_rank = np.zeros(len(ranks))
|
||
|
||
multi_num = ql.shape[0] / len(pos_gt)
|
||
really_pos_gt = pos_gt * int(multi_num)
|
||
|
||
for q_idx, pred in enumerate(predictions):
|
||
for i, n in enumerate(ranks):
|
||
# if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1][:ranks[i]])): # test_40
|
||
if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1])):
|
||
correct_at_rank[i] += 1
|
||
|
||
# 测试是问题，设置一个train小样本，快速迭代
|
||
|
||
|
||
correct_at_rank = correct_at_rank / len(predictions)
|
||
|
||
|
||
return correct_at_rank, predictions,really_pos_gt
|
||
|
||
def evaluate_other(config,
|
||
model,
|
||
query_loader,
|
||
gallery_loader,
|
||
pos_gt,
|
||
ranks=[1, 5, 10],
|
||
name = None,
|
||
cleanup=True,
|
||
LPN=False):
|
||
|
||
|
||
print("Extract Features:")
|
||
# img_features_query = predict(config, model, query_loader)
|
||
# img_features_gallery = predict(config, model, gallery_loader)
|
||
img_features_query = predict_backbone(config, model, query_loader)
|
||
img_features_gallery = predict_backbone(config, model, gallery_loader)
|
||
|
||
gl = img_features_gallery.cpu()
|
||
ql = img_features_query.cpu()
|
||
|
||
# t-sne
|
||
# import numpy as np
|
||
# from sklearn.manifold import TSNE
|
||
# from sklearn.preprocessing import StandardScaler
|
||
# import matplotlib.pyplot as plt
|
||
|
||
# ql_stand = StandardScaler().fit_transform(ql)
|
||
# num = int(ql_stand.shape[0] / 76)
|
||
# t_sne_save = config.dataset_root_dir + '/' + name + '/'
|
||
# y = list(range(0,10))
|
||
# reap_y = np.array([item for item in y for _ in range(num)])
|
||
|
||
# f_1 = ql_stand[::76, :]
|
||
# f_2 = ql_stand[5::76, :]
|
||
# f_3 = ql_stand[10::76, :]
|
||
# f_4 = ql_stand[15::76, :]
|
||
# f_5 = ql_stand[20::76, :]
|
||
# f_6 = ql_stand[25::76, :]
|
||
# f_7 = ql_stand[30::76, :]
|
||
# f_8 = ql_stand[35::76, :]
|
||
# f_9 = ql_stand[40::76, :]
|
||
# f_10 = ql_stand[45::76, :]
|
||
|
||
# x_stand = np.concatenate((f_1, f_2, f_3, f_4, f_5, f_6, f_7, f_8, f_9,f_10), axis=0)
|
||
|
||
# tsne = TSNE(n_components=2, perplexity=num-1, n_iter=5000, n_jobs=-1)
|
||
# X_tsne = tsne.fit_transform(x_stand)
|
||
# plt.figure(figsize=(8, 8))
|
||
# # 归一化颜色值
|
||
# norm = plt.Normalize(reap_y.min(), reap_y.max())
|
||
# # 选择不同的颜色映射
|
||
# cmap = plt.get_cmap('plasma')
|
||
|
||
# # 转换颜色值到[0, 1]区间内
|
||
# colors = cmap(norm(reap_y))
|
||
# scatter = plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=colors, alpha=0.7)
|
||
# plt.colorbar(scatter)
|
||
|
||
# plt.savefig(t_sne_save + 't_sne_' + 'dinov2'+ '.png')
|
||
|
||
|
||
# -------------------------init------------------------------------------
|
||
faiss_index = faiss.IndexFlatL2(gl.shape[1])
|
||
# add references
|
||
faiss_index.add(gl)
|
||
|
||
# search for queries in the index
|
||
_, predictions = faiss_index.search(ql, max(ranks))
|
||
|
||
|
||
|
||
correct_at_rank = np.zeros(len(ranks))
|
||
|
||
|
||
really_pos_gt = pos_gt
|
||
|
||
for q_idx, pred in enumerate(predictions):
|
||
for i, n in enumerate(ranks):
|
||
if np.any(np.in1d(pred[:n], really_pos_gt[q_idx][1])):
|
||
correct_at_rank[i] += 1
|
||
|
||
|
||
|
||
correct_at_rank = correct_at_rank / len(predictions)
|
||
|
||
|
||
return correct_at_rank, predictions,really_pos_gt
|
||
|
||
|
||
|
||
|
||
|
||
|