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)