test_singlescale.py

# Copyright 2022-present NAVER Corp.
# CC BY-NC-SA 4.0
# Available only for non-commercial use

from pdb import set_trace as bb
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from core import functional as myF
from core.pixel_desc import PixelDesc
from tools.common import mkdir_for, todevice, cudnn_benchmark, nparray, image, image_with_trf
from tools.viz import dbgfig, show_correspondences


def arg_parser():
    import argparse
    parser = argparse.ArgumentParser('SingleScalePUMP on GPU with PyTorch')

    parser.add_argument('--img1', required=True, help='path to img1')
    parser.add_argument('--img2', required=True, help='path to img2')
    parser.add_argument('--resize', type=int, default=512, nargs='+', help='prior downsize of img1 and img2')

    parser.add_argument('--output', default=None, help='output path for correspondences')

    parser.add_argument('--levels', type=int, default=99, help='number of pyramid levels')
    parser.add_argument('--min-shape', type=int, default=5, help='minimum size of corr maps')
    parser.add_argument('--nlpow', type=float, default=1.5, help='non-linear activation power in [1,2]')
    parser.add_argument('--border', type=float, default=0.9, help='border invariance level in [0,1]')
    parser.add_argument('--dtype', default='float16', choices='float16 float32 float64'.split())

    parser.add_argument('--desc', default='PUMP-stytrf', help='checkpoint name')
    parser.add_argument('--first-level', choices='torch'.split(), default='torch')
    parser.add_argument('--activation', choices='torch'.split(), default='torch')
    parser.add_argument('--forward', choices='torch cuda cuda-lowmem'.split(), default='cuda-lowmem')
    parser.add_argument('--backward', choices='python torch cuda'.split(), default='cuda')
    parser.add_argument('--reciprocal', choices='cpu cuda'.split(), default='cpu')

    parser.add_argument('--post-filter', default=None, const=True, nargs='?', help='post-filtering (See post_filter.py)')

    parser.add_argument('--verbose', type=int, default=0, help='verbosity')
    parser.add_argument('--device', default='cuda', help='gpu device')
    parser.add_argument('--dbg', nargs='*', default=(), help='debug options')

    return parser


class SingleScalePUMP (nn.Module):
    def __init__(self, levels = 9, nlpow = 1.4, cutoff = 1, 
                 border_inv=0.9, min_shape=5, renorm=(),
                 pixel_desc = None, dtype = torch.float32,
                 verbose = True ):
        super().__init__()
        self.levels = levels
        self.min_shape = min_shape
        self.nlpow = nlpow
        self.border_inv = border_inv
        assert pixel_desc, 'Requires a pixel descriptor'
        self.pixel_desc = pixel_desc.configure(self)
        self.dtype = dtype
        self.verbose = verbose

    @torch.no_grad()
    def forward(self, img1, img2, ret='corres', dbg=()):
        with cudnn_benchmark(False):
            # compute descriptors
            (img1, img2), pixel_descs, trfs = self.extract_descs(img1, img2, dtype=self.dtype)

            # backward and forward passes
            pixel_corr = self.first_level(*pixel_descs, dbg=dbg)
            pixel_corr = self.backward_pass(self.forward_pass(pixel_corr, dbg=dbg), dbg=dbg)

            # recover correspondences
            corres = myF.best_correspondences( pixel_corr )

        if dbgfig('corres', dbg): viz_correspondences(img1[0], img2[0], *corres, fig='last')
        corres = [(myF.affmul(trfs,pos),score) for pos, score in corres] # rectify scaling etc.
        if ret == 'raw': return corres, trfs
        return self.reciprocal(*corres)

    def extract_descs(self, img1, img2, dtype=None):
        img1, sca1 = self.demultiplex_img_trf(img1)
        img2, sca2 = self.demultiplex_img_trf(img2)
        desc1, trf1 = self.pixel_desc(img1)
        desc2, trf2 = self.pixel_desc(img2)
        return (img1, img2), (desc1.type(dtype), desc2.type(dtype)), (sca1@trf1, sca2@trf2)

    def demultiplex_img_trf(self, img, **kw):
        return img if isinstance(img, tuple) else (img, torch.eye(3, device=img.device))

    def forward_pass(self, pixel_corr, dbg=()):
        weights = None
        if isinstance(pixel_corr, tuple):
            pixel_corr, weights = pixel_corr

        # first-level with activation
        if self.verbose: print(f'  Pyramid level {0} shape={tuple(pixel_corr.shape)}')
        pyramid = [ self.activation(0,pixel_corr) ]
        if dbgfig(f'corr0', dbg): viz_correlation_maps(*from_stack('img1','img2'), pyramid[0], fig='last')

        for level in range(1, self.levels+1):
            upper, weights = self.forward_level(level, pyramid[-1], weights)
            if weights.sum() == 0: break # img1 has become too small

            # activation
            pyramid.append( self.activation(level,upper) )

            if self.verbose: print(f'  Pyramid level {level} shape={tuple(upper.shape)}')
            if dbgfig(f'corr{level}', dbg): viz_correlation_maps(*from_stack('img1','img2'), upper, level=level, fig='last')
            if min(upper.shape[-2:]) <= self.min_shape: break # img2 has become too small

        return pyramid

    def forward_level(self, level, corr, weights):
        # max-pooling
        pooled = F.max_pool2d(corr, 3, padding=1, stride=2)

        # sparse conv
        return myF.sparse_conv(level, pooled, weights, norm=self.border_inv)

    def backward_pass(self, pyramid, dbg=()):
        # same than forward in reverse order
        for level in range(len(pyramid)-1, 0, -1):
            lower = self.backward_level(level, pyramid)
            # assert not torch.isnan(lower).any(), bb()
            if self.verbose: print(f'  Pyramid level {level-1} shape={tuple(lower.shape)}')
            del pyramid[-1] # free memory
            if dbgfig(f'corr{level}-bw', dbg): viz_correlation_maps(img1, img2, lower, fig='last')
        return pyramid[0]

    def backward_level(self, level, pyramid):
        # reverse sparse-coonv
        pooled = myF.sparse_conv(level, pyramid[level], reverse=True)

        # reverse max-pool and add to lower level
        return myF.max_unpool(pooled, pyramid[level-1])

    def activation(self, level, corr):
        assert 1 <= self.nlpow <= 3
        corr.clamp_(min=0).pow_(self.nlpow)
        return corr

    def first_level(self, desc1, desc2, dbg=()):
        assert desc1.ndim == desc2.ndim == 4
        assert len(desc1) == len(desc2) == 1, "not implemented"
        H1, W1 = desc1.shape[-2:]
        H2, W2 = desc2.shape[-2:]

        patches = F.unfold(desc1, 4, stride=4) # C*4*4, H1*W1//16
        B, C, N = patches.shape
        # rearrange(patches, 'B (C Kh Kw) H1W1 -> B H1W1 C Kh Kw', Kh=4, Kw=4)
        patches = patches.permute(0, 2, 1).view(B, H1W1, C//16, 4, 4)

        corr, norms = myF.normalized_corr(patches[0], desc2[0], ret_norms=True)
        if dbgfig('ncc',dbg):
            for j in range(0,len(corr),9):
              for i in range(9):
                pl.subplot(3,3,i+1).cla()
                i += j
                pl.imshow(corr[i], vmin=0.9, vmax=1)
                pl.plot(2+(i%16)*4, 2+(i//16)*4,'xr', ms=10)
              bb()
        return corr.view(H1//4, W1//4, H2+1, W2+1), (norms.view(H1//4, W1//4)>0).float()

    def reciprocal(self, corres1, corres2 ):
        corres1, corres2 = todevice(corres1, 'cpu'), todevice(corres2, 'cpu')
        return myF.reciprocal(self, corres1, corres2)


class Main:
    def __init__(self):
        self.post_filtering = False

    def run_from_args(self, args):
        device = args.device
        self.matcher = self.build_matcher(args, device)
        if args.post_filter:
            self.post_filtering = {} if args.post_filter is True else eval(f'dict({args.post_filter})')

        corres = self(*self.load_images(args, device), dbg=set(args.dbg))

        if args.output:
            self.save_output( args.output, corres )

    def run_from_args_with_images(self, img1, img2, args):
        device = args.device
        self.matcher = self.build_matcher(args, device)
        if args.post_filter:
            self.post_filtering = {} if args.post_filter is True else eval(f'dict({args.post_filter})')
        
        if isinstance(args.resize, int): # user can provide 2 separate sizes for each image
            args.resize = (args.resize, args.resize)

        if len(args.resize) == 1: 
            args.resize = 2 * args.resize

        images = []
        for imgx, size in zip([img1, img2], args.resize):
            img = torch.from_numpy(np.array(imgx.convert('RGB'))).permute(2,0,1).to(device)
            img = myF.imresize(img, size)
            images.append( img )
        
        corres = self(*images, dbg=set(args.dbg))
        
        if args.output:
            self.save_output( args.output, corres )
        
        return corres 
        

    @staticmethod
    def get_options( args ):
        # configure the pipeline
        pixel_desc = PixelDesc(path=f'checkpoints/{args.desc}.pt')
        return dict(levels=args.levels, min_shape=args.min_shape, border_inv=args.border, nlpow=args.nlpow,
                    pixel_desc=pixel_desc, dtype=eval(f'torch.{args.dtype}'), verbose=args.verbose)

    @staticmethod
    def tune_matcher( args, matcher, device ):
        if device == 'cpu': 
            matcher.dtype = torch.float32
            args.forward = 'torch'
            args.backward = 'torch'
            args.reciprocal = 'cpu'

        if args.forward == 'cuda':       type(matcher).forward_level = myF.forward_cuda
        if args.forward == 'cuda-lowmem':type(matcher).forward_level = myF.forward_cuda_lowmem
        if args.backward == 'python':    type(matcher).backward_pass = legacy.backward_python
        if args.backward == 'cuda':      type(matcher).backward_level = myF.backward_cuda
        if args.reciprocal == 'cuda':    type(matcher).reciprocal = myF.reciprocal

        return matcher.to(device)

    @staticmethod
    def build_matcher(args, device):
        options = Main.get_options(args)
        matcher = SingleScalePUMP(**options)
        return Main.tune_matcher(args, matcher, device)

    def __call__(self, *imgs, dbg=()):
        corres = self.matcher( *imgs, dbg=dbg).cpu().numpy()
        if self.post_filtering is not False: 
            corres = self.post_filter( imgs, corres )

        if 'print' in dbg: print(corres)
        if dbgfig('viz',dbg):   show_correspondences(*imgs, corres)
        return corres

    @staticmethod
    def load_images( args, device='cpu' ):
        def read_image(impath):
            try:
                from torchvision.io.image import read_image, ImageReadMode
                return read_image(impath, mode=ImageReadMode.RGB)
            except RuntimeError:
                from PIL import Image
                return torch.from_numpy(np.array(Image.open(impath).convert('RGB'))).permute(2,0,1)

        if isinstance(args.resize, int): # user can provide 2 separate sizes for each image
            args.resize = (args.resize, args.resize)

        if len(args.resize) == 1: 
            args.resize = 2 * args.resize

        images = []
        for impath, size in zip([args.img1, args.img2], args.resize):
            img = read_image(impath).to(device)
            img = myF.imresize(img, size)
            images.append( img )
        return images

    def post_filter(self, imgs, corres ):
        from post_filter import filter_corres
        return filter_corres(*map(image_with_trf,imgs), corres, **self.post_filtering)

    def save_output(self, output_path, corres ):
        mkdir_for( output_path )
        np.savez(open(output_path,'wb'), corres=corres)



if __name__ == '__main__':
    Main().run_from_args(arg_parser().parse_args())