-
Notifications
You must be signed in to change notification settings - Fork 2
/
dataset.py
53 lines (39 loc) · 1.54 KB
/
dataset.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import os
import pickle
DATASETS = ['roxford5k', 'rparis6k', 'revisitop1m']
def configdataset(dataset, dir_main):
dataset = dataset.lower()
if dataset not in DATASETS:
raise ValueError('Unknown dataset: {}!'.format(dataset))
if dataset == 'roxford5k' or dataset == 'rparis6k':
# loading imlist, qimlist, and gnd, in cfg as a dict
gnd_fname = os.path.join(dir_main, dataset, 'gnd_{}.pkl'.format(dataset))
with open(gnd_fname, 'rb') as f:
cfg = pickle.load(f)
cfg['gnd_fname'] = gnd_fname
cfg['ext'] = '.jpg'
cfg['qext'] = '.jpg'
elif dataset == 'revisitop1m':
# loading imlist from a .txt file
cfg = {}
cfg['imlist_fname'] = os.path.join(dir_main, dataset, '{}.txt'.format(dataset))
cfg['imlist'] = read_imlist(cfg['imlist_fname'])
cfg['qimlist'] = []
cfg['ext'] = ''
cfg['qext'] = ''
cfg['dir_data'] = os.path.join(dir_main, dataset)
cfg['dir_images'] = os.path.join(cfg['dir_data'], 'jpg')
cfg['n'] = len(cfg['imlist'])
cfg['nq'] = len(cfg['qimlist'])
cfg['im_fname'] = config_imname
cfg['qim_fname'] = config_qimname
cfg['dataset'] = dataset
return cfg
def config_imname(cfg, i):
return os.path.join(cfg['dir_images'], cfg['imlist'][i] + cfg['ext'])
def config_qimname(cfg, i):
return os.path.join(cfg['dir_images'], cfg['qimlist'][i] + cfg['qext'])
def read_imlist(imlist_fn):
with open(imlist_fn, 'r') as file:
imlist = file.read().splitlines()
return imlist