def learning_lw(net):
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])
test_whiten = "retrieval-SfM-30k"
print(">> {}: Learning whitening...".format(test_whiten))
# loading db
db_root = os.path.join(get_data_root(), "train", test_whiten)
ims_root = os.path.join(db_root, "ims")
db_fn = os.path.join(db_root, "{}-whiten.pkl".format(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, images, 1024, transform)
# learning whitening
print(">> {}: Learning...".format(args.test_whiten))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db["qidxs"], db["pidxs"])
Lw = {"m": m, "P": P}
return Lw
def __init__(self,
name,
imsize=None,
transform=None,
loader=default_loader,
random_size=False):
# setting up paths
mode = 'val'
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]
# initializing tuples dataset
self.name = name
self.imsize = imsize
self.images = [
cid2filename(db['cids'][i], ims_root)
for i in range(len(db['cids']))
]
self.clusters = db['cluster']
self.qpool = db['qidxs']
self.transform = transform
self.loader = loader
self.random_size = random_size
def __init__(self, params):
# prepare config structure for the test dataset
self.image_size = params.pop("image_size")
self.dataset = params.pop("dataset")
self.transforms = initialize_transforms(params.pop("transforms"), params.pop("mean_std"))
if isinstance(self.dataset, dict):
# Tsv dataset files
assert self.dataset.keys() == {"name", "queries", "db", "imgdir"}
imgdir = self.dataset['imgdir']
with initialize_file_reader(self.dataset['db'], keys=["identifier"]) as reader:
data = reader.get()
self.images = [path_join(imgdir, x) for x in data["identifier"]]
mapping = {x: i for i, x in enumerate(data["identifier"])}
with initialize_file_reader(self.dataset['queries'], keys=["query", "bbx", "ok", "junk"]) as reader:
data = reader.get()
self.qimages = [path_join(imgdir, x) for x in data["query"]]
self.bbxs = [tuple(x) if x else None for x in data["bbx"]]
self.gnd = [{'ok': [mapping[x] for x in ok], 'junk': [mapping[x] for x in junk]} \
for ok, junk in zip(data["ok"], data["junk"])]
self.dataset = self.dataset['name']
else:
# Official cirtorch files
cfg = configdataset(self.dataset, os.path.join(get_data_root(), 'test'))
self.images = [cfg['im_fname'](cfg,i) for i in range(cfg['n'])]
self.qimages = [cfg['qim_fname'](cfg,i) for i in range(cfg['nq'])]
self.bbxs = [tuple(cfg['gnd'][i]['bbx']) if cfg['gnd'][i]['bbx'] else None for i in range(cfg['nq'])]
self.gnd = cfg['gnd']
assert not params, params.keys()
def _compute_whitening(whitening, net, image_size, transform, ms, msp):
# compute whitening
start = time.time()
print('>> {}: Learning whitening...'.format(whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train', whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(whitening))
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(whitening))
wvecs = extract_vectors(net, images, image_size, transform, ms=ms, msp=msp)
# learning whitening
print('>> {}: Learning...'.format(whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
elapsed = time.time() - start
print('>> {}: elapsed time: {}'.format(whitening, htime(elapsed)))
return Lw, elapsed
def test(datasets, net):
print(">> Evaluating network on test datasets...")
image_size = 1024
# 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])
# compute whitening
# Lw = None
Lw = net.meta["Lw"]["retrieval-SfM-120k"]["ss"]
# 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"])]
# extract database and query vectors
print(">> {}: database images...".format(dataset))
vecs = extract_vectors(net, images, image_size, transform)
print(">> {}: query images...".format(dataset))
qvecs = extract_vectors(net, qimages, image_size, transform, bbxs)
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"])
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"])
print(">> {}: elapsed time: {}".format(dataset,
htime(time.time() - start)))
def __init__(
self,
imsize=None,
nnum=5,
qsize=2000,
poolsize=20000,
transform=None,
loader=default_loader,
norm=None,
filename=None,
random=True,
):
# setting up paths
data_root = get_data_root()
name = "retrieval-SfM-120k"
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)["val"]
# initializing tuples dataset
self.imsize = imsize
self.images = [
cid2filename(db["cids"][i], ims_root)
for i in range(len(db["cids"]))
]
self.clusters = db["cluster"]
self.qpool = db["qidxs"]
# self.ppool = db['pidxs']
# 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.poolvecs = None
self.transform = transform
self.loader = loader
self.pool_clusters_centers = None
self.clustered_pool = []
self.norm = norm
self.kmeans_ = None
if filename is None:
self.filename = FNAME
else:
self.filename = filename
self.loaded_imgs = []
self.random = random
def __init__(self,
name,
mode,
imsize=None,
nnum=5,
qsize=2000,
poolsize=20000,
transform=None,
loader=loader_hashed):
if not (mode == 'train' or mode == 'val'):
raise RuntimeError(
"MODE should be either train or val, passed as string")
# 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']))
]
# initializing tuples dataset
self.name = name
self.mode = mode
self.imsize = imsize
self.clusters = db['cluster']
self.qpool = db['qidxs']
self.ppool = db['pidxs']
## 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 testUkbench(net, eConfig):
#datasets = eConfig['test-datasets'].split(',')
#results = []
#
#for dataset in datasets:
# results.append((dataset, np.random.rand(1)[0]))
#
#return results
print('>> Evaluating network on test datasets...')
# for testing we use image size of max 1024
dataset = 'ukbench'
image_size = 362
# 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])
dbpath = os.path.join(get_data_root(), 'test', 'ukbench', 'full')
images = [
os.path.join(dbpath, 'ukbench{:05d}.jpg'.format(i))
for i in range(10200)
]
labels = np.arange(10200, dtype=np.int) // 4
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
X = extract_vectors(net, images, image_size, transform)
print('>> {}: Evaluating...'.format(dataset))
# rank the similarities
X = X.numpy()
scores = np.dot(X.T, X)
ranks = np.argsort(-scores, axis=1)
ranks = ranks[:, 0:4]
# compute the average accuracy for the first 4 entries
ranksLabel = labels[ranks]
accs = np.sum(ranksLabel == np.repeat(labels[:, np.newaxis], 4, axis=1),
axis=1)
avgAcc = np.mean(accs)
print('avgAcc: {:.6f}'.format(avgAcc))
return [('ukbench', avgAcc)]
def gen_query_bbx_img():
datasets = ['oxford5k', 'paris6k', 'roxford5k', 'rparis6k']
for dataset in datasets:
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
width = 5
for i in range(len(qimages)):
im = Image.open(qimages[i])
draw = ImageDraw.Draw(im)
(x0, y0, x1, y1) = bbxs[i]
for j in range(width):
draw.rectangle([x0-j, y0-j, x1+j, y1+j], outline='yellow')
im.save('_bbx.jpg'.join(qimages[i].split('.jpg')))
print("{} qurery_bbx generate ok".format(dataset))
def eval_datasets(model,
datasets=('oxford5k', 'paris6k', 'roxford5k', 'rparis6k'),
ms=False,
tta_gem_p=1.0,
logger=None):
model = model.eval()
data_root = get_data_root()
scales = [1 / 2**(1 / 2), 1.0, 2**(1 / 2)] if ms else [1.0]
results = dict()
for dataset in datasets:
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(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'])]
tqdm_desc = cfg['dataset']
db_feats = extract_vectors(model,
images=images,
bbxs=None,
scales=scales,
tta_gem_p=tta_gem_p,
tqdm_desc=tqdm_desc)
query_feats = extract_vectors(model,
images=qimages,
bbxs=bbxs,
scales=scales,
tta_gem_p=tta_gem_p,
tqdm_desc=tqdm_desc)
scores = np.dot(db_feats, query_feats.T)
ranks = np.argsort(-scores, axis=0)
results[dataset] = compute_map_and_print(dataset,
ranks,
cfg['gnd'],
kappas=[1, 5, 10],
logger=logger)
return results
def load_network(network_path):
"""Load network from a path."""
print(">> Loading network:\n>>>> '{}'".format(network_path))
if network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[network_path], model_dir=os.path.join(get_data_root(), 'networks'))
else:
state = torch.load(network_path)
net = init_network(model=state['meta']['architecture'], pooling=state['meta']['pooling'], whitening=state['meta']['whitening'],
mean=state['meta']['mean'], std=state['meta']['std'], pretrained=False)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
return net
def main():
args = parser.parse_args()
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# 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
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
if "epoch" in state:
print("Model after {} epochs".format(state["epoch"]))
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters: test both single scale and multiscale
ms_singlescale = [1]
msp_singlescale = 1
ms_multiscale = list(eval(args.multiscale))
msp_multiscale = 1
if len(ms_multiscale
) > 1 and net.meta['pooling'] == 'gem' and not net.meta[
'regional'] and not net.meta['whitening']:
msp_multiscale = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms_multiscale))
print(">>>> msp: {}".format(msp_multiscale))
# 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])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = net.meta['Lw'][args.whitening]
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(
args.whitening)
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = torch.load(whiten_fn)
else:
Lw = {}
for whiten_type, ms, msp in zip(
["ss", "ms"], [ms_singlescale, ms_multiscale],
[msp_singlescale, msp_multiscale]):
print('>> {0}: Learning whitening {1}...'.format(
args.whitening, whiten_type))
# loading db
db_root = os.path.join(get_data_root(), 'train',
args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(
db_root, '{}-whiten.pkl'.format(args.whitening))
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.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw[whiten_type] = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(
args.whitening, htime(time.time() - start)))
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(
args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
else:
Lw = None
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
for whiten_type, ms, msp in zip(["ss", "ms"],
[ms_singlescale, ms_multiscale],
[msp_singlescale, msp_multiscale]):
print('>> Extracting feature on {0}, whitening {1}'.format(
dataset, whiten_type))
# 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'])]
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms,
msp=msp)
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'])
if Lw is not None:
# whiten the vectors
vecs_lw = whitenapply(vecs, Lw[whiten_type]['m'],
Lw[whiten_type]['P'])
qvecs_lw = whitenapply(qvecs, Lw[whiten_type]['m'],
Lw[whiten_type]['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 {}'.format(whiten_type), ranks,
cfg['gnd'])
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
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')
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)
# get output dimensionality size
dim = OUTPUT_DIM[architecture]
# loading network from torchvision
if architecture not in FEATURES:
if architecture == 'w_vgg16':
net_in = vgg16(invariant='W', pretrained=pretrained)
elif architecture == 'w_resnet101':
net_in = resnet101(invariant='W', pretrained=pretrained)
else:
# initialize with network pretrained on imagenet in pytorch
net_in = getattr(torchvision.models, architecture)(pretrained=pretrained)
else:
# initialize with random weights, later on we will fill features with custom pretrained network
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 == 'w_vgg16':
features = [net_in.ciconv] + list(net_in.features.children())[:-1]
elif architecture.startswith('resnet') or architecture == 'w_resnet101':
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](mp=dim)
else:
pool = POOLING[pooling]()
# 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,
}
# create a generic image retrieval network
net = ImageRetrievalNet(features, lwhiten, pool, whiten, meta)
# initialize features with custom pretrained network if needed
if pretrained 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')
net.features.load_state_dict(model_zoo.load_url(FEATURES[architecture], model_dir=model_dir))
return net
def test(datasets, net, wandb_enabled=False, epoch=-1):
global global_step
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.cuda()
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, 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'])]
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net, images, image_size, transform) # implemented with torch.no_grad
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net, qimages, image_size, transform, bbxs) # 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'], wandb_enabled=wandb_enabled, epoch=epoch, global_step=global_step)
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'], wandb_enabled=wandb_enabled, epoch=epoch, global_step=global_step)
print('>> {}: elapsed time: {}'.format(dataset, htime(time.time()-start)))
def main():
global args, min_loss, global_step
args = parser.parse_args()
if args.wandb:
# initialize wandb
wandb.init(
project='cnnimageretrieval-pytorch', name=args.directory,
entity='ronaldseoh')
# save args provided for this experiment to wandb
wandb.config.update(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_train(get_data_root())
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)
if args.local_whitening:
directory += "_lwhiten"
if args.regional:
directory += "_r"
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, 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)
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 cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
# 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['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 = init_network(model_params)
# move network to gpu
model.cuda()
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
elif args.loss == 'triplet':
criterion = TripletLoss(margin=args.loss_margin).cuda()
else:
raise(RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features, pool, whitening
# IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM
parameters = []
# add feature parameters
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*10, '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*10, '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(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
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)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
if args.query_size == -1:
train_query_size = float('Inf')
else:
train_query_size = args.query_size
if args.pool_size == -1:
train_pool_size = float('Inf')
else:
train_pool_size = args.pool_size
train_dataset = TuplesDataset(
name=args.training_dataset,
mode='train',
imsize=args.image_size,
nnum=args.neg_num,
qsize=train_query_size,
poolsize=train_pool_size,
transform=transform,
store_nidxs_others_up_to=args.store_nidxs_others_up_to,
store_nidxs_others_order_by=args.store_nidxs_others_order_by,
totally_random_nidxs=args.totally_random_nidxs,
totally_random_nidxs_others=args.totally_random_nidxs_others,
dense_refresh_batch_and_nearby=args.dense_refresh_batch_and_nearby,
dense_refresh_batch_multi_hop=args.dense_refresh_batch_multi_hop,
dense_refresh_batch_random=args.dense_refresh_batch_random,
)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=0, pin_memory=True, sampler=None,
drop_last=True, collate_fn=collate_tuples
)
if args.val:
val_dataset = TuplesDataset(
name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
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, collate_fn=collate_tuples
)
if args.wandb:
# Start watching 'model' from wandb
wandb.watch(model)
# evaluate the network before starting
# this might not be necessary?
test(args.test_datasets, model, wandb_enabled=args.wandb, epoch=-1)
indexes_to_refresh = []
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
random.seed(args.seed + epoch)
np.random.seed(args.seed + epoch)
torch.manual_seed(args.seed + epoch)
torch.cuda.manual_seed_all(args.seed + epoch)
# train for one epoch on train set
loss, indexes_to_refresh = train(train_loader, model, criterion, optimizer, epoch, indexes_to_refresh)
if args.wandb:
wandb.log({"loss_avg": loss, "epoch": epoch, "global_step": global_step}) ## This is average loss
# adjust learning rate for each epoch
scheduler.step()
# # debug printing to check if everything ok
# lr_feat = optimizer.param_groups[0]['lr']
# lr_pool = optimizer.param_groups[1]['lr']
# print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool))
# evaluate on validation set
if args.val:
with torch.no_grad():
loss = validate(val_loader, model, criterion, epoch)
if args.wandb:
wandb.log({"loss_validation": loss, "epoch": epoch, "global_step": global_step})
# evaluate on test datasets every test_freq epochs
if (epoch + 1) % args.test_freq == 0:
with torch.no_grad():
test(args.test_datasets, model, wandb_enabled=args.wandb, epoch=epoch)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
if is_best:
print("Epoch", str(epoch + 1), "lower 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)
# calculate avg_neg_distance and avg_pos_distance for one last time
print("Training finished. Calculating the final avg_{neg,pos}_distance...")
avg_neg_distance, _ = train_loader.dataset.create_epoch_tuples(
model,
batch_members=[],
refresh_query_selection=False,
refresh_query_vectors=False,
refresh_negative_pool=False,
refresh_negative_pool_vectors=False,
refresh_nidxs=False,
refresh_nidxs_vectors=False)
if args.wandb:
wandb.log({"avg_neg_distance": avg_neg_distance, 'epoch': epoch, "global_step": global_step})
if args.calculate_positive_distance:
avg_pos_distance = train_loader.dataset.calculate_average_positive_distance()
if args.wandb:
wandb.log({"avg_pos_distance": avg_pos_distance, 'epoch': epoch, "global_step": global_step})
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_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# 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
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
if len(ms) > 1 and net.meta['pooling'] == 'gem' and not net.meta[
'regional'] and not net.meta['whitening']:
msp = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms))
print(">>>> msp: {}".format(msp))
else:
msp = 1
# 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])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if len(ms) > 1:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(
args.whitening)
if len(ms) > 1:
whiten_fn += '_ms'
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = torch.load(whiten_fn)
else:
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train',
args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(args.whitening))
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.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(
args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
# 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, msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, msp=msp)
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
print (vecs.shape)
print (qvecs.shape)
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
print (scores.shape)
# to save scores (single query)
# oxford
#f = 'oxf_single.npy'
# paris
#f = 'par_single.npy'
# roxford
#f = 'roxf_single.npy'
# rparis
f = 'rpar_single.npy'
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset, ranks, cfg['gnd'])
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)
# save
np.save(f, scores)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd'])
"""
############################################################
# Test
# 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'])]
print(qimages)
# to load scores
# oxford
#f = 'oxf_single.npy'
#f = 'oxf_mq_avg.npy'
#f = 'oxf_mq_max.npy'
#f = 'oxf_sc_imf.npy'
# paris
#f = 'par_single.npy'
#f = 'par_mq_avg.npy'
#f = 'par_mq_max.npy'
f = 'par_sc_imf.npy'
# roxford
#f = 'roxf_single.npy'
#f = 'roxf_mq_avg.npy'
#f = 'roxf_mq_max.npy'
#f = 'roxf_sc_imf.npy'
# rparis
#f = 'rpar_single.npy'
#f = 'rpar_mq_avg.npy'
#f = 'rpar_mq_max.npy'
#f = 'rpar_sc_imf.npy'
# load
scores = np.load(f)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd'])
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_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network
# pretrained networks (downloaded automatically)
print(">> Loading network:\n>>>> '{}'".format(args.network))
state = load_url(PRETRAINED[args.network],
model_dir=os.path.join(get_data_root(), 'networks'))
# 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
# network initialization
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
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)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms)
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)
top_k = 100
ranks_fnames_qs = []
for q_id in range(len(cfg["qimlist"])):
ranks_q = list(ranks[:top_k, q_id])
ranks_fname_per_q = []
for img_id in ranks_q:
ranks_fname_per_q.append(cfg["imlist"][img_id])
ranks_fnames_qs.append(ranks_fname_per_q)
compute_map_and_print(dataset, ranks, cfg['gnd'])
compute_map_and_print_top_k(dataset, ranks_fnames_qs, cfg['gnd'],
cfg["imlist"])
sys.exit()
with open(dataset + "_gl18_tl_resnet101_gem_w_m.pkl", "wb") as f:
data = {"ranks": ranks, "db_images": images, "q_images": qimages}
pickle.dump(data, f)
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def testHolidays(net, eConfig, dataset, Lw):
print('>> Evaluating network on test dataset: {}'.format(dataset))
# for testing we use image size of max 1024
image_size = 1024
ms = [1]
msp = 1
if (eConfig['multiscale']):
ms = [1, 1. / math.sqrt(2), 1. / 2]
if net.meta['pooling'] == 'gem' and net.whiten is None:
msp = net.pool.p.data.tolist()[0]
# 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])
# read the images and generate paths and queries-positive indexes
dbpath = os.path.join(get_data_root(), 'test', 'holidays')
ext = 'jpg' if dataset == 'holidays' else 'rjpg'
images = sorted(os.listdir(os.path.join(dbpath, ext)))
with open(os.path.join(dbpath, 'straight_gnd_holidays.pkl'), 'rb') as f:
queries = pickle.load(f)
positives = pickle.load(f)
qidx = []
pidx = []
for i in range(len(queries)):
qidx.append(images.index(queries[i]))
aux = []
for j in range(len(positives[i])):
aux.append(images.index(positives[i][j]))
pidx.append(aux)
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
X = extract_vectors(net, [os.path.join(dbpath, ext, n) for n in images],
image_size,
transform,
ms=ms,
msp=msp)
print('>> {}: Evaluating...'.format(dataset))
# rank the similarities
X = X.numpy()
if (Lw is not None):
X = whitenapply(X, Lw['m'], Lw['P'])
scores = np.dot(X.T, X)
ranks = np.argsort(-scores, axis=1)
ranks = ranks[qidx, 1::]
APs = []
for i, r in enumerate(ranks):
trueRanks = np.isin(r, pidx[i])
trueRanks = np.where(trueRanks == True)[0]
APs.append(compute_ap(trueRanks, len(pidx[i])))
mAP = np.mean(APs)
print(">> {}: mAP {:.2f}".format(dataset, mAP * 100))
# return the average mAP
return (dataset + ('+ multiscale' if eConfig['multiscale'] else ''), mAP)
# %%
# check if there are unknown datasets
if dataset not in datasets_names:
raise ValueError('Unsupported or unknown dataset: {}!'.format(dataset))
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
# loading network from path
if network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(network_path))
if network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[network_path],
model_dir=os.path.join(get_data_root(), 'networks'))
else:
# fine-tuned network from path
state = torch.load(network_path)
# 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
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
db_root = '/mnt/fry2/users/datasets/landmarkscvprw18/recognition/'
ims_root = os.path.join(db_root, 'images', 'train')
# 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 = [
os.path.join(ims_root, db['cids'][i] + '.jpg')
for i in range(len(db['cids']))
]
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']
## 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 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')
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', False)
multi_layer_cat = params.get('multi_layer_cat', 1)
# get output dimensionality size
dim = OUTPUT_DIM[architecture]
# # hxq modified, remove resnet50 con5_x for smaller receptive field
# if architecture == 'resnet50':
# dim = 1024
# else:
# 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'):
features = list(net_in.children())[:-2]
# # hxq modified, remove resnet50 con5_x for smaller receptive field
# if architecture == 'resnet50':
# features = list(net_in.children())[:-3]
# else:
# 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())
elif architecture.startswith('resnext101_32x8d'):
features = list(net_in.children())[:-2]
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
pool = POOLING[pooling]()
# pool = POOLING[pooling](p=1.5)
# print('>> GeM pool p: 1.5')
# 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
output_dim = dim
if whitening:
if multi_layer_cat != 1:
whiten = nn.Linear(dim + 1024, output_dim, bias=True)
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_M2_120k_IS1024_MS1_WL.pth'
# # whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_CA_M2_120k_IS1024_MS1_WL.pth'
# # whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/RX101_M2_120k_IS1024_MS1_WL.pth'
# # whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_M4_120k_IS1024_MS1_WL.pth'
# # whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/RX101_M4_120k_IS1024_MS1_WL.pth'
# print('>> load computed whitening \'{}\''.format(whiten_fn.split('/')[-1]))
# Lw = torch.load(whiten_fn)
# P = Lw['P'][:dim, :]
# m = Lw['m']
# P = torch.from_numpy(P).float()
# m = torch.from_numpy(m).float()
# whiten.weight.data = P
# whiten.bias.data = -torch.mm(P, m).squeeze()
else:
# whiten = nn.Linear(dim, dim, bias=True)
output_dim = 512
whiten = nn.Linear(dim, output_dim, bias=True)
nn.init.xavier_normal_(whiten.weight)
nn.init.constant_(whiten.bias, 0)
# hxq added, for whitening test
# print('>> load the parameters of supervised whitening for the FC layer initialization')
# whiten_fn = '/media/iap205/Data/Export/cnnimageretrieval-google_landmark_retrieval/trained_network/' \
# 'R101_O_GL_FC/google-landmarks-dataset-resize_resnet101_gem_whiten_contrastive_m0.85_' \
# 'adam_lr5.0e-07_wd1.0e-04_nnum5_qsize2000_psize22000_bsize5_imsize362/' \
# 'model_epoch114.pth.tar_google-landmarks-dataset_whiten_ms.pth'
# whiten_fn = '/home/iap205/Datasets/google-landmarks-dataset-resize/whiten/imagenet-caffe-resnet101-' \
# 'features-10a101d.pth_google-landmarks-dataset-test_256_whiten_MS1.pth'
# whiten_fn = '/media/iap205/Data4T/Datasets/google-landmarks-dataset-v2/whiten/imagenet-caffe-' \
# 'resnet101-features-10a101d.pth_google-landmarks-dataset-v2-test_256_whiten_MS1.pth'
# whiten_fn = '/media/iap205/Data4T/Datasets/google-landmarks-dataset-v2/whiten/imagenet-caffe-' \
# 'resnet101-features-10a101d.pth_google-landmarks-dataset-v2-test_1024_whiten_MS1.pth'
# whiten_fn = '/media/iap205/Data/Export/cnnimageretrieval-pytorch/trained_network/retrieval-SfM-120k' \
# '_resnet101_gem_contrastive_m0.85_adam_lr5.0e-07_wd1.0e-04_nnum5_qsize2000_psize22000_bsize5' \
# '_imsize362/model_best.pth.tar_retrieval-SfM-120k_whiten_ms.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_120k_IS1024_MS1.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_120k_IS362_MS1.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_120k_IS1024_MS1_WL.pth'
# # whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_120k_IS1024_MS1_WL_WL.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_CA_120k_IS1024_MS1_WL.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_top9_120k_IS1024_MS1_WL.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_top0.25_120k_IS1024_MS1_WL.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_LA_120k_IS1024_MS1_WL.pth'
# whiten_fn = '/home/iap205/Datasets/retrieval-SfM/whiten/R101_top0.25_p1.5_120k_IS1024_MS1_WL.pth'
# print('>> load computed whitening \'{}\''.format(whiten_fn.split('/')[-1]))
# Lw = torch.load(whiten_fn)
# P = Lw['P']
# m = Lw['m']
# P = torch.from_numpy(P).float()
# m = torch.from_numpy(m).float()
# whiten.weight.data = P
# whiten.bias.data = -torch.mm(P, m).squeeze()
# multi PCA test
# whiten = nn.Sequential(nn.Linear(dim, dim, bias=True),
# nn.Linear(dim, dim, bias=True))
# whiten_fn = ['/home/iap205/Datasets/retrieval-SfM/whiten/R101_120k_362_IS1024_MS1_WL.pth',
# '/home/iap205/Datasets/retrieval-SfM/whiten/R101_120k_362_IS1024_MS1_WL_WL.pth']
# for i in range(len(whiten)):
# print('>> load computed whitening \'{}\''.format(whiten_fn[i].split('/')[-1]))
# Lw = torch.load(whiten_fn[i])
# P = Lw['P']
# m = Lw['m']
# P = torch.from_numpy(P).float()
# m = torch.from_numpy(m).float()
# whiten[i].weight.data = P
# whiten[i].bias.data = -torch.mm(P, m).squeeze()
# 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': output_dim,
'multi_layer_cat': multi_layer_cat
}
# create a generic image retrieval network
net = ImageRetrievalNet(features, lwhiten, pool, whiten, meta)
# initialize features with custom pretrained network if needed
if pretrained 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')
# hxq modified, remove resnet50 con5_x for smaller receptive field
if architecture == 'resnet50-3':
state_dict = model_zoo.load_url(FEATURES[architecture],
model_dir=model_dir)
for key in list(state_dict.keys()):
if list(key)[0] == '7':
state_dict.pop(key)
net.features.load_state_dict(state_dict)
else:
# hxq modified
# state_dict = model_zoo.load_url(FEATURES[architecture], model_dir=model_dir)
# state_dict_mGPU = {}
# for key, val in state_dict.items():
# state_dict_mGPU['module.'+key] = val
# net.features.load_state_dict(state_dict_mGPU)
net.features.load_state_dict(
model_zoo.load_url(FEATURES[architecture],
model_dir=model_dir))
return net
def main():
global args, min_loss
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_train(get_data_root())
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)
if args.local_whitening:
directory += "_lwhiten"
if args.regional:
directory += "_r"
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, 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)
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 cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds
# TODO: maybe pass as argument in future implementation?
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.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['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 = init_network(model_params)
# move network to gpu
model.cuda()
if model_params['pooling'] == 'netvlad':
normalize = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
transform = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=(364, 364)),
torchvision.transforms.ToTensor(),
normalize,
])
image_folder = folder.ImageFolder(root="/mnt/m2/dataset", transform=transform)
train_loader = torch.utils.data.DataLoader(
image_folder,
batch_size=64,
num_workers=8,
shuffle=True
)
n_batches = 10
descs_list = []
i = 0
with torch.no_grad():
for x, _ in train_loader:
model.eval()
desc = model.compute_features(x.cuda())
max_pooled_feat_3 = torch.nn.functional.max_pool2d(desc, kernel_size=3, stride=1)
max_pooled_feat_2 = torch.nn.functional.max_pool2d(desc, kernel_size=2, stride=1)
reshaped_pool_3 = make_locals(max_pooled_feat_3)
reshaped_pool_2 = make_locals(max_pooled_feat_2)
desc = torch.cat([reshaped_pool_2, reshaped_pool_3], dim=1)
# N, dim, h, w = desc.shape
# desc = desc.view(N, dim, h*w).permute(0, 2, 1).reshape(N, -1, 512)
desc = desc.cpu().numpy().astype('float32')
descs_list.append(desc)
print(">> Extracted batch {}/{} - NetVLAD initialization -".format(i+1, n_batches))
i+=1
if i == n_batches:
break
descs_list = np.array(descs_list).reshape(-1, 512)
print(descs_list.shape)
print(">> Sampling local features ")
# locals = np.vstack((m[np.random.randint(len(m), size=150)] for m in descs_list)).astype('float32')
locals = descs_list[np.random.randint(len(descs_list), size=len(descs_list)//3)]
np.random.shuffle(locals)
print(">> Locals extracted shape : {}".format(locals.shape))
n_clust = 64
locals = preprocessing.normalize(locals, axis=1)
print(">> Fitting centroids with K-Means")
kmeans = MiniBatchKMeans(n_clusters=n_clust).fit(locals)
centroids = kmeans.cluster_centers_
print(">> Centroids shape: ", centroids.shape)
model.pool.init_params(centroids.T)
print(">> NetVLAD initialized")
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
elif args.loss == 'triplet':
criterion = TripletLoss(margin=args.loss_margin).cuda()
else:
raise(RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features, pool, whitening
# IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM
parameters = []
# add feature parameters
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*10, '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*10, '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(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
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)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
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:
val_dataset = TuplesDataset(
name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
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, collate_fn=collate_tuples
)
# evaluate the network before starting
# this might not be necessary?
test(args.test_datasets, model)
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# adjust learning rate for each epoch
scheduler.step()
# # debug printing to check if everything ok
# lr_feat = optimizer.param_groups[0]['lr']
# lr_pool = optimizer.param_groups[1]['lr']
# print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool))
# train for one epoch on train set
loss = train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
if args.val:
with torch.no_grad():
loss = validate(val_loader, model, criterion, epoch)
# evaluate on test datasets every test_freq epochs
if (epoch + 1) % args.test_freq == 0:
with torch.no_grad():
test(args.test_datasets, model)
# 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)
def __init__(
self,
imsize=None,
nnum=5,
qsize=2000,
poolsize=20000,
transform=None,
loader=default_loader,
filename=None,
q_percent=1,
):
# setting up paths
data_root = get_data_root()
name = "retrieval-SfM-120k"
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)["val"]
# initializing tuples dataset
self.imsize = imsize
self.images = [
cid2filename(db["cids"][i], ims_root)
for i in range(len(db["cids"]))
]
self.clusters = db["cluster"]
self.qpool = db["qidxs"]
# self.ppool = db['pidxs']
# 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 = self.qpool
self.index = np.arange(len(self.qidxs))
if q_percent < 1:
number = int(len(self.qidxs) * q_percent)
self.index = np.random.permutation(self.index)
self.index = self.index[:number]
self.pidxs = []
self.nidxs = []
self.poolvecs = None
self.transform = transform
self.loader = loader
self.filename = filename
self.phase = 1
self.ranks = torch.load(f"{filename}/ranks_362")
if os.path.isfile(f"{filename}/pool_vecs"):
self.pool_vecs = pickle.load(open(f"{filename}/pool_vecs", "rb"))
print(len(self.images))
self.loaded_images = []
if os.path.exists("./images"):
self.loaded_images = pickle.load(open("./images", "rb"))
else:
for i in range(len(self.images)):
try:
img = self.loader(self.images[i])
if self.imsize is not None:
img = imresize(img, self.imsize)
if self.transform is not None:
img_tensor = self.transform(img).unsqueeze(0)
img.close()
self.loaded_images.append(img_tensor)
except:
self.loaded_images.append(None)
pickle.dump(self.loaded_images, open("./images", "wb"))
# whitening = 'google-landmarks-dataset-resize'
#%%
# check if there are unknown datasets
if dataset not in datasets_names:
raise ValueError('Unsupported or unknown dataset: {}!'.format(dataset))
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
# loading network from path
print(">> Loading network:\n>>>> '{}'".format(network_path))
if network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[network_path],
model_dir=os.path.join(get_data_root(), 'networks'))
else:
# fine-tuned network from path
state = torch.load(network_path)
# 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
def main():
args = parser.parse_args()
# check if test dataset are downloaded
# and download if they are not
#download_train(get_data_root())
#download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
state = torch.load(args.network_path)
net = init_network(model=state['meta']['architecture'],
pooling=state['meta']['pooling'],
whitening=state['meta']['whitening'],
mean=state['meta']['mean'],
std=state['meta']['std'],
pretrained=False)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
offtheshelf = args.network_offtheshelf.split('-')
if len(offtheshelf) == 3:
if offtheshelf[2] == 'whiten':
offtheshelf_whiten = True
else:
raise (RuntimeError(
"Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten"
))
else:
offtheshelf_whiten = False
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(model=offtheshelf[0],
pooling=offtheshelf[1],
whitening=offtheshelf_whiten)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = [1]
msp = 1
if args.multiscale:
ms = [1, 1. / math.sqrt(2), 1. / 2]
if net.meta['pooling'] == 'gem' and net.whiten is None:
msp = net.pool.p.data.tolist()[0]
# 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])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if args.multiscale:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train', args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(args.whitening))
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.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
# 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'])]
bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms,
msp=msp)
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'])
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'])
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def gen_train_val_test_pkl(dataset):
if dataset == 'google-landmarks-dataset-v2':
data_table_path = os.path.join(get_data_root(), 'train.csv')
img_check_path = os.path.join(get_data_root(), 'train')
train_val_file_output_path = os.path.join(get_data_root())
test_file_output_path = os.path.join(get_data_root(), 'test',
dataset + '-test')
val_set_size = 10000
test_set_size = 10000
q_index_ratio = 100000. / 700000.
train_qp_pair_num = 12
elif dataset == 'google-landmarks-dataset-resize':
data_table_path = os.path.join(get_data_root(), 'train.csv')
img_check_path = os.path.join(get_data_root(), 'resize_train_image')
train_val_file_output_path = os.path.join(get_data_root())
test_file_output_path = os.path.join(get_data_root(), 'test',
dataset + '-test')
val_set_size = 10000
test_set_size = 10000
q_index_ratio = 118000. / 1100000.
train_qp_pair_num = 50
elif dataset == 'google-landmarks-dataset':
pass
# step0: check image is exist
print(">> load {} file...".format(data_table_path.split('/')[-1]))
csvfile = open(data_table_path, 'r')
csvreader = csv.reader(csvfile)
id_url_class_list = [line[:3] for line in csvreader][1:]
csvfile.close()
if dataset == 'google-landmarks-dataset-v2':
# https://github.com/cvdfoundation/google-landmark
# images = [os.path.join(img_check_path, '/'.join(list(id_url_class_list[i][0])[:3]), id_url_class_list[i][0])
# + '.jpg' for i in range(len(id_url_class_list))]
train_temp = [['1', id_url_class_list[i][0], id_url_class_list[i][2]]
for i in range(len(id_url_class_list))]
elif dataset == 'google-landmarks-dataset-resize':
print(">> check {} file...".format(data_table_path.split('/')[-1]))
train_temp = []
files = os.listdir(img_check_path)
files = {
files[i].split('.jpg')[0]: files[i].split('.jpg')[0]
for i in range(len(files))
}
missfile = open(
os.path.join(train_val_file_output_path, 'train_miss.csv'), 'w')
miss_writer = csv.writer(missfile)
miss_writer.writerow(['id', 'url'])
miss_cnt = 0
for line in id_url_class_list:
if files.__contains__(line[0]):
if line[2] != 'None':
train_temp.append(['1', line[0], line[2]])
else:
train_temp.append(['1', line[0], '-1'])
files.pop(line[0])
else:
if line[2] != 'None':
train_temp.append(['0', line[0], line[2]])
else:
train_temp.append(['0', line[0], '-1'])
miss_writer.writerow(line)
miss_cnt += 1
missfile.close()
print('>>>> train image miss: {}, train_miss.csv save done...'.format(
miss_cnt))
elif dataset == 'google-landmarks-dataset':
pass
np.random.seed(0) # added after the competition
shuffle = np.random.permutation(len(train_temp))
train_shuffle = []
for i in shuffle:
train_shuffle.append(train_temp[i])
train_shuffle = pd.DataFrame(train_shuffle,
columns=['mark', 'id', 'landmark_id'])
train_shuffle['landmark_id'] = train_shuffle['landmark_id'].astype('int')
# train = train.sort_values('landmark_id')
train_shuffle = train_shuffle.to_numpy()
# step1: extract test set
print(">> extract test set...")
landmark_id_shuffle = np.random.permutation(max(train_shuffle[:, 2]))
test_landmark_num = round(test_set_size /
(len(train_shuffle) / max(train_shuffle[:, 2])))
test_sel_landmark_id = landmark_id_shuffle[:test_landmark_num]
landmark_id_shuffle = landmark_id_shuffle[test_landmark_num:]
test = []
for i in range(len(train_shuffle)):
if train_shuffle[i][2] in test_sel_landmark_id:
if train_shuffle[i][0] == '1':
test.append(train_shuffle[i][1:])
train_shuffle[i][0] = '0'
test = pd.DataFrame(test, columns=['id', 'landmark_id'])
test = test.sort_values('landmark_id')
test = test.to_numpy()
print(">> copy test image from train set...")
imlist_all = list(test[:, 0])
if dataset == 'google-landmarks-dataset-v2':
img_path = [
os.path.join(img_check_path, '/'.join(list(imlist_all[i])[:3]),
imlist_all[i] + '.jpg')
for i in range(len(imlist_all))
]
else:
img_path = [
os.path.join(img_check_path, imlist_all[i] + '.jpg')
for i in range(len(imlist_all))
]
for i in range(len(imlist_all)):
shutil.copyfile(
img_path[i],
os.path.join(test_file_output_path, 'jpg', imlist_all[i] + '.jpg'))
test_fold = []
for i in range(test_landmark_num):
test_fold.append([])
landmark_cnt = 0
test_fold[0].append([test[0][0], 0])
for i in range(1, len(test)):
if test[i][1] != test[i - 1][1]:
landmark_cnt += 1
test_fold[landmark_cnt].append([test[i][0], i])
qimlist = []
gnd = []
imlist = []
for i in range(test_landmark_num):
q_num = max(round(q_index_ratio * len(test_fold[i])), 1)
for j in range(q_num):
qimlist.append(test_fold[i][j][0])
gnd.append({'ok': []})
for k in range(len(imlist),
len(imlist) + len(test_fold[i]) - q_num):
gnd[-1]['ok'].append(k)
for j in range(q_num, len(test_fold[i])):
imlist.append(test_fold[i][j][0])
test_pkl = {'imlist': imlist, 'qimlist': qimlist, 'gnd': gnd}
pkl_output_path = os.path.join(test_file_output_path,
'gnd_{}-test.pkl'.format(dataset))
pickle_file = open(pkl_output_path, 'wb')
pickle.dump(test_pkl, pickle_file)
pickle_file.close()
print(">> save gnd_{}-test.pkl done...".format(dataset))
# step2: extract validation set
print('>> extract validation set...')
val_landmark_num = round(val_set_size /
(len(train_shuffle) / max(train_shuffle[:, 2])))
val_sel_landmark_id = landmark_id_shuffle[:val_landmark_num]
val = []
for i in range(len(train_shuffle)):
if train_shuffle[i][2] in val_sel_landmark_id:
if train_shuffle[i][0] == '1':
val.append(train_shuffle[i][1:])
train_shuffle[i][0] = '0'
val = pd.DataFrame(val, columns=['id', 'landmark_id'])
val = val.sort_values('landmark_id')
val = val.to_numpy()
val_dict = {}
val_dict['id'] = list(val[:, 0])
val_dict['landmark_id'] = list(val[:, 1])
pair_cnt = 0
val_qidxs, val_pidxs = [], []
val_qp_pair_num = round(
train_qp_pair_num *
(val_set_size / (len(train_shuffle) - val_set_size - test_set_size)))
for i in range(len(val_dict['landmark_id']) - 1):
if val_dict['landmark_id'][i] != -1:
if val_dict['landmark_id'][i] == val_dict['landmark_id'][i + 1]:
if pair_cnt < max(1, val_qp_pair_num):
val_qidxs.append(i)
val_pidxs.append(i + 1)
pair_cnt += 1
else:
pair_cnt = 0
val_dict['qidxs'] = val_qidxs
val_dict['pidxs'] = val_pidxs
if dataset == 'google-landmarks-dataset-v2':
print(">> copy val image from train set...")
img_path = [
os.path.join(img_check_path, '/'.join(list(val_dict['id'][i])[:3]),
val_dict['id'][i] + '.jpg')
for i in range(len(val_dict['id']))
]
for i in range(len(val_dict['id'])):
shutil.copyfile(
img_path[i],
os.path.join(train_val_file_output_path, 'val',
val_dict['id'][i] + '.jpg'))
# step3: extract train set
print('>> extract train set...')
train_dict = {}
train = train_shuffle
train = pd.DataFrame(train, columns=['mark', 'id', 'landmark_id'])
train = train.sort_values('landmark_id')
train = train.to_numpy()
train_id, train_landmark_id = [], []
for i in range(len(train)):
if train[i, 0] == '1':
train_id.append(train[i, 1])
train_landmark_id.append(train[i, 2])
train_dict['id'] = train_id
train_dict['landmark_id'] = train_landmark_id
pair_cnt = 0
train_qidxs, train_pidxs = [], []
for i in range(len(train_dict['landmark_id']) - 1):
if train_dict['landmark_id'][i] != -1:
if train_dict['landmark_id'][i] == train_dict['landmark_id'][i +
1]:
if pair_cnt < train_qp_pair_num:
train_qidxs.append(i)
train_pidxs.append(i + 1)
pair_cnt += 1
else:
pair_cnt = 0
train_dict['qidxs'] = train_qidxs
train_dict['pidxs'] = train_pidxs
# warp all dict to one pickle file
all_dict = {'train': train_dict, 'val': val_dict}
pickle_file = open(
os.path.join(train_val_file_output_path, '{}.pkl'.format(dataset)),
'wb')
pickle.dump(all_dict, pickle_file)
pickle_file.close()
print(">> save {}.pkl done...".format(dataset))
def init_network(model='resnet101', pooling='gem', whitening=False, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], pretrained=True):
# loading network from torchvision
if pretrained:
if model not in FEATURES:
# initialize with network pretrained on imagenet in pytorch
net_in = getattr(torchvision.models, model)(pretrained=True)
else:
# initialize with random weights, later on we will fill features with custom pretrained network
net_in = getattr(torchvision.models, model)(pretrained=False)
else:
# initialize with random weights
net_in = getattr(torchvision.models, model)(pretrained=False)
# initialize features
# take only convolutions for features,
# always ends with ReLU to make last activations non-negative
if model.startswith('alexnet'):
features = list(net_in.features.children())[:-1]
elif model.startswith('vgg'):
features = list(net_in.features.children())[:-1]
elif model.startswith('resnet'):
features = list(net_in.children())[:-2]
elif model.startswith('densenet'):
features = list(net_in.features.children())
features.append(nn.ReLU(inplace=True))
elif model.startswith('squeezenet'):
features = list(net_in.features.children())
else:
raise ValueError('Unsupported or unknown model: {}!'.format(model))
# initialize pooling
pool = POOLING[pooling]()
# get output dimensionality size
dim = OUTPUT_DIM[model]
# initialize whitening
if whitening:
w = '{}-{}'.format(model, pooling)
whiten = nn.Linear(dim, dim, bias=True)
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':model, 'pooling':pooling, 'whitening':whitening, 'outputdim':dim, 'mean':mean, 'std':std}
# create a generic image retrieval network
net = ImageRetrievalNet(features, pool, whiten, meta)
# initialize features with custom pretrained network if needed
if pretrained and model in FEATURES:
print(">> {}: for '{}' custom pretrained features '{}' are used"
.format(os.path.basename(__file__), model, os.path.basename(FEATURES[model])))
model_dir = os.path.join(get_data_root(), 'networks')
net.features.load_state_dict(model_zoo.load_url(FEATURES[model], model_dir=model_dir))
return net
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_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# 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
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
if len(ms) > 1 and net.meta['pooling'] == 'gem' and not net.meta[
'regional'] and not net.meta['whitening']:
msp = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms))
print(">>>> msp: {}".format(msp))
else:
msp = 1
# 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])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if len(ms) > 1:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(
args.whitening)
if len(ms) > 1:
whiten_fn += '_ms'
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = torch.load(whiten_fn)
else:
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train',
args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(args.whitening))
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.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(
args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
images = get_imlist("E:\\PycharmProjects\\image-retrieval\\holiday2\\")
names = []
for i, img_path in enumerate(images):
img_name = os.path.split(img_path)[1]
print(img_name)
names.append(img_name)
# 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'])]
# try:
# bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
# except:
# bbxs = None # for holidaysmanrot and copydays
# names = []
# for i, img_path in enumerate(images):
# img_name = os.path.split(img_path)[1]
# print(img_name)
# names.append(img_name)
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# convert to numpy
vecs = vecs.numpy()
vecs = vecs.T
print("--------------------------------------------------")
print(" writing feature extraction results ...")
print("--------------------------------------------------")
output = "gem_res_holiday_3.h5"
h5f = h5py.File(output, 'w')
h5f.create_dataset('dataset_1', data=vecs)
h5f.create_dataset('dataset_2', data=np.string_(names))
h5f.close()
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def testOxfordParisHolidays(net, eConfig):
#datasets = eConfig['test-datasets'].split(',')
#results = []
#
#for dataset in datasets:
# results.append((dataset, np.random.rand(1)[0]))
#
#return results
print('>> Evaluating network on test datasets...')
# for testing we use image size of max 1024
image_size = 1024
# setting up the multi-scale parameters
ms = [1]
msp = 1
if (eConfig['multiscale']):
ms = [1, 1. / math.sqrt(2), 1. / 2]
if net.meta['pooling'] == 'gem' and net.whiten is None:
msp = net.pool.p.data.tolist()[0]
# 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])
# compute whitening
if eConfig['whitening']:
start = time.time()
print('>> {}: Learning whitening...'.format(eConfig['test-whiten']))
# loading db
db_root = os.path.join(get_data_root(), 'train',
eConfig['test-whiten'])
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(eConfig['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(eConfig['test-whiten']))
wvecs = extract_vectors(net,
images,
image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(eConfig['test-whiten']))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(eConfig['test-whiten'],
htime(time.time() - start)))
else:
Lw = None
# evaluate on test datasets
datasets = eConfig['test-datasets'].split(',')
results = []
for dataset in datasets:
start = time.time()
if (dataset != 'holidays' and dataset != 'rholidays'):
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'])]
if (dataset == 'oxford105k' or dataset == 'paris106k'):
images.extend(cfg['distractors'])
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
image_size,
transform,
ms=ms,
msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
image_size,
transform,
bbxs,
ms=ms,
msp=msp)
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)
results.append(
compute_map_and_print(
dataset +
('+ multiscale' if eConfig['multiscale'] else ''), ranks,
cfg['gnd']))
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)
results.append(
compute_map_and_print(dataset + ' + whiten', ranks,
cfg['gnd']))
else:
results.append(testHolidays(net, eConfig, dataset, Lw))
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
return results
def main():
args = parser.parse_args()
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
state = torch.load(args.network_path)
net = init_network(model=state['meta']['architecture'],
pooling=state['meta']['pooling'],
whitening=state['meta']['whitening'],
mean=state['meta']['mean'],
std=state['meta']['std'],
pretrained=False)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
offtheshelf = args.network_offtheshelf.split('-')
if len(offtheshelf) == 3:
if offtheshelf[2] == 'whiten':
offtheshelf_whiten = True
else:
raise (RuntimeError(
"Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten"
))
else:
offtheshelf_whiten = False
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(model=offtheshelf[0],
pooling=offtheshelf[1],
whitening=offtheshelf_whiten)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = [1]
msp = 1
if args.multiscale:
ms = [1, 1. / math.sqrt(2), 1. / 2]
if net.meta['pooling'] == 'gem' and net.whiten is None:
msp = net.pool.p.data.tolist()[0]
# 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])
# compute whitening
if args.whitening is not None:
start = time.time()
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train', args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.whitening))
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.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
if dataset == 'reco':
images, qimages = landmark_recognition_dataset()
bbxs = [None for x in qimages]
elif dataset == 'retr':
images, _ = landmark_retrieval_dataset()
qimages = []
bbxs = [None for x in qimages]
else:
# 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'])]
with open('%s_fnames.pkl' % dataset, 'wb') as f:
pickle.dump([images, qimages], f)
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
vecs = vecs.numpy()
print('>> saving')
np.save('{}_vecs.npy'.format(dataset), vecs)
if len(qimages) > 0:
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms,
msp=msp)
qvecs = qvecs.numpy()
np.save('{}_qvecs.npy'.format(dataset), qvecs)
if Lw is not None:
# whiten the vectors
vecs_lw = whitenapply(vecs, Lw['m'], Lw['P'])
qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P'])
# TODO
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))