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 AerialDatasetEvalGroup, AerialDatasetEvalVanilia from models import model import glob import matplotlib.pyplot as plt import matplotlib.patches as patches import numpy as np import argparse from tqdm import tqdm 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_root_dir', type=str, default='/media/guan/新加卷/EdgeBing/TestData/test_40_midref_rot0/', 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') parser.add_argument('--save_dir_path', type=str, default='/media/Shen/Data/RingoData/DenseUAV/test/db.txt', help='Root directory of the dataset') # 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' 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() # 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_root_dir": args.dataset_root_dir, "checkpoint_path": args.checkpoint_path, "save_dir_path":args.save_dir_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 config #-------------------------------------------------------------------------------------------# # Test Config #-------------------------------------------------------------------------------------------# config = parse_config() if not os.path.exists(config['save_dir_path']): os.mkdir(config['save_dir_path']) # test angle angle_list = list(range(0, 1)) 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']) else: model = model.GrounpDinoGlobal(config['group']['group_arch'], config['agg']['agg_arch'], config['agg']['agg_config']) for angle in tqdm(angle_list): if config["mode"] == "vanilia": eva_dataset_query = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'], mode='query', angle=angle, 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 = AerialDatasetEvalVanilia(data_dir=config['dataset_root_dir'], 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) else: eva_dataset_query = AerialDatasetEvalGroup(data_dir=config["dataset_root_dir"], mode='query', angle=angle, 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 = AerialDatasetEvalGroup(data_dir=config["dataset_root_dir"], 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) # model = model.GrounpGlobal(config.group['group_arch'], # config.agg['agg_arch'], # config.agg['agg_config']) model_state_dict = torch.load(config['checkpoint_path'], map_location='cuda:0') model.load_state_dict(model_state_dict, strict=False) model = model.to(config['device']) # pos_gt = eval_dataloader_db.dataset.get_gt_npy() 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 save_result_txt = config['save_dir_path'] + '/' + str(angle) + '.txt' with open(save_result_txt, 'w') as f_w: info = 'top 1: '+ str(round(result[0]*100,2)) + ' top 5: ' +str(round(result[1]*100,2)) + ' top 10: ' + str(round(result[2]*100,2)) f_w.write(info + '\n') f_w.close() # 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)