Initial import: World-UAV prepro
Add dataloaders (v1/v2), analysis scripts, and documentation for working with UAV-GeoLoc (World-UAV). Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
209
GeoLoc-UAV-main/eval_simidataset_parser_game4loc.py
Normal file
209
GeoLoc-UAV-main/eval_simidataset_parser_game4loc.py
Normal file
@@ -0,0 +1,209 @@
|
||||
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 WorldDatasetEvalVanilia, WorldDatasetEvalGroup
|
||||
from models import model
|
||||
import glob
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
import numpy as np
|
||||
import argparse
|
||||
|
||||
from models.game4loc import DesModel
|
||||
|
||||
def get_parser():
|
||||
parser = argparse.ArgumentParser(description="Configuration for training the model")
|
||||
|
||||
# Model Configurations
|
||||
parser.add_argument('--mode', type=str, default='vit_base_patch16_rope_reg1_gap_256.sbb_in1k', help='Model architecture')
|
||||
parser.add_argument('--model', type=str, default='vanilia', help='Path to save model checkpoints')
|
||||
|
||||
|
||||
# Dataset Paths
|
||||
parser.add_argument('--dataset_root', type=str, default='/media/guan/新加卷/EdgeBing/WorldLoc/', help='Root directory of the dataset')
|
||||
parser.add_argument('--test_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset')
|
||||
parser.add_argument('--save_txt', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.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="/media/guan/新加卷/Code/Code/vit_base_eva_gta_same_area.pth", 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=1, help='Batch size')
|
||||
parser.add_argument('--verbose', type=bool, default=True, help='Verbose output during training')
|
||||
parser.add_argument('--gpu_ids', type=tuple, default=(1,), 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()
|
||||
|
||||
|
||||
config = {
|
||||
"mode": args.mode,
|
||||
"model": args.model,
|
||||
"dataset_root_dir": args.dataset_root,
|
||||
"test_index_txt": args.test_txt,
|
||||
"save_txt":args.save_txt,
|
||||
"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((384, 384), interpolation=T.InterpolationMode.BILINEAR),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=IMAGENET_MEAN_STD["mean"], std=IMAGENET_MEAN_STD["std"]),
|
||||
])
|
||||
|
||||
model = DesModel(model_name='vit_base_patch16_rope_reg1_gap_256.sbb_in1k',
|
||||
pretrained=True,
|
||||
img_size=384,
|
||||
share_weights=True)
|
||||
|
||||
if config["checkpoint_path"] is not None:
|
||||
print("Start from:", config["checkpoint_path"])
|
||||
model_state_dict = torch.load(config["checkpoint_path"])
|
||||
model.load_state_dict(model_state_dict, strict=True)
|
||||
|
||||
# Data parallel
|
||||
print("GPUs available:", torch.cuda.device_count())
|
||||
# if torch.cuda.device_count() > 1 and len(config.gpu_ids) > 1:
|
||||
# model = torch.nn.DataParallel(model, device_ids=config.gpu_ids)
|
||||
|
||||
# Model to device
|
||||
model = model.to(config["device"])
|
||||
|
||||
#------------------------------------------------------------Eval---------------------------------------------------------------------#
|
||||
result_list_recall = []
|
||||
result_list_precision = []
|
||||
with open(config['save_txt'], 'w') as f_w:
|
||||
with open(config["test_index_txt"],"r") as val_test:
|
||||
for line in val_test:
|
||||
if config["model"] == 'vanilia':
|
||||
eva_dataset_query = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
|
||||
name=line.strip('\n'),
|
||||
mode='query',
|
||||
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 = WorldDatasetEvalVanilia(data_dir=config["dataset_root_dir"],
|
||||
name=line.strip('\n'),
|
||||
mode='DB',
|
||||
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)
|
||||
|
||||
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["model"],LPN=False)
|
||||
|
||||
print('top 1: ', round(result[0]*100,2), 'top 5: ', round(result[1]*100,2), 'top 10: ', round(result[2]*100,2)) #vanilia
|
||||
f_w.write(line + ' ' + str(round(result[0]*100,2)) + ' ' + str(round(result[1]*100,2)) + '\n')
|
||||
|
||||
|
||||
# ap@5
|
||||
ap_list = []
|
||||
for i in range(predictions.shape[0]):
|
||||
ex = np.isin(predictions[i, 5:], really_pos_gt[i][1])
|
||||
num_all = np.sum(ex) / 5 * 100
|
||||
ap_list.append(num_all)
|
||||
average_ap = np.mean(np.array(ap_list))
|
||||
|
||||
result_list_recall.append(result)
|
||||
result_list_precision.append(average_ap)
|
||||
|
||||
|
||||
result_array = np.array(result_list_recall)
|
||||
average_result = np.mean(result_array, axis=0)
|
||||
print('Average', 'top 1: ', round(average_result[0]*100,2), 'top 5: ', round(average_result[1]*100,2), 'top 10: ', round(average_result[2]*100,2))
|
||||
|
||||
|
||||
result_precision = np.array(result_list_precision)
|
||||
av_p = np.mean(result_precision)
|
||||
print('AP@5 is', round(av_p,2))
|
||||
|
||||
info1 = 'Average' + 'top 1: ' + str(round(average_result[0]*100,2)) + 'top 5: ' + str(round(average_result[1]*100,2)) + 'top 10' + str(round(average_result[2]*100,2)) + '\n'
|
||||
f_w.write(info1)
|
||||
f_w.write('AP@5 is'+str(round(av_p,2)))
|
||||
|
||||
|
||||
|
||||
# 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 = eval_dataloader_db.dataset.samples[predictions[i,0]]
|
||||
info = query_path + ' ' + pred_path + ' ' + str(num) + '\n'
|
||||
f_w.write(info)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user