temp state

LICENSE

@@ -0,0 +1,85 @@
+FIRe, Copyright (c) 2021-2022 Naver Corporation, is licensed under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 license.
+
+A summary of the CC BY-NC-SA 4.0 license is located here:
+	https://creativecommons.org/licenses/by-nc-sa/4.0/
+
+The CC BY-NC-SA 4.0 license is located here:
+	https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+	
+	
+SEE NOTICE BELOW WITH RESPECT TO THE FILE: train.py and evaluate.py
+SEE NOTICE BELOW WITH RESPECT TO THE FILES in folder how/
+
+**********************************
+
+
+NOTICE WITH RESPECT TO THE FILE: train.py and evaluate.py
+
+
+This software is being redistributed in a modifiled form. The original form is available here:
+
+https://github.com/gtolias/how
+
+
+ORIGINAL COPYRIGHT NOTICE AND PERMISSION NOTICE AVAILABLE HERE IS REPRODUCE BELOW:
+
+https://github.com/gtolias/how/blob/master/LICENSE
+
+
+MIT License
+
+Copyright (c) 2020 Giorgos Tolias, Tomas Jenicek
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
+
+**********************************
+
+SEE NOTICE BELOW WITH RESPECT TO THE FILES in folder how/
+
+This project contains subcomponents with separate copyright notices and license terms. 
+Your use of the source code for these subcomponents is subject to the terms and conditions of the following licenses.
+
+====
+
+gtolias/how
+https://github.com/gtolias/how
+
+MIT License
+
+Copyright (c) 2020 Giorgos Tolias, Tomas Jenicek
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+====

README.md

@@ -1,13 +1 @@
----
-title: Superfeatures
-emoji: 🏢
-colorFrom: blue
-colorTo: indigo
-sdk: gradio
-sdk_version: 2.9.1
-app_file: app.py
-pinned: false
-license: cc-by-nc-sa-4.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
+TBD

fire.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ddeb04ebdd5ca3e7a9d86ce6a5dec5dabfbb23a70a6f3d0907b17e484474202
+size 52765649

how/__init__.py

@@ -0,0 +1,4 @@
+"""
+Official Python implementation of HOW method for ECCV 2020 paper "Learning and aggregating deep
+local descriptors for instance-level recognition"
+"""

how/layers/__init__.py

@@ -0,0 +1,5 @@
+"""
+Modules implementing layers in pytorch by inheriting from torch.nn.Module
+"""
+
+from . import attention, dim_reduction, pooling

how/layers/attention.py

@@ -0,0 +1,10 @@
+"""Layers producing a 2D attention map from a feature map"""
+
+from torch import nn
+
+
+class L2Attention(nn.Module):
+    """Compute the attention as L2-norm of local descriptors"""
+
+    def forward(self, x):
+        return (x.pow(2.0).sum(1) + 1e-10).sqrt().squeeze(0)

how/layers/dim_reduction.py

@@ -0,0 +1,29 @@
+"""Layers implementing dimensionality reduction of a feature map"""
+
+import torch
+from torch import nn
+
+from ..utils import whitening
+
+
+class ConvDimReduction(nn.Conv2d):
+    """Dimensionality reduction as a convolutional layer
+
+    :param int input_dim: Network out_channels
+    :param in dim: Whitening out_channels, for dimensionality reduction
+    """
+
+    def __init__(self, input_dim, dim):
+        super().__init__(input_dim, dim, (1, 1), padding=0, bias=True)
+
+    def initialize_pca_whitening(self, des):
+        """Initialize PCA whitening from given descriptors. Return tuple of shift and projection."""
+        m, P = whitening.pcawhitenlearn_shrinkage(des)
+        m, P = m.T, P.T
+
+        projection = torch.Tensor(P[:self.weight.shape[0], :]).unsqueeze(-1).unsqueeze(-1)
+        self.weight.data = projection.to(self.weight.device)
+
+        projected_shift = -torch.mm(torch.FloatTensor(P), torch.FloatTensor(m)).squeeze()
+        self.bias.data = projected_shift[:self.weight.shape[0]].to(self.bias.device)
+        return m.T, P.T

how/layers/functional.py

@@ -0,0 +1,73 @@
+"""Layer functions"""
+
+import torch
+import torch.nn.functional as F
+
+import cirtorch.layers.functional as CF
+
+
+def smoothing_avg_pooling(feats, kernel_size):
+    """Smoothing average pooling
+
+    :param torch.Tensor feats: Feature map
+    :param int kernel_size: kernel size of pooling
+    :return torch.Tensor: Smoothend feature map
+    """
+    pad = kernel_size // 2
+    return F.avg_pool2d(feats, (kernel_size, kernel_size), stride=1, padding=pad,
+                        count_include_pad=False)
+
+
+def weighted_spoc(ms_feats, ms_weights):
+    """Weighted SPoC pooling, summed over scales.
+
+    :param list ms_feats: A list of feature maps, each at a different scale
+    :param list ms_weights: A list of weights, each at a different scale
+    :return torch.Tensor: L2-normalized global descriptor
+    """
+    desc = torch.zeros((1, ms_feats[0].shape[1]), dtype=torch.float32, device=ms_feats[0].device)
+    for feats, weights in zip(ms_feats, ms_weights):
+        desc += (feats * weights).sum((-2, -1)).squeeze()
+    return CF.l2n(desc)
+
+
+def how_select_local(ms_feats, ms_masks, *, scales, features_num):
+    """Convert multi-scale feature maps with attentions to a list of local descriptors
+
+    :param list ms_feats: A list of feature maps, each at a different scale
+    :param list ms_masks: A list of attentions, each at a different scale
+    :param list scales: A list of scales (floats)
+    :param int features_num: Number of features to be returned (sorted by attenions)
+    :return tuple: A list of descriptors, attentions, locations (x_coor, y_coor) and scales where
+            elements from each list correspond to each other
+    """
+    device = ms_feats[0].device
+    size = sum(x.shape[0] * x.shape[1] for x in ms_masks)
+
+    desc = torch.zeros(size, ms_feats[0].shape[1], dtype=torch.float32, device=device)
+    atts = torch.zeros(size, dtype=torch.float32, device=device)
+    locs = torch.zeros(size, 2, dtype=torch.int16, device=device)
+    scls = torch.zeros(size, dtype=torch.float16, device=device)
+
+    pointer = 0
+    for sc, vs, ms in zip(scales, ms_feats, ms_masks):
+        if len(ms.shape) == 0:
+            continue
+
+        height, width = ms.shape
+        numel = torch.numel(ms)
+        slc = slice(pointer, pointer+numel)
+        pointer += numel
+
+        desc[slc] = vs.squeeze(0).reshape(vs.shape[1], -1).T
+        atts[slc] = ms.reshape(-1)
+        width_arr = torch.arange(width, dtype=torch.int16)
+        locs[slc, 0] = width_arr.repeat(height).to(device) # x axis
+        height_arr = torch.arange(height, dtype=torch.int16)
+        locs[slc, 1] = height_arr.view(-1, 1).repeat(1, width).reshape(-1).to(device) # y axis
+        scls[slc] = sc
+
+    keep_n = min(features_num, atts.shape[0]) if features_num is not None else atts.shape[0]
+    idx = atts.sort(descending=True)[1][:keep_n]
+
+    return desc[idx], atts[idx], locs[idx], scls[idx]

how/layers/pooling.py

@@ -0,0 +1,19 @@
+"""Spatial pooling layers"""
+
+from torch import nn
+
+from . import functional as LF
+
+
+class SmoothingAvgPooling(nn.Module):
+    """Average pooling that smoothens the feature map, keeping its size
+
+    :param int kernel_size: Kernel size of given pooling (e.g. 3)
+    """
+
+    def __init__(self, kernel_size):
+        super().__init__()
+        self.kernel_size = kernel_size
+
+    def forward(self, x):
+        return LF.smoothing_avg_pooling(x, kernel_size=self.kernel_size)

how/networks/__init__.py

@@ -0,0 +1,5 @@
+"""
+Pytorch networks
+"""
+
+from . import how_net

how/networks/how_net.py

@@ -0,0 +1,221 @@
+"""Module of the HOW method"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision
+
+from cirtorch.networks import imageretrievalnet
+
+from .. import layers
+from ..layers import functional as HF
+from ..utils import io_helpers
+
+NUM_WORKERS = 6
+
+CORERCF_SIZE = {
+    'resnet18': 32,
+    'resnet50': 32,
+    'resnet101': 32,
+}
+
+
+class HOWNet(nn.Module):
+    """Network for the HOW method
+
+    :param list features: A list of torch.nn.Module which act as feature extractor
+    :param torch.nn.Module attention: Attention layer
+    :param torch.nn.Module smoothing: Smoothing layer
+    :param torch.nn.Module dim_reduction: Dimensionality reduction layer
+    :param dict meta: Metadata that are stored with the network
+    :param dict runtime: Runtime options that can be used as default for e.g. inference
+    """
+
+    def __init__(self, features, attention, smoothing, dim_reduction, meta, runtime):
+        super().__init__()
+
+        self.features = features
+        self.attention = attention
+        self.smoothing = smoothing
+        self.dim_reduction = dim_reduction
+
+        self.meta = meta
+        self.runtime = runtime
+
+
+    def copy_excluding_dim_reduction(self):
+        """Return a copy of this network without the dim_reduction layer"""
+        meta = {**self.meta, "outputdim": self.meta['backbone_dim']}
+        return self.__class__(self.features, self.attention, self.smoothing, None, meta, self.runtime)
+
+    def copy_with_runtime(self, runtime):
+        """Return a copy of this network with a different runtime dict"""
+        return self.__class__(self.features, self.attention, self.smoothing, self.dim_reduction, self.meta, runtime)
+
+
+    # Methods of nn.Module
+
+    @staticmethod
+    def _set_batchnorm_eval(mod):
+        if mod.__class__.__name__.find('BatchNorm') != -1:
+            # freeze running mean and std
+            mod.eval()
+
+    def train(self, mode=True):
+        res = super().train(mode)
+        if mode:
+            self.apply(HOWNet._set_batchnorm_eval)
+        return res
+
+    def parameter_groups(self, optimizer_opts):
+        """Return torch parameter groups"""
+        layers = [self.features, self.attention, self.smoothing]
+        parameters = [{'params': x.parameters()} for x in layers if x is not None]
+        if self.dim_reduction:
+            # Do not update dimensionality reduction layer
+            parameters.append({'params': self.dim_reduction.parameters(), 'lr': 0.0})
+        return parameters
+
+
+    # Forward
+
+    def features_attentions(self, x, *, scales):
+        """Return a tuple (features, attentions) where each is a list containing requested scales"""
+        feats = []
+        masks = []
+        for s in scales:
+            xs = nn.functional.interpolate(x, scale_factor=s, mode='bilinear', align_corners=False)
+            o = self.features(xs)
+            m = self.attention(o)
+            if self.smoothing:
+                o = self.smoothing(o)
+            if self.dim_reduction:
+                o = self.dim_reduction(o)
+            feats.append(o)
+            masks.append(m)
+
+        # Normalize max weight to 1
+        mx = max(x.max() for x in masks)
+        masks = [x/mx for x in masks]
+
+        return feats, masks
+
+    def forward(self, x):
+        return self.forward_global(x, scales=self.runtime['training_scales'])
+
+    def forward_global(self, x, *, scales):
+        """Return global descriptor"""
+        feats, masks = self.features_attentions(x, scales=scales)
+        return HF.weighted_spoc(feats, masks)
+
+    def forward_local(self, x, *, features_num, scales):
+        """Return local descriptors"""
+        feats, masks = self.features_attentions(x, scales=scales)
+        return HF.how_select_local(feats, masks, scales=scales, features_num=features_num)
+
+
+    # String conversion
+
+    def __repr__(self):
+        meta_str = "\n".join("    %s: %s" % x for x in self.meta.items())
+        return "%s(meta={\n%s\n})" % (self.__class__.__name__, meta_str)
+
+    def meta_repr(self):
+        """Return meta representation"""
+        return str(self)
+
+
+def init_network(architecture, pretrained, skip_layer, dim_reduction, smoothing, runtime):
+    """Initialize HOW network
+
+    :param str architecture: Network backbone architecture (e.g. resnet18)
+    :param bool pretrained: Whether to start with a network pretrained on ImageNet
+    :param int skip_layer: How many layers of blocks should be skipped (from the end)
+    :param dict dim_reduction: Options for the dimensionality reduction layer
+    :param dict smoothing: Options for the smoothing layer
+    :param dict runtime: Runtime options to be stored in the network
+    :return HOWNet: Initialized network
+    """
+    # Take convolutional layers as features, always ends with ReLU to make last activations non-negative
+    net_in = getattr(torchvision.models, architecture)(pretrained=pretrained)
+    if architecture.startswith('alexnet') or architecture.startswith('vgg'):
+        features = list(net_in.features.children())[:-1]
+    elif architecture.startswith('resnet'):
+        features = list(net_in.children())[:-2]
+    elif architecture.startswith('densenet'):
+        features = list(net_in.features.children()) + [nn.ReLU(inplace=True)]
+    elif architecture.startswith('squeezenet'):
+        features = list(net_in.features.children())
+    else:
+        raise ValueError('Unsupported or unknown architecture: {}!'.format(architecture))
+
+    if skip_layer > 0:
+        features = features[:-skip_layer]
+    backbone_dim = imageretrievalnet.OUTPUT_DIM[architecture] // (2 ** skip_layer)
+
+    att_layer = layers.attention.L2Attention()
+    smooth_layer = None
+    if smoothing:
+        smooth_layer = layers.pooling.SmoothingAvgPooling(**smoothing)
+    reduction_layer = None
+    if dim_reduction:
+        reduction_layer = layers.dim_reduction.ConvDimReduction(**dim_reduction, input_dim=backbone_dim)
+
+    meta = {
+        "architecture": architecture,
+        "backbone_dim": backbone_dim,
+        "outputdim": reduction_layer.out_channels if dim_reduction else backbone_dim,
+        "corercf_size": CORERCF_SIZE[architecture] // (2 ** skip_layer),
+    }
+    return HOWNet(nn.Sequential(*features), att_layer, smooth_layer, reduction_layer, meta, runtime)
+
+
+def extract_vectors(net, dataset, device, *, scales):
+    """Return global descriptors in torch.Tensor"""
+    net.eval()
+    loader = torch.utils.data.DataLoader(dataset, shuffle=False, pin_memory=True, num_workers=NUM_WORKERS)
+
+    with torch.no_grad():
+        vecs = torch.zeros(len(loader), net.meta['outputdim'])
+        for i, inp in io_helpers.progress(enumerate(loader), size=len(loader), print_freq=100):
+            vecs[i] = net.forward_global(inp.to(device), scales=scales).cpu().squeeze()
+
+    return vecs
+
+
+def extract_vectors_local(net, dataset, device, *, features_num, scales):
+    """Return tuple (local descriptors, image ids, strenghts, locations and scales) where locations
+        consists of (coor_x, coor_y, scale) and elements of each list correspond to each other"""
+    net.eval()
+    loader = torch.utils.data.DataLoader(dataset, shuffle=False, pin_memory=True, num_workers=NUM_WORKERS)
+
+    with torch.no_grad():
+        vecs, strengths, locs, scls, imids = [], [], [], [], []
+        for imid, inp in io_helpers.progress(enumerate(loader), size=len(loader), print_freq=100):
+            output = net.forward_local(inp.to(device), features_num=features_num, scales=scales)
+
+            vecs.append(output[0].cpu().numpy())
+            strengths.append(output[1].cpu().numpy())
+            locs.append(output[2].cpu().numpy())
+            scls.append(output[3].cpu().numpy())
+            imids.append(np.full((output[0].shape[0],), imid))
+
+    return np.vstack(vecs), np.hstack(imids), np.hstack(strengths), np.vstack(locs), np.hstack(scls)
+
+
+
+def extract_vectors_all(net, dataset, device, *, features_num, scales):
+    """Return tuple (local descriptors, image ids, strenghts, locations and scales) where locations
+        consists of (coor_x, coor_y, scale) and elements of each list correspond to each other"""
+    net.eval()
+    loader = torch.utils.data.DataLoader(dataset, shuffle=False, pin_memory=True, num_workers=NUM_WORKERS)
+
+    with torch.no_grad():
+        feats, attns, strenghts = [], [], []
+        for imid, inp in io_helpers.progress(enumerate(loader), size=len(loader), print_freq=100):
+            output = net.get_superfeatures(inp.to(device), scales=scales)
+            feats.append(output[0])
+            attns.append(output[1])
+            strenghts.append(output[2])
+
+    return feats, attns, strenghts

how/stages/__init__.py

@@ -0,0 +1,5 @@
+"""
+Implementation of different network stages, such as training and evaluation
+"""
+
+from . import evaluate, train

how/stages/evaluate.py

@@ -0,0 +1,314 @@
+"""Implements evaluation of trained models"""
+
+import time
+import warnings
+from pathlib import Path
+import pickle
+import numpy as np
+import torch
+from torchvision import transforms
+from PIL import ImageFile
+
+from cirtorch.datasets.genericdataset import ImagesFromList
+
+from asmk import asmk_method, kernel as kern_pkg
+from ..networks import how_net
+from ..utils import score_helpers, data_helpers, logging
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+warnings.filterwarnings("ignore", r"^Possibly corrupt EXIF data", category=UserWarning)
+
+
+def evaluate_demo(demo_eval, evaluation, globals):
+    """Demo evaluating a trained network
+
+    :param dict demo_eval: Demo-related options
+    :param dict evaluation: Evaluation-related options
+    :param dict globals: Global options
+    """
+    globals["device"] = torch.device("cpu")
+    if demo_eval['gpu_id'] is not None:
+        globals["device"] = torch.device(("cuda:%s" % demo_eval['gpu_id']))
+
+    # Handle net_path when directory
+    net_path = Path(demo_eval['exp_folder']) / demo_eval['net_path']
+    if net_path.is_dir() and (net_path / "epochs/model_best.pth").exists():
+        net_path = net_path / "epochs/model_best.pth"
+
+    # Load net
+    state = _convert_checkpoint(torch.load(net_path, map_location='cpu'))
+    net = how_net.init_network(**state['net_params']).to(globals['device'])
+    net.load_state_dict(state['state_dict'])
+    globals["transform"] = transforms.Compose([transforms.ToTensor(), \
+                transforms.Normalize(**dict(zip(["mean", "std"], net.runtime['mean_std'])))])
+
+    # Eval
+    if evaluation['global_descriptor']['datasets']:
+        eval_global(net, evaluation['inference'], globals, **evaluation['global_descriptor'])
+
+    if evaluation['multistep']:
+        eval_asmk_multistep(net, evaluation['inference'], evaluation['multistep'], globals, **evaluation['local_descriptor'])
+    elif evaluation['local_descriptor']['datasets']:
+        eval_asmk(net, evaluation['inference'], globals, **evaluation['local_descriptor'])
+
+
+def eval_global(net, inference, globals, *, datasets):
+    """Evaluate global descriptors"""
+    net.eval()
+    time0 = time.time()
+    logger = globals["logger"]
+    logger.info("Starting global evaluation")
+
+    results = {}
+    for dataset in datasets:
+        images, qimages, bbxs, gnd = data_helpers.load_dataset(dataset, data_root=globals['root_path'])
+        logger.info(f"Evaluating {dataset}")
+
+        with logging.LoggingStopwatch("extracting database images", logger.info, logger.debug):
+            dset = ImagesFromList(root='', images=images, imsize=inference['image_size'], bbxs=None,
+                                  transform=globals['transform'])
+            vecs = how_net.extract_vectors(net, dset, globals["device"], scales=inference['scales'])
+        with logging.LoggingStopwatch("extracting query images", logger.info, logger.debug):
+            qdset = ImagesFromList(root='', images=qimages, imsize=inference['image_size'], bbxs=bbxs,
+                                   transform=globals['transform'])
+            qvecs = how_net.extract_vectors(net, qdset, globals["device"], scales=inference['scales'])
+
+        vecs, qvecs = vecs.numpy(), qvecs.numpy()
+        ranks = np.argsort(-np.dot(vecs, qvecs.T), axis=0)
+        results[dataset] = score_helpers.compute_map_and_log(dataset, ranks, gnd, logger=logger)
+
+    logger.info(f"Finished global evaluation in {int(time.time()-time0) // 60} min")
+    return results
+
+
+def eval_asmk(net, inference, globals, *, datasets, codebook_training, asmk):
+    """Evaluate local descriptors with ASMK"""
+    net.eval()
+    time0 = time.time()
+    logger = globals["logger"]
+    logger.info("Starting asmk evaluation")
+
+    asmk = asmk_method.ASMKMethod.initialize_untrained(asmk)
+    asmk = asmk_train_codebook(net, inference, globals, logger, codebook_training=codebook_training,
+                               asmk=asmk, cache_path=None)
+
+    results = {}
+    for dataset in datasets:
+        dataset_name = dataset if isinstance(dataset, str) else dataset['name']
+        images, qimages, bbxs, gnd = data_helpers.load_dataset(dataset, data_root=globals['root_path'])
+        logger.info(f"Evaluating '{dataset_name}'")
+
+        asmk_dataset = asmk_index_database(net, inference, globals, logger, asmk=asmk, images=images)
+        asmk_query_ivf(net, inference, globals, logger, dataset=dataset, asmk_dataset=asmk_dataset,
+                       qimages=qimages, bbxs=bbxs, gnd=gnd, results=results,
+                       cache_path=globals["exp_path"] / "query_results.pkl")
+
+    logger.info(f"Finished asmk evaluation in {int(time.time()-time0) // 60} min")
+    return results
+
+
+def eval_asmk_multistep(net, inference, multistep, globals, *, datasets, codebook_training, asmk):
+    """Evaluate local descriptors with ASMK"""
+    valid_steps = ["train_codebook", "aggregate_database", "build_ivf", "query_ivf", "aggregate_build_query"]
+    assert multistep['step'] in valid_steps, multistep['step']
+
+    net.eval()
+    time0 = time.time()
+    logger = globals["logger"]
+    (globals["exp_path"] / "eval").mkdir(exist_ok=True)
+    logger.info(f"Starting asmk evaluation step '{multistep['step']}'")
+
+    # Handle partitioning
+    partition = {"suffix": "", "norm_start": 0, "norm_end": 1}
+    if multistep.get("partition"):
+        total, index = multistep['partition']
+        partition = {"suffix": f":{total}_{str(index).zfill(len(str(total-1)))}",
+                     "norm_start": index / total,
+                     "norm_end": (index+1) / total}
+        if multistep['step'] == "aggregate_database" or multistep['step'] == "query_ivf":
+            logger.info(f"Processing partition '{total}_{index}'")
+
+    # Handle distractors
+    distractors_path = None
+    distractors = multistep.get("distractors")
+    if distractors:
+        distractors_path = globals["exp_path"] / f"eval/{distractors}.ivf.pkl"
+
+    # Train codebook
+    asmk = asmk_method.ASMKMethod.initialize_untrained(asmk)
+    cdb_path = globals["exp_path"] / "eval/codebook.pkl"
+    if multistep['step'] == "train_codebook":
+        asmk_train_codebook(net, inference, globals, logger, codebook_training=codebook_training,
+                            asmk=asmk, cache_path=cdb_path)
+        return None
+
+    asmk = asmk.train_codebook(None, cache_path=cdb_path)
+
+    results = {}
+    for dataset in datasets:
+        dataset_name = database_name = dataset if isinstance(dataset, str) else dataset['name']
+        if distractors and multistep['step'] != "aggregate_database":
+            dataset_name = f"{distractors}_{database_name}"
+        images, qimages, bbxs, gnd = data_helpers.load_dataset(dataset, data_root=globals['root_path'])
+        logger.info(f"Processing dataset '{dataset_name}'")
+
+        # Infer database
+        if multistep['step'] == "aggregate_database":
+            agg_path = globals["exp_path"] / f"eval/{database_name}.agg{partition['suffix']}.pkl"
+            asmk_aggregate_database(net, inference, globals, logger, asmk=asmk, images=images,
+                                    partition=partition, cache_path=agg_path)
+
+        # Build ivf
+        elif multistep['step'] == "build_ivf":
+            ivf_path = globals["exp_path"] / f"eval/{dataset_name}.ivf.pkl"
+            asmk_build_ivf(globals, logger, asmk=asmk, cache_path=ivf_path, database_name=database_name,
+                           distractors=distractors, distractors_path=distractors_path)
+
+        # Query ivf
+        elif multistep['step'] == "query_ivf":
+            asmk_dataset = asmk.build_ivf(None, None, cache_path=globals["exp_path"] / f"eval/{dataset_name}.ivf.pkl")
+            start, end = int(len(qimages)*partition['norm_start']), int(len(qimages)*partition['norm_end'])
+            bbxs = bbxs[start:end] if bbxs is not None else None
+            results_path = globals["exp_path"] / f"eval/{dataset_name}.results{partition['suffix']}.pkl"
+            asmk_query_ivf(net, inference, globals, logger, dataset=dataset, asmk_dataset=asmk_dataset,
+                           qimages=qimages[start:end], bbxs=bbxs, gnd=gnd, results=results,
+                           cache_path=results_path, imid_offset=start)
+
+        # All 3 dataset steps
+        elif multistep['step'] == "aggregate_build_query":
+            if multistep.get("partition"):
+                raise NotImplementedError("Partitions within step 'aggregate_build_query' are not" \
+                                          " supported, use separate steps")
+            results_path = globals["exp_path"] / "query_results.pkl"
+            if gnd is None and results_path.exists():
+                logger.debug("Step results already exist")
+                continue
+            asmk_dataset = asmk_index_database(net, inference, globals, logger, asmk=asmk, images=images,
+                                               distractors_path=distractors_path)
+            asmk_query_ivf(net, inference, globals, logger, dataset=dataset, asmk_dataset=asmk_dataset,
+                           qimages=qimages, bbxs=bbxs, gnd=gnd, results=results, cache_path=results_path)
+
+    logger.info(f"Finished asmk evaluation step '{multistep['step']}' in {int(time.time()-time0) // 60} min")
+    return results
+
+#
+# Separate steps
+#
+
+def asmk_train_codebook(net, inference, globals, logger, *, codebook_training, asmk, cache_path):
+    """Asmk evaluation step 'train_codebook'"""
+    if cache_path and cache_path.exists():
+        return asmk.train_codebook(None, cache_path=cache_path)
+
+    images = data_helpers.load_dataset('train', data_root=globals['root_path'])[0]
+    images = images[:codebook_training['images']]
+    dset = ImagesFromList(root='', images=images, imsize=inference['image_size'], bbxs=None,
+                          transform=globals['transform'])
+    infer_opts = {"scales": codebook_training['scales'], "features_num": inference['features_num']}
+    des_train = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts)[0]
+    asmk = asmk.train_codebook(des_train, cache_path=cache_path)
+    logger.info(f"Codebook trained in {asmk.metadata['train_codebook']['train_time']:.1f}s")
+    return asmk
+
+def asmk_aggregate_database(net, inference, globals, logger, *, asmk, images, partition, cache_path):
+    """Asmk evaluation step 'aggregate_database'"""
+    if cache_path.exists():
+        logger.debug("Step results already exist")
+        return
+    codebook = asmk.codebook
+    kernel = kern_pkg.ASMKKernel(codebook, **asmk.params['build_ivf']['kernel'])
+    start, end = int(len(images)*partition['norm_start']), int(len(images)*partition['norm_end'])
+    data_opts = {"imsize": inference['image_size'], "transform": globals['transform']}
+    infer_opts = {"scales": inference['scales'], "features_num": inference['features_num']}
+    # Aggregate database
+    dset = ImagesFromList(root='', images=images[start:end], bbxs=None, **data_opts)
+    vecs, imids, *_ = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts)
+    imids += start
+    quantized = codebook.quantize(vecs, imids, **asmk.params["build_ivf"]["quantize"])
+    aggregated = kernel.aggregate(*quantized, **asmk.params["build_ivf"]["aggregate"])
+    with cache_path.open("wb") as handle:
+        pickle.dump(dict(zip(["des", "word_ids", "image_ids"], aggregated)), handle)
+
+def asmk_build_ivf(globals, logger, *, asmk, cache_path, database_name, distractors, distractors_path):
+    """Asmk evaluation step 'build_ivf'"""
+    if cache_path.exists():
+        logger.debug("Step results already exist")
+        return asmk.build_ivf(None, None, cache_path=cache_path)
+    builder = asmk.create_ivf_builder(cache_path=cache_path)
+    # Build ivf
+    if not builder.loaded_from_cache:
+        if distractors:
+            builder.initialize_with_distractors(distractors_path)
+            logger.debug(f"Loaded ivf with distractors '{distractors}'")
+        for path in sorted(globals["exp_path"].glob(f"eval/{database_name}.agg*.pkl")):
+            with path.open("rb") as handle:
+                des = pickle.load(handle)
+            builder.ivf.add(des['des'], des['word_ids'], des['image_ids'])
+            logger.info(f"Indexed '{path.name}'")
+    asmk_dataset = asmk.add_ivf_builder(builder)
+    logger.debug(f"IVF stats: {asmk_dataset.metadata['build_ivf']['ivf_stats']}")
+    return asmk_dataset
+
+def asmk_index_database(net, inference, globals, logger, *, asmk, images, distractors_path=None):
+    """Asmk evaluation step 'aggregate_database' and 'build_ivf'"""
+    data_opts = {"imsize": inference['image_size'], "transform": globals['transform']}
+    infer_opts = {"scales": inference['scales'], "features_num": inference['features_num']}
+    # Index database vectors
+    dset = ImagesFromList(root='', images=images, bbxs=None, **data_opts)
+    vecs, imids, *_ = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts)
+    asmk_dataset = asmk.build_ivf(vecs, imids, distractors_path=distractors_path)
+    logger.info(f"Indexed images in {asmk_dataset.metadata['build_ivf']['index_time']:.2f}s")
+    logger.debug(f"IVF stats: {asmk_dataset.metadata['build_ivf']['ivf_stats']}")
+    return asmk_dataset
+
+def asmk_query_ivf(net, inference, globals, logger, *, dataset, asmk_dataset, qimages, bbxs, gnd,
+                   results, cache_path, imid_offset=0):
+    """Asmk evaluation step 'query_ivf'"""
+    if gnd is None and cache_path and cache_path.exists():
+        logger.debug("Step results already exist")
+        return
+    data_opts = {"imsize": inference['image_size'], "transform": globals['transform']}
+    infer_opts = {"scales": inference['scales'], "features_num": inference['features_num']}
+    # Query vectors
+    qdset = ImagesFromList(root='', images=qimages, bbxs=bbxs, **data_opts)
+    qvecs, qimids, *_ = how_net.extract_vectors_local(net, qdset, globals["device"], **infer_opts)
+    qimids += imid_offset
+    metadata, query_ids, ranks, scores = asmk_dataset.query_ivf(qvecs, qimids)
+    logger.debug(f"Average query time (quant+aggr+search) is {metadata['query_avg_time']:.3f}s")
+    # Evaluate
+    if gnd is not None:
+        results[dataset] = score_helpers.compute_map_and_log(dataset, ranks.T, gnd, logger=logger)
+    with cache_path.open("wb") as handle:
+        pickle.dump({"metadata": metadata, "query_ids": query_ids, "ranks": ranks, "scores": scores}, handle)
+
+#
+# Helpers
+#
+
+def _convert_checkpoint(state):
+    """Enable loading checkpoints in the old format"""
+    if "_version" not in state:
+        # Old checkpoint format
+        meta = state['meta']
+        state['net_params'] = {
+            "architecture": meta['architecture'],
+            "pretrained": True,
+            "skip_layer": meta['skip_layer'],
+            "dim_reduction": {"dim": meta["dim"]},
+            "smoothing": {"kernel_size": meta["feat_pool_k"]},
+            "runtime": {
+                "mean_std": [meta['mean'], meta['std']],
+                "image_size": 1024,
+                "features_num": 1000,
+                "scales": [2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25],
+                "training_scales": [1],
+            },
+        }
+
+        state_dict = state['state_dict']
+        state_dict['dim_reduction.weight'] = state_dict.pop("whiten.weight")
+        state_dict['dim_reduction.bias'] = state_dict.pop("whiten.bias")
+
+        state['_version'] = "how/2020"
+
+    return state

how/stages/train.py

@@ -0,0 +1,241 @@
+"""Implements training new models"""
+
+import time
+import copy
+from collections import defaultdict
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+
+from cirtorch.layers.loss import ContrastiveLoss
+from cirtorch.datasets.datahelpers import collate_tuples
+from cirtorch.datasets.traindataset import TuplesDataset
+from cirtorch.datasets.genericdataset import ImagesFromList
+
+from ..networks import how_net
+from ..utils import data_helpers, io_helpers, logging, plots
+from . import evaluate
+
+
+def train(demo_train, training, validation, model, globals):
+    """Demo training a network
+
+    :param dict demo_train: Demo-related options
+    :param dict training: Training options
+    :param dict validation: Validation options
+    :param dict model: Model options
+    :param dict globals: Global options
+    """
+    logger = globals["logger"]
+    (globals["exp_path"] / "epochs").mkdir(exist_ok=True)
+    if (globals["exp_path"] / f"epochs/model_epoch{training['epochs']}.pth").exists():
+        logger.info("Skipping network training, already trained")
+        return
+
+    # Global setup
+    set_seed(0)
+    globals["device"] = torch.device("cpu")
+    if demo_train['gpu_id'] is not None:
+        globals["device"] = torch.device(("cuda:%s" % demo_train['gpu_id']))
+
+    # Initialize network
+    net = how_net.init_network(**model).to(globals["device"])
+    globals["transform"] = transforms.Compose([transforms.ToTensor(), \
+                transforms.Normalize(**dict(zip(["mean", "std"], net.runtime['mean_std'])))])
+    with logging.LoggingStopwatch("initializing network whitening", logger.info, logger.debug):
+        initialize_dim_reduction(net, globals, **training['initialize_dim_reduction'])
+
+    # Initialize training
+    optimizer, scheduler, criterion, train_loader = \
+            initialize_training(net.parameter_groups(training["optimizer"]), training, globals)
+    validation = Validation(validation, globals)
+
+    for epoch in range(training['epochs']):
+        epoch1 = epoch + 1
+        set_seed(epoch1)
+
+        time0 = time.time()
+        train_loss = train_epoch(train_loader, net, globals, criterion, optimizer, epoch1)
+
+        validation.add_train_loss(train_loss, epoch1)
+        validation.validate(net, epoch1)
+
+        scheduler.step()
+
+        io_helpers.save_checkpoint({
+            'epoch': epoch1, 'meta': net.meta, 'state_dict': net.state_dict(),
+            'optimizer' : optimizer.state_dict(), 'best_score': validation.best_score[1],
+            'scores': validation.scores, 'net_params': model, '_version': 'how/2020',
+        }, validation.best_score[0] == epoch1, epoch1 == training['epochs'], globals["exp_path"] / "epochs")
+
+        logger.info(f"Epoch {epoch1} finished in {time.time() - time0:.1f}s")
+
+
+def train_epoch(train_loader, net, globals, criterion, optimizer, epoch1):
+    """Train for one epoch"""
+    logger = globals['logger']
+    batch_time = data_helpers.AverageMeter()
+    data_time = data_helpers.AverageMeter()
+    losses = data_helpers.AverageMeter()
+
+    # Prepare epoch
+    train_loader.dataset.create_epoch_tuples(net)
+    net.train()
+
+    end = time.time()
+    for i, (input, target) in enumerate(train_loader):
+        data_time.update(time.time() - end)
+        optimizer.zero_grad()
+
+        num_images = len(input[0]) # number of images per tuple
+        for inp, trg in zip(input, target):
+            output = torch.zeros(net.meta['outputdim'], num_images).to(globals["device"])
+            for imi in range(num_images):
+                output[:, imi] = net(inp[imi].to(globals["device"])).squeeze()
+            loss = criterion(output, trg.to(globals["device"]))
+            loss.backward()
+            losses.update(loss.item())
+
+        optimizer.step()
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if (i+1) % 20 == 0 or i == 0 or (i+1) == len(train_loader):
+            logger.info(f'>> Train: [{epoch1}][{i+1}/{len(train_loader)}]\t' \
+                        f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
+                        f'Data {data_time.val:.3f} ({data_time.avg:.3f})\t' \
+                        f'Loss {losses.val:.4f} ({losses.avg:.4f})')
+
+    return losses.avg
+
+
+def set_seed(seed):
+    """Sets given seed globally in used libraries"""
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+
+
+def initialize_training(net_parameters, training, globals):
+    """Initialize classes necessary for training"""
+    # Need to check for keys because of defaults
+    assert training['optimizer'].keys() == {"lr", "weight_decay"}
+    assert training['lr_scheduler'].keys() == {"gamma"}
+    assert training['loss'].keys() == {"margin"}
+    assert training['dataset'].keys() == {"name", "mode", "imsize", "nnum", "qsize", "poolsize"}
+    assert training['loader'].keys() == {"batch_size"}
+
+    optimizer = torch.optim.Adam(net_parameters, **training["optimizer"])
+    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, **training["lr_scheduler"])
+    criterion = ContrastiveLoss(**training["loss"]).to(globals["device"])
+    train_dataset = TuplesDataset(**training['dataset'], transform=globals["transform"])
+    train_loader = torch.utils.data.DataLoader(train_dataset, **training['loader'], \
+            pin_memory=True, drop_last=True, shuffle=True, collate_fn=collate_tuples, \
+            num_workers=how_net.NUM_WORKERS)
+    return optimizer, scheduler, criterion, train_loader
+
+
+
+def extract_train_descriptors(net, globals, *, images, features_num):
+    """Extract descriptors for a given number of images from the train set"""
+    if features_num is None:
+        features_num = net.runtime['features_num']
+
+    images = data_helpers.load_dataset('train', data_root=globals['root_path'])[0][:images]
+    dataset = ImagesFromList(root='', images=images, imsize=net.runtime['image_size'], bbxs=None,
+                             transform=globals["transform"])
+    des_train = how_net.extract_vectors_local(net, dataset, globals["device"],
+                                              scales=net.runtime['training_scales'],
+                                              features_num=features_num)[0]
+    return des_train
+
+
+def initialize_dim_reduction(net, globals, **kwargs):
+    """Initialize dimensionality reduction by PCA whitening from 'images' number of descriptors"""
+    if not net.dim_reduction:
+        return
+
+    print(">> Initializing dim reduction")
+    des_train = extract_train_descriptors(net.copy_excluding_dim_reduction(), globals, **kwargs)
+    net.dim_reduction.initialize_pca_whitening(des_train)
+
+
+class Validation:
+    """A convenient interface to validation, keeping historical values and plotting continuously
+
+    :param dict validations: Options for each validation type (e.g. local_descriptor)
+    :param dict globals: Global options
+    """
+
+    methods = {
+        "global_descriptor": evaluate.eval_global,
+        "local_descriptor": evaluate.eval_asmk,
+    }
+
+    def __init__(self, validations, globals):
+        validations = copy.deepcopy(validations)
+        self.frequencies = {x: y.pop("frequency") for x, y in validations.items()}
+        self.validations = validations
+        self.globals = globals
+        self.scores = {x: defaultdict(list) for x in validations}
+        self.scores["train_loss"] = []
+
+    def add_train_loss(self, loss, epoch):
+        """Store training loss for given epoch"""
+        self.scores['train_loss'].append((epoch, loss))
+
+        fig = plots.EpochFigure("train set", ylabel="loss")
+        fig.plot(*list(zip(*self.scores["train_loss"])), 'o-', label='train')
+        fig.save(self.globals['exp_path'] / "fig_train.jpg")
+
+    def validate(self, net, epoch):
+        """Perform validation of the network and store the resulting score for given epoch"""
+        for name, frequency in self.frequencies.items():
+            if frequency and epoch % frequency == 0:
+                scores = self.methods[name](net, net.runtime, self.globals, **self.validations[name])
+                for dataset, values in scores.items():
+                    value = values['map_medium'] if "map_medium" in values else values['map']
+                    self.scores[name][dataset].append((epoch, value))
+
+                if "val_eccv20" in scores:
+                    fig = plots.EpochFigure(f"val set - {name}", ylabel="mAP")
+                    fig.plot(*list(zip(*self.scores[name]['val_eccv20'])), 'o-', label='val')
+                    fig.save(self.globals['exp_path'] / f"fig_val_{name}.jpg")
+
+                if scores.keys() - {"val_eccv20"}:
+                    fig = plots.EpochFigure(f"test set - {name}", ylabel="mAP")
+                    for dataset, value in self.scores[name].items():
+                        if dataset != "val_eccv20":
+                            fig.plot(*list(zip(*value)), 'o-', label=dataset)
+                    fig.save(self.globals['exp_path'] / f"fig_test_{name}.jpg")
+
+    @property
+    def decisive_scores(self):
+        """List of pairs (epoch, score) where score is decisive for comparing epochs"""
+        for name in ["local_descriptor", "global_descriptor"]:
+            if self.frequencies[name] and "val_eccv20" in self.scores[name]:
+                return self.scores[name]['val_eccv20']
+        return self.scores["train_loss"]
+
+    @property
+    def last_epoch(self):
+        """Tuple (last epoch, last score) or (None, None) before decisive score is computed"""
+        decisive_scores = self.decisive_scores
+        if not decisive_scores:
+            return None, None
+
+        return decisive_scores[-1]
+
+    @property
+    def best_score(self):
+        """Tuple (best epoch, best score) or (None, None) before decisive score is computed"""
+        decisive_scores = self.decisive_scores
+        if not decisive_scores:
+            return None, None
+
+        aggr = min
+        for name in ["local_descriptor", "global_descriptor"]:
+            if self.frequencies[name] and "val_eccv20" in self.scores[name]:
+                aggr = max
+        return aggr(decisive_scores, key=lambda x: x[1])

how/utils/__init__.py

@@ -0,0 +1,3 @@
+"""
+Standalone utilities, mainly helper functions
+"""

how/utils/data_helpers.py

@@ -0,0 +1,90 @@
+"""Data manipulation helpers"""
+
+import os.path
+import pickle
+
+from cirtorch.datasets.datahelpers import cid2filename
+from cirtorch.datasets.testdataset import configdataset
+
+
+def load_dataset(dataset, data_root=''):
+    """Return tuple (image list, query list, bounding boxes, gnd dictionary)"""
+
+    if isinstance(dataset, dict):
+        root = os.path.join(data_root, dataset['image_root'])
+        images, qimages = None, None
+        if dataset['database_list'] is not None:
+            images = [path_join(root, x.strip("\n")) for x in open(dataset['database_list']).readlines()]
+        if dataset['query_list'] is not None:
+            qimages = [path_join(root, x.strip("\n")) for x in open(dataset['query_list']).readlines()]
+        bbxs = None
+        gnd = None
+
+    elif dataset == 'train':
+        training_set = 'retrieval-SfM-120k'
+        db_root = os.path.join(data_root, 'train', training_set)
+        ims_root = os.path.join(db_root, 'ims')
+        db_fn = os.path.join(db_root, '{}.pkl'.format(training_set))
+        with open(db_fn, 'rb') as f:
+            db = pickle.load(f)['train']
+        images = [cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids']))]
+        qimages = []
+        bbxs = None
+        gnd = None
+
+    elif dataset == 'val_eccv20':
+        db_root = os.path.join(data_root, 'train', 'retrieval-SfM-120k')
+        fn_val_proper = db_root+'/retrieval-SfM-120k-val-eccv2020.pkl' # pos are all with #inl >=3 & <= 10
+        with open(fn_val_proper, 'rb') as f:
+            db = pickle.load(f)
+        ims_root = os.path.join(db_root, 'ims')
+        images = [cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids']))]
+        gnd = db['gnd']
+        qidx = db['qidx']
+        qimages = [images[x] for x in qidx]
+        bbxs = None
+
+    elif "/" in dataset:
+        with open(dataset, 'rb') as handle:
+            db = pickle.load(handle)
+        images, qimages, bbxs, gnd = db['imlist'], db['qimlist'], None, db['gnd']
+
+    else:
+        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'])]
+        if 'bbx' in cfg['gnd'][0].keys():
+            bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
+        else:
+            bbxs = None
+        gnd = cfg['gnd']
+
+    return images, qimages, bbxs, gnd
+
+
+def path_join(root, name):
+    """Perform os.path.join by default; if asterisk is present in root, substitute with the name.
+
+    >>> path_join('/data/img_*.jpg', '001')
+    '/data/img_001.jpg'
+    """
+    if "*" in root.rsplit("/", 1)[-1]:
+        return root.replace("*", name)
+    return os.path.join(root, name)
+
+
+class AverageMeter:
+    """Compute and store the average and last value"""
+
+    def __init__(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        """Update the counter by a new value"""
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count

how/utils/download.py

@@ -0,0 +1,44 @@
+"""Functions for downloading files necessary for training and evaluation"""
+
+import os.path
+from cirtorch.utils.download import download_train, download_test
+from . import io_helpers
+
+
+def download_for_eval(evaluation, demo_eval, dataset_url, globals):
+    """Download datasets for evaluation and network if given by url"""
+    # Datasets
+    datasets = evaluation['global_descriptor']['datasets'] \
+                + evaluation['local_descriptor']['datasets']
+    download_datasets(datasets, dataset_url, globals)
+    # Network
+    if demo_eval and (demo_eval['net_path'].startswith("http://") \
+                        or demo_eval['net_path'].startswith("https://")):
+        net_name = os.path.basename(demo_eval['net_path'])
+        io_helpers.download_files([net_name], globals['root_path'] / "models",
+                                  os.path.dirname(demo_eval['net_path']) + "/",
+                                  logfunc=globals["logger"].info)
+        demo_eval['net_path'] = globals['root_path'] / "models" / net_name
+
+
+def download_for_train(validation, dataset_url, globals):
+    """Download datasets for training"""
+
+    datasets = ["train"] + validation['global_descriptor']['datasets'] \
+                + validation['local_descriptor']['datasets']
+    download_datasets(datasets, dataset_url, globals)
+
+
+def download_datasets(datasets, dataset_url, globals):
+    """Download data associated with each required dataset"""
+
+    if "val_eccv20" in datasets:
+        download_train(globals['root_path'])
+        io_helpers.download_files(["retrieval-SfM-120k-val-eccv2020.pkl"],
+                                  globals['root_path'] / "train/retrieval-SfM-120k",
+                                  dataset_url, logfunc=globals["logger"].info)
+    elif "train" in datasets:
+        download_train(globals['root_path'])
+
+    if "roxford5k" in datasets or "rparis6k" in datasets:
+        download_test(globals['root_path'])

how/utils/html.py

@@ -0,0 +1,252 @@
+from __future__ import print_function
+import os
+from PIL import Image
+
+# see help for ocmmon HTML tags at http://www.mountaindragon.com/html/text.htm
+
+
+class Node:
+  def __init__(self,tag,text='',props=dict()):
+    self.children = []
+    self.tag=tag
+    self.text=text
+    self.props=props
+  def add(self, node):
+    self.children.append(node)
+    return node
+  def tostr(self):
+    s = ""
+    if not self.props:
+      s+= "<%s>%s"%(self.tag,self.text)
+    else:
+      s+= "<%s %s>%s"%(self.tag,' '.join(["%s='%s'"%(k,v) for k,v in self.props.items() if v!=None]),self.text)
+    for child in self.children:
+      s += child.tostr()
+    s += "</%s>"%self.tag
+    return s
+  def write(self,fout):
+    if not self.props:
+      print("<%s>%s"%(self.tag,self.text), file=fout)
+    else:
+      print("<%s %s>%s"%(self.tag,' '.join(["%s='%s'"%(k,v) for k,v in self.props.items() if v!=None]),self.text), file=fout)
+    for child in self.children:
+      child.write(fout)
+    print("</%s>"%self.tag, file=fout)
+  def first(self,tag,order=1):
+    if self.tag==tag: return self
+    for c in self.children[::order]:
+      res = c.first(tag,order)
+      if res: return res
+    return None
+  def last(self,tag):
+    return self.first(tag,-1)
+
+
+
+class HTML (Node):
+  def __init__(self):
+    Node.__init__(self,'html')
+  def header(self,**kw):
+    return self.add(Header(**kw))
+  def body(self,**kw):
+    return self.add(BodyNode('body',props=kw))
+  def save(self,fname):
+    fout = open(fname,'w') if type(fname)==str else fname
+    for e in self.children:
+      e.write(fout)
+  def show(self,fname=''):
+    if not fname:  fname = os.tmpname()+'.html' 
+    self.save(fname)
+    os.system('/opt/google/chrome/google-chrome '+fname)
+
+
+
+class Header (HTML):
+  def __init__(self, **kw):
+    Node.__init__(self,'header',props=kw)
+  def title(self,text):
+    return self.add(Node('title',text=text))
+  def script(self, text="", **kw):
+    return self.add(Node('script',text=text, props=kw))
+  def link(self, **kw):
+    return self.add(Node('link', props=kw))
+  def meta(self):
+    return self.add(Node('meta', props={"http-equiv":"Content-Type", "content": "charset=iso-8859-1"}))
+
+class BodyNode (Node):
+  # title of section
+  def h(self, strength, text='', **kw):
+    return self.add(BodyNode('h%d'%strength, text=text, props=kw))
+  # paragraph
+  def p(self, text='', **kw):
+    return self.add(BodyNode('p',text=text, props=kw))
+  # bold
+  def bold(self, text='', **kw):
+    return self.add(BodyNode('b',text=text, props=kw))
+  def b(self, text='', **kw):
+    return self.add(BodyNode('b',text=text, props=kw))
+  # italic
+  def italic(self, text='', **kw):
+    return self.add(BodyNode('i',text=text, props=kw))
+  def i(self, text='', **kw):
+    return self.add(BodyNode('i',text=text, props=kw))
+  # span/text
+  def span(self, text='', **kw):
+    return self.add(BodyNode('span',text=text, props=kw))
+  # font
+  def font(self,text='',color=None,face=None,size=None):
+    return self.add(BodyNode('font',text=text, props={'color':color,'face':face,'size':size}))
+  # small
+  def small(self, text='', **kw):
+    return self.add(BodyNode('small',text=text, props=kw))
+  def big(self, text='', **kw):
+    return self.add(BodyNode('big',text=text, props=kw))
+  # centered 
+  def center(self, text='', **kw):
+    return self.add(BodyNode('center',text=text, props=kw))
+  # div
+  def div(self, text='', **kw):
+    return self.add(BodyNode('div', text=text, props=kw))
+  # unordered list
+  def unordlist(self, text='', **kw):
+    return self.add(BodyNode('ul', text=text, props=kw))
+  # ordered list
+  def ordlist(self, text='', **kw):
+    return self.add(BodyNode('ol', text=text, props=kw))
+  def item(self, text='', type=None, **kw):
+    kw['type'] = type # non-ord {'circle', 'square', 'disc'}, ord  {'1', 'A', 'a', 'I', 'i'}
+    return self.add(BodyNode('li', text=text, props=kw))
+  # line break
+  def br(self):
+    self.add(Node('br'))
+  # horizontal line
+  def hr(self):
+    self.add(Node('hr'))
+  # table
+  def table(self, **kw):
+    return self.add(Table(**kw))
+  # image
+  def image(self, img, **kw):
+    return self.add(Image(img,**kw))
+  # link
+  def a(self, href, text='', **kw):
+    kw['href'] = href
+    return self.add(BodyNode('a', text=text, props=kw))
+  def hidden(self, text, **kw):
+    kw['type'] = 'hidden'
+    kw['value'] = text
+    return self.add(BodyNode('input',props=kw))
+  def imagelink(self, img, **kw):
+    return self.add( BodyNode('a', text=Image(img,**kw).tostr(), props={"href":img}) )
+  
+class Table (Node):
+  def __init__(self,**kw):
+    Node.__init__(self,'table',props=kw)
+  def row(self,elems=[],header=False,**kw):
+    r=TableRow(header, **kw)
+    for e in elems:
+      if issubclass(e.__class__,Node):
+        r.add(e)
+      else: 
+        r.cell(str(e))
+    return self.add(r)
+  def fromlist(self, elems, header=None):
+    if header and type(header)!=bool: elems=[header]+elems; header=True
+    for row in elems:
+      self.row(row,header=header)
+      header=False  # only once
+
+
+class TableRow (Node):
+  def __init__(self, isheader=False, **kw):
+    Node.__init__(self,'tr',props=kw)
+    self.isheader=isheader
+  def cell(self, text='', **kw):
+    return self.add(BodyNode(self.isheader and 'th' or 'td',text=text,props=kw))
+
+
+
+class Image (Node):
+  def __init__(self, img, dir='', name='', width=None, height=None, alt=None):
+    if type(img)==str:  loc = img
+    else:
+      if name:
+        img.save(os.path.join(dir,name))
+        loc = name
+      else:
+        loc =  os.tmpnam()+".png"
+        img.save(loc)
+    Node.__init__(self,'img',props={'src':loc,'width':width,'height':height,'alt':alt,'title':alt})
+
+
+def htmlspace(n):
+    return "&nbsp;".join(["" for i in range(n)])
+def htmloptions(l):
+    return "".join(["<option>"+s+"</option>" for s in l])
+
+
+if __name__=='__main__':
+  import pdb
+  
+  doc = HTML()
+  doc.header().title('test of python-generated HTML page')
+  body=doc.body()
+  body.h(1,"1. Title of page")
+  body.p('a paragraph of text')
+  body.h(2,"2.1 second title")
+  p=body.p()
+  p.italic('another')
+  p.font(color='red').bold('paragraph')
+  p.span('of text')
+  body.h(3,'2.1.1. sub-sub-title')
+  body.p("Here is a list:")
+  ls=body.unordlist()
+  ls.item("first item")
+  ls.item("second item")
+  ls.item("final item")
+  body.hr()
+  body.table(border=1).fromlist([[1,2],[3,4]],header=['col1','col2'])
+  body.br()
+  body.center().image(img='/home/lear/revaud/coca-cola.jpg',width=500,height=300)
+  body.hr()
+  tab=body.table(border=0)
+  tab.row(['coca-cola']*5,header=True)
+  for i in range(3):
+    r = body.last('table').row()
+    for j in range(5):
+      r.cell(bgcolor=['#00FF00','red'][(i+j)%2]).image('/home/lear/revaud/coca-cola2.jpg',width=200)
+  
+  doc.show('/tmp/test.html')
+  print('result stored in /tmp/test.html')
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

how/utils/io_helpers.py

@@ -0,0 +1,105 @@
+"""Helper functions related to io"""
+
+import os.path
+import sys
+import shutil
+import urllib.request
+from pathlib import Path
+import yaml
+import torch
+
+
+def progress(iterable, *, size=None, print_freq=1, handle=sys.stdout):
+    """Generator wrapping an iterable to print progress"""
+    for i, element in enumerate(iterable):
+        yield element
+
+        if i == 0 or (i+1) % print_freq == 0 or (i+1) == size:
+            if size:
+                handle.write(f'\r>>>> {i+1}/{size} done...')
+            else:
+                handle.write(f'\r>>>> {i+1} done...')
+
+    handle.write("\n")
+
+
+# Params
+
+def load_params(path):
+    """Return loaded parameters from a yaml file"""
+    with open(path, "r") as handle:
+        content = yaml.safe_load(handle)
+    return load_nested_templates(content, os.path.dirname(path))
+
+def save_params(path, params):
+    """Save given parameters to a yaml file"""
+    with open(path, "w") as handle:
+        yaml.safe_dump(params, handle, default_flow_style=False)
+
+def load_nested_templates(params, root_path):
+    """Find keys '__template__' in nested dictionary and replace corresponding value with loaded
+        yaml file"""
+    if not isinstance(params, dict):
+        return params
+
+    if "__template__" in params:
+        template_path = os.path.expanduser(params.pop("__template__"))
+        path = os.path.join(root_path, template_path)
+        root_path = os.path.dirname(path)
+        # Treat template as defaults
+        params = dict_deep_overlay(load_params(path), params)
+
+    for key, value in params.items():
+        params[key] = load_nested_templates(value, root_path)
+
+    return params
+
+def dict_deep_overlay(defaults, params):
+    """If defaults and params are both dictionaries, perform deep overlay (use params value for
+        keys defined in params), otherwise use defaults value"""
+    if isinstance(defaults, dict) and isinstance(params, dict):
+        for key in params:
+            defaults[key] = dict_deep_overlay(defaults.get(key, None), params[key])
+        return defaults
+
+    return params
+
+def dict_deep_set(dct, key, value):
+    """Set key to value for a nested dictionary where the key is a sequence (e.g. list)"""
+    if len(key) == 1:
+        dct[key[0]] = value
+        return
+
+    if not isinstance(dct[key[0]], dict) or key[0] not in dct:
+        dct[key[0]] = {}
+    dict_deep_set(dct[key[0]], key[1:], value)
+
+
+# Download
+
+def download_files(names, root_path, base_url, logfunc=None):
+    """Download file names from given url to given directory path. If logfunc given, use it to log
+        status."""
+    root_path = Path(root_path)
+    for name in names:
+        path = root_path / name
+        if path.exists():
+            continue
+        if logfunc:
+            logfunc(f"Downloading file '{name}'")
+        path.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(base_url + name, path)
+
+
+# Checkpoints
+
+def save_checkpoint(state, is_best, keep_epoch, directory):
+    """Save state dictionary to the directory providing whether the corresponding epoch is the best
+        and whether to keep it anyway"""
+    filename = os.path.join(directory, 'model_epoch%d.pth' % state['epoch'])
+    filename_best = os.path.join(directory, 'model_best.pth')
+    if is_best and keep_epoch:
+        torch.save(state, filename)
+        shutil.copyfile(filename, filename_best)
+    elif is_best or keep_epoch:
+        torch.save(state, filename_best if is_best else filename)

how/utils/logging.py

@@ -0,0 +1,63 @@
+"""Logging-related functionality"""
+
+import time
+import logging
+
+# Logging
+
+def init_logger(log_path):
+    """Return a logger instance which logs to stdout and, if log_path is not None, also to a file"""
+    logger = logging.getLogger("HOW")
+    logger.setLevel(logging.DEBUG)
+
+    stdout_handler = logging.StreamHandler()
+    stdout_handler.setLevel(logging.INFO)
+    stdout_handler.setFormatter(logging.Formatter('%(name)s %(levelname)s: %(message)s'))
+    logger.addHandler(stdout_handler)
+
+    if log_path:
+        file_handler = logging.FileHandler(log_path)
+        file_handler.setLevel(logging.DEBUG)
+        formatter = logging.Formatter('%(asctime)s %(name)s %(levelname)s: %(message)s')
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+    return logger
+
+
+# Stopwatch
+
+class LoggingStopwatch:
+    """Stopwatch context that produces one message when entered and another one when exited,
+        with the time spent in the context embedded in the exiting message.
+
+    :param str message: Message to be logged at the start and finish. If the first word
+            of the message ends with 'ing', convert to passive for finish message.
+    :param callable log_start: Will be called with given message at the start
+    :param callable log_finish: Will be called with built message at the finish. If None, use
+            log_start
+    """
+
+    def __init__(self, message, log_start, log_finish=None):
+        self.message = message
+        self.log_start = log_start
+        self.log_finish = log_finish if log_finish is not None else log_start
+        self.time0 = None
+
+    def __enter__(self):
+        self.time0 = time.time()
+        if self.log_start:
+            self.log_start(self.message.capitalize())
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Build message
+        words = self.message.split(" ")
+        secs = "%.1fs" % (time.time() - self.time0)
+        if words[0].endswith("ing"):
+            words += [words.pop(0).replace("ing", "ed"), "in", secs]
+        else:
+            words += ["(%.1f)" % secs]
+
+        # Log message
+        if self.log_finish:
+            self.log_finish(" ".join(words).capitalize())

how/utils/plots.py

@@ -0,0 +1,37 @@
+"""Plotting classes"""
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+
+class EpochFigure:
+    """Basic figure for plotting scores across epochs
+
+    :param str title: Figure title
+    :param str ylabel: Plot's y label
+    """
+
+    def __init__(self, title, *, ylabel):
+        self.fig = plt.figure()
+        self.axes = self.fig.add_subplot(1, 1, 1)
+        self.title = title
+        self.ylabel = ylabel
+
+    def __del__(self):
+        plt.close(self.fig)
+
+    def __getattr__(self, name):
+        # Delegate method calls on self.axes
+        return getattr(self.axes, name)
+
+    def save(self, path):
+        """Save figure to given path"""
+        self.axes.grid(b=True, which='major', color='k', linestyle='-')
+        self.axes.grid(b=True, which='minor', color='r', linestyle='-', alpha=0.2)
+        self.axes.minorticks_on()
+        self.axes.legend()
+        self.axes.set_xlabel('epoch')
+        self.axes.set_ylabel(self.ylabel)
+        self.axes.set_title(self.title)
+        self.fig.savefig(path)

how/utils/score_helpers.py

@@ -0,0 +1,59 @@
+"""Helper functions for computing evaluation scores"""
+
+import numpy as np
+
+from cirtorch.utils.evaluate import compute_map
+
+
+def compute_map_and_log(dataset, ranks, gnd, kappas=(1, 5, 10), logger=None):
+    """Computed mAP and log it
+
+    :param str dataset: Dataset to compute the mAP on (e.g. roxford5k)
+    :param np.ndarray ranks: 2D matrix of ints corresponding to previously computed ranks
+    :param dict gnd: Ground-truth dataset structure
+    :param list kappas: Compute mean precision at each kappa
+    :param logging.Logger logger: If not None, use it to log mAP and all mP@kappa
+    :return tuple: mAP and mP@kappa (medium difficulty for roxford5k and rparis6k)
+    """
+    # new evaluation protocol
+    if dataset.startswith('roxford5k') or dataset.startswith('rparis6k'):
+        gnd_t = []
+        for gndi in gnd:
+            g = {}
+            g['ok'] = np.concatenate([gndi['easy']])
+            g['junk'] = np.concatenate([gndi['junk'], gndi['hard']])
+            gnd_t.append(g)
+        mapE, apsE, mprE, prsE = compute_map(ranks, gnd_t, kappas)
+
+        gnd_t = []
+        for gndi in gnd:
+            g = {}
+            g['ok'] = np.concatenate([gndi['easy'], gndi['hard']])
+            g['junk'] = np.concatenate([gndi['junk']])
+            gnd_t.append(g)
+        mapM, apsM, mprM, prsM = compute_map(ranks, gnd_t, kappas)
+
+        gnd_t = []
+        for gndi in gnd:
+            g = {}
+            g['ok'] = np.concatenate([gndi['hard']])
+            g['junk'] = np.concatenate([gndi['junk'], gndi['easy']])
+            gnd_t.append(g)
+        mapH, apsH, mprH, prsH = compute_map(ranks, gnd_t, kappas)
+
+        if logger:
+            fmap = lambda x: np.around(x*100, decimals=2)
+            logger.info(f"Evaluated {dataset}: mAP E: {fmap(mapE)}, M: {fmap(mapM)}, H: {fmap(mapH)}")
+            logger.info(f"Evaluated {dataset}: mP@k{kappas} E: {fmap(mprE)}, M: {fmap(mprM)}, H: {fmap(mprH)}")
+
+        scores = {"map_easy": mapE.item(), "mp@k_easy": mprE, "ap_easy": apsE, "p@k_easy": prsE,
+                  "map_medium": mapM.item(), "mp@k_medium": mprM, "ap_medium": apsM, "p@k_medium": prsM,
+                  "map_hard": mapH.item(), "mp@k_hard": mprH, "ap_hard": apsH, "p@k_hard": prsH}
+        return scores
+
+    # old evaluation protocol
+    map_score, ap_scores, prk, pr_scores = compute_map(ranks, gnd, kappas=kappas)
+    if logger:
+        fmap = lambda x: np.around(x*100, decimals=2)
+        logger.info(f"Evaluated {dataset}: mAP {fmap(map_score)}, mP@k {fmap(prk)}")
+    return {"map": map_score, "mp@k": prk, "ap": ap_scores, "p@k": pr_scores}

how/utils/visualize.py

@@ -0,0 +1,99 @@
+import os
+import numpy as np
+import cv2
+
+
+from how.utils.html import HTML
+
+def visualize_attention_map(dataset_name, imgpaths, attentions, scales, outdir):
+  assert len(imgpaths) == len(attentions)
+  os.makedirs(outdir, exist_ok=True)
+  for i, imgpath in enumerate(imgpaths): # for each image
+    img_basename = os.path.splitext(os.path.basename(imgpath))[0]
+    atts = attentions[i]
+    # load image
+    img = cv2.imread(imgpath)
+    # generate the visu for each scale independently
+    for j,s in enumerate(scales):
+      a = atts[j]
+      img_s = cv2.resize(img, None, fx=s, fy=s)
+      heatmap_s = cv2.applyColorMap( (255*cv2.resize(a, (img_s.shape[1],img_s.shape[0]))).astype(np.uint8), cv2.COLORMAP_JET)
+      overlay = cv2.addWeighted(heatmap_s, 0.5, img_s, 0.5, 0)
+      cv2.imwrite(outdir+'{:s}_scale{:g}.jpg'.format(img_basename, s), overlay)
+    # generate the visu for the aggregation over scales
+    agg_atts = sum([cv2.resize(a, (img.shape[1],img.shape[0])) for a in atts]) / len(atts)
+    heatmap_s = cv2.applyColorMap( (255*agg_atts).astype(np.uint8), cv2.COLORMAP_JET)
+    overlay = cv2.addWeighted(heatmap_s, 0.5, img, 0.5, 0)
+    cv2.imwrite(outdir+'{:s}_aggregated.jpg'.format(img_basename), overlay)
+  # generate a html webpage for visualization
+  doc = HTML()
+  doc.header().title(dataset_name)
+  b = doc.body()
+  b.h(1, dataset_name+' (attention map)')
+  t = b.table(cellpadding=2, border=1)
+  for i, imgpath in enumerate(imgpaths):
+    img_basename = os.path.splitext(os.path.basename(imgpath))[0]
+    if i%3==0: t.row(['info','image','agg','scale 1']+['scale '+str(s) for s in scales if s!=1], header=True)
+    r = t.row()
+    r.cell(str(i)+': '+img_basename)
+    r.cell('<a href="{img:s}"><img src="{img:s}"/></a>'.format(img=imgpath))
+    r.cell('<a href="{img:s}"><img src="{img:s}"/></a>'.format(img='{:s}_aggregated.jpg'.format(img_basename)))
+    r.cell('<a href="{img:s}"><img src="{img:s}"/></a>'.format(img='{:s}_scale1.jpg'.format(img_basename)))
+    for s in scales:
+      if s==1: continue
+      r.cell('<a href="{img:s}"><img src="{img:s}"/></a>'.format(img='{:s}_scale{:g}.jpg'.format(img_basename,s)))
+  doc.save(outdir+'index.html')
+
+  
+def visualize_region_maps(dataset_name, imgpaths, attentions, regions, scales, outdir, topk=10):  
+  assert len(imgpaths) == len(attentions)
+  assert len(attentions) == len(regions)
+  assert 1 in scales # we display the regions only for scale 1 (at least so far)
+  os.makedirs(outdir, exist_ok=True)
+  # generate visualization of each region
+  for i, imgpath in enumerate(imgpaths): # for each image
+    img_basename = os.path.splitext(os.path.basename(imgpath))[0]
+    regs = regions[i]
+    # load image
+    img = cv2.imread(imgpath)
+    # for each scale
+    for j,s in enumerate(scales):
+      if s!=1: continue    # just consider scale 1
+      r = regs[j][-1]
+      img_s = cv2.resize(img, None, fx=s, fy=s)
+      for ir in range(r.shape[0]):
+        heatmap_s = cv2.applyColorMap( (255*cv2.resize(np.minimum(1,100*r[ir,:,:]), (img_s.shape[1],img_s.shape[0]))).astype(np.uint8), cv2.COLORMAP_JET) # factor 10 for easier visualization
+        overlay = cv2.addWeighted(heatmap_s, 0.5, img_s, 0.5, 0)
+        cv2.imwrite(outdir+'{:s}_region{:d}_scale{:g}.jpg'.format(img_basename, ir, s), overlay)
+    # generate a html webpage for visualization
+    doc = HTML()
+    doc.header().title(dataset_name)
+    b = doc.body()
+    b.h(1, dataset_name+' (region maps)')
+    t = b.table(cellpadding=2, border=1)
+    for i, imgpath in enumerate(imgpaths):
+      atts = attentions[i]
+      regs = regions[i]
+      for j,s in enumerate(scales):
+        a = atts[j]
+        rr = regs[j][-1] # -1 because it is a list of the history of regions
+        if s==1: break
+      argsort = np.argsort(-a)
+      img_basename = os.path.splitext(os.path.basename(imgpath))[0]
+      if i%3==0: t.row(['info','image']+['scale 1 - region {:d}'.format(ir) for ir in range(topk)], header=True)
+      r = t.row()
+      r.cell(str(i)+': '+img_basename)
+      r.cell('<a href="{img:s}"><img src="{img:s}"/></a>'.format(img=imgpath))
+      for ir in range(topk):
+        index = argsort[ir]
+        r.cell('<a href="{img:s}"><img src="{img:s}"/></a><br>index: {index:d}, att: {att:g}, rmax: {rmax:g}'.format(img='{:s}_region{:d}_scale{:g}.jpg'.format(img_basename,index,s), index=index, att=a[index], rmax=rr[index,:,:].max()))
+    doc.save(outdir+'index.html')
+  
+if __name__=='__main__':
+  dataset = 'roxford5k'
+  from how.utils import data_helpers
+  images, qimages, bbxs, gnd = data_helpers.load_dataset(dataset, data_root="/tmp-network/user/pweinzae/CNNImageRetrieval/data/")
+  import pickle
+  with open('/tmp-network/user/pweinzae/roxford5k_features_attentions.pkl', 'rb') as fid:
+    features, attentions = pickle.load(fid)
+  visualize_attention_maps(qimages, attentions, scales=[2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25], outdir='/tmp-network/user/pweinzae/tmp/visu_attention_maps/'+dataset)

how/utils/whitening.py

@@ -0,0 +1,36 @@
+"""Functions for training and applying whitening"""
+
+import numpy as np
+
+
+def l2_normalize_vec(X):
+    """L2-normalize given descriptors"""
+    return X / (np.linalg.norm(X, ord=2, axis=1, keepdims=True) + 1e-6)
+
+
+def whitenapply(X, m, P, dimensions=None):
+    """Apply whitening (m, P) on descriptors X. If dimensions not None, perform dim reduction."""
+    if not dimensions:
+        dimensions = P.shape[1]
+
+    X = np.dot(X-m, P[:, :dimensions])
+    return l2_normalize_vec(X)
+
+
+def pcawhitenlearn_shrinkage(X, s=1.0):
+    """Learn PCA whitening with shrinkage from given descriptors"""
+    N = X.shape[0]
+
+    # Learning PCA w/o annotations
+    m = X.mean(axis=0, keepdims=True)
+    Xc = X - m
+    Xcov = np.dot(Xc.T, Xc)
+    Xcov = (Xcov + Xcov.T) / (2*N)
+    eigval, eigvec = np.linalg.eig(Xcov)
+    order = eigval.argsort()[::-1]
+    eigval = eigval[order]
+    eigvec = eigvec[:, order]
+
+    P = np.dot(np.linalg.inv(np.diag(np.power(eigval, 0.5*s))), eigvec.T)
+
+    return m, P.T

requirements.txt

@@ -0,0 +1,5 @@
+numpy
+pyaml
+matplotlib
+torch==1.3.1
+torchvision==0.4.2