app building

app.py

@@ -1,6 +1,6 @@
 import spaces
 import subprocess
-import sys
+import sys, os
 from pathlib import Path
 
 ''' loading modules '''
@@ -13,21 +13,28 @@ if not CKPT.exists():
     subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "./sam2-src[notebooks]"], cwd=ROOT)
     subprocess.check_call(["bash", "download_ckpts.sh"], cwd=SAM2 / "checkpoints")
 try:
-    import asmk  # noqa: F401
-except ImportError:
+    import asmk.index  # noqa: F401
+except Exception as e:
     subprocess.check_call(
-        [sys.executable, "-m", "cythonize", "*.pyx"], cwd='./asmk/cython'
+        ["cythonize", "*.pyx"], cwd='./asmk-src/cython'
     )
     subprocess.check_call(
-        [sys.executable, "-m", "pip", "install", './asmk', "--no-build-isolation"]
+        [sys.executable, "-m", "pip", "install", './asmk-src', "--no-build-isolation"]
+    )
+if not os.path.exists('./private'):
+    from huggingface_hub import snapshot_download
+    local_dir = snapshot_download(
+        repo_id="nycu-cplab/3AM",
+        local_dir="./private",
+        repo_type="model",
     )
+
 import gradio as gr
 import torch
 torch.no_grad().__enter__()
 import numpy as np
 from PIL import Image, ImageDraw
 import cv2
-import os
 import copy
 import json
 import logging

asmk-src/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Tomas Jenicek, Giorgos Tolias
+
+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.

asmk-src/README.md

@@ -0,0 +1,65 @@
+# Python ASMK (Aggregated Selective Match Kernels)
+
+This is a Python implementation of the ASMK approach published in [ICCV 2013](http://hal.inria.fr/docs/00/86/46/84/PDF/iccv13_tolias.pdf):
+
+```
+@InProceedings{TAJ13,
+  author       = "Giorgos Tolias and Yannis Avrithis and Herv\'e J\'egou",
+  title        = "To aggregate or not to aggregate: Selective match kernels for image search",
+  booktitle    = "IEEE International Conference on Computer Vision",
+  year         = "2013"
+}
+```
+
+This package is provided to support image retrieval with local descriptors and to reproduce the results of our [ECCV 2020 paper](https://arxiv.org/abs/2007.13172) with HOW deep local descriptors:
+
+```
+@InProceedings{TJ20,
+  author      = "Giorgos Tolias and Tomas Jenicek and Ond\v{r}ej Chum}",
+  title       = "Learning and aggregating deep local descriptors for instance-level recognition",
+  booktitle   = "European Conference on Computer Vision",
+  year        = "2020"
+}
+```
+
+There are minor differences compared to the original ASMK approach (ICCV'13) and [implementation](https://github.com/gtolias/asmk), which are described in our ECCV'20 paper. Using the provided package to run ASMK with other local descriptors is straightforward.
+
+
+
+## Running the Code
+
+1. Install the requirements (`faiss-cpu` for cpu-only setup)
+
+```
+pip3 install pyaml numpy faiss-gpu
+```
+
+
+2. Build C library for your Python version
+
+```
+python3 setup.py build_ext --inplace
+rm -r build
+```
+
+
+3. Download `cirtorch` and add it to your `PYTHONPATH`
+
+```
+wget "https://github.com/filipradenovic/cnnimageretrieval-pytorch/archive/master.zip"
+unzip master.zip
+rm master.zip
+export PYTHONPATH=${PYTHONPATH}:$(realpath cnnimageretrieval-pytorch-master)
+```
+
+
+4. Run `examples/demo_how.py` giving it any `.yaml` parameter file from `examples/params/*.yml`
+
+
+### Reproducing ECCV 2020 results with HOW local descriptors
+
+Reproducing results from **Table 2.**
+
+- R18<sub>how</sub> (n = 1000): &nbsp; `examples/demo_how.py eccv20_how_r18_1000` &ensp; _ROxf (M): 75.1, RPar (M): 79.4_
+- -R50<sub>how</sub> (n = 1000): &nbsp; `examples/demo_how.py eccv20_how_r50-_1000` &ensp; _ROxf (M): 78.3, RPar (M): 80.1_
+- -R50<sub>how</sub> (n = 2000): &nbsp; `examples/demo_how.py eccv20_how_r50-_2000` &ensp; _ROxf (M): 79.4, RPar (M): 81.6_

asmk-src/asmk.egg-info/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.1
+Name: asmk
+Version: 0.1
+Summary: ASMK Python implementation for ECCV'20 paper "Learning and aggregating deep local descriptors for instance-level recognition"
+Author: Tomas Jenicek, Giorgos Tolias
+License-File: LICENSE

asmk-src/asmk.egg-info/SOURCES.txt

@@ -0,0 +1,19 @@
+LICENSE
+README.md
+setup.py
+asmk/__init__.py
+asmk/asmk_method.py
+asmk/codebook.py
+asmk/functional.py
+asmk/index.py
+asmk/inverted_file.py
+asmk/io_helpers.py
+asmk/kernel.py
+asmk.egg-info/PKG-INFO
+asmk.egg-info/SOURCES.txt
+asmk.egg-info/dependency_links.txt
+asmk.egg-info/requires.txt
+asmk.egg-info/top_level.txt
+asmk.egg-info/zip-safe
+cython/hamming.c
+test/test_hamming.py

asmk-src/asmk.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

asmk-src/asmk.egg-info/requires.txt

@@ -0,0 +1,2 @@
+numpy
+pyaml

asmk-src/asmk.egg-info/top_level.txt

@@ -0,0 +1 @@
+asmk

asmk-src/asmk.egg-info/zip-safe

@@ -0,0 +1 @@
+

asmk-src/asmk/__init__.py

@@ -0,0 +1,6 @@
+"""
+Python implementation of ASMK for ECCV 2020 paper "Learning and aggregating deep local descriptors
+for instance-level recognition"
+"""
+
+from .asmk_method import ASMKMethod

asmk-src/asmk/asmk_method.py

@@ -0,0 +1,232 @@
+"""A wrapper around all asmk-related packages for convenient use"""
+
+import os.path
+import time
+import numpy as np
+
+from . import kernel as kern_pkg, codebook as cdb_pkg, index as idx_pkg, inverted_file as ivf_pkg
+from . import io_helpers
+
+
+class ASMKMethod:
+    """
+    Class to keep necessary objects and provide easy access to asmk method's steps. Each step
+    of asmk method corresponds to one method. Use initialize_untrained() class method instead
+    of directly calling the constructor.
+
+    :param dict params: contains keys index, train_codebook, build_ivf and query_ivf, each
+        containing the corresponding step parameters
+    :param dict metadata: only stored by this object, never changed
+    :param Codebook codebook: object from the codebook module
+    :param ASMKKernel kernel: object from the kernel module
+    :param IVF inverted_file: object from the inverted_file module
+    """
+
+    def __init__(self, params, metadata, *, codebook=None, kernel=None, inverted_file=None):
+        self.params = params
+        self.metadata = metadata
+
+        self.codebook = codebook
+        self.kernel = kernel
+        self.inverted_file = inverted_file
+
+
+    @classmethod
+    def initialize_untrained(cls, params):
+        """Initialize asmk method before training, provided only params (see constructor docs)"""
+        return cls(params, {})
+
+    #
+    # Method steps
+    #
+
+    def train_codebook(self, *columns, cache_path=None, step_params=None):
+        """The first step of the method - training the codebook (or loading from cache)
+
+        :param ndarray vecs: 2D numpy array, rows are vectors for codebook training
+        :param str cache_path: trained codebook will be stored under given file path and loaded
+            next time without training (None to turn off)
+        :param dict step_params: parameters that will override stored parameters for this step
+            (self.params['train_codebook'])
+        :return: new ASMKMethod object (containing metadata of this step), do not change self
+        """
+        assert not self.codebook, "Codebook already trained"
+        index_factory = idx_pkg.initialize_index(**self.params['index'])
+        step_params = step_params or self.params.get("train_codebook")
+
+        if cache_path and os.path.exists(cache_path):
+            time0 = time.time()
+            cdb = cdb_pkg.Codebook.initialize_from_state(io_helpers.load_pickle(cache_path),
+                                                         index_factory=index_factory)
+            cdb.index()
+            assert cdb.params == step_params['codebook']
+            metadata = {"load_time": time.time() - time0}
+        else:
+            cdb = cdb_pkg.Codebook(**step_params['codebook'], index_factory=index_factory)
+            metadata = cdb.train(*columns)
+            if cache_path:
+                io_helpers.save_pickle(cache_path, cdb.state_dict())
+
+        metadata["index_class"] = index_factory.__class__.__name__
+        return self.__class__({**self.params, "train_codebook": step_params},
+                              {**self.metadata, "train_codebook": metadata},
+                              codebook=cdb)
+
+
+    def build_ivf(self, *columns, distractors_path=None, cache_path=None, step_params=None):
+        """The second step of the method - building the ivf (or loading from cache)
+
+        :param ndarray vecs: 2D numpy array, rows are vectors to be indexed by the ivf
+        :param ndarray imids: 1D numpy array of image ids corresponding to 'vecs'
+        :param str distractors_path: ivf will be initialized with given distractors ivf path
+        :param str cache_path: built ivf will be stored under given file path and loaded
+            next time without training (None to turn off)
+        :param dict step_params: parameters that will override stored parameters for this step
+            (self.params['build_ivf'])
+        :return: new ASMKMethod object (containing metadata of this step), do not change self
+        """
+
+        builder = self.create_ivf_builder(cache_path=cache_path, step_params=step_params)
+
+        # Skip if loaded, otherwise add everything at once
+        if not builder.loaded_from_cache:
+            if distractors_path:
+                builder.initialize_with_distractors(distractors_path)
+            builder.add(*columns)
+
+        return self.add_ivf_builder(builder)
+
+
+    def create_ivf_builder(self, *, cache_path=None, step_params=None):
+        """Part of the second step of the method, see build_ivf() method implementation for usage
+
+        :param str cache_path: built ivf will be stored under given file path and loaded
+            next time without training (None to turn off)
+        :param dict step_params: parameters that will override stored parameters for this step
+            (self.params['build_ivf'])
+        :return: IvfBuilder object
+        """
+        assert not self.kernel and not self.inverted_file, "Inverted file already built"
+        step_params = step_params or self.params.get("build_ivf")
+        kern = kern_pkg.ASMKKernel(self.codebook, **step_params['kernel'])
+
+        return IvfBuilder(step_params, self.codebook, kern, cache_path=cache_path)
+
+
+    def add_ivf_builder(self, ivf_builder):
+        """Part of the second step of the method, see build_ivf() method implementation for usage
+
+        :param IvfBuilder ivf_builder: Builder with vectors added
+        :return: new ASMKMethod object (containing metadata from the builder), do not change self
+        """
+        ivf_metadata = ivf_builder.save()
+
+        return self.__class__({**self.params, "build_ivf": ivf_builder.step_params},
+                              {**self.metadata, "build_ivf": ivf_metadata},
+                              codebook=self.codebook, kernel=ivf_builder.kernel,
+                              inverted_file=ivf_builder.ivf)
+
+
+    def query_ivf(self, *columns, step_params=None, progress=None):
+        """The last step of the method - querying the ivf
+
+        :param ndarray qvecs: 2D numpy array, rows are vectors, each acting as a query for the ivf
+        :param ndarray qimids: 1D numpy array of image ids corresponding to 'qvecs'
+        :param dict step_params: parameters that will override stored parameters for this step
+            (self.params['query_ivf'])
+        :param bool progress: step at which update progress printing (None to disable)
+        :return: tuple (dict metadata, ndarray images, 2D ndarray ranks, 2D ndarray scores), do not
+            change self
+        """
+
+        step_params = step_params or self.params.get("query_ivf")
+
+        time0 = time.time()
+        images, ranks, scores = self.accumulate_scores(self.codebook, self.kernel, \
+                                    self.inverted_file, *columns, params=step_params, progress=progress)
+        metadata = {"query_avg_time": (time.time()-time0)/len(ranks)}
+        return metadata, images, ranks, scores
+
+
+    #
+    # Helper functions
+    #
+
+    @staticmethod
+    def accumulate_scores(cdb, kern, ivf, qvecs, qimids, *cols, params, progress=None):
+        """Accumulate scores for every query image (qvecs, qimids) given codebook, kernel,
+            inverted_file and parameters."""
+        similarity_func = lambda *x: kern.similarity(*x, **params["similarity"])
+
+        acc = []
+        slices = list(io_helpers.slice_unique(qimids))
+        for imid, seq in io_helpers.progress(slices, frequency=progress, header="Query"):
+            quantized = cdb.quantize(qvecs[seq], *(x[seq] for x in cols), **params["quantize"])
+            aggregated = kern.aggregate_image(*quantized, **params["aggregate"])
+            ranks, scores = ivf.search(*aggregated, **params["search"], similarity_func=similarity_func)
+            acc.append((imid, ranks, scores))
+
+        imids_all, ranks_all, scores_all = zip(*acc)
+        return np.array(imids_all), np.vstack(ranks_all), np.vstack(scores_all)
+
+
+class IvfBuilder:
+    """Inverted file (IVF) wrapper simplifying vector addition
+
+    :param dict step_params: contains parameters for build_ivf step
+    :param Codebook codebook: object from the codebook module
+    :param ASMKKernel kernel: object from the kernel module
+    :param str cache_path: built ivf will be stored under given file path and loaded
+        next time without training (None to turn off)
+    """
+
+    def __init__(self, step_params, codebook, kernel, *, cache_path):
+        self.step_params = step_params
+        self.codebook = codebook
+        self.kernel = kernel
+
+        if cache_path and os.path.exists(cache_path):
+            time0 = time.time()
+            self.ivf = ivf_pkg.IVF.initialize_from_state(io_helpers.load_pickle(cache_path))
+            self.metadata = {"load_time": time.time() - time0}
+            self.cache_path = None
+        else:
+            self.ivf = ivf_pkg.IVF.initialize_empty(**step_params['ivf'],
+                                                    codebook_size=codebook.size)
+            self.metadata = {"index_time": 0}
+            self.cache_path = cache_path
+
+    @property
+    def loaded_from_cache(self):
+        """If the contained IVF was loaded (otherwise, it is empty after initialization)"""
+        return "load_time" in self.metadata
+
+    def initialize_with_distractors(self, path):
+        """Initialize with distractors ivf at given path"""
+        self.ivf = ivf_pkg.IVF.initialize_from_state(io_helpers.load_pickle(path))
+        self.ivf.imid_offset = self.ivf.n_images
+
+    def add(self, *columns, progress=None):
+        """Add descriptors and cooresponding image ids to the IVF
+
+        :param np.ndarray vecs: 2D array of local descriptors
+        :param np.ndarray imids: 1D array of image ids
+        :param bool progress: step at which update progress printing (None to disable)
+        """
+        time0 = time.time()
+        quantized = self.codebook.quantize(*columns, **self.step_params["quantize"])
+        if progress:
+            print(">> Descriptors quantized")
+        aggregated = self.kernel.aggregate(*quantized, **self.step_params["aggregate"], progress=progress)
+        self.ivf.add(*aggregated, progress=200*progress if progress else None)
+        self.metadata['index_time'] += time.time() - time0
+
+    def save(self):
+        """Save to cache path if defined
+
+        :return: dict metadata with ivf stats
+        """
+        if self.cache_path:
+            io_helpers.save_pickle(self.cache_path, self.ivf.state_dict())
+
+        return {**self.metadata, "ivf_stats": self.ivf.stats}

asmk-src/asmk/codebook.py

@@ -0,0 +1,85 @@
+"""Codebook implementations for quantization of features into visual words"""
+
+import time
+import warnings
+import numpy as np
+
+
+class Codebook:
+    """Codebook of a fixed size for feature quantization"""
+
+    def __init__(self, index_factory, *, size):
+        self.params = {
+            "size": size,
+        }
+        if isinstance(size, str) and size[-1] in "kM":
+            size = int(size[:-1]) * {"k": 1024, "M": 1024**2}[size[-1]]
+        self.size = size
+        assert isinstance(self.size, int), self.size
+
+        self.index_factory = index_factory
+        self.search_index = None
+        self.centroids = None
+
+    #
+    # Create index
+    #
+
+    def index(self, centroids=None):
+        """Index either provided or stored centroids (when centroids=None). Return a dictionary
+            with 'index' key where value is how long it took to index the centroids."""
+        if centroids is not None:
+            assert self.size == centroids.shape[0], (self.size, centroids.shape[0])
+            self.centroids = centroids
+        time0 = time.time()
+        self.search_index = self.index_factory.create_index(self.centroids)
+        return {"index_time": time.time() - time0}
+
+    def train(self, des):
+        """Cluster descriptors and index resulting cluster centers. Return a dictionary with
+            'cluster', 'index' and 'train' keys where the value is how long it took to cluster,
+            index or train (sum of all)."""
+        time0 = time.time()
+        centroids = self.index_factory.cluster(des, self.size)
+        time_taken = time.time() - time0
+        meta = self.index(centroids)
+        return {**meta, "cluster_time": time_taken, "train_time": sum(meta.values()) + time_taken}
+
+    #
+    # Search in index
+    #
+
+    def quantize(self, des, *cols, multiple_assignment):
+        """Quantize given descriptors. Additional cols can be given, and this function will make
+            sure, that they still correspond to returned descriptors. Multiple assignment can be
+            applied when multiple_assignment > 1."""
+        _, centroid_ids = self.search_index.search(np.ascontiguousarray(des, dtype=np.float32),
+                                                   multiple_assignment)
+
+        return (des, centroid_ids) + cols
+
+    #
+    # Load and save
+    #
+
+    def state_dict(self):
+        """Return state dict which is a checkpoint of current state for future recovery"""
+        if self.centroids is None:
+            warnings.warn("Returning state_dict when codebook is not indexed (meaningless)")
+
+        return {
+            "type": self.__class__.__name__,
+            "params": self.params,
+            "state": {
+                "centroids": self.centroids,
+            }
+        }
+
+    @classmethod
+    def initialize_from_state(cls, state, index_factory):
+        """Initialize from a previously stored state_dict given an index factory"""
+        assert state["type"] == cls.__name__
+        codebook = cls(**state["params"], index_factory=index_factory)
+        if state["state"]["centroids"] is not None:
+            codebook.index(state["state"]["centroids"])
+        return codebook

asmk-src/asmk/functional.py

@@ -0,0 +1,15 @@
+"""Mathematical functions operating on arrays"""
+
+import numpy as np
+
+
+def normalize_vec_l2(vecs):
+    """Perform l2 normalization on each vector in a given matrix (axis 1)"""
+    norm = np.linalg.norm(vecs, ord=2, axis=1, keepdims=True) + 1e-6
+    return vecs / norm
+
+def asmk_kernel(sim, image_ids, *, alpha, similarity_threshold):
+    """Compute scores for visual words"""
+    mask = (sim>=similarity_threshold)
+    sim = np.power(sim[mask], alpha) # monomial kernel
+    return image_ids[mask], sim

asmk-src/asmk/index.py

@@ -0,0 +1,51 @@
+"""Index factories for efficient search (clustering)"""
+
+import numpy as np
+import faiss
+
+
+class FaissL2Index:
+    """Faiss factory for indexes on cpu"""
+
+    @staticmethod
+    def _faiss_index_flat(dim):
+        """Return initialized faiss.IndexFlatL2"""
+        return faiss.IndexFlatL2(dim)
+
+    def cluster(self, points, k, **index_kwargs):
+        """Clustering given points into k clusters"""
+        index = self._faiss_index_flat(points.shape[1], **index_kwargs)
+        clus = faiss.Clustering(points.shape[1], k)
+        clus.verbose = False
+        clus.niter = 10
+        clus.train(np.ascontiguousarray(points, dtype=np.float32), index)
+        return faiss.vector_float_to_array(clus.centroids).reshape(clus.k, clus.d)
+
+    def create_index(self, points, **index_kwargs):
+        """Return faiss index with given points"""
+        index = self._faiss_index_flat(points.shape[1], **index_kwargs)
+        index.add(points)
+        return index
+
+
+class FaissGpuL2Index(FaissL2Index):
+    """Faiss factory for indexes on gpu"""
+
+    def __init__(self, gpu_id):
+        super().__init__()
+        self.gpu_id = gpu_id
+
+    def _faiss_index_flat(self, dim):
+        """Return initialized faiss.GpuIndexFlatL2"""
+        res = faiss.StandardGpuResources()
+        flat_config = faiss.GpuIndexFlatConfig()
+        flat_config.device = self.gpu_id
+        return faiss.GpuIndexFlatL2(res, dim, flat_config)
+
+
+def initialize_index(gpu_id):
+    """Return either gpu faiss factory or cpu faiss factory (gpu_id is None)"""
+    if gpu_id is not None:
+        return FaissGpuL2Index(gpu_id)
+
+    return FaissL2Index()

asmk-src/asmk/inverted_file.py

@@ -0,0 +1,154 @@
+"""Inverted file implementations for efficient search in a set of visual words"""
+
+import numpy as np
+
+from . import io_helpers
+
+
+class IVF:
+    """Inverted File for efficient feature indexation with idf support. Can be updated."""
+
+    def __init__(self, norm_factor, n_images, ivf_vecs, ivf_image_ids, counts, idf, imid_offset, *, use_idf):
+        self.params = {
+            "use_idf": use_idf,
+        }
+
+        self.norm_factor = norm_factor
+        self.n_images = n_images
+        self.ivf_vecs = ivf_vecs
+        self.ivf_image_ids = ivf_image_ids
+        self.counts = counts
+        self.idf = idf
+        self.imid_offset = imid_offset
+
+
+    @classmethod
+    def initialize_empty(cls, *, codebook_size, **params):
+        """Return an empty IVF object given codebook params (size) and IVF params."""
+        ivf_vecs = [None for _ in range(codebook_size)]
+        ivf_image_ids = [None for _ in range(codebook_size)]
+        counts = np.zeros(codebook_size, dtype=np.int32)
+        idf = np.ones(counts.shape, dtype=np.float32)
+
+        return cls(**params, norm_factor=[], n_images=0, ivf_vecs=ivf_vecs,
+                   ivf_image_ids=ivf_image_ids, counts=counts, idf=idf, imid_offset=0)
+
+    #
+    # Index and search
+    #
+
+    @staticmethod
+    def _append_to_np_array(arr, size, item):
+        initial_size, increase_ratio = 10, 1.5
+
+        if arr is None:
+            # Initialization
+            arr = np.zeros((initial_size,) + item.shape, dtype=item.dtype)
+        elif size >= arr.shape[0]:
+            # Extension
+            new_size = int(np.ceil(arr.shape[0] * increase_ratio))
+            arr = np.resize(arr, (new_size,) + arr.shape[1:])
+
+        arr[size] = item
+        return arr
+
+
+    def add(self, des, word_ids, image_ids, *, progress=None):
+        """Add descriptors with corresponding visual word ids and image ids to this ivf"""
+        image_ids += self.imid_offset
+        min_imid, max_imid = image_ids.min(), image_ids.max()
+        assert min_imid >= self.n_images # The next chunk must be consequtive
+
+        norm_append = np.zeros(max_imid + 1 - len(self.norm_factor))
+        self.norm_factor = np.concatenate((self.norm_factor, norm_append))
+        self.n_images = np.max((self.n_images, max_imid + 1))
+
+        for i, word in io_helpers.progress(enumerate(word_ids), size=len(word_ids),
+                                           frequency=progress, header="Index"):
+            self.ivf_vecs[word] = self._append_to_np_array(self.ivf_vecs[word], self.counts[word],
+                                                           des[i])
+            self.ivf_image_ids[word] = self._append_to_np_array(self.ivf_image_ids[word],
+                                                                self.counts[word], image_ids[i])
+            self.counts[word] += 1
+            self.norm_factor[image_ids[i]] += 1
+
+            if self.params["use_idf"]:
+                self.idf[word] = np.log(self.n_images / self.counts[word])**2
+
+        if self.params["use_idf"]:
+            # Re-compute norm_factor to use idf
+            self.norm_factor[:] = 0
+            for word, imids in enumerate(self.ivf_image_ids):
+                if imids is not None:
+                    for imid in imids[:self.counts[word]]:
+                        self.norm_factor[imid] += self.idf[word]
+
+    def search(self, des, word_ids, *, similarity_func, topk):
+        """Search in this ivf with given descriptors and corresponding visual word ids. Return
+            similarity computed by provided function downweighted by idf and accumulated for all
+            visual words. Return topk results per query."""
+        scores = np.zeros(self.n_images)
+        q_norm_factor = 0
+
+        for qvec, word in zip(des, word_ids):
+            q_norm_factor += self.idf[word]
+            if self.ivf_image_ids[word] is None:
+                # Empty visual word
+                continue
+
+            image_ids, sim = similarity_func(qvec, self.ivf_vecs[word][:self.counts[word]],
+                                             self.ivf_image_ids[word][:self.counts[word]])
+
+            sim *= self.idf[word] # apply idf
+            sim /= np.sqrt(self.norm_factor[image_ids]) # normalize
+            scores[image_ids] += sim
+
+        scores = scores / np.sqrt(q_norm_factor)
+        ranks = np.argsort(-scores)[:topk]
+        return ranks - self.imid_offset, scores[ranks]
+
+    #
+    # Load, save and stats
+    #
+
+    @property
+    def stats(self):
+        """Return a shallow dictionary with stats of the ivf"""
+        sum_counts = self.counts.sum()
+        imbalance_factor = self.counts.shape[0] * np.power(self.counts, 2).sum() / (sum_counts**2 or 1)
+        return {
+            "images": self.n_images,
+            "vectors_per_image": sum_counts / (self.n_images or 1),
+            "mean_entries_per_vw": self.counts.mean(),
+            "empty_vw": sum(1 for x in self.counts if x == 0),
+            "min_entries_per_vw": self.counts.min(),
+            "max_entries_per_vw": self.counts.max(),
+            "std_of_entries_per_vw": self.counts.std(),
+            "imbalance_factor_of_vw": imbalance_factor,
+        }
+
+
+    def state_dict(self):
+        """Return state dict which is a checkpoint of current state for future recovery"""
+        return {
+            "type": self.__class__.__name__,
+            "params": self.params,
+            "state": {
+                "norm_factor": self.norm_factor,
+                "n_images": self.n_images,
+                "ivf_vecs": self.ivf_vecs,
+                "ivf_image_ids": self.ivf_image_ids,
+                "counts": self.counts,
+                "idf": self.idf,
+                "imid_offset": self.imid_offset,
+            }
+        }
+
+    @classmethod
+    def initialize_from_state(cls, state):
+        """Initialize from a previously stored state_dict given an index factory"""
+        assert state["type"] == cls.__name__
+        if "imid_offset" not in state['state']:
+            # For backwards compatibility
+            state['state']['imid_offset'] = 0
+        return cls(**state["params"], **state["state"])

asmk-src/asmk/io_helpers.py

@@ -0,0 +1,128 @@
+"""Helper functions related to io"""
+
+import os.path
+import time
+import sys
+import logging
+import pickle
+import urllib.request
+from io import StringIO
+from pathlib import Path
+import yaml
+import numpy as np
+
+# Params
+
+def load_params(path):
+    """Return loaded parameters from a yaml file"""
+    with open(path, "r") as handle:
+        content = yaml.safe_load(handle)
+    if "__template__" in content:
+        # Treat template as defaults
+        template_path = os.path.expanduser(content.pop("__template__"))
+        template = load_params(os.path.join(os.path.dirname(path), template_path))
+        content = dict_deep_overlay(template, content)
+    return content
+
+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
+
+
+# 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("ASMK")
+    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
+
+def progress(iterable, *, size=None, frequency=1, header=""):
+    """Generator that wraps an iterable and prints progress"""
+    if size is None:
+        size = len(iterable)
+    header = f"{header.capitalize()}: " if header else ""
+    charsize = len(str(size))
+    if frequency:
+        print(f"{header}[{'0'.rjust(charsize)}/{size}]", end="  ")
+        sys.stdout.flush()
+    time0 = time.time()
+    for i, element in enumerate(iterable):
+        yield element
+        i1 = i+1
+        if frequency and (i1 % frequency == 0 or i1 == size):
+            avg_time = (time.time() - time0) / i1
+            print(f"\r{header}[{str(i1).rjust(charsize)}/{size}] " \
+                    f"elapsed {int(avg_time*i1/60):02d}m/{int(avg_time*size/60):02d}m", end="  ")
+            sys.stdout.flush()
+    if frequency:
+        print()
+
+def capture_stdout(func, logger):
+    """Redirect stdout to logger"""
+    sys.stdout, stdout = StringIO(), sys.stdout
+    func()
+    sys.stdout, out_text = stdout, sys.stdout.getvalue()
+    for line in out_text.strip().split("\n"):
+        logger.info(line)
+
+
+# Load and save state dicts
+
+def load_pickle(path):
+    """Load pickled data from path"""
+    with open(path, 'rb') as handle:
+        return pickle.load(handle)
+
+def save_pickle(path, data):
+    """Save data to path using pickle"""
+    with open(path, 'wb') as handle:
+        pickle.dump(data, handle)
+
+
+# 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)
+
+
+# Iteration
+
+def slice_unique(ids):
+    """Generate slices that mark a sequence of identical values in a given array of ids. The
+        sequence must be uninterrupted (compact)."""
+    pointer = 0
+    for i, counts in zip(*np.unique(ids, return_counts=True)):
+        seq = slice(pointer, pointer+counts)
+        assert (ids[seq] == i).all()
+        yield i, seq
+        pointer += counts

asmk-src/asmk/kernel.py

@@ -0,0 +1,85 @@
+"""Kernel functionality implementation - aggregation and similarity computation"""
+
+import numpy as np
+
+from . import functional, hamming, io_helpers
+
+
+class ASMKKernel:
+    """Kernel for ASMK with the option of binarization."""
+
+    binary_shortcuts = {"bin": True, "nobin": False}
+
+    def __init__(self, codebook, *, binary):
+        self.params = {
+            "binary": binary,
+        }
+        self.binary = self.binary_shortcuts.get(binary, binary)
+        assert self.binary in self.binary_shortcuts.values()
+
+        self.codebook = codebook
+
+    #
+    # Aggregation
+    #
+
+    def aggregate_image(self, des, word_ids):
+        """Aggregate descriptors (with corresponding visual word ids) for a single image"""
+        unique_ids = np.unique(word_ids)
+        ades = np.empty((unique_ids.shape[0], des.shape[1]), dtype=np.float32)
+
+        for i, word in enumerate(unique_ids):
+            ades[i] = (des[(word_ids==word).any(axis=1)] - self.codebook.centroids[word]).sum(0)
+
+        if self.binary:
+            ades = hamming.binarize_and_pack_2D(ades)
+        else:
+            ades = functional.normalize_vec_l2(ades)
+
+        return ades, unique_ids
+
+    def aggregate(self, des, word_ids, image_ids, *, progress=None, **kwargs):
+        """Aggregate descriptors with corresponding visual word ids for corresponding image ids"""
+        acc = []
+        slices = list(io_helpers.slice_unique(image_ids))
+        for imid, seq in io_helpers.progress(slices, frequency=progress, header="Aggregate"):
+            ades, ids = self.aggregate_image(des[seq], word_ids[seq], **kwargs)
+            acc.append((ades, ids, np.full(ids.shape[0], imid)))
+
+        agg_des, agg_words, agg_imids = zip(*acc)
+        return np.vstack(agg_des), np.hstack(agg_words), np.hstack(agg_imids)
+
+    #
+    # Similarity
+    #
+
+    def similarity(self, qvec, vecs, image_ids, *, alpha, similarity_threshold):
+        """Compute similarity between given query vector and database feature vectors with their
+            corresponding image ids. Alpha is the similarity exponent after the similarity
+            threshold is applied."""
+        # Compute similarity with vw residuals for all other images
+        if self.binary:
+            norm_hdist = hamming.hamming_cdist_packed(qvec.reshape(1, -1), vecs)
+            sim = -2*norm_hdist.squeeze(0) + 1 # normalized hamming dist -> similarity in [-1, 1]
+        else:
+            sim = np.matmul(vecs, qvec)
+
+        return functional.asmk_kernel(sim, image_ids, alpha=alpha,
+                                      similarity_threshold=similarity_threshold)
+
+    #
+    # Load and save
+    #
+
+    def state_dict(self):
+        """Return state dict which is a checkpoint of current state for future recovery"""
+        return {
+            "type": self.__class__.__name__,
+            "params": self.params,
+        }
+
+    @classmethod
+    def initialize_from_state(cls, state, codebook):
+        """Initialize from a previously stored state_dict given a codebook"""
+        assert state["type"] == cls.__name__
+        return cls(**state["params"], codebook=codebook)

asmk-src/cython/build.sh

@@ -0,0 +1 @@
+cythonize *.pyx

asmk-src/cython/hamming.pyx

@@ -0,0 +1,152 @@
+# cython: language_level=3, boundscheck=False, wraparound=False
+
+import numpy as np
+cimport cython
+from libc.math cimport ceil
+
+
+cdef unsigned int BIT_MASK_1 = 0x55555555
+cdef unsigned int BIT_MASK_2 = 0x33333333
+cdef unsigned int BIT_MASK_4 = 0x0f0f0f0f
+cdef unsigned int BIT_MASK_8 = 0x00ff00ff
+cdef unsigned int BIT_MASK_16 = 0x0000ffff
+
+
+cdef int c_count_bits(unsigned int n) nogil:
+    n = (n & BIT_MASK_1) + ((n >> 1) & BIT_MASK_1)
+    n = (n & BIT_MASK_2) + ((n >> 2) & BIT_MASK_2)
+    n = (n & BIT_MASK_4) + ((n >> 4) & BIT_MASK_4)
+    n = (n & BIT_MASK_8) + ((n >> 8) & BIT_MASK_8)
+    n = (n & BIT_MASK_16) + ((n >> 16) & BIT_MASK_16)
+    return n
+
+
+cdef unsigned int c_binarize_and_pack_uint32(float[::1] arr, Py_ssize_t length, int threshold) nogil:
+    cdef unsigned int tmp = 0
+    cdef Py_ssize_t i
+
+    for i in range(length):
+        tmp = (tmp << 1) + (arr[i] > threshold)
+    return tmp
+
+
+@cython.cdivision(True)
+cdef float c_hamming_dist_uint32_arr(unsigned int[::1] n1, unsigned int[::1] n2, float normalization) nogil:
+    cdef Py_ssize_t length = n1.shape[0]
+    if normalization == 0:
+        normalization = length * 32
+
+    cdef int sum = 0
+    for i in range(length):
+        sum += c_count_bits(n1[i] ^ n2[i])
+    return sum / normalization
+
+
+#
+# Python API
+#
+
+def binarize_and_pack(float[::1] arr, int threshold = 0):
+    """
+    binarize_and_pack(float[::1] arr, int threshold = 0)
+
+    Binarizes given 1D numpy array by 'arr = arr > threshold' and packs its elements into bits
+    in uint32 array. Returns a 1D uint32 array where each element corresponds to a started set
+    of 32 bits.
+
+    >> binarize_and_pack((np.random.rand(10) - 0.5).astype(np.float32))
+    array([2206203904], dtype=uint32)
+    """
+    cdef Py_ssize_t dim_orig = arr.shape[0]
+    cdef Py_ssize_t dim = int(ceil(dim_orig / 32.0))
+    result = np.zeros(dim, dtype=np.uint32)
+    cdef unsigned int[::1] result_view = result
+
+    cdef Py_ssize_t i, offset
+    cdef unsigned int tmp
+    offset = 0
+    for i in range(dim-1):
+        result_view[i] = c_binarize_and_pack_uint32(arr[offset:], 32, threshold)
+        offset += 32
+
+    # Last iteration
+    tmp = c_binarize_and_pack_uint32(arr[offset:], dim_orig - offset, threshold)
+    result_view[dim-1] = tmp << (offset + 32 - dim_orig)
+
+    return result
+
+
+def binarize_and_pack_2D(float[:,::1] arr, int threshold = 0):
+    """
+    binarize_and_pack_2D(float[:,::1] arr, int threshold = 0)
+
+    Binarizes given 2D numpy array by 'arr = arr > threshold' and packs its elements into bits
+    in uint32 array. Returns a 2D uint32 array where each row corresponds to row in the original
+    array and each element to a started set of 32 bits.
+
+    >> binarize_and_pack_2D((np.random.rand(2, 10) - 0.5).astype(np.float32))
+    array([[1786773504]
+           [1509949440]], dtype=uint32)
+    """
+    cdef Py_ssize_t dim0 = arr.shape[0]
+    cdef Py_ssize_t dim1_orig = arr.shape[1]
+    cdef Py_ssize_t dim1 = int(ceil(dim1_orig / 32.0))
+    result = np.zeros((dim0, dim1), dtype=np.uint32)
+    cdef unsigned int[:, ::1] result_view = result
+
+    cdef Py_ssize_t i, j, offset
+    cdef unsigned int tmp
+    for i in range(dim0):
+        offset = 0
+        for j in range(dim1-1):
+            result_view[i,j] = c_binarize_and_pack_uint32(arr[i][offset:], 32, threshold)
+            offset += 32
+
+        # Last iteration
+        tmp = c_binarize_and_pack_uint32(arr[i][offset:], dim1_orig - offset, threshold)
+        result_view[i,dim1-1] = tmp << (offset + 32 - dim1_orig)
+
+    return result
+
+
+def hamming_dist_packed(unsigned int[::1] n1, unsigned int[::1] n2, float normalization = 0):
+    """
+    hamming_dist_packed(unsigned int[::1] n1, unsigned int[::1] n2, float normalization = 0)
+
+    Computes a hamming distance between two bit arrays packed into uint32 arrays and divides
+    it by normalization, if provided, otherwise by the number of bits in an array (always
+    a multiplication of 32).
+
+    >> hamming_dist_packed(np.array([3], dtype=np.uint32), np.array([1], dtype=np.uint32), 2)
+    0.5
+    """
+    assert n1 is not None and n2 is not None
+    assert n1.shape[0] == n2.shape[0]
+    return c_hamming_dist_uint32_arr(n1, n2, normalization)
+
+
+def hamming_cdist_packed(unsigned int[:,::1] arr1, unsigned int[:,::1] arr2, float normalization = 0):
+    """
+    hamming_cdist_packed(unsigned int[:,::1] arr1, unsigned int[:,::1] arr2, float normalization = 0)
+
+    Computes a hamming distance between two sets of bit arrays packed into uint32 using
+    hamming_dist_packed. Returns an array of size (arr1.shape[0], arr2.shape[0]).
+
+    >> hamming_cdist_packed(np.array([[3], [1]], dtype=np.uint32), np.array([[1], [2]], dtype=np.uint32), 2)
+    array([[0.5, 0.5],
+           [0. , 1. ]], dtype=float32)
+    """
+    assert arr1 is not None and arr2 is not None
+    assert arr1.shape[1] == arr2.shape[1]
+
+    cdef Py_ssize_t dim0 = arr1.shape[0]
+    cdef Py_ssize_t dim1 = arr2.shape[0]
+    result = np.zeros((dim0, dim1), dtype=np.float32)
+    cdef float[:, ::1] result_view = result
+
+    cdef Py_ssize_t i, j
+    for i in range(dim0):
+        for j in range(dim1):
+            result_view[i, j] = c_hamming_dist_uint32_arr(arr1[i], arr2[j], normalization)
+
+    return result

asmk-src/examples/demo_how.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+import os.path
+import sys
+import time
+import argparse
+from pathlib import Path
+import numpy as np
+
+# Add package root to pythonpath
+sys.path.append(os.path.realpath(f"{__file__}/../../"))
+
+from cirtorch.datasets.testdataset import configdataset
+from cirtorch.utils.evaluate import compute_map_and_print
+from asmk import io_helpers, ASMKMethod
+
+FEATURES_URL = "http://ptak.felk.cvut.cz/personal/toliageo/share/how/features/"
+DATASETS_URL = "http://cmp.felk.cvut.cz/cnnimageretrieval/data/test/"
+
+
+def initialize(params, demo_params, globals, logger):
+    """Download necessary files and initialize structures"""
+    logger.info(f"ECCV20 demo with parameters '{globals['exp_path'].name}'")
+
+    # Download featues
+    features = ["%s_%s.pkl" % (x, demo_params['eval_features']) \
+                    for x in demo_params['eval_datasets']]
+    features.append("%s_%s.pkl" % (demo_params['codebook_dataset'],
+                                   demo_params['codebook_features']))
+    io_helpers.download_files(features, globals['root_path'] / "features", FEATURES_URL,
+                              logfunc=logger.info)
+
+    # Download test datasets
+    pkls = ["%s/gnd_%s.pkl" % (x, x) for x in demo_params['eval_datasets']]
+    io_helpers.download_files(pkls, globals['root_path'] / "test", DATASETS_URL,
+                              logfunc=logger.info)
+
+    # Initialize asmk method wrapper
+    return ASMKMethod.initialize_untrained(params)
+
+
+def train_codebook(asmk, demo_params, globals, logger):
+    """The first step of asmk method - training the codebook"""
+    codebook_path = f"{globals['exp_path']}/codebook.pkl"
+    features_path = f"{globals['root_path']}/features/{demo_params['codebook_dataset']}_" \
+                    f"{demo_params['codebook_features']}.pkl"
+
+    desc = io_helpers.load_pickle(features_path)
+    logger.info(f"Loaded descriptors for codebook")
+    asmk = asmk.train_codebook(desc['vecs'], cache_path=codebook_path)
+
+    metadata = asmk.metadata['train_codebook']
+    logger.debug(f"Using {metadata['index_class']} index")
+    if "load_time" in metadata:
+        logger.info("Loaded pre-trained codebook")
+    else:
+        logger.info(f"Codebook trained in {metadata['train_time']:.1f}s")
+        logger.debug(f"Vectors for codebook clustered in {metadata['cluster_time']:.1f}s " \
+                     f"and indexed in {metadata['index_time']:.1f}s")
+    return asmk
+
+
+def build_ivf(asmk, dataset, desc, globals, logger):
+    """The second step of asmk method - building the ivf"""
+    ivf_path = f"{globals['exp_path']}/ivf_{dataset}.pkl"
+
+    asmk = asmk.build_ivf(desc['vecs'], desc['imids'], cache_path=ivf_path)
+
+    metadata = asmk.metadata['build_ivf']
+    if "load_time" in metadata:
+        logger.info("Loaded indexed ivf")
+    else:
+        logger.info(f"Indexed descriptors in {metadata['index_time']:.2f}s")
+
+    logger.debug(f"IVF stats: {metadata['ivf_stats']}")
+    return asmk
+
+
+def query_ivf(asmk, dataset, desc, globals, logger):
+    """The last step of asmk method - querying the ivf"""
+    metadata, _images, ranks, _scores = asmk.query_ivf(desc['qvecs'], desc['qimids'])
+    logger.debug(f"Average query time (quant+aggr+search) is {metadata['query_avg_time']:.3f}s")
+    gnd = configdataset(dataset, f"{globals['root_path']}/test/")['gnd']
+    io_helpers.capture_stdout(lambda: compute_map_and_print(dataset, ranks.T, gnd), logger)
+
+
+def demo_how(params, globals, logger):
+    """Demo where asmk is applied to the HOW descriptors from eccv'20 paper, replicating reported
+        results. Params is a dictionary with parameters for each step."""
+    demo_params = params.pop("demo_how")
+    asmk = initialize(params, demo_params, globals, logger)
+
+    asmk = train_codebook(asmk, demo_params, globals, logger)
+
+    # Create db and evaluate datasets
+    for dataset in demo_params['eval_datasets']:
+        desc = io_helpers.load_pickle(f"{globals['root_path']}/features/{dataset}_" \
+                                      f"{demo_params['eval_features']}.pkl")
+        logger.info(f"Loaded DB and query descriptors for {dataset}")
+
+        asmk_dataset = build_ivf(asmk, dataset, desc, globals, logger)
+
+        query_ivf(asmk_dataset, dataset, desc, globals, logger)
+
+
+def main(args):
+    """Argument parsing and parameter preparation for the demo"""
+    # Arguments
+    parser = argparse.ArgumentParser(description="ASMK demo replicating results for HOW " \
+                                                 "descriptors from ECCV 2020.")
+    parser.add_argument('parameters', nargs='+', type=str,
+                        help="Relative path to a yaml file that contains parameters.")
+    args = parser.parse_args(args)
+
+    package_root = Path(__file__).resolve().parent.parent
+    for parameters_path in args.parameters:
+        # Load yaml params
+        if not parameters_path.endswith(".yml"):
+            parameters_path = package_root / "examples" / ("params/%s.yml" % parameters_path)
+        params = io_helpers.load_params(parameters_path)
+
+        # Resolve data folders
+        globals = {}
+        globals["root_path"] = (package_root / params['demo_how']['data_folder'])
+        globals["root_path"].mkdir(parents=True, exist_ok=True)
+        exp_name = Path(parameters_path).name[:-len(".yml")]
+        globals["exp_path"] = (package_root / params['demo_how']['exp_folder']) / exp_name
+        globals["exp_path"].mkdir(parents=True, exist_ok=True)
+
+        # Setup logging
+        logger = io_helpers.init_logger(globals["exp_path"] / "output.log")
+
+        # Run demo
+        demo_how(params, globals, logger)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

asmk-src/examples/params/_eccv20_how.yml

@@ -0,0 +1,40 @@
+# Common parameters for HOW descriptors
+
+demo_how:
+  data_folder: data # Data path relative to the package root (can be absolute path)
+  exp_folder: data/experiments # Experiment path relative to the package root (can be absolute path)
+  eval_datasets: [roxford5k, rparis6k] # Datasets to evaluate on
+  eval_features: null # Name of the local features for evaluation (will be downloaded for each dataset)
+  codebook_dataset: sfm120k_subset20k # Dataset to train the codebook on
+  codebook_features: null # Name of the local features to train the codebook on (will be downloaded for codebook_dataset)
+
+# General properties
+
+index:
+  gpu_id: 0 # Gpu id to use, None to run on cpu
+
+# Steps
+
+train_codebook:
+  codebook:
+    size: "64k" # Number of visual words (i.e. clusters of local features) in the codebook
+
+build_ivf:
+  kernel:
+    binary: True # Binarization option in the asmk
+  ivf:
+    use_idf: False # IDF weighting in the inverted file
+
+  quantize:
+    multiple_assignment: 1 # Number of assignments for the db vectors
+  aggregate: {}
+
+query_ivf:
+  quantize:
+    multiple_assignment: 5 # Number of assignments for the query vectors
+  aggregate: {}
+  search:
+    topk: null # Limit the number of results (db images) returned for each query (null means unlimited)
+  similarity:
+    similarity_threshold: 0.0 # Feature similarity threshold
+    alpha: 3.0 # Feature similarity exponent

asmk-src/examples/params/eccv20_how_r18_1000.yml

@@ -0,0 +1,7 @@
+# Demo scenario for HOW descriptors with Resnet 18 backbone
+
+__template__: _eccv20_how.yml
+
+demo_how:
+  eval_features: how_r18_1000
+  codebook_features: how_r18_1000

asmk-src/examples/params/eccv20_how_r50-_1000.yml

@@ -0,0 +1,7 @@
+# Demo scenario for HOW descriptors with Resnet 50 backbone, without the last block
+
+__template__: _eccv20_how.yml
+
+demo_how:
+  eval_features: how_r50-_1000
+  codebook_features: how_r50-_1000

asmk-src/examples/params/eccv20_how_r50-_2000.yml

@@ -0,0 +1,7 @@
+# Demo scenario for HOW descriptors with Resnet 50 backbone, without the last block
+
+__template__: _eccv20_how.yml
+
+demo_how:
+  eval_features: how_r50-_2000
+  codebook_features: how_r50-_1000

asmk-src/setup.py

@@ -0,0 +1,35 @@
+import sys
+from setuptools import setup, Extension
+from setuptools.command.install import install
+
+class InstallWrapper(install):
+
+    def run(self):
+        try:
+            import faiss
+        except ImportError:
+            sys.stderr.write("\nERROR: faiss package not installed (install either faiss-cpu or " \
+                             "faiss-gpu before installing this package.).\n\n")
+            sys.exit(1)
+
+        install.run(self)
+
+
+setup(
+    name="asmk",
+    version="0.1",
+    description="ASMK Python implementation for ECCV'20 paper \"Learning and aggregating deep " \
+                "local descriptors for instance-level recognition\"",
+    author="Tomas Jenicek, Giorgos Tolias",
+    packages=[
+        "asmk",
+    ],
+    ext_modules=[Extension("asmk.hamming", ["cython/hamming.c"])],
+    install_requires=[
+        "numpy",
+        "pyaml",
+    ],
+    cmdclass={
+        "install": InstallWrapper,
+    },
+    zip_safe=True)

asmk-src/test/test_hamming.py

@@ -0,0 +1,54 @@
+"""Unit tests of asmk.hamming"""
+
+import unittest
+from asmk import hamming
+import numpy as np
+from scipy.spatial.distance import cdist
+
+
+class TestFunctions(unittest.TestCase):
+    """Unit test of functions"""
+
+    @staticmethod
+    def _numpy_pack_uint32(arr):
+        res = np.empty((arr.shape[0], int(np.ceil(arr.shape[1] / 32))), dtype=np.uint32)
+        packed = np.packbits(arr, axis=1).astype(np.uint32)
+        packed = np.pad(packed, ((0, 0), (0, 4 - packed.shape[1] % 4)), 'constant')
+        for i in range(res.shape[1]):
+            res[:,i] = (packed[:,4*i+0] << 24) + (packed[:,4*i+1] << 16) + (packed[:,4*i+2] << 8) + packed[:,4*i+3]
+        return res
+
+    def test_binarize_and_pack(self):
+        for dim1 in range(1, 40):
+            arr = (np.random.rand(dim1) - 0.5).astype(np.float32)
+            self.assertTrue(np.allclose(self._numpy_pack_uint32(np.expand_dims(arr, 0) > 0).squeeze(),
+                                        hamming.binarize_and_pack(arr)))
+
+    def test_binarize_and_pack_2D(self):
+        for dim1 in range(1, 40):
+            arr = (np.random.rand(10, dim1) - 0.5).astype(np.float32)
+            self.assertTrue(np.allclose(self._numpy_pack_uint32(arr > 0), hamming.binarize_and_pack_2D(arr)))
+
+    def test_hamming_dist_packed(self):
+        for dim1 in range(100, 140):
+            arr1 = (np.random.rand(1, dim1) - 0.5).astype(np.float32)
+            arr2 = (np.random.rand(1, dim1) - 0.5).astype(np.float32)
+            res = hamming.hamming_dist_packed(hamming.binarize_and_pack_2D(arr1).squeeze(0),
+                                              hamming.binarize_and_pack_2D(arr2).squeeze(0), dim1)
+            if dim1 % 32 == 0:
+                # Test default behaviour
+                res = hamming.hamming_dist_packed(hamming.binarize_and_pack_2D(arr1).squeeze(0),
+                                                  hamming.binarize_and_pack_2D(arr2).squeeze(0))
+            self.assertTrue(np.allclose(res, cdist(arr1 > 0, arr2 > 0, 'hamming').squeeze()))
+
+    def test_hamming_cdist_packed(self):
+        for dim1 in range(100, 140):
+            arr1 = (np.random.rand(10, dim1) - 0.5).astype(np.float32)
+            arr2 = (np.random.rand(10, dim1) - 0.5).astype(np.float32)
+            if dim1 % 32 == 0:
+                # Test default behaviour
+                res = hamming.hamming_cdist_packed(hamming.binarize_and_pack_2D(arr1),
+                                                   hamming.binarize_and_pack_2D(arr2))
+            res = hamming.hamming_cdist_packed(hamming.binarize_and_pack_2D(arr1),
+                                               hamming.binarize_and_pack_2D(arr2), dim1)
+            self.assertTrue(np.allclose(res, cdist(arr1 > 0, arr2 > 0, 'hamming')))