vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/imatch-toolbox/third_party/r2d2/tools/dataloader.py

@@ -0,0 +1,367 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+import pdb
+from PIL import Image
+import numpy as np
+
+import torch
+import torchvision.transforms as tvf
+
+from tools.transforms import instanciate_transformation
+from tools.transforms_tools import persp_apply
+
+
+RGB_mean = [0.485, 0.456, 0.406]
+RGB_std  = [0.229, 0.224, 0.225]
+
+norm_RGB = tvf.Compose([tvf.ToTensor(), tvf.Normalize(mean=RGB_mean, std=RGB_std)])
+
+
+class PairLoader:
+    """ On-the-fly jittering of pairs of image with dense pixel ground-truth correspondences.
+    
+    crop:   random crop applied to both images
+    scale:  random scaling applied to img2
+    distort: random ditorsion applied to img2
+    
+    self[idx] returns a dictionary with keys: img1, img2, aflow, mask
+     - img1: cropped original
+     - img2: distorted cropped original
+     - aflow: 'absolute' optical flow = (x,y) position of each pixel from img1 in img2
+     - mask: (binary image) valid pixels of img1
+    """
+    def __init__(self, dataset, crop='', scale='', distort='', norm = norm_RGB, 
+                       what = 'aflow mask', idx_as_rng_seed = False):
+        assert hasattr(dataset, 'npairs')
+        assert hasattr(dataset, 'get_pair')
+        self.dataset = dataset
+        self.distort = instanciate_transformation(distort)
+        self.crop = instanciate_transformation(crop)
+        self.norm = instanciate_transformation(norm)
+        self.scale = instanciate_transformation(scale)
+        self.idx_as_rng_seed = idx_as_rng_seed # to remove randomness
+        self.what = what.split() if isinstance(what, str) else what
+        self.n_samples = 5 # number of random trials per image
+
+    def __len__(self):
+        assert len(self.dataset) == self.dataset.npairs, pdb.set_trace() # and not nimg
+        return len(self.dataset)
+
+    def __repr__(self):
+        fmt_str = 'PairLoader\n'
+        fmt_str += repr(self.dataset)
+        fmt_str += '  npairs: %d\n' % self.dataset.npairs
+        short_repr = lambda s: repr(s).strip().replace('\n',', ')[14:-1].replace('    ',' ')
+        fmt_str += '  Distort: %s\n' % short_repr(self.distort)
+        fmt_str += '  Crop: %s\n' % short_repr(self.crop)
+        fmt_str += '  Norm: %s\n' % short_repr(self.norm)
+        return fmt_str
+
+    def __getitem__(self, i):
+        #from time import time as now; t0 = now()
+        if self.idx_as_rng_seed:
+            import random
+            random.seed(i)
+            np.random.seed(i)
+
+        # Retrieve an image pair and their absolute flow
+        img_a, img_b, metadata = self.dataset.get_pair(i, self.what)
+        
+        # aflow contains pixel coordinates indicating where each 
+        # pixel from the left image ended up in the right image
+        # as (x,y) pairs, but its shape is (H,W,2)
+        aflow = np.float32(metadata['aflow'])
+        mask = metadata.get('mask', np.ones(aflow.shape[:2],np.uint8))
+
+        # apply transformations to the second image
+        img_b = {'img': img_b, 'persp':(1,0,0,0,1,0,0,0)}
+        if self.scale:
+            img_b = self.scale(img_b)
+        if self.distort:
+            img_b = self.distort(img_b)
+        
+        # apply the same transformation to the flow
+        aflow[:] = persp_apply(img_b['persp'], aflow.reshape(-1,2)).reshape(aflow.shape)
+        corres = None
+        if 'corres' in metadata:
+            corres = np.float32(metadata['corres'])
+            corres[:,1] = persp_apply(img_b['persp'], corres[:,1])
+        
+        # apply the same transformation to the homography
+        homography = None
+        if 'homography' in metadata:
+            homography = np.float32(metadata['homography'])
+            # p_b = homography * p_a
+            persp = np.float32(img_b['persp']+(1,)).reshape(3,3)
+            homography = persp @ homography
+
+        # determine crop size
+        img_b = img_b['img']
+        crop_size = self.crop({'imsize':(10000,10000)})['imsize']
+        output_size_a = min(img_a.size, crop_size)
+        output_size_b = min(img_b.size, crop_size)
+        img_a = np.array(img_a)
+        img_b = np.array(img_b)
+
+        ah,aw,p1 = img_a.shape
+        bh,bw,p2 = img_b.shape
+        assert p1 == 3
+        assert p2 == 3
+        assert aflow.shape == (ah, aw, 2)
+        assert mask.shape == (ah, aw)
+
+        # Let's start by computing the scale of the
+        # optical flow and applying a median filter:
+        dx = np.gradient(aflow[:,:,0])
+        dy = np.gradient(aflow[:,:,1])
+        scale = np.sqrt(np.clip(np.abs(dx[1]*dy[0] - dx[0]*dy[1]), 1e-16, 1e16))
+
+        accu2 = np.zeros((16,16), bool)
+        Q = lambda x, w: np.int32(16 * (x - w.start) / (w.stop - w.start))
+        
+        def window1(x, size, w):
+            l = x - int(0.5 + size / 2)
+            r = l + int(0.5 + size)
+            if l < 0: l,r = (0, r - l)
+            if r > w: l,r = (l + w - r, w)
+            if l < 0: l,r = 0,w # larger than width
+            return slice(l,r)
+        def window(cx, cy, win_size, scale, img_shape):
+            return (window1(cy, win_size[1]*scale, img_shape[0]), 
+                    window1(cx, win_size[0]*scale, img_shape[1]))
+
+        n_valid_pixel = mask.sum()
+        sample_w = mask / (1e-16 + n_valid_pixel)
+        def sample_valid_pixel():
+            n = np.random.choice(sample_w.size, p=sample_w.ravel())
+            y, x = np.unravel_index(n, sample_w.shape)
+            return x, y
+        
+        # Find suitable left and right windows
+        trials = 0 # take the best out of few trials
+        best = -np.inf, None
+        for _ in range(50*self.n_samples):
+            if trials >= self.n_samples: break # finished!
+
+            # pick a random valid point from the first image
+            if n_valid_pixel == 0: break
+            c1x, c1y = sample_valid_pixel()
+            
+            # Find in which position the center of the left
+            # window ended up being placed in the right image
+            c2x, c2y = (aflow[c1y, c1x] + 0.5).astype(np.int32)
+            if not(0 <= c2x < bw and 0 <= c2y < bh): continue
+
+            # Get the flow scale
+            sigma = scale[c1y, c1x]
+
+            # Determine sampling windows
+            if 0.2 < sigma < 1: 
+                win1 = window(c1x, c1y, output_size_a, 1/sigma, img_a.shape)
+                win2 = window(c2x, c2y, output_size_b, 1, img_b.shape)
+            elif 1 <= sigma < 5:
+                win1 = window(c1x, c1y, output_size_a, 1, img_a.shape)
+                win2 = window(c2x, c2y, output_size_b, sigma, img_b.shape)
+            else:
+                continue # bad scale
+
+            # compute a score based on the flow
+            x2,y2 = aflow[win1].reshape(-1, 2).T.astype(np.int32)
+            # Check the proportion of valid flow vectors
+            valid = (win2[1].start <= x2) & (x2 < win2[1].stop) \
+                  & (win2[0].start <= y2) & (y2 < win2[0].stop)
+            score1 = (valid * mask[win1].ravel()).mean()
+            # check the coverage of the second window
+            accu2[:] = False
+            accu2[Q(y2[valid],win2[0]), Q(x2[valid],win2[1])] = True
+            score2 = accu2.mean()
+            # Check how many hits we got
+            score = min(score1, score2)
+
+            trials += 1
+            if score > best[0]:
+                best = score, win1, win2
+        
+        if None in best: # counldn't find a good window
+            img_a = np.zeros(output_size_a[::-1]+(3,), dtype=np.uint8)
+            img_b = np.zeros(output_size_b[::-1]+(3,), dtype=np.uint8)
+            aflow = np.nan * np.ones((2,)+output_size_a[::-1], dtype=np.float32)
+            homography = np.nan * np.ones((3,3), dtype=np.float32)
+
+        else:
+            win1, win2 = best[1:]
+            img_a = img_a[win1]
+            img_b = img_b[win2]
+            aflow = aflow[win1] - np.float32([[[win2[1].start, win2[0].start]]])
+            mask = mask[win1]
+            aflow[~mask.view(bool)] = np.nan # mask bad pixels!
+            aflow = aflow.transpose(2,0,1) # --> (2,H,W)
+            
+            if corres is not None:
+                corres[:,0] -= (win1[1].start, win1[0].start)
+                corres[:,1] -= (win2[1].start, win2[0].start)
+            
+            if homography is not None:
+                trans1 = np.eye(3, dtype=np.float32)
+                trans1[:2,2] = (win1[1].start, win1[0].start)
+                trans2 = np.eye(3, dtype=np.float32)
+                trans2[:2,2] = (-win2[1].start, -win2[0].start)
+                homography = trans2 @ homography @ trans1
+                homography /= homography[2,2]
+            
+            # rescale if necessary
+            if img_a.shape[:2][::-1] != output_size_a:
+                sx, sy = (np.float32(output_size_a)-1)/(np.float32(img_a.shape[:2][::-1])-1)
+                img_a = np.asarray(Image.fromarray(img_a).resize(output_size_a, Image.ANTIALIAS))
+                mask = np.asarray(Image.fromarray(mask).resize(output_size_a, Image.NEAREST))
+                afx = Image.fromarray(aflow[0]).resize(output_size_a, Image.NEAREST)
+                afy = Image.fromarray(aflow[1]).resize(output_size_a, Image.NEAREST)
+                aflow = np.stack((np.float32(afx), np.float32(afy)))
+                
+                if corres is not None:
+                    corres[:,0] *= (sx, sy)
+                
+                if homography is not None:
+                    homography = homography @ np.diag(np.float32([1/sx,1/sy,1]))
+                    homography /= homography[2,2]
+
+            if img_b.shape[:2][::-1] != output_size_b:
+                sx, sy = (np.float32(output_size_b)-1)/(np.float32(img_b.shape[:2][::-1])-1)
+                img_b = np.asarray(Image.fromarray(img_b).resize(output_size_b, Image.ANTIALIAS))
+                aflow *= [[[sx]], [[sy]]]
+                
+                if corres is not None:
+                    corres[:,1] *= (sx, sy)
+                
+                if homography is not None:
+                    homography = np.diag(np.float32([sx,sy,1])) @ homography
+                    homography /= homography[2,2]
+    
+        assert aflow.dtype == np.float32, pdb.set_trace()
+        assert homography is None or homography.dtype == np.float32, pdb.set_trace()
+        if 'flow' in self.what:
+            H, W = img_a.shape[:2]
+            mgrid = np.mgrid[0:H, 0:W][::-1].astype(np.float32)
+            flow = aflow - mgrid
+        
+        result = dict(img1=self.norm(img_a), img2=self.norm(img_b))
+        for what in self.what:
+            try: result[what] = eval(what)
+            except NameError: pass
+        return result
+
+
+
+def threaded_loader( loader, iscuda, threads, batch_size=1, shuffle=True):
+    """ Get a data loader, given the dataset and some parameters.
+    
+    Parameters
+    ----------
+    loader : object[i] returns the i-th training example.
+    
+    iscuda : bool
+        
+    batch_size : int
+    
+    threads : int
+    
+    shuffle : int
+    
+    Returns
+    -------
+        a multi-threaded pytorch loader.
+    """
+    return torch.utils.data.DataLoader(
+        loader,
+        batch_size = batch_size,
+        shuffle = shuffle,
+        sampler = None,
+        num_workers = threads,
+        pin_memory = iscuda,
+        collate_fn=collate)
+
+
+
+def collate(batch, _use_shared_memory=True):
+    """Puts each data field into a tensor with outer dimension batch size.
+    Copied from https://github.com/pytorch in torch/utils/data/_utils/collate.py
+    """
+    import re
+    error_msg = "batch must contain tensors, numbers, dicts or lists; found {}"
+    elem_type = type(batch[0])
+    if isinstance(batch[0], torch.Tensor):
+        out = None
+        if _use_shared_memory:
+            # If we're in a background process, concatenate directly into a
+            # shared memory tensor to avoid an extra copy
+            numel = sum([x.numel() for x in batch])
+            storage = batch[0].storage()._new_shared(numel)
+            out = batch[0].new(storage)
+        return torch.stack(batch, 0, out=out)
+    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
+            and elem_type.__name__ != 'string_':
+        elem = batch[0]
+        assert elem_type.__name__ == 'ndarray'
+        # array of string classes and object
+        if re.search('[SaUO]', elem.dtype.str) is not None:
+            raise TypeError(error_msg.format(elem.dtype))
+        batch = [torch.from_numpy(b) for b in batch]
+        try:
+            return torch.stack(batch, 0)
+        except RuntimeError:
+            return batch
+    elif batch[0] is None:
+        return list(batch)
+    elif isinstance(batch[0], int):
+        return torch.LongTensor(batch)
+    elif isinstance(batch[0], float):
+        return torch.DoubleTensor(batch)
+    elif isinstance(batch[0], str):
+        return batch
+    elif isinstance(batch[0], dict):
+        return {key: collate([d[key] for d in batch]) for key in batch[0]}
+    elif isinstance(batch[0], (tuple,list)):
+        transposed = zip(*batch)
+        return [collate(samples) for samples in transposed]
+
+    raise TypeError((error_msg.format(type(batch[0]))))
+
+
+
+def tensor2img(tensor, model=None):
+    """ convert back a torch/numpy tensor to a PIL Image
+        by undoing the ToTensor() and Normalize() transforms.
+    """
+    mean = norm_RGB.transforms[1].mean
+    std =  norm_RGB.transforms[1].std
+    if isinstance(tensor, torch.Tensor):
+        tensor = tensor.detach().cpu().numpy()
+    
+    res = np.uint8(np.clip(255*((tensor.transpose(1,2,0) * std) + mean), 0, 255))
+    from PIL import Image
+    return Image.fromarray(res)
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser("Tool to debug/visualize the data loader")
+    parser.add_argument("dataloader", type=str, help="command to create the data loader")
+    args = parser.parse_args()
+
+    from datasets import *
+    auto_pairs = lambda db: SyntheticPairDataset(db,
+        'RandomScale(256,1024,can_upscale=True)', 
+        'RandomTilting(0.5), PixelNoise(25)')
+        
+    loader = eval(args.dataloader)
+    print("Data loader =", loader)
+
+    from tools.viz import show_flow
+    for data in loader:
+        aflow = data['aflow']
+        H, W = aflow.shape[-2:]
+        flow = (aflow - np.mgrid[:H, :W][::-1]).transpose(1,2,0)
+        show_flow(tensor2img(data['img1']), tensor2img(data['img2']), flow)
+

vismatch/third_party/imatch-toolbox/third_party/r2d2/tools/trainer.py

@@ -0,0 +1,76 @@
+# Copyright 2019-present NAVER Corp.
+# CC BY-NC-SA 3.0
+# Available only for non-commercial use
+
+import pdb
+from tqdm import tqdm
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+
+
+class Trainer (nn.Module):
+    """ Helper class to train a deep network.
+        Overload this class `forward_backward` for your actual needs.
+    
+    Usage: 
+        train = Trainer(net, loader, loss, optimizer)
+        for epoch in range(n_epochs):
+            train()
+    """
+    def __init__(self, net, loader, loss, optimizer):
+        nn.Module.__init__(self)
+        self.net = net
+        self.loader = loader
+        self.loss_func = loss
+        self.optimizer = optimizer
+
+    def iscuda(self):
+        return next(self.net.parameters()).device != torch.device('cpu')
+
+    def todevice(self, x):
+        if isinstance(x, dict):
+            return {k:self.todevice(v) for k,v in x.items()}
+        if isinstance(x, (tuple,list)):
+            return [self.todevice(v)  for v in x]
+        
+        if self.iscuda(): 
+            return x.contiguous().cuda(non_blocking=True)
+        else:
+            return x.cpu()
+
+    def __call__(self):
+        self.net.train()
+        
+        stats = defaultdict(list)
+        
+        for iter,inputs in enumerate(tqdm(self.loader)):
+            inputs = self.todevice(inputs)
+            
+            # compute gradient and do model update
+            self.optimizer.zero_grad()
+            
+            loss, details = self.forward_backward(inputs)
+            if torch.isnan(loss):
+                raise RuntimeError('Loss is NaN')
+            
+            self.optimizer.step()
+            
+            for key, val in details.items():
+                stats[key].append( val )
+        
+        print(" Summary of losses during this epoch:")
+        mean = lambda lis: sum(lis) / len(lis)
+        for loss_name, vals in stats.items():
+            N = 1 + len(vals)//10
+            print(f"  - {loss_name:20}:", end='')
+            print(f" {mean(vals[:N]):.3f} --> {mean(vals[-N:]):.3f} (avg: {mean(vals):.3f})")
+        return mean(stats['loss']) # return average loss
+
+    def forward_backward(self, inputs):
+        raise NotImplementedError()
+
+
+
+