naver/PUMP

1# Copyright 2022-present NAVER Corp.2# CC BY-NC-SA 4.03# Available only for non-commercial use4 5from pdb import set_trace as bb6from tqdm import tqdm7import numpy as np8import torch9 10import test_singlescale as tss11import core.functional as myF12from tools.viz import dbgfig, show_correspondences13 14 15def arg_parser(parser = None):16    parser = parser or tss.arg_parser()17 18    parser.add_argument('--rec-overlap', type=float, default=0.5, help='overlap between tiles in [0,0.5]')19    parser.add_argument('--rec-score-thr', type=float, default=1, help='corres score threshold to guide fine levels')20    parser.add_argument('--rec-fast-thr', type=float, default=0.1, help='prune block if less than `fast` corres fall in it')21 22    return parser23 24 25class RecursivePUMP (tss.SingleScalePUMP):26    """ Recursive PUMP: 27        1) find initial correspondences at a coarse scale, 28        2) refine them at a selection of finer scales29    """30    def __init__(self, coarse_size=512, fine_size=512, rec_overlap=0.5, rec_score_thr=1.0, 31                       rec_fast_thr = 0.1, **other_options ):32        super().__init__(**other_options)33        assert 10 < coarse_size < 102434        assert 10 < fine_size < 102435        assert 0 <= rec_overlap < 136        assert 0 < rec_fast_thr < 137        self.coarse_size = coarse_size38        self.fine_size = fine_size39        self.overlap = rec_overlap40        self.score_thr = rec_score_thr41        self.fast_thr = rec_fast_thr42 43    @torch.no_grad()44    def forward(self, img1, img2, ret='corres', dbg=()):45        img1, sca1 = self.demultiplex_img_trf(img1, force=True)46        img2, sca2 = self.demultiplex_img_trf(img2, force=True)47        input_trfs = (sca1, sca2)48 49        # coarse first level with low-res images50        corres = self.coarse_correspondences(img1, img2)51 52        # fine level: iterate on HQ blocks53        accu1, accu2 = (self._make_accu(img1), self._make_accu(img2))54        for block1, block2 in tqdm(list(self._enumerate_blocks(img1, img2, corres))):55            # print(f"img1[{block1[}:{}, {}:{}]"56            accus, trfs = tss.SingleScalePUMP.forward(self, block1, block2, ret='raw', dbg=dbg)57            self._update_accu( accu1, accus[0], trfs[0][:2,2] )58            self._update_accu( accu2, accus[1], trfs[1][:2,2] )59 60        demul = lambda accu: (accu[:,:,:4].reshape(-1,4).clone(), accu[:,:,4].clone())61        corres = demul(accu1), demul(accu2)62        if dbgfig('corres', dbg): viz_correspondences(img1, img2, *corres, fig='last')63        corres = [(myF.affmul(input_trfs,pos),score) for pos, score in corres] # rectify scaling etc.64        if ret == 'raw': return corres, input_trfs65        return self.reciprocal(*corres)66 67    def coarse_correspondences(self, img1, img2, **kw):68        # joint image resize, because relative size is important (multiscale)69        shape1, shape2 = img1.shape[-2:], img2.shape[-2:]70        if max(shape1 + shape2) > self.coarse_size:71            f1 = self.coarse_size / max(shape1)72            f2 = self.coarse_size / max(shape2)73            f = min(f1, f2)74            img1 = myF.imresize( img1, int(0.5+f*max(shape1)) )75            img2 = myF.imresize( img2, int(0.5+f*max(shape2)) )76        else:77            f = 178 79        init_corres = tss.SingleScalePUMP.forward(self, img1, img2, **kw)80        # show_correspondences(img1, img2, init_corres, fig='last')81        corres = init_corres[init_corres[:,4] > self.score_thr]82        print(f"  keeping {len(corres)}/{len(init_corres)} corres with score > {self.score_thr} ...")83        return corres84 85    def _update_accu(self, accu, update, offset ):86        pos, scores = update87        H, W = scores.shape88        offx, offy = map(lambda i: int(i/4), offset)89        accu = accu[offy:offy+H, offx:offx+W]90        better = accu[:,:,4] < scores91        accu[:,:,4][better] = scores[better].float()92        accu[:,:,0:4][better] = pos.reshape(H,W,4)[better]93 94    def _enumerate_blocks(self, img1, img2, corres):95        H1, W1, H2, W2 = img1.shape[1:] + img2.shape[1:]96        size, step = self.fine_size, int(self.overlap * self.fine_size)97        def regular_steps(size): 98            if size <= self.fine_size: return [0]99            nb = int(np.ceil(size / step)) - 1 # garranted >= 1100            return (np.linspace(0, size-self.fine_size, nb) / 4 + 0.5).astype(int) * 4101        def translation(x,y):102            res = torch.eye(3, device=img1.device)103            res[0,2] = x104            res[1,2] = y105            return res106        def block2(x2,y2):107            return img2[:,y2:y2+size,x2:x2+size], translation(x2,y2)108        cx1, cy1 = corres[:,0:2].T109 110        for y1 in regular_steps(H1):111          for x1 in regular_steps(W1):112            block1 = (img1[:,y1:y1+size,x1:x1+size], translation(x1,y1))113            c2 = corres[(y1<=cy1) & (cy1<y1+size) & (x1<=cx1) & (cx1<x1+size)]114            nb_init = len(c2)115            while len(c2):116                cx2, cy2 = c2[:,2:4].T117                x2, y2 = (int(max(0,min(W2-size,cx2.median()-size//2)) / 4 + 0.5) * 4, 118                          int(max(0,min(H2-size,cy2.median()-size//2)) / 4 + 0.5) * 4)119                inside = (y2<=cy2) & (cy2<y2+size) & (x2<=cx2) & (cx2<x2+size)120                if not inside.any(): 121                    x2, y2 = c2[np.random.choice(len(c2)),2:4]122                    x2 = int(max(0,min(W2-size,x2-size//2)) / 4 + 0.5) * 4123                    y2 = int(max(0,min(H2-size,y2-size//2)) / 4 + 0.5) * 4124                    inside = (y2<=cy2) & (cy2<y2+size) & (x2<=cx2) & (cx2<x2+size)125 126                if inside.sum()/nb_init >= self.fast_thr:127                    yield block1, block2(x2,y2)128 129                c2 = c2[~inside] # remove130 131    def _make_accu(self, img):132        C, H, W = img.shape133        return img.new_zeros(((H+3)//4, (W+3)//4, 5), dtype=torch.float32)134 135 136 137class Main (tss.Main):138    @staticmethod139    def build_matcher(args, device):140        # set coarse and fine size based on now obsolete --resize argument141        if isinstance(args.resize, int): args.resize = [args.resize]142        if len(args.resize) == 1: args.resize *= 2143        args.rec_coarse_size, args.rec_fine_size = args.resize144        args.resize = 0 # disable it so that image loading does not downsize images145 146        options = Main.get_options( args )147 148        matcher = RecursivePUMP( coarse_size=args.rec_coarse_size, fine_size=args.rec_fine_size, 149            rec_overlap=args.rec_overlap, rec_score_thr=args.rec_score_thr, rec_fast_thr=args.rec_fast_thr,150            **options)151 152        return tss.Main.tune_matcher(matcher, **vars(args) ).to(device)153 154 155if __name__ == '__main__':156    Main().run_from_args(arg_parser().parse_args())157