Source code for aitlas.clustering.pic

import logging
import time

from .utils import make_graph, preprocess_features, run_pic


logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")


class PIC:
    """Class to perform Power Iteration Clustering on a graph of nearest neighbors. Arguments for consistency with k-means init:

    :param sigma: bandwith of the Gaussian kernel (default 0.2)
    :type sigma: float
    :param nnn: number of nearest neighbors (default 5)
    :type nnn: int
    :param alpha: parameter in PIC (default 0.001)
    :type alpha: float
    :param distribute_singletons: If True, reassign each singleton to the cluster of its closest nonsingleton nearest neighbors (up to nnn nearest neighbors).
    :type distribute_singletons: bool
    :param images_lists: for each cluster, the list of image indexes belonging to this cluster
    :type images_lists: list of lists of ints
    """

    def __init__(
        self, args=None, sigma=0.2, nnn=5, alpha=0.001, distribute_singletons=True
    ):
        self.sigma = sigma
        self.alpha = alpha
        self.nnn = nnn
        self.distribute_singletons = distribute_singletons

    def cluster(self, data, verbose=False):
        start = time.time()

        # preprocess the data
        xb = preprocess_features(data)

        # construct nnn graph
        I, D = make_graph(xb, self.nnn)

        # run PIC
        clust = run_pic(I, D, self.sigma, self.alpha)
        images_lists = {}
        for h in set(clust):
            images_lists[h] = []
        for data, c in enumerate(clust):
            images_lists[c].append(data)

        # allocate singletons to clusters of their closest NN not singleton
        if self.distribute_singletons:
            clust_NN = {}
            for i in images_lists:
                # if singleton
                if len(images_lists[i]) == 1:
                    s = images_lists[i][0]
                    # for NN
                    for n in I[s, 1:]:
                        # if NN is not a singleton
                        if not len(images_lists[clust[n]]) == 1:
                            clust_NN[s] = n
                            break
            for s in clust_NN:
                del images_lists[clust[s]]
                clust[s] = clust[clust_NN[s]]
                images_lists[clust[s]].append(s)

        self.images_lists = []
        for c in images_lists:
            self.images_lists.append(images_lists[c])

        if verbose:
            logging.info("pic time: {0:.0f} s".format(time.time() - start))

        return 0