def __init__(self, architecture='resnet101', pretrained_type='gl18', soa_layers='45', mode='train'): super(ResNetSOAs, self).__init__() base_model = vars(models)[architecture](pretrained=True) last_feat_in = base_model.inplanes base_model = nn.Sequential(*list(base_model.children())[:-2]) if pretrained_type=='caffenet' and architecture in FEATURES: print(">> {}: for '{}' custom pretrained features '{}' are used" .format(os.path.basename(__file__), architecture, os.path.basename(FEATURES[architecture]))) model_dir = os.path.join(get_data_root(), 'networks') base_model.load_state_dict(model_zoo.load_url(FEATURES[architecture], model_dir=model_dir)) elif pretrained_type in ['SfM120k', 'gl18'] and architecture in FEATURES: pretrained_name = pretrained_type + '-tl-' + architecture + '-gem-w' print(">> {}: for '{}' custom pretrained features '{}' are used" .format(os.path.basename(__file__), architecture, os.path.basename(PRETRAINED[pretrained_name]))) model_dir = os.path.join(get_data_root(), 'networks') base_model.load_state_dict(extract_features_from_e2e(model_zoo.load_url(PRETRAINED[pretrained_name], model_dir=model_dir))) res_blocks = list(base_model.children()) self.conv1 = nn.Sequential(*res_blocks[0:2]) self.conv2_x = nn.Sequential(*res_blocks[2:5]) self.conv3_x = res_blocks[5] self.conv4_x = res_blocks[6] self.conv5_x = res_blocks[7] self.soa_layers = soa_layers if '4' in self.soa_layers: print("SOA_4:") self.soa4 = SOABlock(in_ch=last_feat_in // 2, k=4) if '5' in self.soa_layers: print("SOA_5:") self.soa5 = SOABlock(in_ch=last_feat_in, k=2)
def main(): args = parser.parse_args() # check if test dataset are downloaded # and download if they are not download_distractors(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network net = load_network(network_name=args.network) net.mode = 'test' print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = list(eval(args.multiscale)) print(">>>> Evaluating scales: {}".format(ms)) # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # evaluate on test datasets dataset = 'revisitop1m' start = time.time() print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])] try: bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] except: bbxs = None # for holidaysmanrot and copydays # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, mode='test') torch.save(vecs, args.network + '_vecs_' + dataset + '.pt') print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
def load_network(network_name='resnet101-solar-best.pth'): # loading network # pretrained networks (downloaded automatically) print(">> Loading network:\n>>>> '{}'".format(network_name)) state = torch.load(os.path.join(get_data_root(), 'networks', network_name)) # parsing net params from meta # architecture, pooling, mean, std required # the rest has default values, in case that is doesnt exist net_params = {} net_params['architecture'] = state['meta']['architecture'] net_params['pooling'] = state['meta']['pooling'] net_params['local_whitening'] = state['meta'].get('local_whitening', False) net_params['regional'] = state['meta'].get('regional', False) net_params['whitening'] = state['meta'].get('whitening', False) net_params['mean'] = state['meta']['mean'] net_params['std'] = state['meta']['std'] net_params['pretrained'] = False net_params['pretrained_type'] = None net_params['soa'] = state['meta']['soa'] net_params['soa_layers'] = state['meta']['soa_layers'] net = init_network(net_params) net.load_state_dict(state['state_dict']) return net
def test(datasets, net, epoch, summary): #ms = [1] ms = [1, 2**(1 / 2), 1 / 2**(1 / 2)] print('>> Evaluating network on test datasets...') # for testing we use image size of max 1024 image_size = 1024 # moving network to gpu and eval mode net.mode = 'test' net.to(device) net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # compute whitening if args.test_whiten: start = time.time() print('>> {}: Learning whitening...'.format(args.test_whiten)) # loading db db_root = os.path.join(get_data_root(), 'train', args.test_whiten) ims_root = os.path.join(db_root, 'ims') db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.test_whiten)) with open(db_fn, 'rb') as f: db = pickle.load(f) images = [ cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids'])) ] # extract whitening vectors print('>> {}: Extracting...'.format(args.test_whiten)) wvecs = extract_vectors(net, None, images, image_size, transform) # implemented with torch.no_grad # learning whitening print('>> {}: Learning...'.format(args.test_whiten)) wvecs = wvecs.numpy() m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs']) Lw = {'m': m, 'P': P} print('>> {}: elapsed time: {}'.format(args.test_whiten, htime(time.time() - start))) else: Lw = None # evaluate on test datasets datasets = args.test_datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])] bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, image_size, transform, summary=summary, mode='test', ms=ms) # implemented with torch.no_grad print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, image_size, transform, bbxs, summary=summary, mode='test', ms=ms) # implemented with torch.no_grad print('>> {}: Evaluating...'.format(dataset)) # convert to numpy vecs = vecs.numpy() qvecs = qvecs.numpy() # search, rank, and print scores = np.dot(vecs.T, qvecs) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset, ranks, cfg['gnd'], summary=summary, epoch=epoch) # for protocol in ['hard']: #'easy', 'medium', 'hard']: # plot_ranks(qimages, images, ranks, cfg['gnd'], bbxs, summary, dataset, epoch, 20, protocol) if Lw is not None: # whiten the vectors vecs_lw = whitenapply(vecs, Lw['m'], Lw['P']) qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P']) # search, rank, and print scores = np.dot(vecs_lw.T, qvecs_lw) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd'], summary=summary, epoch=epoch) print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
def main(): global args, min_loss, device args = parser.parse_args() # manually check if there are unknown test datasets for dataset in args.test_datasets.split(','): if dataset not in test_datasets_names: raise ValueError( 'Unsupported or unknown test dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_test(get_data_root()) # create export dir if it doesnt exist directory = "{}".format(args.training_dataset) directory += "_{}".format(args.arch) directory += "_{}".format(args.pool) + str(args.p) if args.local_whitening: directory += "_lwhiten" if args.regional: directory += "_r" if args.whitening: directory += "_whiten" if not args.pretrained: directory += "_notpretrained" if args.soa: directory += "_soa_" directory += args.soa_layers directory = os.path.join('second_order_attn', directory) if args.unfreeze_last: directory += "_unfreeze_last" if args.sos: directory += "_SOS_lambda{:.2f}".format(args._lambda) directory += "_{}".format(args.pretrained_type) directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin) directory += "_{}_lr{:.1e}_lrd{:.1e}_wd{:.1e}".format( args.optimizer, args.lr, args.lr_decay, args.weight_decay) directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size) directory += "_bsize{}_uevery{}_imsize{}".format(args.batch_size, args.update_every, args.image_size) device = torch.device('cuda:' + args.gpu_id) args.directory = os.path.join(args.directory, directory) print(">> Creating directory if it does not exist:\n>> '{}'".format( args.directory)) if not os.path.exists(args.directory): os.makedirs(args.directory) # set random seeds torch.manual_seed(0) torch.cuda.manual_seed_all(0) np.random.seed(0) random.seed(0) # initialize model if args.pretrained: print(">> Using pre-trained model '{}'".format(args.arch)) else: print(">> Using model from scratch (random weights) '{}'".format( args.arch)) model_params = {} model_params['architecture'] = args.arch model_params['pooling'] = args.pool model_params['p'] = args.p model_params['local_whitening'] = args.local_whitening model_params['regional'] = args.regional model_params['whitening'] = args.whitening # model_params['mean'] = ... # will use default # model_params['std'] = ... # will use default model_params['pretrained'] = args.pretrained model_params['pretrained_type'] = args.pretrained_type model_params['flatten_desc'] = args.flatten_desc model_params['soa'] = args.soa model_params['soa_layers'] = args.soa_layers model = init_network(model_params) # move network to gpu model.to(device) # define loss function (criterion) and optimizer if args.loss == 'contrastive': criterion = ContrastiveLoss(margin=args.loss_margin).to(device) elif args.loss == 'triplet': criterion = TripletLoss(margin=args.loss_margin).to(device) else: raise (RuntimeError("Loss {} not available!".format(args.loss))) if args.sos: criterionB = SOSLoss().to(device) else: criterionB = None # parameters split into features, pool, whitening # IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM parameters = [] # add feature parameters if args.soa: for p in model.features.conv1.parameters(): p.requires_grad = False for p in model.features.conv2_x.parameters(): p.requires_grad = False for p in model.features.conv3_x.parameters(): p.requires_grad = False for p in model.features.conv4_x.parameters(): p.requires_grad = False if args.unfreeze_last: parameters.append({ 'params': model.features.conv5_x.parameters(), 'lr': args.lr * 0.0 }) #, 'weight_decay': 0}) if '4' in args.soa_layers: parameters.append({'params': model.features.soa4.parameters() }) #, 'lr': args.lr*10}) #, 'weight_decay': 0}) if '5' in args.soa_layers: parameters.append({'params': model.features.soa5.parameters() }) #, 'lr': args.lr*10}) #, 'weight_decay': 0}) else: parameters.append({'params': model.features.parameters()}) # add local whitening if exists if model.lwhiten is not None: parameters.append({'params': model.lwhiten.parameters()}) # add pooling parameters (or regional whitening which is part of the pooling layer!) if not args.regional: # global, only pooling parameter p weight decay should be 0 if args.pool == 'gem': parameters.append({ 'params': model.pool.parameters(), 'lr': args.lr * 100, 'weight_decay': 0 }) elif args.pool == 'gemmp': parameters.append({ 'params': model.pool.parameters(), 'lr': args.lr * 100, 'weight_decay': 0 }) else: # regional, pooling parameter p weight decay should be 0, # and we want to add regional whitening if it is there if args.pool == 'gem': parameters.append({ 'params': model.pool.rpool.parameters(), 'lr': args.lr * 1, 'weight_decay': 0 }) elif args.pool == 'gemmp': parameters.append({ 'params': model.pool.rpool.parameters(), 'lr': args.lr * 100, 'weight_decay': 0 }) if model.pool.whiten is not None: parameters.append({'params': model.pool.whiten.parameters()}) # add final whitening if exists if model.whiten is not None: parameters.append({'params': model.whiten.parameters()}) # define optimizer if args.optimizer == 'sgd': optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay) # define learning rate decay schedule # TODO: maybe pass as argument in future implementation? exp_decay = math.exp(-args.lr_decay) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay) # optionally resume from a checkpoint start_epoch = 0 checkpoint_exists = False if args.resume: args.resume = os.path.join(args.directory, args.resume) if os.path.isfile(args.resume): # load checkpoint weights and update model and optimizer print(">> Loading checkpoint:\n>> '{}'".format(args.resume)) checkpoint = torch.load(args.resume) start_epoch = checkpoint['epoch'] min_loss = checkpoint['min_loss'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # important not to forget scheduler updating scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch'] - 1) checkpoint_exists = True else: print(">> No checkpoint found at '{}'".format(args.resume)) # Data loading code #if args.training_dataset.startswith('megadepth'): # model.meta['mean'] += sum(mean for mean in model.meta['mean']) / 3 # model.meta['std'] += sum(std for std in model.meta['std']) / 3 normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std']) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) if args.training_dataset.startswith('gl18'): train_dataset = TuplesBatchedDataset(name=args.training_dataset, mode='train', imsize=args.image_size, nnum=args.neg_num, qsize=args.query_size, poolsize=args.pool_size, transform=transform) else: train_dataset = TuplesDataset(name=args.training_dataset, mode='train', imsize=args.image_size, nnum=args.neg_num, qsize=args.query_size, poolsize=args.pool_size, transform=transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, sampler=None, drop_last=True #, collate_fn=collate_tuples ) if args.val: if args.training_dataset.startswith('gl18'): val_dataset = TuplesBatchedDataset(name=args.training_dataset, mode='val', imsize=args.image_size, nnum=args.neg_num, qsize=args.val_query_size, poolsize=args.val_pool_size, transform=transform) else: val_dataset = TuplesDataset(name=args.training_dataset, mode='val', imsize=args.image_size, nnum=args.neg_num, qsize=args.val_query_size, poolsize=args.val_pool_size, transform=transform) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=True) # set up tensorboad summary = tb_setup(args.directory, checkpoint_exists) ################################################### evaluate the network before starting # this might not be necessary? # test(args.test_datasets, model, start_epoch, summary) for epoch in range(start_epoch, args.epochs): # set manual seeds per epoch random.seed(epoch) np.random.seed(epoch) torch.manual_seed(epoch) torch.cuda.manual_seed_all(epoch) loss = train(train_loader, model, criterion, criterionB, optimizer, epoch, summary) # evaluate on test datasets every test_freq epochs with torch.no_grad(): if args.val: loss = validate(val_loader, model, criterion, criterionB, epoch, summary) if (epoch + 1) % args.test_freq == 0: # evaluate on validation set test(args.test_datasets, model, epoch + 1, summary) # remember best loss and save checkpoint is_best = loss < min_loss min_loss = min(loss, min_loss) save_checkpoint( { 'epoch': epoch + 1, 'meta': model.meta, 'state_dict': model.state_dict(), 'min_loss': min_loss, 'optimizer': optimizer.state_dict(), }, is_best, args.directory) # train for one epoch on train set # adjust learning rate for each epoch scheduler.step()
def init_network(params): # parse params with default values architecture = params.get('architecture', 'resnet101') local_whitening = params.get('local_whitening', False) pooling = params.get('pooling', 'gem') p = params.get('p', 3.) regional = params.get('regional', False) whitening = params.get('whitening', False) mean = params.get('mean', [0.485, 0.456, 0.406]) std = params.get('std', [0.229, 0.224, 0.225]) pretrained = params.get('pretrained', True) soa = params.get('soa', False) soa_layers = params.get('soa_layers', '45') pretrained_type = params.get('pretrained_type', 'SfM120k') flatten_desc = params.get('flatten_desc', False) mode = params.get('mode', 'train') # get output dimensionality size dim = OUTPUT_DIM[architecture] # loading network from torchvision if pretrained: if architecture not in FEATURES: # initialize with network pretrained on imagenet in pytorch net_in = getattr(torchvision.models, architecture)(pretrained=True) else: # initialize with random weights, later on we will fill features with custom pretrained network net_in = getattr(torchvision.models, architecture)(pretrained=False) else: # initialize with random weights net_in = getattr(torchvision.models, architecture)(pretrained=False) # initialize features # take only convolutions for features, # always ends with ReLU to make last activations non-negative if architecture.startswith('alexnet'): features = list(net_in.features.children())[:-1] elif architecture.startswith('vgg'): features = list(net_in.features.children())[:-1] elif architecture.startswith('resnet'): #feat0.weight.data = torch.nn.init.kaiming_normal_(feat0.weight.data) features = list(net_in.children())[:-2] elif architecture.startswith('densenet'): features = list(net_in.features.children()) features.append(nn.ReLU(inplace=True)) elif architecture.startswith('squeezenet'): features = list(net_in.features.children()) else: raise ValueError( 'Unsupported or unknown architecture: {}!'.format(architecture)) # initialize local whitening if local_whitening: lwhiten = nn.Linear(dim, dim, bias=True) # TODO: lwhiten with possible dimensionality reduce if pretrained: lw = architecture if lw in L_WHITENING: print( ">> {}: for '{}' custom computed local whitening '{}' is used" .format(os.path.basename(__file__), lw, os.path.basename(L_WHITENING[lw]))) whiten_dir = os.path.join(get_data_root(), 'whiten') lwhiten.load_state_dict( model_zoo.load_url(L_WHITENING[lw], model_dir=whiten_dir)) else: print( ">> {}: for '{}' there is no local whitening computed, random weights are used" .format(os.path.basename(__file__), lw)) else: lwhiten = None # initialize pooling if pooling == 'gemmp': pool = POOLING[pooling](p=p, mp=dim) else: pool = POOLING[pooling](p=p) # initialize regional pooling if regional: rpool = pool rwhiten = nn.Linear(dim, dim, bias=True) # TODO: rwhiten with possible dimensionality reduce if pretrained: rw = '{}-{}-r'.format(architecture, pooling) if rw in R_WHITENING: print( ">> {}: for '{}' custom computed regional whitening '{}' is used" .format(os.path.basename(__file__), rw, os.path.basename(R_WHITENING[rw]))) whiten_dir = os.path.join(get_data_root(), 'whiten') rwhiten.load_state_dict( model_zoo.load_url(R_WHITENING[rw], model_dir=whiten_dir)) else: print( ">> {}: for '{}' there is no regional whitening computed, random weights are used" .format(os.path.basename(__file__), rw)) pool = Rpool(rpool, rwhiten) # initialize whitening if whitening: whiten = nn.Linear(dim, dim, bias=True) # TODO: whiten with possible dimensionality reduce if pretrained: w = architecture if local_whitening: w += '-lw' w += '-' + pooling if regional: w += '-r' if w in WHITENING: print(">> {}: for '{}' custom computed whitening '{}' is used". format(os.path.basename(__file__), w, os.path.basename(WHITENING[w]))) whiten_dir = os.path.join(get_data_root(), 'whiten') whiten.load_state_dict( model_zoo.load_url(WHITENING[w], model_dir=whiten_dir)) else: print( ">> {}: for '{}' there is no whitening computed, random weights are used" .format(os.path.basename(__file__), w)) else: whiten = None # create meta information to be stored in the network meta = { 'architecture': architecture, 'local_whitening': local_whitening, 'pooling': pooling, 'regional': regional, 'whitening': whitening, 'mean': mean, 'std': std, 'outputdim': dim, 'soa': soa, 'soa_layers': soa_layers, } # create a generic image retrieval network net = SOLAR_Global_Retrieval(architecture, features, lwhiten, pool, whiten, meta, pretrained_type=pretrained_type, soa_layers=soa_layers, mode=mode) return net
def __init__(self, name, mode, imsize=None, nnum=5, qsize=2000, poolsize=20000, transform=None, loader=default_loader): if not (mode == 'train' or mode == 'val'): raise(RuntimeError("MODE should be either train or val, passed as string")) if name.startswith('retrieval-SfM'): # setting up paths data_root = get_data_root() db_root = os.path.join(data_root, 'train', name) ims_root = os.path.join(db_root, 'ims') # loading db db_fn = os.path.join(db_root, '{}.pkl'.format(name)) with open(db_fn, 'rb') as f: db = pickle.load(f)[mode] # setting fullpath for images self.images = [cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids']))] elif name.startswith('gl'): ## TODO: NOT IMPLEMENTED YET PROPOERLY (WITH AUTOMATIC DOWNLOAD) # setting up paths data_root = get_data_root() db_root = os.path.join(data_root, 'train', name) self.ims_root = os.path.join('/media/ssd/datasets/gl18', 'train') # loading db db_fn = os.path.join(db_root, 'db_{}.pkl'.format(name)) with open(db_fn, 'rb') as f: db = pickle.load(f)[mode] # setting fullpath for images self.images = [os.path.join(self.ims_root, db['cids'][i]+'.jpg') for i in range(len(db['cids']))] self.pfns = db['pidxs'] else: raise(RuntimeError("Unknown dataset name!")) # initializing tuples dataset self.name = name self.mode = mode self.imsize = imsize self.clusters = db['cluster'] self.qpool = db['qidxs'] self.ppool = db['pidxs'] self.bbxs = db['bbxs'] if mode == 'train': print('__'*50) print('Dataset:', name) print('__'*50) print('__'*50) print('Total number of', mode, 'samples: ', len(self.qpool)) print('__'*50) ## If we want to keep only unique q-p pairs ## However, ordering of pairs will change, although that is not important # qpidxs = list(set([(self.qidxs[i], self.pidxs[i]) for i in range(len(self.qidxs))])) # self.qidxs = [qpidxs[i][0] for i in range(len(qpidxs))] # self.pidxs = [qpidxs[i][1] for i in range(len(qpidxs))] # size of training subset for an epoch self.nnum = nnum self.qsize = min(qsize, len(self.qpool)) self.poolsize = min(poolsize, len(self.images)) self.qidxs = None self.pidxs = None self.nidxs = None self.transform = transform self.loader = loader self.print_freq = 10
def main(): args = parser.parse_args() # check if there are unknown datasets for dataset in args.datasets.split(','): if dataset not in datasets_names: raise ValueError( 'Unsupported or unknown dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network net = load_network(network_name=args.network) net.mode = 'test' print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = list(eval(args.multiscale)) print(">>>> Evaluating scales: {}".format(ms)) # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: summary_ranks = tb_setup( os.path.join('specs/ranks/', dataset, args.network)) summary_embeddings = tb_setup( os.path.join('specs/embeddings/', dataset, args.network)) start = time.time() print('') print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])] try: bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] except: bbxs = None # for holidaysmanrot and copydays # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, mode='test') vecs = vecs.numpy() print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, mode='test') qvecs = qvecs.numpy() print('>> {}: Evaluating...'.format(dataset)) # search, rank, and print scores = np.dot(vecs.T, qvecs) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset, ranks, cfg['gnd']) print('') # plot retrieval rankings and save to tensorboard summary for protocol in ['easy', 'medium', 'hard']: plot_ranks(qimages, images, ranks, cfg['gnd'], bbxs, summary_ranks, dataset, 'solar-best: ', 20, protocol) print('') # plot embeddings for cluster visualisation in tensorboard/projector plot_embeddings(images, vecs, summary_embeddings, imsize=64, sample_freq=1) print('') print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
def main(): args = parser.parse_args() # check if there are unknown datasets for dataset in args.datasets.split(','): if dataset not in datasets_names: raise ValueError( 'Unsupported or unknown dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network net = load_network(network_name=args.network) net.mode = 'test' print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = list(eval(args.multiscale)) print(">>>> Evaluating scales: {}".format(ms)) # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])] try: bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] except: bbxs = None # for holidaysmanrot and copydays # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, mode='test') print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, mode='test') qvecs = qvecs.numpy() print('>> {}: Evaluating...'.format(dataset)) vecs_1m = torch.load(args.network + '_vecs_' + 'revisitop1m' + '.pt') vecs = torch.cat([vecs, vecs_1m], dim=1) vecs = vecs.numpy() # search, rank, and print scores = np.dot(vecs.T, qvecs) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset, ranks, cfg['gnd']) print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))