vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/imatch-toolbox/third_party/r2d2/extract_kapture.py

@@ -0,0 +1,194 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+
+from PIL import Image
+
+from tools import common
+from tools.dataloader import norm_RGB
+from nets.patchnet import *
+from os import path
+
+from extract import load_network, NonMaxSuppression, extract_multiscale
+
+# Kapture is a pivot file format, based on text and binary files, used to describe SfM (Structure From Motion)
+# and more generally sensor-acquired data
+# it can be installed with
+# pip install kapture
+# for more information check out https://github.com/naver/kapture
+import kapture
+from kapture.io.records import get_image_fullpath
+from kapture.io.csv import kapture_from_dir
+from kapture.io.csv import get_feature_csv_fullpath, keypoints_to_file, descriptors_to_file
+from kapture.io.features import get_keypoints_fullpath, keypoints_check_dir, image_keypoints_to_file
+from kapture.io.features import get_descriptors_fullpath, descriptors_check_dir, image_descriptors_to_file
+from kapture.io.csv import get_all_tar_handlers
+
+
+def extract_kapture_keypoints(args):
+    """
+    Extract r2d2 keypoints and descritors to the kapture format directly
+    """
+    print('extract_kapture_keypoints...')
+    with get_all_tar_handlers(args.kapture_root,
+                              mode={kapture.Keypoints: 'a',
+                                    kapture.Descriptors: 'a',
+                                    kapture.GlobalFeatures: 'r',
+                                    kapture.Matches: 'r'}) as tar_handlers:
+        kdata = kapture_from_dir(args.kapture_root, None,
+                                 skip_list=[kapture.GlobalFeatures,
+                                            kapture.Matches,
+                                            kapture.Points3d,
+                                            kapture.Observations],
+                                 tar_handlers=tar_handlers)
+
+        assert kdata.records_camera is not None
+        image_list = [filename for _, _, filename in kapture.flatten(kdata.records_camera)]
+        if args.keypoints_type is None:
+            args.keypoints_type = path.splitext(path.basename(args.model))[0]
+            print(f'keypoints_type set to {args.keypoints_type}')
+        if args.descriptors_type is None:
+            args.descriptors_type = path.splitext(path.basename(args.model))[0]
+            print(f'descriptors_type set to {args.descriptors_type}')
+
+        if kdata.keypoints is not None and args.keypoints_type in kdata.keypoints \
+                and kdata.descriptors is not None and args.descriptors_type in kdata.descriptors:
+            print('detected already computed features of same keypoints_type/descriptors_type, resuming extraction...')
+            image_list = [name
+                          for name in image_list
+                          if name not in kdata.keypoints[args.keypoints_type] or
+                          name not in kdata.descriptors[args.descriptors_type]]
+
+        if len(image_list) == 0:
+            print('All features were already extracted')
+            return
+        else:
+            print(f'Extracting r2d2 features for {len(image_list)} images')
+
+        iscuda = common.torch_set_gpu(args.gpu)
+
+        # load the network...
+        net = load_network(args.model)
+        if iscuda:
+            net = net.cuda()
+
+        # create the non-maxima detector
+        detector = NonMaxSuppression(
+            rel_thr=args.reliability_thr,
+            rep_thr=args.repeatability_thr)
+
+        if kdata.keypoints is None:
+            kdata.keypoints = {}
+        if kdata.descriptors is None:
+            kdata.descriptors = {}
+
+        if args.keypoints_type not in kdata.keypoints:
+            keypoints_dtype = None
+            keypoints_dsize = None
+        else:
+            keypoints_dtype = kdata.keypoints[args.keypoints_type].dtype
+            keypoints_dsize = kdata.keypoints[args.keypoints_type].dsize
+        if args.descriptors_type not in kdata.descriptors:
+            descriptors_dtype = None
+            descriptors_dsize = None
+        else:
+            descriptors_dtype = kdata.descriptors[args.descriptors_type].dtype
+            descriptors_dsize = kdata.descriptors[args.descriptors_type].dsize
+
+        for image_name in image_list:
+            img_path = get_image_fullpath(args.kapture_root, image_name)
+            print(f"\nExtracting features for {img_path}")
+            img = Image.open(img_path).convert('RGB')
+            W, H = img.size
+            img = norm_RGB(img)[None]
+            if iscuda:
+                img = img.cuda()
+
+            # extract keypoints/descriptors for a single image
+            xys, desc, scores = extract_multiscale(net, img, detector,
+                                                   scale_f=args.scale_f,
+                                                   min_scale=args.min_scale,
+                                                   max_scale=args.max_scale,
+                                                   min_size=args.min_size,
+                                                   max_size=args.max_size,
+                                                   verbose=True)
+
+            xys = xys.cpu().numpy()
+            desc = desc.cpu().numpy()
+            scores = scores.cpu().numpy()
+            idxs = scores.argsort()[-args.top_k or None:]
+
+            xys = xys[idxs]
+            desc = desc[idxs]
+            if keypoints_dtype is None or descriptors_dtype is None:
+                keypoints_dtype = xys.dtype
+                descriptors_dtype = desc.dtype
+
+                keypoints_dsize = xys.shape[1]
+                descriptors_dsize = desc.shape[1]
+
+                kdata.keypoints[args.keypoints_type] = kapture.Keypoints('r2d2', keypoints_dtype, keypoints_dsize)
+                kdata.descriptors[args.descriptors_type] = kapture.Descriptors('r2d2', descriptors_dtype,
+                                                                               descriptors_dsize,
+                                                                               args.keypoints_type, 'L2')
+                keypoints_config_absolute_path = get_feature_csv_fullpath(kapture.Keypoints,
+                                                                          args.keypoints_type,
+                                                                          args.kapture_root)
+                descriptors_config_absolute_path = get_feature_csv_fullpath(kapture.Descriptors,
+                                                                            args.descriptors_type,
+                                                                            args.kapture_root)
+                keypoints_to_file(keypoints_config_absolute_path, kdata.keypoints[args.keypoints_type])
+                descriptors_to_file(descriptors_config_absolute_path, kdata.descriptors[args.descriptors_type])
+            else:
+                assert kdata.keypoints[args.keypoints_type].dtype == xys.dtype
+                assert kdata.descriptors[args.descriptors_type].dtype == desc.dtype
+                assert kdata.keypoints[args.keypoints_type].dsize == xys.shape[1]
+                assert kdata.descriptors[args.descriptors_type].dsize == desc.shape[1]
+                assert kdata.descriptors[args.descriptors_type].keypoints_type == args.keypoints_type
+                assert kdata.descriptors[args.descriptors_type].metric_type == 'L2'
+
+            keypoints_fullpath = get_keypoints_fullpath(args.keypoints_type, args.kapture_root,
+                                                        image_name, tar_handlers)
+            print(f"Saving {xys.shape[0]} keypoints to {keypoints_fullpath}")
+            image_keypoints_to_file(keypoints_fullpath, xys)
+            kdata.keypoints[args.keypoints_type].add(image_name)
+
+            descriptors_fullpath = get_descriptors_fullpath(args.descriptors_type, args.kapture_root,
+                                                            image_name, tar_handlers)
+            print(f"Saving {desc.shape[0]} descriptors to {descriptors_fullpath}")
+            image_descriptors_to_file(descriptors_fullpath, desc)
+            kdata.descriptors[args.descriptors_type].add(image_name)
+
+        if not keypoints_check_dir(kdata.keypoints[args.keypoints_type], args.keypoints_type,
+                                   args.kapture_root, tar_handlers) or \
+                not descriptors_check_dir(kdata.descriptors[args.descriptors_type], args.descriptors_type,
+                                          args.kapture_root, tar_handlers):
+            print('local feature extraction ended successfully but not all files were saved')
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser(
+        "Extract r2d2 local features for all images in a dataset stored in the kapture format")
+    parser.add_argument("--model", type=str, required=True, help='model path')
+    parser.add_argument('--keypoints-type', default=None,  help='keypoint type_name, default is filename of model')
+    parser.add_argument('--descriptors-type', default=None,  help='descriptors type_name, default is filename of model')
+
+    parser.add_argument("--kapture-root", type=str, required=True, help='path to kapture root directory')
+
+    parser.add_argument("--top-k", type=int, default=5000, help='number of keypoints')
+
+    parser.add_argument("--scale-f", type=float, default=2**0.25)
+    parser.add_argument("--min-size", type=int, default=256)
+    parser.add_argument("--max-size", type=int, default=1024)
+    parser.add_argument("--min-scale", type=float, default=0)
+    parser.add_argument("--max-scale", type=float, default=1)
+
+    parser.add_argument("--reliability-thr", type=float, default=0.7)
+    parser.add_argument("--repeatability-thr", type=float, default=0.7)
+
+    parser.add_argument("--gpu", type=int, nargs='+', default=[0], help='use -1 for CPU')
+    args = parser.parse_args()
+
+    extract_kapture_keypoints(args)

vismatch/third_party/imatch-toolbox/third_party/r2d2/nets/ap_loss.py

@@ -0,0 +1,67 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+import pdb
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class APLoss (nn.Module):
+    """ differentiable AP loss, through quantization.
+        
+        Input: (N, M)   values in [min, max]
+        label: (N, M)   values in {0, 1}
+        
+        Returns: list of query AP (for each n in {1..N})
+                 Note: typically, you want to minimize 1 - mean(AP)
+    """
+    def __init__(self, nq=25, min=0, max=1, euc=False):
+        nn.Module.__init__(self)
+        assert isinstance(nq, int) and 2 <= nq <= 100
+        self.nq = nq
+        self.min = min
+        self.max = max
+        self.euc = euc
+        gap = max - min
+        assert gap > 0
+        
+        # init quantizer = non-learnable (fixed) convolution
+        self.quantizer = q = nn.Conv1d(1, 2*nq, kernel_size=1, bias=True)
+        a = (nq-1) / gap
+        #1st half = lines passing to (min+x,1) and (min+x+1/a,0) with x = {nq-1..0}*gap/(nq-1)
+        q.weight.data[:nq] = -a
+        q.bias.data[:nq] = torch.from_numpy(a*min + np.arange(nq, 0, -1)) # b = 1 + a*(min+x)
+        #2nd half = lines passing to (min+x,1) and (min+x-1/a,0) with x = {nq-1..0}*gap/(nq-1)
+        q.weight.data[nq:] = a
+        q.bias.data[nq:] = torch.from_numpy(np.arange(2-nq, 2, 1) - a*min) # b = 1 - a*(min+x)
+        # first and last one are special: just horizontal straight line
+        q.weight.data[0] = q.weight.data[-1] = 0
+        q.bias.data[0] = q.bias.data[-1] = 1
+
+    def compute_AP(self, x, label):
+        N, M = x.shape
+        if self.euc:  # euclidean distance in same range than similarities
+            x = 1 - torch.sqrt(2.001 - 2*x)
+
+        # quantize all predictions
+        q = self.quantizer(x.unsqueeze(1))
+        q = torch.min(q[:,:self.nq], q[:,self.nq:]).clamp(min=0) # N x Q x M
+
+        nbs = q.sum(dim=-1) # number of samples  N x Q = c
+        rec = (q * label.view(N,1,M).float()).sum(dim=-1) # nb of correct samples = c+ N x Q
+        prec = rec.cumsum(dim=-1) / (1e-16 + nbs.cumsum(dim=-1)) # precision
+        rec /= rec.sum(dim=-1).unsqueeze(1) # norm in [0,1]
+
+        ap = (prec * rec).sum(dim=-1) # per-image AP
+        return ap
+
+    def forward(self, x, label):
+        assert x.shape == label.shape # N x M
+        return self.compute_AP(x, label)
+
+
+
+
+

vismatch/third_party/imatch-toolbox/third_party/r2d2/nets/losses.py

@@ -0,0 +1,56 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+import pdb
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nets.sampler import *
+from nets.repeatability_loss import *
+from nets.reliability_loss import *
+
+
+class MultiLoss (nn.Module):
+    """ Combines several loss functions for convenience.
+    *args: [loss weight (float), loss creator, ... ]
+    
+    Example:
+        loss = MultiLoss( 1, MyFirstLoss(), 0.5, MySecondLoss() )
+    """
+    def __init__(self, *args, dbg=()):
+        nn.Module.__init__(self)
+        assert len(args) % 2 == 0, 'args must be a list of (float, loss)'
+        self.weights = []
+        self.losses = nn.ModuleList()
+        for i in range(len(args)//2):
+            weight = float(args[2*i+0])
+            loss = args[2*i+1]
+            assert isinstance(loss, nn.Module), "%s is not a loss!" % loss
+            self.weights.append(weight)
+            self.losses.append(loss)
+
+    def forward(self, select=None, **variables):
+        assert not select or all(1<=n<=len(self.losses) for n in select)
+        d = dict()
+        cum_loss = 0
+        for num, (weight, loss_func) in enumerate(zip(self.weights, self.losses),1):
+            if select is not None and num not in select: continue
+            l = loss_func(**{k:v for k,v in variables.items()})
+            if isinstance(l, tuple):
+                assert len(l) == 2 and isinstance(l[1], dict)
+            else:
+                l = l, {loss_func.name:l}
+            cum_loss = cum_loss + weight * l[0]
+            for key,val in l[1].items():
+                d['loss_'+key] = float(val)
+        d['loss'] = float(cum_loss)
+        return cum_loss, d
+
+
+
+
+
+

vismatch/third_party/imatch-toolbox/third_party/r2d2/nets/patchnet.py

@@ -0,0 +1,134 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+import pdb
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BaseNet (nn.Module):
+    """ Takes a list of images as input, and returns for each image:
+        - a pixelwise descriptor
+        - a pixelwise confidence
+    """
+    def softmax(self, ux):
+        if ux.shape[1] == 1:
+            x = F.softplus(ux)
+            return x / (1 + x)  # for sure in [0,1], much less plateaus than softmax
+        elif ux.shape[1] == 2:
+            return F.softmax(ux, dim=1)[:,1:2]
+
+    def normalize(self, x, ureliability, urepeatability):
+        return dict(descriptors = F.normalize(x, p=2, dim=1),
+                    repeatability = self.softmax( urepeatability ),
+                    reliability = self.softmax( ureliability ))
+
+    def forward_one(self, x):
+        raise NotImplementedError()
+
+    def forward(self, imgs, **kw):
+        res = [self.forward_one(img) for img in imgs]
+        # merge all dictionaries into one
+        res = {k:[r[k] for r in res if k in r] for k in {k for r in res for k in r}}
+        return dict(res, imgs=imgs, **kw)
+
+
+
+class PatchNet (BaseNet):
+    """ Helper class to construct a fully-convolutional network that
+        extract a l2-normalized patch descriptor.
+    """
+    def __init__(self, inchan=3, dilated=True, dilation=1, bn=True, bn_affine=False):
+        BaseNet.__init__(self)
+        self.inchan = inchan
+        self.curchan = inchan
+        self.dilated = dilated
+        self.dilation = dilation
+        self.bn = bn
+        self.bn_affine = bn_affine
+        self.ops = nn.ModuleList([])
+
+    def _make_bn(self, outd):
+        return nn.BatchNorm2d(outd, affine=self.bn_affine)
+
+    def _add_conv(self, outd, k=3, stride=1, dilation=1, bn=True, relu=True):
+        d = self.dilation * dilation
+        if self.dilated: 
+            conv_params = dict(padding=((k-1)*d)//2, dilation=d, stride=1)
+            self.dilation *= stride
+        else:
+            conv_params = dict(padding=((k-1)*d)//2, dilation=d, stride=stride)
+        self.ops.append( nn.Conv2d(self.curchan, outd, kernel_size=k, **conv_params) )
+        if bn and self.bn: self.ops.append( self._make_bn(outd) )
+        if relu: self.ops.append( nn.ReLU(inplace=True) )
+        self.curchan = outd
+    
+    def forward_one(self, x):
+        assert self.ops, "You need to add convolutions first"
+        for n,op in enumerate(self.ops):
+            x = op(x)
+        return self.normalize(x)
+
+
+class L2_Net (PatchNet):
+    """ Compute a 128D descriptor for all overlapping 32x32 patches.
+        From the L2Net paper (CVPR'17).
+    """
+    def __init__(self, dim=128, **kw ):
+        PatchNet.__init__(self, **kw)
+        add_conv = lambda n,**kw: self._add_conv((n*dim)//128,**kw)
+        add_conv(32)
+        add_conv(32)
+        add_conv(64, stride=2)
+        add_conv(64)
+        add_conv(128, stride=2)
+        add_conv(128)
+        add_conv(128, k=7, stride=8, bn=False, relu=False)
+        self.out_dim = dim
+
+
+class Quad_L2Net (PatchNet):
+    """ Same than L2_Net, but replace the final 8x8 conv by 3 successive 2x2 convs.
+    """
+    def __init__(self, dim=128, mchan=4, relu22=False, **kw ):
+        PatchNet.__init__(self, **kw)
+        self._add_conv(  8*mchan)
+        self._add_conv(  8*mchan)
+        self._add_conv( 16*mchan, stride=2)
+        self._add_conv( 16*mchan)
+        self._add_conv( 32*mchan, stride=2)
+        self._add_conv( 32*mchan)
+        # replace last 8x8 convolution with 3 2x2 convolutions
+        self._add_conv( 32*mchan, k=2, stride=2, relu=relu22)
+        self._add_conv( 32*mchan, k=2, stride=2, relu=relu22)
+        self._add_conv(dim, k=2, stride=2, bn=False, relu=False)
+        self.out_dim = dim
+
+
+
+class Quad_L2Net_ConfCFS (Quad_L2Net):
+    """ Same than Quad_L2Net, with 2 confidence maps for repeatability and reliability.
+    """
+    def __init__(self, **kw ):
+        Quad_L2Net.__init__(self, **kw)
+        # reliability classifier
+        self.clf = nn.Conv2d(self.out_dim, 2, kernel_size=1)
+        # repeatability classifier: for some reasons it's a softplus, not a softmax!
+        # Why? I guess it's a mistake that was left unnoticed in the code for a long time...
+        self.sal = nn.Conv2d(self.out_dim, 1, kernel_size=1) 
+
+    def forward_one(self, x):
+        assert self.ops, "You need to add convolutions first"
+        for op in self.ops:
+            x = op(x)
+        # compute the confidence maps
+        ureliability = self.clf(x**2)
+        urepeatability = self.sal(x**2)
+        return self.normalize(x, ureliability, urepeatability)
+
+
+
+
+

vismatch/third_party/imatch-toolbox/third_party/r2d2/nets/reliability_loss.py

@@ -0,0 +1,59 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+import pdb
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nets.ap_loss import APLoss
+
+
+class PixelAPLoss (nn.Module):
+    """ Computes the pixel-wise AP loss:
+        Given two images and ground-truth optical flow, computes the AP per pixel.
+        
+        feat1:  (B, C, H, W)   pixel-wise features extracted from img1
+        feat2:  (B, C, H, W)   pixel-wise features extracted from img2
+        aflow:  (B, 2, H, W)   absolute flow: aflow[...,y1,x1] = x2,y2
+    """
+    def __init__(self, sampler, nq=20):
+        nn.Module.__init__(self)
+        self.aploss = APLoss(nq, min=0, max=1, euc=False)
+        self.name = 'pixAP'
+        self.sampler = sampler
+
+    def loss_from_ap(self, ap, rel):
+        return 1 - ap
+
+    def forward(self, descriptors, aflow, **kw):
+        # subsample things
+        scores, gt, msk, qconf = self.sampler(descriptors, kw.get('reliability'), aflow)
+        
+        # compute pixel-wise AP
+        n = qconf.numel()
+        if n == 0: return 0
+        scores, gt = scores.view(n,-1), gt.view(n,-1)
+        ap = self.aploss(scores, gt).view(msk.shape)
+
+        pixel_loss = self.loss_from_ap(ap, qconf)
+        
+        loss = pixel_loss[msk].mean()
+        return loss
+
+
+class ReliabilityLoss (PixelAPLoss):
+    """ same than PixelAPLoss, but also train a pixel-wise confidence
+        that this pixel is going to have a good AP.
+    """
+    def __init__(self, sampler, base=0.5, **kw):
+        PixelAPLoss.__init__(self, sampler, **kw)
+        assert 0 <= base < 1
+        self.base = base
+        self.name = 'reliability'
+
+    def loss_from_ap(self, ap, rel):
+        return 1 - ap*rel - (1-rel)*self.base
+
+
+

vismatch/third_party/imatch-toolbox/third_party/r2d2/nets/repeatability_loss.py

@@ -0,0 +1,66 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+import pdb
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nets.sampler import FullSampler
+
+class CosimLoss (nn.Module):
+    """ Try to make the repeatability repeatable from one image to the other.
+    """
+    def __init__(self, N=16):
+        nn.Module.__init__(self)
+        self.name = f'cosim{N}'
+        self.patches = nn.Unfold(N, padding=0, stride=N//2)
+
+    def extract_patches(self, sal):
+        patches = self.patches(sal).transpose(1,2) # flatten
+        patches = F.normalize(patches, p=2, dim=2) # norm
+        return patches
+        
+    def forward(self, repeatability, aflow, **kw):
+        B,two,H,W = aflow.shape
+        assert two == 2
+
+        # normalize
+        sali1, sali2 = repeatability
+        grid = FullSampler._aflow_to_grid(aflow)
+        sali2 = F.grid_sample(sali2, grid, mode='bilinear', padding_mode='border')
+
+        patches1 = self.extract_patches(sali1)
+        patches2 = self.extract_patches(sali2)
+        cosim = (patches1 * patches2).sum(dim=2)
+        return 1 - cosim.mean()
+
+
+class PeakyLoss (nn.Module):
+    """ Try to make the repeatability locally peaky.
+
+    Mechanism: we maximize, for each pixel, the difference between the local mean
+               and the local max.
+    """
+    def __init__(self, N=16):
+        nn.Module.__init__(self)
+        self.name = f'peaky{N}'
+        assert N % 2 == 0, 'N must be pair'
+        self.preproc = nn.AvgPool2d(3, stride=1, padding=1)
+        self.maxpool = nn.MaxPool2d(N+1, stride=1, padding=N//2)
+        self.avgpool = nn.AvgPool2d(N+1, stride=1, padding=N//2)
+
+    def forward_one(self, sali):
+        sali = self.preproc(sali) # remove super high frequency
+        return 1 - (self.maxpool(sali) - self.avgpool(sali)).mean()
+
+    def forward(self, repeatability, **kw):
+        sali1, sali2 = repeatability
+        return (self.forward_one(sali1) + self.forward_one(sali2)) /2
+
+
+
+
+