naver/PUMP

1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5import pdb, sys, os6import argparse7import numpy as np8from scipy.sparse import coo_matrix, csr_matrix, triu, csgraph9 10import core.functional as myF11from tools.common import image, image_with_trf12from tools.viz import dbgfig, show_correspondences13 14 15def arg_parser():16    parser = argparse.ArgumentParser("Post-filtering of Deep matching correspondences")17 18    parser.add_argument("--img1", required=True, help="path to first image")19    parser.add_argument("--img2", required=True, help="path to second image")20    parser.add_argument("--resize", default=0, type=int, help="prior image downsize (0 if recursive)")21    parser.add_argument("--corres", required=True, help="input path")22    parser.add_argument("--output", default="", help="filtered corres output")23 24    parser.add_argument("--locality", type=float, default=2, help="tolerance to deformation")25    parser.add_argument("--min-cc-size", type=int, default=50, help="min connex-component size")26    parser.add_argument("--densify", default='no', choices=['no','full','cc','convex'], help="output pixel-dense corres field")27    parser.add_argument("--dense-side", default='left', choices=['left','right'], help="img to densify")28 29    parser.add_argument("--verbose", "-v", type=int, default=0, help="verbosity level")30    parser.add_argument("--dbg", type=str, nargs='+', default=(), help="debug options")31    return parser32 33 34def main(args):35    import test_singlescale as pump36    corres = np.load(args.corres)['corres']37    imgs = tuple(map(image, pump.Main.load_images(args)))38 39    if dbgfig('raw',args.dbg):40        show_correspondences(*imgs, corres)41 42    corres = filter_corres( *imgs, corres, 43        locality=args.locality, min_cc_size=args.min_cc_size, 44        densify=args.densify, dense_side=args.dense_side,45        verbose=args.verbose, dbg=args.dbg)46 47    if dbgfig('viz',args.dbg):48        show_correspondences(*imgs, corres)49 50    return pump.save_output( args, corres )51 52 53def filter_corres( img0, img1, corres, 54            locality = None, # graph edge locality55            min_cc_size = None, # min CC size56            densify = None, 57            dense_side = None,58            verbose = 0, dbg=()):59 60    if None in (locality, min_cc_size, densify, dense_side):61        default_params = arg_parser()62        locality = locality or default_params.get_default('locality')63        min_cc_size = min_cc_size or default_params.get_default('min_cc_size')64        densify = densify or default_params.get_default('densify')65        dense_side = dense_side or default_params.get_default('dense_side')66 67    img0, trf0 = img0 if isinstance(img0,tuple) else (img0, np.eye(3))68    img1, trf1 = img1 if isinstance(img1,tuple) else (img1, np.eye(3))69    assert isinstance(img0, np.ndarray) and isinstance(img1, np.ndarray)70 71    corres = myF.affmul((np.linalg.inv(trf0),np.linalg.inv(trf1)), corres)72    n_corres = len(corres)73    if verbose: print(f'>> input: {len(corres)} correspondences')74 75    graph = compute_graph(corres, max_dis=locality*4)76    if verbose: print(f'>> {locality=}: {graph.nnz} nodes in graph')77 78    cc_sizes = measure_connected_components(graph)79    corres[:,4] += np.log2(cc_sizes)80    corres = corres[cc_sizes > min_cc_size]81    if verbose: print(f'>> {min_cc_size=}: remaining {len(corres)} correspondences')82 83    final = myF.affmul((trf0,trf1), corres)84 85    if densify != 'no':86        # densify correspondences87        if dense_side == 'right': # temporary swap88            final = final[:,[2,3,0,1]]89            H = round(img1.shape[0] / trf1[1,1])90            W = round(img1.shape[1] / trf1[0,0])91        else:92            H = round(img0.shape[0] / trf0[1,1])93            W = round(img0.shape[1] / trf0[0,0])94        95        if densify == 'cc':96            assert False, 'todo'97        elif densify in (True, 'full', 'convex'):98            # recover true image0's shape99            final = densify_corres( final, (H, W), full=(densify!='convex') )100        else:101            raise ValueError(f'Bad mode for {densify=}')102 103        if dense_side == 'right': # undo temporary swap104            final = final[:,[2,3,0,1]]105 106    return final107 108 109def compute_graph(corres, max_dis=10, min_ang=90):110    """ 4D distances (corres can only be connected to same scale)111        using sparse matrices for efficiency112 113    step1: build horizontal and vertical binning, binsize = max_dis114           add in each bin all neighbor bins115    step2: for each corres, we can intersect 2 bins to get a short list of candidates116    step3: verify euclidean distance < maxdis (optional?)117    """118    def bin_positions(pos):119        # every corres goes into a single bin120        bin_indices = np.int32(pos.clip(min=0) // max_dis) + 1121        cols = np.arange(len(pos))122        123        # add the cell before and the cell after, to handle border effects124        res = csr_matrix((np.ones(len(bin_indices)*3,dtype=np.float32), 125            (np.r_[bin_indices-1, bin_indices, bin_indices+1], np.r_[cols,cols,cols])),126            shape=(bin_indices.max()+2 if bin_indices.size else 1, len(pos)))127 128        return res, bin_indices129 130    # 1-hot matrices of shape = nbins x n_corres131    x1_bins = bin_positions(corres[:,0])132    y1_bins = bin_positions(corres[:,1])133    x2_bins = bin_positions(corres[:,2])134    y2_bins = bin_positions(corres[:,3])    135 136    def row_indices(ngh):137        res = np.bincount(ngh.indptr[1:-1], minlength=ngh.indptr[-1])[:-1]138        return res.cumsum()139 140    def compute_dist( ngh, pts, scale=None ):141        # pos from the second point142        x_pos = pts[ngh.indices,0]143        y_pos = pts[ngh.indices,1]144 145        # subtract pos from the 1st point146        rows = row_indices(ngh)147        x_pos -= pts[rows, 0]148        y_pos -= pts[rows, 1]149        dis = np.sqrt(np.square(x_pos) + np.square(y_pos))150        if scale is not None: 151            # there is a scale for each of the 2 pts, we encline to choose the worst one152            dis *= (scale[rows] + scale[ngh.indices]) / 2 # so we use arithmetic instead of geometric mean153 154        return normed(np.c_[x_pos, y_pos]), dis155 156    def Rot( ngh, degrees ):157        rows = row_indices(ngh)158        rad = degrees * np.pi / 180159        rad = (rad[rows] + rad[ngh.indices]) / 2 # average angle between 2 corres160        cos, sin = np.cos(rad), np.sin(rad)161        return np.float32(((cos, -sin), (sin,cos))).transpose(2,0,1)162 163    def match(xbins, ybins, pt1, pt2, way):164        xb, ixb = xbins165        yb, iyb = ybins166        167        # gets for each corres a list of potential matches168        ngh = xb[ixb].multiply( yb[iyb] ) # shape = n_corres x n_corres    169        ngh = triu(ngh, k=1).tocsr() # remove mirrored matches170        # ngh = matches of matches, shape = n_corres x n_corres171 172        # verify locality and flow173        vec1, d1 = compute_dist(ngh, pt1) # for each match, distance and orientation in img1174        # assert d1.max()**0.5 < 2*max_dis*1.415, 'cannot be larger than 2 cells in diagonals, or there is a bug'+bb()175        scale, rot = myF.decode_scale_rot(corres[:,5])176        vec2, d2 = compute_dist(ngh, pt2, scale=scale**(-way))177        ang = np.einsum('ik,ik->i', (vec1[:,None] @ Rot(ngh,way*rot))[:,0], vec2)178 179        valid = (d1 <= max_dis) & (d2 <= max_dis) & (ang >= np.cos(min_ang*np.pi/180))180        res = csr_matrix((valid, ngh.indices, ngh.indptr), shape=ngh.shape)181        res.eliminate_zeros()182        return res183 184    # find all neihbors within each xy bin185    ngh1 = match(x1_bins, y1_bins, corres[:,0:2], corres[:,2:4], way=+1)186    ngh2 = match(x2_bins, y2_bins, corres[:,2:4], corres[:,0:2], way=-1).T187 188    return ngh1 + ngh2 # union189    190 191def measure_connected_components(graph, dbg=()):192    # compute connected components193    nc, labels = csgraph.connected_components(graph, directed=False)194 195    # filter and remove all small components196    count = np.bincount(labels)197    198    return count[labels]199 200def normed( mat ):201    return mat / np.linalg.norm(mat, axis=-1, keepdims=True).clip(min=1e-16)202 203 204def densify_corres( corres, shape, full=True ):205    from scipy.interpolate import LinearNDInterpolator206    from scipy.spatial import cKDTree as KDTree207 208    assert len(corres) > 3, 'Not enough corres for densification'209    H, W = shape210 211    interp = LinearNDInterpolator(corres[:,0:2], corres[:,2:4])212    X, Y = np.mgrid[0:H, 0:W][::-1] # H x W, H x W213    p1 = np.c_[X.ravel(), Y.ravel()]214    p2 = interp(X, Y) # H x W x 2215 216    p2 = p2.reshape(-1,2)217    invalid = np.isnan(p2).any(axis=1)218 219    if full:220        # interpolate pixels outside of the convex hull221        badp = p1[invalid]222        tree = KDTree(corres[:,0:2])223        _, nn = tree.query(badp, 3) # find 3 closest neighbors224        corflow = corres[:,2:4] - corres[:,0:2]225        p2.reshape(-1,2)[invalid] = corflow[nn].mean(axis=1) + p1[invalid]226    else:227        # remove nans, i.e. remove points outside of convex hull228        p1, p2 = p1[~invalid], p2[~invalid]229 230    # return correspondence field231    return np.c_[p1, p2]232 233 234if __name__ == '__main__':235    main(arg_parser().parse_args())236