vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/imatch-toolbox/third_party/patch2pix/networks/utils.py

@@ -0,0 +1,111 @@
+import torch
+import numpy as np
+
+def select_local_patch_feats(feats1, feats2, ibatch, imatches, 
+                             feat_idx=[1, 2, 3, 4],
+                             feats_downsample=[1, 2, 2, 2, 2],
+                             psize=16, ptype='center'):
+    dy, dx = torch.meshgrid(torch.arange(psize), torch.arange(psize))
+    dx = dx.flatten().view(1, -1).to(imatches.device)
+    dy = dy.flatten().view(1, -1).to(imatches.device)
+
+    if ptype == 'center':
+        shift = psize // 2
+        dy -= shift
+        dx -= shift    
+    
+    _, _, h1, w1 = feats1[0].shape
+    _, _, h2, w2 = feats2[0].shape
+    x1, y1, x2, y2 = imatches.permute(1, 0).long()
+
+    # Ids for local patch
+    get_x_pids = lambda x, w, ds: ((x.view(-1, 1) + dx).view(-1) // ds).long().clamp(min=0, max=w//ds-1)
+    get_y_pids = lambda y, h, ds: ((y.view(-1, 1) + dy).view(-1) // ds).long().clamp(min=0, max=h//ds-1)
+
+    # Collect features for local matches
+    f1s, f2s = [], []
+    for j, (fmap1, fmap2) in enumerate(zip(feats1, feats2)):
+        if j not in feat_idx:
+            continue 
+        ds = np.prod(feats_downsample[0:j+1])
+        f1s.append(fmap1[ibatch, :, get_y_pids(y1, h1, ds), get_x_pids(x1, w1, ds)])
+        f2s.append(fmap2[ibatch, :, get_y_pids(y2, h2, ds), get_x_pids(x2, w2, ds)])
+        
+    f1s = torch.cat(f1s, dim=0) # D, N*16
+    f2s = torch.cat(f2s, dim=0) # D, N*16
+    return f1s, f2s, dx.squeeze(), dy.squeeze()
+
+def filter_coarse(coarse_matches, match_scores, ncn_thres=0.0, mutual=True, ptmax=None):
+    matches = []
+    scores = []
+    for imatches, iscores in  zip(coarse_matches, match_scores):
+        _, ids, counts = np.unique(imatches.cpu().data.numpy(), axis=0, return_index=True, return_counts=True)
+        if mutual:
+            # Consider only if they are multual consistant 
+            ids = ids[counts > 1]
+            #print(len(imatches), len(ids))
+            
+        if len(ids) > 0:
+            iscores = iscores[ids]
+            imatches = imatches[ids]
+
+        # NC score filtering
+        ids = torch.nonzero(iscores.flatten() > ncn_thres, as_tuple=False).flatten()
+        
+        # Cut or fill upto ptmax for memory control
+        if ptmax: 
+            if len(ids) == 0:
+                # insert a random match
+                ids = torch.tensor([0, 0, 0, 0]).long()
+            iids = np.arange(len(ids))
+            np.random.shuffle(iids)
+            iids = np.tile(iids, (ptmax // len(ids) + 1))[:ptmax]
+            ids = ids[iids]
+            
+        if len(ids) > 0: 
+            iscores = iscores[ids]
+            imatches = imatches[ids]
+            
+        matches.append(imatches)
+        scores.append(iscores)
+        
+    return matches, scores
+
+def sym_epi_dist(matches, F, sqrt=True, eps=1e-8):
+    # matches: Nx4
+    # F: 3x3
+    N = matches.shape[0]
+    matches = matches.to(F)
+    ones = torch.ones((N,1)).to(F)
+    p1 = torch.cat([matches[:, 0:2] , ones], dim=1) 
+    p2 = torch.cat([matches[:, 2:4] , ones], dim=1)
+
+    # l2=F*x1, l1=F^T*x2
+    l2 = F.matmul(p1.transpose(1, 0)) # 3,N
+    l1 = F.transpose(1, 0).matmul(p2.transpose(1, 0))
+    dd = (l2.transpose(1, 0) * p2).sum(dim=1)
+
+    sqrt = False
+    if sqrt:
+        d = dd.abs() * (1.0 / (eps + l1[0, :] ** 2 + l1[1, :] ** 2).sqrt() + 1.0 / (eps + l2[0, :] ** 2 + l2[1, :] ** 2).sqrt())    
+    else:
+        d = dd ** 2 * (1.0 / (eps + l1[0, :] ** 2 + l1[1, :] ** 2) + 1.0 / (eps + l2[0, :] ** 2 + l2[1, :] ** 2))
+    return d.float()
+
+def sampson_dist(matches, F, eps=1e-8):
+    # First-order approximation to reprojection error
+    # matches: Nx4
+    # F: 3x3
+    N = matches.shape[0]
+    matches = matches.to(F)
+    ones = torch.ones((N,1)).to(F)
+    p1 = torch.cat([matches[:, 0:2] , ones], dim=1) 
+    p2 = torch.cat([matches[:, 2:4] , ones], dim=1)
+
+    # l2=F*x1, l1=F^T*x2
+    l2 = F.matmul(p1.transpose(1, 0)) # 3,N
+    l1 = F.transpose(1, 0).matmul(p2.transpose(1, 0))
+    dd = (l2.transpose(1, 0) * p2).sum(dim=1)
+    d = dd ** 2 / (eps + l1[0, :] ** 2 + l1[1, :] ** 2 + l2[0, :] ** 2 + l2[1, :] ** 2)   
+    return d.float()
+

vismatch/third_party/imatch-toolbox/third_party/patch2pix/train_patch2pix.py

@@ -0,0 +1,374 @@
+import argparse
+from argparse import Namespace
+import os
+import time
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.utils.data as data
+
+from utils.datasets import ImMatchDatasetMega
+from utils.train.helper import *
+from utils.train.eval_epoch_immatch import eval_immatch_val_sets
+from utils.common.setup_helper import *
+from networks.utils import sampson_dist, filter_coarse
+from networks.patch2pix import Patch2Pix
+
+def parse_agrs():
+    parser = argparse.ArgumentParser(description='Train Patch2Pix Matching Network')
+    parser.add_argument('--gpu', '-gpu', type=int, default=0) 
+    parser.add_argument('--seed',  type=int, default=1)     
+    parser.add_argument('--epochs', type=int, default=100) 
+    parser.add_argument('--save_step', type=int, default=1) 
+    parser.add_argument('--plot_counts', type=int, default=5)    
+    parser.add_argument('--batch', type=int, default=8) 
+    parser.add_argument('--regr_batch', type=int, default=1200)
+    parser.add_argument('--visdom_host', '-vh',  type=str, default=None)  
+    parser.add_argument('--visdom_port', '-vp',  type=str, default=None)      
+    parser.add_argument('--prefix', type=str, default='')    
+    parser.add_argument('--out_dir', '-o', type=str, default='output/patch2pix')  
+    
+    # Data loading config
+    parser.add_argument('--dataset', type=str, default='MegaDepth')
+    parser.add_argument('--data_root', type=str, default='data')
+    parser.add_argument('--pair_root', type=str, default='data_pairs')
+    parser.add_argument(
+        '--match_npy', type=str,                     
+        default='megadepth_pairs.ov0.35_imrat1.5.pair500.excl_test.npy'
+    )
+    
+    # Model architecture
+    parser.add_argument('--backbone', type=str, default='ResNet34')
+    parser.add_argument('--change_stride', action='store_true')
+    parser.add_argument('--ksize', type=int, default=2)
+    parser.add_argument('--freeze_feat', type=int, default=87)    
+    parser.add_argument('--feat_idx', type=int, nargs='*', default=[0, 1, 2, 3])    
+    parser.add_argument('--feat_comb', type=str, default='pre')
+    parser.add_argument('--conv_kers', type=int, nargs='*', default=[3, 3])
+    parser.add_argument('--conv_dims', type=int, nargs='*', default=[512, 512])
+    parser.add_argument('--conv_strs', type=int, nargs='*', default=[2, 1])
+    parser.add_argument('--fc_dims', type=int, nargs='*', default=[512, 256])    
+    parser.add_argument('--psize', type=int, nargs=2, default=[16, 16])
+    parser.add_argument('--pshift', type=int, default=8)
+    parser.add_argument('--panc', type=int, choices=[8, 1], default=8)    
+    parser.add_argument('--ptmax', type=int, default=400)    
+    parser.add_argument('--shared', action='store_true')    
+    
+    # Matching thresholds
+    parser.add_argument('--cthres', type=float, default=0.5)
+    parser.add_argument('--cls_dthres', type=int, nargs=2, default=[50, 5])
+    parser.add_argument('--epi_dthres', type=int, nargs=2, default=[50, 5])
+
+    # Model intialize
+    parser.add_argument('--pretrain', type=str, default=None)   
+    parser.add_argument('--ckpt', type=str, default=None)                            
+    parser.add_argument('--resume', action='store_true')  # Auto load last cpkt
+        
+    # Optimization
+    parser.add_argument('--lr_init', '-lr', metavar='%f', type=float, default=5e-4) 
+    parser.add_argument('--lr_decay', '-lrd', metavar='%s[type] %f[*factor] %d[*step]', nargs='*', default=None) # Opt: 'step' 'multistep'
+    parser.add_argument('--weight_decay', '-wd', metavar='%f', type=float, default=0) 
+    parser.add_argument('--weight_cls', '-wcls', metavar='%f', type=float, default=10.0)
+    parser.add_argument('--weight_epi', '-wepi', metavar='%f[fine] %f[mid]', type=float, nargs='*', default=[1, 1])
+
+    args = parser.parse_args()
+    return args
+    
+def train_epoch(epoch, net, train_loader, train_vis, args, lprint_):    
+    net.train()
+    train_vis.clear()  # Clearn visdom plot data per epoch    
+    plot_step = len(train_loader) // args.plot_counts
+    
+    # Setup threshold params
+    net.panc = args.panc
+    ksize = args.ksize
+    cthres, cls_dthres, epi_dthres = args.cthres, args.cls_dthres, args.epi_dthres
+    cls_loss_weight = args.weight_cls
+    efine_weight, emid_weight = args.weight_epi
+
+    # Start training
+    skipped = 0
+    lprint_(f'ksize={ksize} cthres={cthres} cls_dthres={cls_dthres} '
+            f'epi_dthre={epi_dthres} ptmax={args.ptmax} panc={net.panc}')
+    for i, batch in enumerate(train_loader):
+        im_src, im_pos, Fs = net.load_batch_(batch, dtype='pair')
+        
+        # Estimate patch-level matches 
+        corr4d, delta4d, feats1, feats2 = net.forward(im_src, im_pos, ksize=ksize, return_feats=True)
+        coarse_matches, match_scores = net.cal_coarse_matches(corr4d, delta4d, ksize=ksize, upsample=net.upsample, center=True)
+        
+        if net.panc > 1 and args.ptmax > 0:
+            coarse_matches, match_scores = filter_coarse(coarse_matches, match_scores, 0.0, True, ptmax=args.ptmax)        
+
+        # Coarse matches to locate anchors    
+        coarse_matches = net.shift_to_anchors(coarse_matches)
+
+        # Mid level matching on positive pairs
+        mid_matches, mid_probs = net.forward_fine_match(feats1, feats2, 
+                                                        coarse_matches, 
+                                                        psize=net.psize[0],
+                                                        ptype=net.ptype[0],
+                                                        regressor=net.regress_mid)
+        
+        # Fine level matching based on mid matches
+        fine_matches, fine_probs = net.forward_fine_match(feats1, feats2, 
+                                                          mid_matches,               
+                                                          psize=net.psize[1],  
+                                                          ptype=net.ptype[1],
+                                                          regressor=net.regress_fine)
+        # Calculate per pair losses
+        cls_batch_lss = []
+        epi_batch_lss = []
+        for F, cmat, mmat, fmat, mcls_pred, fcls_pred in zip(Fs, coarse_matches, 
+                                                             mid_matches, fine_matches, 
+                                                             mid_probs, fine_probs):
+            N = len(cmat)
+            
+            # Classification gt based on coarse matches             
+            cdist = net.geo_dist_fn(cmat, F)
+            mdist = net.geo_dist_fn(mmat, F)
+            fdist = net.geo_dist_fn(fmat, F)
+            ones = torch.ones_like(cdist)
+            zeros = torch.zeros_like(cdist)
+
+            # Classification loss
+            mcls_pos = torch.where(cdist < cls_dthres[0], ones, zeros)
+            fcls_pos = torch.where(mdist < cls_dthres[1], ones, zeros)
+            mcls_neg = 1 - mcls_pos
+            fcls_neg = 1 - fcls_pos
+            
+            if mcls_pos.sum() == 0 or fcls_pos.sum() == 0:
+                skipped += 1
+                continue
+                        
+            mcls_weights = mcls_neg.sum() / mcls_pos.sum() * mcls_pos + mcls_neg 
+            mcls_lss = nn.functional.binary_cross_entropy(mcls_pred, mcls_pos, reduction='none')
+            mcls_lss = (mcls_weights * mcls_lss).mean()
+
+            fcls_weights = fcls_neg.sum() / fcls_pos.sum() * fcls_pos + fcls_neg 
+            fcls_lss = nn.functional.binary_cross_entropy(fcls_pred, fcls_pos, reduction='none')
+            fcls_lss = (fcls_weights * fcls_lss).mean()            
+                
+            cls_lss = mcls_lss + fcls_lss    
+            cls_batch_lss.append(cls_lss)
+            
+            # Plot cls metric
+            plot_cls_metric(mcls_pred, mcls_pos, cthres, train_vis.plots.cls_mid)
+            plot_cls_metric(fcls_pred, fcls_pos, cthres, train_vis.plots.cls_fine)
+            
+            # Plot statis
+            train_vis.plots.cls_ratios.mpos_gt.append(mcls_pos.sum().item() / N)
+            train_vis.plots.cls_ratios.fpos_gt.append(fcls_pos.sum().item() / N)
+            train_vis.plots.loss.cls_mid.append(mcls_lss.item())
+            train_vis.plots.loss.cls_fine.append(fcls_lss.item())
+                        
+            # Epipolar loss
+            mids_gt = torch.where(cdist < epi_dthres[0], ones, zeros).nonzero(as_tuple=False).flatten()
+            fids_gt = torch.where(mdist < epi_dthres[1], ones, zeros).nonzero(as_tuple=False).flatten()
+            #lprint_(f'{len(mdist)} {len(mids_gt)} {len(fdist)} {len(fids_gt)}')
+            
+            if len(fids_gt) == 0 and len(mids_gt) == 0: 
+                skipped += 1
+                continue
+
+            epi_mid = mdist[mids_gt].mean() if len(mids_gt) > 0 else torch.tensor(0).to(mdist)
+            epi_fine = fdist[fids_gt].mean() if len(fids_gt) > 0 else torch.tensor(0).to(fdist)
+            epi_lss = emid_weight * epi_mid + efine_weight * epi_fine
+            epi_batch_lss.append(epi_lss)                        
+
+            # Plot epi dists
+            if len(mids_gt) > 0:  
+                train_vis.plots.loss.epi_mid.append(epi_mid.item())
+                train_vis.plots.match_dist.mmid_gt.append(epi_mid.item())
+                train_vis.plots.match_dist.cmid_gt.append(cdist[mids_gt].mean().item())                
+                        
+            if len(fids_gt) > 0:
+                train_vis.plots.loss.epi_fine.append(epi_fine.item())                
+                train_vis.plots.match_dist.ffid_gt.append(epi_fine.item())
+                train_vis.plots.match_dist.mfid_gt.append(mdist[fids_gt].mean().item())
+
+        # Total loss
+        cls_loss = torch.stack(cls_batch_lss).mean() if len(cls_batch_lss) > 0 else torch.tensor(0.0, requires_grad=True).to(net.device)
+        epi_loss = torch.stack(epi_batch_lss).mean() if len(epi_batch_lss) > 0 else torch.tensor(0.0, requires_grad=True).to(net.device)        
+        loss = cls_loss_weight * cls_loss + epi_loss
+        train_vis.plots.loss.pair.append(loss.item())
+
+        # Optimize
+        net.optim_step_(loss)
+
+        # Monitor memory usage    
+        rss, vms = get_sys_mem()
+        train_vis.plots.mem.rss.append(rss)
+        train_vis.plots.mem.vms.append(vms)
+        
+        gpu_maloc, gpu_mres = get_gpu_mem()
+        train_vis.plots.mem.gpu_maloc.append(gpu_maloc)
+        train_vis.plots.mem.gpu_mres.append(gpu_mres)
+        torch.cuda.empty_cache()
+        
+        # Update plots periocially
+        if i % plot_step == 0 and i > 0:                        
+            train_vis.plot(epoch=epoch + (i / len(train_loader)))
+            lprint_('Batch:{} Loss:{}'.format(i, train_vis.get_plot_print(train_vis.plots.loss)))
+            lprint_('Cls_mid:{}'.format(train_vis.get_plot_print(train_vis.plots.cls_mid)))
+            lprint_('Cls_fine:{}'.format(train_vis.get_plot_print(train_vis.plots.cls_fine)))            
+            lprint_('Match:{}\n'.format(train_vis.get_plot_print(train_vis.plots.match_dist)))
+    
+    # Always update plots in the end of an epoch            
+    train_vis.plot(epoch=epoch + 1)                
+    lprint_('>Epoch:{} Skipped:{} Loss:{}'.format(epoch + 1, skipped, train_vis.get_plot_print(train_vis.plots.loss)))
+    lprint_('Cls_mid:{}'.format(train_vis.get_plot_print(train_vis.plots.cls_mid)))
+    lprint_('Cls_fine:{}'.format(train_vis.get_plot_print(train_vis.plots.cls_fine)))            
+    lprint_('Match:{}'.format(train_vis.get_plot_print(train_vis.plots.match_dist)))
+    
+def main():
+    np.set_printoptions(precision=3)   
+    args = parse_agrs()
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
+    make_deterministic(args.seed)
+    print(args)
+    
+    # Init data loader
+    match_npy_name = args.match_npy 
+    match_npy = os.path.join(args.pair_root, match_npy_name)    
+    pair_type = match_npy_name.replace('megadepth_pairs.', '').replace('_imrat1.5','').replace('.npy','')
+    data_tag = 'Mega.' + pair_type
+    train_set = ImMatchDatasetMega(args.data_root, match_npy, wt=480, ht=320)
+    train_loader = data.DataLoader(train_set, batch_size=args.batch, shuffle=True)
+    
+    # Init output dir names
+    if args.prefix is not '':
+        odir_tag = args.prefix + '.' + data_tag
+    else:
+        odir_tag = data_tag
+    odir_tag += '.freeze{}'.format(args.freeze_feat)
+    if args.change_stride:
+        odir_tag += '.cs'    
+    if args.pretrain:
+        odir_tag += '.pretrain'
+    
+    # fe1234nc0.9ep50-100_cls{}lr5e-4..
+    feat_tag = 'ks{}fe{}'.format(args.ksize, ''.join([str(v) for v in args.feat_idx]))
+    thres_tag = 'ep{}-{}cls{}-{}'.format(args.epi_dthres[0], args.epi_dthres[1],
+                                              args.cls_dthres[0], args.cls_dthres[1])
+    train_tag = '_wcls{}wepi{}-{}.lr{}'.format(args.weight_cls, args.weight_epi[0], 
+                                               args.weight_epi[1], args.lr_init)
+    
+    if args.weight_decay > 0:
+        train_tag += 'wd{}'.format(args.weight_decay)
+    if args.lr_decay:
+        decay_type = args.lr_decay[0]
+        if decay_type == 'step':
+            train_tag = '{}lrst{}-{}'.format(train_tag, args.lr_decay[1], args.lr_decay[2])
+        elif decay_type == 'multistep':
+            train_tag = '{}lrms{}-{}'.format(train_tag, args.lr_decay[1], args.lr_decay[2])
+
+    exp_tag = '{}{}{}'.format(feat_tag, thres_tag, train_tag)   
+    
+    # Regressor
+    regress_tag = '{}{}_conv{}dim{}str{}fc{}_psz{}-{}a{}'.format(args.feat_comb, args.ptmax,
+                                                                ''.join(map(str, args.conv_kers)),
+                                                                '-'.join(map(str, args.conv_dims)), 
+                                                                '-'.join(map(str, args.conv_strs)),
+                                                                '-'.join(map(str, args.fc_dims)),
+                                                                args.psize[0], args.psize[1], 
+                                                                args.panc)        
+    if args.shared:
+        regress_tag += '.shared'
+        
+    # Create output dirs and log file
+    out_dir = os.path.join(args.out_dir, odir_tag, exp_tag, regress_tag)
+    args.out_dir = out_dir
+    if not os.path.exists(out_dir):
+        os.makedirs(out_dir)
+    log = open(os.path.join(out_dir, 'log.txt'), 'a')
+    lprint_ = lambda ms: lprint(ms, log)
+    lprint_(config2str(args))
+    lprint_('Log dir {}'.format(out_dir))
+    lprint_(f'>>>Load dataset:{data_tag}, train:{len(train_loader.dataset)}')
+            
+    
+    # Initialize visdom
+    env = '{}.{}_{}{}_{}'.format(odir_tag, feat_tag, thres_tag, train_tag, regress_tag)
+    server = args.visdom_host
+    port = args.visdom_port
+    lprint_('>>Visdom server: {} port: {} env: {}'.format(server, port, env))    
+    train_vis = get_visdom_plots(prefix='train', env=env, server=server, port=port)
+    test_vis = get_visdom_plots(prefix='test', env=env, server=server, port=port)
+    
+    # Initialize model
+    config, best_vals = init_model_config(args, lprint_)
+    config.freeze_nc = True                
+    net = Patch2Pix(config)
+    if args.weight_epi[0] == 0:
+        lprint_('Freeze regressor_mid ...')        
+        for param in net.regress_mid.parameters():
+            param.requires_grad = False
+            
+    lprint_('Params backboone={} ncn={} regress_mid={} regress_fine={}'.format(
+        count_parameters(net.extract),
+        count_parameters(net.ncn),
+        count_parameters(net.regress_mid), 
+        count_parameters(net.regress_fine)        
+    ))
+    lprint_('Set geo dist: sampson distance')            
+    net.geo_dist_fn = sampson_dist
+
+    
+    # Training and validation
+    t0 = time.time()    
+    lprint_('Start training from {} to {} ..'.format(config.start_epoch, args.epochs))  
+    for epoch in range(config.start_epoch, args.epochs):        
+                
+        # Always train on normally matching pairs
+        lprint_('\n>>>Epoch {} training...'.format(epoch+1))        
+        lprint_('>>>Current_lr={}\n'.format(net.optimizer.param_groups[0]['lr']))
+        t1 = time.time()        
+        train_epoch(epoch, net, train_loader, train_vis, args, lprint_)
+        lprint_('Epoch training time: {:.2f}s'.format(time.time() - t1))        
+        
+        # Validation 
+        net.panc = 1  # Hard set topk to 1   
+        lprint_(f'Validation setting: panc={net.panc}')
+        
+        # Always save last ckpt 
+        save_ckpt(net, epoch, out_dir, best_vals=best_vals, last_ckpt=True)    
+
+        # Save model periodically
+        if (epoch + 1) % args.save_step == 0:
+            save_ckpt(net, epoch, out_dir, best_vals=best_vals)                                       
+            
+
+        # Eval immatch 
+        try:
+            res = eval_immatch_val_sets(net, 
+                                        data_root=f'{args.data_root}/immatch_benchmark/val_dense',
+                                        ksize=2, imsize=1024,
+                                        eval_type='fine', io_thres=0.5, 
+                                        sample_max=150, lprint_=lprint_)
+
+            # Save the best model based on immatch
+            qt_err, pass_rate = res    
+            rate = 0.34 * pass_rate[0] + 0.33 * pass_rate[4] + 0.33 * pass_rate[9]  # % < 1/5/10 px
+            if qt_err < best_vals[2] or rate > best_vals[3]:
+                if qt_err < best_vals[2]:
+                    best_vals[2] = qt_err
+                if rate > best_vals[3]:
+                    best_vals[3] = rate
+                save_ckpt(net, epoch, out_dir, best_vals=best_vals, name='immatch_best_ckpt')
+                lprint_('>>Save best immatch model: epoch={} qt={:.3f} rate={:.2f}%'.format(epoch+1, qt_err, rate))
+        
+        except:
+            lprint_('Failed to eval immatch')
+            res = None
+                    
+        # Update the learning rate    
+        if net.lr_scheduler:           
+            net.lr_scheduler.step()
+    lprint_('Finished, time:{:.4f}s'.format(time.time() - t0))
+    log.close()
+    
+    
+if __name__ == '__main__':
+    main()
+        

vismatch/third_party/imatch-toolbox/third_party/patch2pix/utils/colmap/data_loading.py

@@ -0,0 +1,169 @@
+import os
+from argparse import Namespace
+import numpy as np
+from utils.colmap.read_database import COLMAPDataLoader
+from utils.colmap.read_write_model import *
+
+def sav_model_multi_ov_pairs(model_dir, overlaps):
+    sav_file_path = os.path.join(model_dir, 'ov_pairs.npy')
+    if os.path.exists(sav_file_path):
+        ov_pair_dict = np.load(sav_file_path, allow_pickle=True).item()
+        all_exists = True        
+        for k in overlaps:
+            if k not in ov_pair_dict:
+                all_exists = False
+        if all_exists:
+            print('All overlaps have been computed.')
+            return ov_pair_dict
+
+    ov_pair_dict = {}        
+    images = read_images_binary(os.path.join(model_dir, 'images.bin'))
+    im_ids = list(images.keys())
+    overlap_scores, _ = cal_overlap_scores(im_ids, images)               
+
+    for min_overlap in overlaps:
+        if min_overlap in ov_pair_dict:
+            print(f'ov>{min_overlap} exists, skip.')
+            continue
+        valid_scores = np.logical_and(overlap_scores >= min_overlap, overlap_scores < 1)
+        pair_ids = np.vstack(np.where(valid_scores)).T
+        pair_names = []
+        for id1, id2 in pair_ids:
+            im1 = images[im_ids[id1]].name
+            im2 = images[im_ids[id2]].name
+            pair_names.append((max(im1, im2), min(im1, im2)))
+        print(f'ov>{min_overlap} pairs: {len(pair_names)}')
+        ov_pair_dict[min_overlap] = pair_names        
+    np.save(sav_file_path, ov_pair_dict)
+    return ov_pair_dict    
+    
+def load_model_ov_pairs(model_dir, min_overlap=0.3):    
+    images = read_images_binary(os.path.join(model_dir, 'images.bin'))
+    im_ids = list(images.keys())
+    overlap_scores, _ = cal_overlap_scores(im_ids, images)               
+    valid_scores = np.logical_and(overlap_scores >= min_overlap, overlap_scores < 1)
+    pair_ids = np.vstack(np.where(valid_scores)).T
+    pair_names = []
+    for id1, id2 in pair_ids:
+        im1 = images[im_ids[id1]].name
+        im2 = images[im_ids[id2]].name
+        pair_names.append((max(im1, im2), min(im1, im2)))
+    print('Loaded ov>{} pairs: {}'.format(min_overlap, len(pair_names)))
+    return pair_names
+
+def cal_overlap_scores(im_ids, images):
+    N = len(im_ids)
+    overlap_scores = np.eye(N)
+    im_3ds = [np.where(images[i].point3D_ids > 0)[0] for i in im_ids]
+
+    for i in range(N):
+        im1 = images[im_ids[i]]
+        pts1 = im_3ds[i]    
+        for j in range(N):
+            if j <= i :
+                continue
+            im2 = images[im_ids[j]]
+            pts2 = im_3ds[j]        
+            ov = len(np.intersect1d(pts1, pts2)) / max(len(pts1), len(pts2))
+            overlap_scores[i, j] = ov
+    nums_3d = np.array([len(v) for v in im_3ds])
+    return overlap_scores, nums_3d
+
+def load_model_ims(model_dir):
+    cameras = read_cameras_binary(os.path.join(model_dir, 'cameras.bin'))
+    images = read_images_binary(os.path.join(model_dir, 'images.bin'))
+    #print(len(cameras), len(images))   
+    imdict = {}
+    for i in images:
+        cid = images[i].camera_id
+        if cid not in cameras:
+            continue        
+        data = parse_data(images[i], cameras[cid])
+        imdict[data.name] = data
+    return imdict
+
+def cam_params_to_matrix(params, model):
+    if model == 'SIMPLE_PINHOLE':
+        f, ox, oy = params
+        K = np.array([[f, 0, ox], [0, f, oy], [0, 0, 1]])
+    elif model == 'PINHOLE':
+        f1, f2, ox, oy = params
+        K = np.array([[f1, 0, ox], [0, f2, oy], [0, 0, 1]])
+    elif model == 'SIMPLE_RADIAL':
+        f, ox, oy, _ = params
+        K = np.array([[f, 0, ox], [0, f, oy], [0, 0, 1]])
+    elif model == 'RADIAL':
+        f, ox, oy, _, _ = params
+        K = np.array([[f, 0, ox], [0, f, oy], [0, 0, 1]])
+    return K
+
+def parse_data(im, cam):
+    # Extract information from Image&Camera objects
+    K = cam_params_to_matrix(cam.params, cam.model)
+    q = im.qvec
+    R = qvec2rotmat(q)
+    t = im.tvec
+    c = - R.T.dot(t)
+    return Namespace(name=im.name, K=K, c=c, q=q, R=R, id=im.id)
+
+def load_colmap_matches(db_path, pair_names):
+    # Loading data from colmap database
+    db_loader = COLMAPDataLoader(db_path)
+    keypoints = db_loader.load_keypoints(key_len=6)
+    images = db_loader.load_images(name_based=True)   
+    pair_ids = [(images[im1][0], images[im2][0]) for im1, im2 in pair_names]
+    db_matches = db_loader.load_pair_matches(pair_ids)
+    match_dict = {}
+    for pname, pid in zip(pair_names, pair_ids):
+        (im1, im2) = pid
+        kpts1 = keypoints[im1]
+        kpts2 = keypoints[im2]
+        if pid not in db_matches:
+            matches = None
+        else:
+            key_ids = db_matches[pid]
+            N = key_ids.shape[0]
+            matches = np.zeros((N, 4))
+            for j in range(N):
+                k1, k2 = key_ids[j,:]
+                x1, y1 = kpts1[k1][0:2]
+                x2, y2 = kpts2[k2][0:2]
+                matches[j, :] = [x1, y1, x2, y2]    
+        match_dict[pname] = matches
+    return match_dict
+
+def export_intrinsics_txt(model_dir, sav_path): 
+    cameras = read_cameras_binary(os.path.join(model_dir, 'cameras.bin'))
+    images = read_images_binary(os.path.join(model_dir, 'images.bin'))
+    with open(sav_path, 'w') as f:
+        for imid in images:
+            im = images[imid]
+            cid = im.camera_id
+            if cid not in cameras:
+                continue
+            cam = cameras[cid]
+            model = cam.model
+            w, h = cam.width, cam.height
+            params = cam.params
+            
+            # Line format: im SIMPLE_RADIAL 1600 1200 1199.91 800 600 -0.0324314
+            line = f'{im.name} {model} {w} {h} '
+            for p in params:
+                line += f'{p} '
+            line += '\n'
+            f.write(line)
+            f.flush()
+    print('Finished, save to ', sav_path)
+        
+def parse_camera_matrices(intrinsic_txt):
+    with open(intrinsic_txt) as f:
+        # Line format: im SIMPLE_RADIAL 1600 1200 1199.91 800 600 -0.0324314
+        intrinsic_lines = f.readlines()
+        camera_matrices = {}
+        for line in intrinsic_lines:
+            cur = line.split()
+            name, model = cur[0:2]
+            params = [float(v) for v in cur[4::]]
+            K = cam_params_to_matrix(params, model)
+            camera_matrices[name] = K
+    return camera_matrices