def get_data(args): root = osp.join(args.data_dir, args.dataset) dataset = datasets.create(args.dataset, root, scale=args.scale) test_transformer_db = get_transformer_test(args.height, args.width) test_transformer_q = get_transformer_test(args.height, args.width, tokyo=(args.dataset=='tokyo')) pitts = datasets.create('pitts', osp.join(args.data_dir, 'pitts'), scale='30k', verbose=False) pitts_train = sorted(list(set(pitts.q_train) | set(pitts.db_train))) train_extract_loader = DataLoader( Preprocessor(pitts_train, root=pitts.images_dir, transform=test_transformer_db), batch_size=args.test_batch_size, num_workers=args.workers, sampler=DistributedSliceSampler(pitts_train), shuffle=False, pin_memory=True) test_loader_q = DataLoader( Preprocessor(dataset.q_test, root=dataset.images_dir, transform=test_transformer_q), batch_size=(1 if args.dataset=='tokyo' else args.test_batch_size), num_workers=args.workers, sampler=DistributedSliceSampler(dataset.q_test), shuffle=False, pin_memory=True) test_loader_db = DataLoader( Preprocessor(dataset.db_test, root=dataset.images_dir, transform=test_transformer_db), batch_size=args.test_batch_size, num_workers=args.workers, sampler=DistributedSliceSampler(dataset.db_test), shuffle=False, pin_memory=True) return dataset, pitts_train, train_extract_loader, test_loader_q, test_loader_db
def get_data(args, iters): root = osp.join(args.data_dir, args.dataset) dataset = datasets.create(args.dataset, root, scale=args.scale) train_transformer = get_transformer_train(args.height, args.width) test_transformer = get_transformer_test(args.height, args.width) sampler = DistributedRandomTupleSampler(dataset.q_train, dataset.db_train, dataset.train_pos, dataset.train_neg, neg_num=args.neg_num, neg_pool=args.neg_pool) train_loader = IterLoader(DataLoader(Preprocessor( dataset.q_train + dataset.db_train, root=dataset.images_dir, transform=train_transformer), batch_size=args.tuple_size, num_workers=args.workers, sampler=sampler, shuffle=False, pin_memory=True, drop_last=True), length=iters) train_extract_loader = DataLoader( Preprocessor(sorted(list(set(dataset.q_train) | set(dataset.db_train))), root=dataset.images_dir, transform=test_transformer), batch_size=args.test_batch_size, num_workers=args.workers, sampler=DistributedSliceSampler( sorted(list(set(dataset.q_train) | set(dataset.db_train)))), shuffle=False, pin_memory=True) val_loader = DataLoader( Preprocessor(sorted(list(set(dataset.q_val) | set(dataset.db_val))), root=dataset.images_dir, transform=test_transformer), batch_size=args.test_batch_size, num_workers=args.workers, sampler=DistributedSliceSampler( sorted(list(set(dataset.q_val) | set(dataset.db_val)))), shuffle=False, pin_memory=True) test_loader = DataLoader( Preprocessor(sorted(list(set(dataset.q_test) | set(dataset.db_test))), root=dataset.images_dir, transform=test_transformer), batch_size=args.test_batch_size, num_workers=args.workers, sampler=DistributedSliceSampler( sorted(list(set(dataset.q_test) | set(dataset.db_test)))), shuffle=False, pin_memory=True) return dataset, train_loader, val_loader, test_loader, sampler, train_extract_loader
def get_data(args, nIm): root = osp.join(args.data_dir, args.dataset) dataset = datasets.create(args.dataset, root, scale='30k') cluster_set = list(set(dataset.q_train) | set(dataset.db_train)) transformer = get_transformer_test(args.height, args.width) sampler = SubsetRandomSampler( np.random.choice(len(cluster_set), nIm, replace=False)) cluster_loader = DataLoader(Preprocessor(cluster_set, root=dataset.images_dir, transform=transformer), batch_size=args.batch_size, num_workers=args.workers, sampler=sampler, shuffle=False, pin_memory=True) return dataset, cluster_loader
def main_worker(args): init_dist(args.launcher, args) synchronize() print("Use GPU: {} for inference, rank no.{} of world_size {}".format( args.gpu, args.rank, args.world_size)) if (args.rank == 0): print("==========\nArgs:{}\n==========".format(args)) if dist.get_rank() == 0: print("inference on '%s'" % args.img_path) if os.path.exists(result_dir): os.system("rm -rf %s" % result_dir) os.mkdir(result_dir) # print("Loading model on GPU-%d" % args.gpu) model = vgg16_netvlad(args, pretrained=True) # read image img = Image.open(args.img_path).convert('RGB') transformer = transforms.Compose([ transforms.Resize((480, 640)), transforms.ToTensor(), transforms.Normalize( mean=[0.48501960784313836, 0.4579568627450961, 0.4076039215686255], std=[ 0.00392156862745098, 0.00392156862745098, 0.00392156862745098 ]) ]) img = transformer(img) img = img.cuda(args.gpu) # use GPU (optional) # model = model.cuda() # img = img.cuda() # extract descriptor (4096-dim) with torch.no_grad(): outputs = model(img.unsqueeze(0)) descriptor = outputs[0].cpu() if (isinstance(outputs, list) or isinstance(outputs, tuple)): x_pool, x_vlad = outputs outputs = F.normalize(x_vlad, p=2, dim=-1) else: outputs = F.normalize(outputs, p=2, dim=-1) pca_parameters_path = osp.join( osp.dirname(args.resume), 'pca_params_' + osp.basename(args.resume).split('.')[0] + '.h5') pca = PCA(4096, True, pca_parameters_path) pca.load(gpu=args.gpu) if dist.get_rank() == 0: outputs = pca.infer(outputs) outputs = outputs.data.cpu() features = OrderedDict() features[args.img_path] = outputs[0] root = osp.join(args.data_dir, args.dataset) dataset = datasets.create(args.dataset, root, scale="30k") query = dataset.q_test gallery = dataset.db_test # distmat, _, _ = pairwise_distance(features, query, gallery) test_transformer_db = get_transformer_test(480, 640) test_loader_db = DataLoader(Preprocessor(dataset.db_test, root=dataset.images_dir, transform=test_transformer_db), batch_size=args.test_batch_size, num_workers=args.workers, sampler=DistributedSliceSampler( dataset.db_test), shuffle=False, pin_memory=True) features_db = extract_features(model, test_loader_db, gallery, vlad=True, pca=pca, gpu=args.gpu, sync_gather=args.sync_gather) synchronize() if (dist.get_rank() == 0): x = torch.cat([features[args.img_path].unsqueeze(0)], 0) y = torch.cat([features_db[f].unsqueeze(0) for f, _, _, _ in gallery], 0) m, n = x.size(0), y.size(0) x = x.view(m, -1) y = y.view(n, -1) distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n) + \ torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n, m).t() distmat.addmm_(1, -2, x, y.t()) sort_idx = np.argsort(distmat, axis=1) del distmat db_ids = [db[1] for db in gallery] for qIx, pred in enumerate(sort_idx): pred = spatial_nms(pred.tolist(), db_ids, max(recall_topk) * 12) # # comput recall # query_id = # correct_at_n = np.zeros(len(recall_topk)) # gt = dataset.test_pos # for i, n in enumerate(recall_topk): # # if in top N then also in top NN, where NN > N # if np.any(np.in1d(pred[:n], gt[query_id])): # correct_at_n[i:] += 1 # break # recalls = correct_at_n / len(gt) # print('Recall Scores:') # for i, k in enumerate(recall_topk): # print(' top-{:<4}{:12.1%}'.format(k, recalls[i])) # save images for i, n in enumerate(recall_topk): # if in top N then also in top NN, where NN > N result = np.array(gallery)[pred[:n]][:, 0] # print("top-%d: " % n, result) img_save_dir = osp.join(result_dir, "top-%d" % n) if not osp.exists(img_save_dir): os.mkdir(img_save_dir) for result_img_path in result: os.system("cp %s %s" % (osp.join( dataset.images_dir, result_img_path), img_save_dir)) print("top-%d image%s saved under '%s'" % (n, "s are" if n > 1 else " is", img_save_dir)) descriptor = descriptor.numpy() pd.DataFrame(descriptor).to_csv( os.path.join(result_dir, "descriptor.csv")) print("Saved features on CSV") synchronize()