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()