<commit_before>import drawing_functions2 as df2
import chip_compute2, feature_compute2, load_data2
import report_results2
import sys
import matplotlib.pyplot as plt
import numpy as np
from hotspotter.Parallelize import parallel_compute
from hotspotter.other.ConcretePrintable import DynStruct
from hotspotter.helpers import Timer, get_exec_src, check_path, tic, toc
from drawing_functions2 import draw_matches, draw_kpts, tile_all_figures
from hotspotter.tpl.pyflann import FLANN
import hotspotter.tpl.cv2  as cv2
from itertools import chain
from numpy import linalg
from cvransac2 import H_homog_from_RANSAC, H_homog_from_DELSAC, H_homog_from_PCVSAC, H_homog_from_CV2SAC
import params2

__K__ = params2.__K__
__NUM_RERANK__   = params2.__NUM_RERANK__
__RATIO_THRESH__ = params2.__RATIO_THRESH__
__FLANN_PARAMS__ = params2.__FLANN_PARAMS__
__FEAT_TYPE__    = params2.__FEAT_TYPE__ 
__XY_THRESH__    = params2.__XY_THRESH__ = .05

def printDBG(msg):
    pass

def runall_match(hs):
    with Timer(msg=None):
        flann_1vM = precompute_index_1vM(hs)
    #functools.partial
    cx2_res_1vM = __run_matching(hs, assign_matches_1vM, flann_1vM)
    cx2_res_1v1 = __run_matching(hs, assign_matches_1v1)

class HotspotterQueryResult(DynStruct):
    def __init__(self):
        super(HotspotterQueryResult, self).__init__()
        self.qcx    = -1
        self.cx2_fm = []
        self.cx2_fs = []
        self.cx2_fm_V = []
        self.cx2_fs_V = []

# Work function that is the basic matching pipeline. 
# specific functions need to be bound: 
# 
# fn_assign - assign feature matches
def __run_matching(hs, fn_assign, *args):
    cx2_cid    = hs.tables.cx2_cid
    cx2_kpts, cx2_desc = hs.get_feats(__FEAT_TYPE__)
    cx2_desc   = [d for (k,d) in cx2_feats]
    cx2_kpts   = [k for (k,d) in cx2_feats]
    cx2_res = [HotspotterQueryResult() for _ in xrange(len(cx2_cid))]
    tt_ALL = tic('all queries')
    assign_times = []
    verify_times = []
    skip_list = []
    for qcx, qcid in enumerate(cx2_cid):
        if qcid == 0: 
            skip_list.append(qcx)
            continue
        tt_A = tic('query(qcx=%d)' % qcx)
        # Assign matches with the chosen function (1v1) or (1vM)
        (cx2_fm, cx2_fs) = fn_assign(qcx, cx2_cid, cx2_desc, *args)
        assign_times.append(toc(tt_A))
        # Spatially verify the assigned matches
        tt_V = tic('verify(qcx=%d)' % qcx)
        (cx2_fm_V, cx2_fs_V) = spatially_verify_1vX(qcx, cx2_kpts, cx2_fm, cx2_fs)
        verify_times.append(toc(tt_V))
        # Assign output to a query result
        res = cx2_res[qcx]
        res.qcx = qcx
        res.cx2_fm    = cx2_fm
        res.cx2_fs    = cx2_fs
        res.cx2_fm_V = cx2_fm_V
        res.cx2_fs_V = cx2_fs_V
    if len(skip_list) > 0:
        print('Skipped more queries than you should have: %r ' % skip_list)
    total_time = toc(tt_ALL)
    report_results2.report_results(cx2_res, hs.tables)
    return cx2_res

#@profile
def __spatially_verify(func_homog, kpts1, kpts2, fm, fs, DBG=None):
    '''1) compute a robust transform from img2 -> img1
       2) keep feature matches which are inliers '''
    # ugg transpose, I like row first, but ransac seems not to
    if len(fm) == 0: 
        return (np.empty((0,2)),np.empty((0,1)), np.eye(3))
    kpts1_m = kpts1[fm[:,0],:].T
    kpts2_m = kpts2[fm[:,1],:].T
    # -----------------------------------------------
    # TODO: SHOULD THIS HAPPEN HERE? (ISSUE XY_THRESH)
    # Get match threshold 10% of matching keypoint extent diagonal
    img1_extent = (kpts1_m[0:2,:].max(1) - kpts1_m[0:2,:].min(1))[0:2]
    xy_thresh1_sqrd = np.sum(img1_extent**2) * (__xy_thresh_percent__**2)
    if not DBG is None:
        print('---------------------------------------')
        print('INFO: spatially_verify xy threshold:')
        print(' * Threshold is %.1f%% of diagonal length' % (__xy_thresh_percent__*100))
        print(' * img1_extent = %r '     % img1_extent)
        print(' * img1_diag_len = %.2f ' % np.sqrt(np.sum(img1_extent**2)))
        print(' * xy_thresh1_sqrd=%.2f'  % np.sqrt(xy_thresh1_sqrd))
        print('---------------------------------------')
    # -----------------------------------------------
    H, inliers = func_homog(kpts2_m, kpts1_m, xy_thresh1_sqrd) 
    fm_V = fm[inliers,:]
    fs_V = fs[inliers,:]
    return fm_V, fs_V, H

def spatially_verify(kpts1, kpts2, fm, fs, DBG=None):
    ''' Concrete implementation of spatial verification
        using the deterministic ellipse based sample conensus'''
    return __spatially_verify(H_homog_from_DELSAC, kpts1, kpts2, fm, fs, DBG)
spatially_verify.__doc__ += '\n'+__spatially_verify.__doc__

#@profile
def spatially_verify_1vX(qcx, cx2_kpts, cx2_fm, cx2_fs):
    kpts1     = cx2_kpts[qcx]
    cx2_cscore = np.array([np.sum(fs) for fs in cx2_fs])
    top_cx     = cx2_cscore.argsort()[::-1]
    num_rerank = min(len(top_cx), __NUM_RERANK__)
    # -----------------------------------------------
    # TODO: SHOULD THIS HAPPEN HERE? (ISSUE XY_THRESH)
    #img1_extent = (kpts1[:,0:2].max(0) - qkpts1[:,0:2].min(0))[0:2]
    #xy_thresh1_sqrd = np.sum(img1_extent**2) * __xy_thresh_percent__
    # -----------------------------------------------
    # Precompute output container
    cx2_fm_V = [[] for _ in xrange(len(cx2_fm))]
    cx2_fs_V = [[] for _ in xrange(len(cx2_fs))]
    # spatially verify the top __NUM_RERANK__ results
    for topx in xrange(num_rerank):
        cx    = top_cx[topx]
        kpts2 = cx2_kpts[cx]
        fm    = cx2_fm[cx]
        fs    = cx2_fs[cx]
        fm_V, fs_V, H = spatially_verify(kpts1, kpts2, fm, fs)
        cx2_fm_V[cx] = fm_V
        cx2_fs_V[cx] = fs_V
    return cx2_fm_V, cx2_fs_V


#@profile
def precompute_index_1vM(hs):
    cx2_cid   = hs.tables.cx2_cid
    cx2_desc  = hs.feats.cx2_desc
    feat_dir  = hs.dirs.feat_dir
    feat_type = hs.feats.feat_type
    print('Precomputing one vs many information')
    cx2_nFeats = [len(k) for k in cx2_desc]
    _ax2_cx = [[cx_]*nFeats for (cx_, nFeats) in iter(zip(range(len(cx2_cid)), cx2_nFeats))]
    _ax2_fx = [range(nFeats) for nFeats in iter(cx2_nFeats)]
    ax2_cx  = np.array(list(chain.from_iterable(_ax2_cx)))
    ax2_fx  = np.array(list(chain.from_iterable(_ax2_fx)))
    ax2_desc   = np.vstack(cx2_desc)
    # Build (or reload) one vs many flann index
    flann_1vM = FLANN()
    flann_1vM_path = feat_dir + '/flann_1vM_'+feat_type+'.index'
    load_success = False
    if check_path(flann_1vM_path):
        try:
            print('Attempting to load flann index')
            flann_1vM.load_index(flann_1vM_path, ax2_desc)
            print('...successfully loaded flann index')
            load_success = True
        except Exception as ex:
            print('Cannot load FLANN index'+repr(ex))
    if not load_success:
        with Timer(msg='rebuilding FLANN index'):
            flann_1vM.build_index(ax2_desc, **__FLANN_PARAMS__)
            flann_1vM.save_index(flann_1vM_path)
    # Keep relevant data in the flann object. 
    # as to prevent them from loosing scope
    flann_1vM.ax2_desc = ax2_desc 
    flann_1vM.ax2_cx   = ax2_cx 
    flann_1vM.ax2_fx   = ax2_fx 
    return flann_1vM

# Feature scoring functions
def LNRAT_fn(vdist, ndist): return np.log(np.divide(ndist, vdist+1E-8)+1) 
def RATIO_fn(vdist, ndist): return np.divide(ndist, vdist+1E-8)
def LNBNN_fn(vdist, ndist): return ndist - vdist 
score_fn = RATIO_fn

#@profile
def assign_matches_1vM(qcx, cx2_cid, cx2_desc, flann_1vM):
    '''
    Matches desc1 vs all database descriptors using 
    Input:
        qcx       - query chip index
        cx2_cid   - chip ID lookup table (for removing self matches)
        cx2_desc  - chip descriptor lookup table
        flann_1vM - prebuild FLANN index of aggregated database descriptors<commit_msg>started matcher class<commit_after>import drawing_functions2 as df2
import chip_compute2, feature_compute2, load_data2
import report_results2
import sys
import matplotlib.pyplot as plt
import numpy as np
from hotspotter.Parallelize import parallel_compute
from hotspotter.other.ConcretePrintable import DynStruct
from hotspotter.helpers import Timer, get_exec_src, check_path, tic, toc
from drawing_functions2 import draw_matches, draw_kpts, tile_all_figures
from hotspotter.tpl.pyflann import FLANN
import hotspotter.tpl.cv2  as cv2
from itertools import chain
from numpy import linalg
from cvransac2 import H_homog_from_RANSAC, H_homog_from_DELSAC, H_homog_from_PCVSAC, H_homog_from_CV2SAC
import params2

__K__ = params2.__K__
__NUM_RERANK__   = params2.__NUM_RERANK__
__RATIO_THRESH__ = params2.__RATIO_THRESH__
__FLANN_PARAMS__ = params2.__FLANN_PARAMS__
__FEAT_TYPE__    = params2.__FEAT_TYPE__ 
__XY_THRESH__    = params2.__XY_THRESH__ = .05

def printDBG(msg):
    pass


class Matcher(object):
    def __init__(self, hs):
        if hs.feats.feat_type == 'FREAK':
            pass
            self.__func_nn_1vM = self.__nn_1vM_flann
        else:
            self.flann_1vM = precompute_index_1vM(hs)
            self.__func_nn_1vM = self.__nn_1vM_flann
        pass

    def nn_1vM(self, desc1, K):
        return self.__func_nn_1vM(desc1, K)

    def __nn_1vM_flann(self, desc1, K):
        (qfx2_ax, qfx2_dists) = flann_1vM.nn_index(desc1, K+1, **__FLANN_PARAMS__)

    def __nn_1vM_freak(self, desc1, K): 
        pass
    def precompute_1vM_index(self, hs):
        pass

#@profile
def precompute_index_1vM(hs):
    cx2_cid   = hs.tables.cx2_cid
    cx2_desc  = hs.feats.cx2_desc
    feat_dir  = hs.dirs.feat_dir
    feat_type = hs.feats.feat_type
    print('Precomputing one vs many information')
    cx2_nFeats = [len(k) for k in cx2_desc]
    _ax2_cx = [[cx_]*nFeats for (cx_, nFeats) in iter(zip(range(len(cx2_cid)), cx2_nFeats))]
    _ax2_fx = [range(nFeats) for nFeats in iter(cx2_nFeats)]
    ax2_cx  = np.array(list(chain.from_iterable(_ax2_cx)))
    ax2_fx  = np.array(list(chain.from_iterable(_ax2_fx)))
    ax2_desc   = np.vstack(cx2_desc)
    # Build (or reload) one vs many flann index
    flann_1vM = FLANN()
    flann_1vM_path = feat_dir + '/flann_1vM_'+feat_type+'.index'
    load_success = False
    if check_path(flann_1vM_path):
        try:
            print('Attempting to load flann index')
            flann_1vM.load_index(flann_1vM_path, ax2_desc)
            print('...successfully loaded flann index')
            load_success = True
        except Exception as ex:
            print('Cannot load FLANN index'+repr(ex))
    if not load_success:
        with Timer(msg='rebuilding FLANN index'):
            flann_1vM.build_index(ax2_desc, **__FLANN_PARAMS__)
            flann_1vM.save_index(flann_1vM_path)
    # Keep relevant data in the flann object. 
    # as to prevent them from loosing scope
    flann_1vM.ax2_desc = ax2_desc 
    flann_1vM.ax2_cx   = ax2_cx 
    flann_1vM.ax2_fx   = ax2_fx 
    return flann_1vM

def runall_match(hs):
    with Timer(msg=None):
        flann_1vM = precompute_index_1vM(hs)
    #functools.partial

    matcher_1vM = Matcher('1vM', hs)

    cx2_res_1vM = __run_matching(hs, assign_matches_1vM, flann_1vM)
    cx2_res_1v1 = __run_matching(hs, assign_matches_1v1)

class HotspotterQueryResult(DynStruct):
    def __init__(self):
        super(HotspotterQueryResult, self).__init__()
        self.qcx    = -1
        self.cx2_fm = []
        self.cx2_fs = []
        self.cx2_fm_V = []
        self.cx2_fs_V = []

# Work function that is the basic matching pipeline. 
# specific functions need to be bound: 
# 
# fn_assign - assign feature matches
def __run_matching(hs, fn_assign, *args):
    cx2_cid    = hs.tables.cx2_cid
    cx2_kpts, cx2_desc = hs.get_feats(__FEAT_TYPE__)
    cx2_desc   = [d for (k,d) in cx2_feats]
    cx2_kpts   = [k for (k,d) in cx2_feats]
    cx2_res = [HotspotterQueryResult() for _ in xrange(len(cx2_cid))]
    tt_ALL = tic('all queries')
    assign_times = []
    verify_times = []
    skip_list = []
    for qcx, qcid in enumerate(cx2_cid):
        if qcid == 0: 
            skip_list.append(qcx)
            continue
        tt_A = tic('query(qcx=%d)' % qcx)
        # Assign matches with the chosen function (1v1) or (1vM)
        (cx2_fm, cx2_fs) = fn_assign(qcx, cx2_cid, cx2_desc, *args)
        assign_times.append(toc(tt_A))
        # Spatially verify the assigned matches
        tt_V = tic('verify(qcx=%d)' % qcx)
        (cx2_fm_V, cx2_fs_V) = spatially_verify_1vX(qcx, cx2_kpts, cx2_fm, cx2_fs)
        verify_times.append(toc(tt_V))
        # Assign output to a query result
        res = cx2_res[qcx]
        res.qcx = qcx
        res.cx2_fm    = cx2_fm
        res.cx2_fs    = cx2_fs
        res.cx2_fm_V = cx2_fm_V
        res.cx2_fs_V = cx2_fs_V
    if len(skip_list) > 0:
        print('Skipped more queries than you should have: %r ' % skip_list)
    total_time = toc(tt_ALL)
    report_results2.report_results(cx2_res, hs.tables)
    return cx2_res

#@profile
def __spatially_verify(func_homog, kpts1, kpts2, fm, fs, DBG=None):
    '''1) compute a robust transform from img2 -> img1
       2) keep feature matches which are inliers '''
    # ugg transpose, I like row first, but ransac seems not to
    if len(fm) == 0: 
        return (np.empty((0,2)),np.empty((0,1)), np.eye(3))
    kpts1_m = kpts1[fm[:,0],:].T
    kpts2_m = kpts2[fm[:,1],:].T
    # -----------------------------------------------
    # TODO: SHOULD THIS HAPPEN HERE? (ISSUE XY_THRESH)
    # Get match threshold 10% of matching keypoint extent diagonal
    img1_extent = (kpts1_m[0:2,:].max(1) - kpts1_m[0:2,:].min(1))[0:2]
    xy_thresh1_sqrd = np.sum(img1_extent**2) * (__xy_thresh_percent__**2)
    if not DBG is None:
        print('---------------------------------------')
        print('INFO: spatially_verify xy threshold:')
        print(' * Threshold is %.1f%% of diagonal length' % (__xy_thresh_percent__*100))
        print(' * img1_extent = %r '     % img1_extent)
        print(' * img1_diag_len = %.2f ' % np.sqrt(np.sum(img1_extent**2)))
        print(' * xy_thresh1_sqrd=%.2f'  % np.sqrt(xy_thresh1_sqrd))
        print('---------------------------------------')
    # -----------------------------------------------
    H, inliers = func_homog(kpts2_m, kpts1_m, xy_thresh1_sqrd) 
    fm_V = fm[inliers,:]
    fs_V = fs[inliers,:]
    return fm_V, fs_V, H

def spatially_verify(kpts1, kpts2, fm, fs, DBG=None):
    ''' Concrete implementation of spatial verification
        using the deterministic ellipse based sample conensus'''
    return __spatially_verify(H_homog_from_DELSAC, kpts1, kpts2, fm, fs, DBG)
spatially_verify.__doc__ += '\n'+__spatially_verify.__doc__

#@profile
def spatially_verify_1vX(qcx, cx2_kpts, cx2_fm, cx2_fs):
    kpts1     = cx2_kpts[qcx]
    cx2_cscore = np.array([np.sum(fs) for fs in cx2_fs])
    top_cx     = cx2_cscore.argsort()[::-1]
    num_rerank = min(len(top_cx), __NUM_RERANK__)
    # -----------------------------------------------
    # TODO: SHOULD THIS HAPPEN HERE? (ISSUE XY_THRESH)
    #img1_extent = (kpts1[:,0:2].max(0) - qkpts1[:,0:2].min(0))[0:2]
    #xy_thresh1_sqrd = np.sum(img1_extent**2) * __xy_thresh_percent__
    # -----------------------------------------------
    # Precompute output container
    cx2_fm_V = [[] for _ in xrange(len(cx2_fm))]
    cx2_fs_V = [[] for _ in xrange(len(cx2_fs))]
    # spatially verify the top __NUM_RERANK__ results
    for topx in xrange(num_rerank):
        cx    = top_cx[topx]
        kpts2 = cx2_kpts[cx]
        fm    = cx2_fm[cx]
        fs    = cx2_fs[cx]
        fm_V, fs_V, H = spatially_verify(kpts1, kpts2, fm, fs)
        cx2_fm_V[cx] = fm_V
        cx2_fs_V[cx] = fs_V
    return cx2_fm_V, cx2_fs_V




# Feature scoring functions
def LNRAT_fn(vdist, ndist): return np.log(np.divide(ndist, vdist+1E-8)+1) 
def RATIO_fn(vdist, ndist): return np.divide(ndist, vdist+1E-8)
def LNBNN_fn(vdist, ndist): return ndist - vdist 
score_fn = RATIO_fn

#@profile
def assign_matches_1vM(qcx, cx2_cid, cx2_desc, flann_1vM):
    '''
    Matches desc1 vs all database descriptors using 
    Input:
        qcx       - query chip index
        cx2_cid   - chip ID lookup table (for removing self matches)
        cx2_desc  - chip descriptor lookup table
        flann_1vM - prebuild FLANN index of aggregated database descriptors
{'fraction_non_alphanumeric': 0.13375285335307546, 'fraction_numerical': 0.028379295453143317, 'mean_word_length': 2.7864984816631626, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1625902', 'n_tokens_mistral': 6381, 'n_tokens_neox': 6236, 'n_words': 1298}