nns_matching.py

Go to the documentation of this file.

#!/usr/bin/env python
import numpy as np
from scipy.spatial import distance
 
 
class NearestNeighborsMatching(object):
    """Nearest Neighbor matching of description vectors
    """
 
    def __init__(self, dim=None):
        """Initialization
 
        Args:
            dim (int, optional): Global descriptor size. Defaults to None.
        """
        self.nn = 0
        self.dimdim = dim
        self.itemsitems = dict()
        self.datadata = []
        if dim is not None:
            self.datadata = np.zeros((1000, dim), dtype='float32')
 
    def add_item(self, vector, item):
        """Add item to the matching list
 
        Args:
            vector (np.array): descriptor
            item: identification info (e.g., int)
        """
        assert vector.ndim == 1
        if self.nn >= len(self.datadata):
            if self.dimdim is None:
                self.dimdim = len(vector)
                self.datadata = np.zeros((1000, self.dimdim), dtype='float32')
            else:
                self.datadata.resize((2 * len(self.datadata), self.dimdim),
                                 refcheck=False)
        self.itemsitems[self.nn] = item
        self.datadata[self.nn] = vector
        self.nn += 1
 
    def search(self, query, k):  # searching from 100000 items consume 30ms
        """Search for nearest neighbors
 
        Args:
            query (np.array): descriptor to match
            k (int): number of best matches to return
 
        Returns:
            list(int, np.array): best matches
        """
        if len(self.datadata) == 0:
            return [], []
 
        similarities = np.zeros(self.nn)
 
        for i in range(self.nn):
            similarities[i] = 1 - distance.cosine(query, self.datadata[i,:].squeeze())
 
        ns = np.argsort(similarities)[::-1][:k]
        return [self.itemsitems[n] for n in ns], similarities[ns]
 
    def search_best(self, query):
        """Search for the nearest neighbor
 
        Args:
            query (np.array): descriptor to match
 
        Returns:
            int, np.array: best match
        """
        if len(self.datadata) == 0:
            return None, None
 
        items, similarities = self.searchsearch(query, 1)
        return items[0], similarities[0]