Pikaliov
4ff36ce188
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
287 lines
13 KiB
Python
287 lines
13 KiB
Python
from torch.utils.data import DataLoader
|
|
from dataclasses import dataclass,field
|
|
from eval import eval
|
|
import os
|
|
import torch
|
|
from torchvision import transforms as T
|
|
from dataset.World import DenseUAVDatasetEvalVanilia,DenseUAVDatasetEvalGroup
|
|
from models import model
|
|
import glob
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as patches
|
|
import numpy as np
|
|
import argparse
|
|
|
|
def get_parser():
|
|
parser = argparse.ArgumentParser(description="Configuration for training the model")
|
|
|
|
# Model Configurations
|
|
parser.add_argument('--mode', type=str, default='group', help='Model architecture')
|
|
parser.add_argument('--model_path', type=str, default='./world', help='Path to save model checkpoints')
|
|
|
|
# Group Config
|
|
parser.add_argument('--group_arch', type=str, default='groupdinonet', help='Group architecture')
|
|
parser.add_argument('--group_config', type=str, default='none', help='Group configuration')
|
|
|
|
# Backbone Config
|
|
parser.add_argument('--backbone_arch', type=str, default='dinov2_vits14', help='Backbone architecture')
|
|
parser.add_argument('--pretrain_flag', type=bool, default=True, help='Flag to use pre-trained weights')
|
|
|
|
# Agg Config
|
|
parser.add_argument('--agg_arch', type=str, default='multiconvap', help='Aggregation architecture')
|
|
parser.add_argument('--agg_in_channels', type=int, default=384, help='Input channels for aggregation')
|
|
parser.add_argument('--agg_out_channels', type=int, default=384, help='Output channels for aggregation')
|
|
parser.add_argument('--agg_s1', type=int, default=1, help='Aggregation s1 parameter')
|
|
parser.add_argument('--agg_s2', type=int, default=1, help='Aggregation s2 parameter')
|
|
parser.add_argument('--agg_LPN', type=bool, default=False, help='Use LPN for aggregation')
|
|
|
|
# Dataset Paths
|
|
parser.add_argument('--dataset_query', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/query.txt', help='Root directory of the dataset')
|
|
parser.add_argument('--dataset_db', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
|
|
parser.add_argument('--dataset_gt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/gt.txt', help='Root directory of the dataset')
|
|
|
|
#'/media/Shen/Data/RingoData/WorldLoc/TestData/vpair test_40_midref_rot0'
|
|
# Checkpoint Config
|
|
parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to start from a checkpoint')
|
|
|
|
# Training Parameters
|
|
parser.add_argument('--num_workers', type=int, default=0 if os.name == 'nt' else 4, help='Number of workers for data loading')
|
|
parser.add_argument('--device', type=str, default='cuda:0' if torch.cuda.is_available() else 'cpu', help='Device for training')
|
|
parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Use cudnn benchmark for performance')
|
|
parser.add_argument('--cudnn_deterministic', type=bool, default=False, help='Make cudnn deterministic')
|
|
|
|
# Training Settings
|
|
parser.add_argument('--mixed_precision', type=bool, default=True, help='Use mixed precision training')
|
|
parser.add_argument('--custom_sampling', type=bool, default=True, help='Use custom sampling')
|
|
parser.add_argument('--seed', type=int, default=1, help='Random seed')
|
|
parser.add_argument('--epochs', type=int, default=30, help='Number of epochs to train')
|
|
parser.add_argument('--batch_size', type=int, default=5, help='Batch size')
|
|
parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
|
|
parser.add_argument('--gpu_ids', type=tuple, default=(0,), help='GPU IDs for training')
|
|
|
|
# Optimizer Config
|
|
parser.add_argument('--clip_grad', type=float, default=100.0, help='Clip gradients (None or float)')
|
|
parser.add_argument('--decay_exclude_bias', type=bool, default=False, help='Exclude bias from decay')
|
|
parser.add_argument('--grad_checkpointing', type=bool, default=False, help='Use gradient checkpointing')
|
|
|
|
# Loss Config
|
|
parser.add_argument('--label_smoothing', type=float, default=0.1, help='Label smoothing factor')
|
|
|
|
# Learning Rate
|
|
parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
|
|
parser.add_argument('--scheduler', type=str, default='cosine', help='Learning rate scheduler')
|
|
parser.add_argument('--warmup_epochs', type=float, default=0.1, help='Warmup epochs for learning rate')
|
|
parser.add_argument('--lr_end', type=float, default=0.0001, help='End learning rate for polynomial scheduler')
|
|
|
|
return parser
|
|
|
|
def parse_config():
|
|
parser = get_parser()
|
|
args = parser.parse_args()
|
|
|
|
# Build the config dictionaries dynamically based on parsed args
|
|
group_config = {
|
|
"group_arch": args.group_arch,
|
|
"group_config": {args.group_config}
|
|
}
|
|
|
|
backbone_config = {
|
|
"backbone_arch": args.backbone_arch,
|
|
"pretrain_flag": args.pretrain_flag
|
|
}
|
|
|
|
agg_config = {
|
|
"agg_arch": args.agg_arch,
|
|
"agg_config": {
|
|
"in_channels": args.agg_in_channels,
|
|
"out_channels": args.agg_out_channels,
|
|
"s1": args.agg_s1,
|
|
"s2": args.agg_s2,
|
|
"LPN": args.agg_LPN
|
|
}
|
|
}
|
|
|
|
config = {
|
|
"mode": args.mode,
|
|
"model_path": args.model_path,
|
|
"group": group_config,
|
|
"backbone": backbone_config,
|
|
"agg": agg_config,
|
|
"dataset_query": args.dataset_query,
|
|
"dataset_db": args.dataset_db,
|
|
"dataset_gt": args.dataset_gt,
|
|
"checkpoint_path": args.checkpoint_path,
|
|
"num_workers": args.num_workers,
|
|
"device": args.device,
|
|
"cudnn_benchmark": args.cudnn_benchmark,
|
|
"cudnn_deterministic": args.cudnn_deterministic,
|
|
"mixed_precision": args.mixed_precision,
|
|
"custom_sampling": args.custom_sampling,
|
|
"seed": args.seed,
|
|
"epochs": args.epochs,
|
|
"batch_size": args.batch_size,
|
|
"verbose": args.verbose,
|
|
"gpu_ids": args.gpu_ids,
|
|
"clip_grad": args.clip_grad,
|
|
"decay_exclude_bias": args.decay_exclude_bias,
|
|
"grad_checkpointing": args.grad_checkpointing,
|
|
"label_smoothing": args.label_smoothing,
|
|
"lr": args.lr,
|
|
"scheduler": args.scheduler,
|
|
"warmup_epochs": args.warmup_epochs,
|
|
"lr_end": args.lr_end
|
|
}
|
|
|
|
return args, config
|
|
|
|
|
|
#-------------------------------------------------------------------------------------------#
|
|
# Train Config
|
|
#-------------------------------------------------------------------------------------------#
|
|
args, config = parse_config()
|
|
|
|
|
|
IMAGENET_MEAN_STD = {'mean': [0.485, 0.456, 0.406],
|
|
'std': [0.229, 0.224, 0.225]}
|
|
eval_transform = T.Compose([
|
|
T.Resize((224, 224), interpolation=T.InterpolationMode.BILINEAR),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
|
|
])
|
|
|
|
if config["mode"] == "vanilia":
|
|
model = model.BackboneGlobal(config['backbone']['backbone_arch'],
|
|
config['backbone']['pretrain_flag'],
|
|
config['agg']['agg_arch'],
|
|
config['agg']['agg_config'])
|
|
|
|
eva_dataset_query = DenseUAVDatasetEvalVanilia(txt=config['dataset_query'],
|
|
mode='query',
|
|
gt_txt=config["dataset_gt"],
|
|
transforms=eval_transform)
|
|
|
|
eval_dataloader_query = DataLoader(eva_dataset_query,
|
|
batch_size=config['batch_size'],
|
|
num_workers=config['num_workers'],
|
|
shuffle=not config['custom_sampling'],
|
|
pin_memory=True)
|
|
|
|
eva_dataset_db = DenseUAVDatasetEvalVanilia(txt=config['dataset_db'],
|
|
mode='DB',
|
|
gt_txt=config["dataset_gt"],
|
|
transforms=eval_transform)
|
|
|
|
eval_dataloader_db = DataLoader(eva_dataset_db,
|
|
batch_size=config['batch_size'],
|
|
num_workers=config['num_workers'],
|
|
shuffle=not config['custom_sampling'],
|
|
pin_memory=True)
|
|
else:
|
|
model = model.GrounpDinoGlobal(config['group']['group_arch'],
|
|
config['agg']['agg_arch'],
|
|
config['agg']['agg_config'])
|
|
eva_dataset_query = DenseUAVDatasetEvalGroup(txt=config["dataset_query"],
|
|
mode='query',
|
|
gt_txt=config['dataset_gt'],
|
|
transforms=eval_transform)
|
|
|
|
eval_dataloader_query = DataLoader(eva_dataset_query,
|
|
batch_size=config['batch_size'],
|
|
num_workers=config['num_workers'],
|
|
shuffle=not config['custom_sampling'],
|
|
pin_memory=True)
|
|
|
|
eva_dataset_db = DenseUAVDatasetEvalGroup(txt=config["dataset_db"],
|
|
mode='DB',
|
|
gt_txt=config['dataset_gt'],
|
|
transforms=eval_transform)
|
|
|
|
eval_dataloader_db = DataLoader(eva_dataset_db,
|
|
batch_size=config['batch_size'],
|
|
num_workers=config['num_workers'],
|
|
shuffle=not config['custom_sampling'],
|
|
pin_memory=True)
|
|
|
|
|
|
|
|
model_state_dict = torch.load(config['checkpoint_path'], map_location=config['device'])
|
|
model.load_state_dict(model_state_dict, strict=False)
|
|
|
|
model = model.to(config['device'])
|
|
|
|
|
|
pos_gt = eval_dataloader_db.dataset.get_gt()
|
|
result, predictions, really_pos_gt = eval.evaluate(config, model, eval_dataloader_query, eval_dataloader_db, pos_gt, mode=config["mode"],LPN=config['agg']['agg_config']['LPN'])
|
|
print(config['checkpoint_path'])
|
|
print('top 1: ', round(result[0]*100,2), 'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
|
|
|
|
|
|
|
|
# vis and save retrieval results
|
|
# save_vis_dir = config.dataset_root_dir + '/' + 'vis' + '/'
|
|
# if not os.path.exists(save_vis_dir):
|
|
# os.makedirs(save_vis_dir)
|
|
|
|
# temp_path = os.path.join(config.dataset_root_dir, 'reference_images')
|
|
# DB_path = sorted(glob.glob(f'{temp_path}/{"*.png"}'))
|
|
|
|
# # save top 1 flase or wrong
|
|
# with open(config.save_pred_txt, 'w') as f:
|
|
# for i in range(predictions.shape[0]):
|
|
# query_path = eval_dataloader_query.dataset.getitem(i)
|
|
# if np.any(np.in1d(predictions[i,0], really_pos_gt[i][1])):
|
|
# num = 1
|
|
# else:
|
|
# num = 0
|
|
# pred_path = DB_path[predictions[i,0]]
|
|
# info = query_path + ' ' + pred_path + ' ' + str(num) + '\n'
|
|
# f.write(info)
|
|
|
|
|
|
# for i in range(predictions.shape[0]):
|
|
# query_path = eval_dataloader_query.dataset.getitem(i)
|
|
# fig, axs = plt.subplots(2, 6, figsize=(15, 5))
|
|
# query_img = plt.imread(query_path)
|
|
# for j in range(2):
|
|
# for k in range(6):
|
|
# if j == 0 and k == 0:
|
|
# axs[j, k].imshow(query_img)
|
|
# axs[j, k].axis('off') # 不显示坐标轴
|
|
# elif j==0 and k != 0:
|
|
# if np.any(np.in1d(predictions[i,k], really_pos_gt[i][1] )):
|
|
|
|
# db_img_path = DB_path[predictions[i,k]]
|
|
# db_img = plt.imread(db_img_path)
|
|
# axs[j, k].imshow(db_img)
|
|
# # 创建一个矩形框
|
|
# rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='blue', facecolor='none')
|
|
|
|
# # 将矩形框添加到图像上，根据图像尺寸调整框的大小
|
|
# rect.set_transform(axs[j, k].transData) # 将框的坐标系设置为数据坐标系
|
|
# axs[j, k].add_patch(rect)
|
|
# axs[j,k].axis('off') # 不显示坐标轴
|
|
# else:
|
|
# db_img_path = DB_path[predictions[i,k]]
|
|
# db_img = plt.imread(db_img_path)
|
|
# axs[j, k].imshow(db_img)
|
|
# # 创建一个矩形框
|
|
# rect = patches.Rectangle((10, 10), 2, 2, linewidth=10, edgecolor='red', facecolor='none')
|
|
|
|
# # 将矩形框添加到图像上，根据图像尺寸调整框的大小
|
|
# rect.set_transform(axs[j, k].transData) # 将框的坐标系设置为数据坐标系
|
|
# axs[j, k].add_patch(rect)
|
|
# axs[j, k].axis('off') # 不显示坐标轴
|
|
# if j ==1:
|
|
# try:
|
|
# db_img_path = DB_path[really_pos_gt[i][1][k]]
|
|
# db_img = plt.imread(db_img_path)
|
|
# axs[j, k].imshow(db_img)
|
|
# axs[j, k].axis('off') # 不显示坐标轴
|
|
# except:
|
|
# break
|
|
|
|
# save_one_path = save_vis_dir + str(i) + '.png'
|
|
# plt.savefig(save_one_path, dpi=300)
|
|
|
|
|