nlc.py 文件源码

def compute_nn(regions, frameEnd, F=15, L=4, redirect=False):
    """
    Compute transition matrix using nearest neighbors
    Input:
        regions: (k,d): k regions with d-dim feature
        frameEnd: (n,): indices of regions where frame ends: 0-indexed, included
        F: temporal radius: nn to be searched in (2F+1) frames around curr frame
        L: nearest neighbors to be found per frame on an average
    Output: transM: (k,k)
    """
    sTime = time.time()
    M = L * (2 * F + 1)
    k, _ = regions.shape
    n = frameEnd.shape[0]
    transM = np.zeros((k, k), dtype=np.float)

    # Build 0-1 nn graph based on L2 distance using KDTree
    for i in range(n):
        # build KDTree with 2F+1 frames around frame i
        startF = max(0, i - F)
        startF = 1 + frameEnd[startF - 1] if startF > 0 else 0
        endF = frameEnd[min(n - 1, i + F)]
        tree = KDTree(regions[startF:1 + endF], leafsize=100)

        # find nn for regions in frame i
        currStartF = 1 + frameEnd[i - 1] if i > 0 else 0
        currEndF = frameEnd[i]
        distNN, nnInd = tree.query(regions[currStartF:1 + currEndF], M)
        nnInd += startF
        currInd = np.mgrid[currStartF:1 + currEndF, 0:M][0]
        transM[currInd, nnInd] = distNN
        if not redirect:
            sys.stdout.write('NearestNeighbor computation: [% 5.1f%%]\r' %
                                (100.0 * float((i + 1) / n)))
            sys.stdout.flush()

    eTime = time.time()
    print('NearestNeighbor computation finished: %.2f s' % (eTime - sTime))

    return transM