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 PIL import Image7import numpy as np8 9from core import functional as myF10from tools.common import todevice11from .transforms import instanciate_transforms12from .utils import *13 14 15class FastPairLoader (DatasetWithRng):16    """ On-the-fly generation of related image pairs17    crop:   random crop applied to both images18    scale:  random scaling applied to img219    distort: random ditorsion applied to img220    21    self[idx] returns: (img1, img2), dict(homography=)22        (homography: 3x3 array, can be nan)23    """24    def __init__(self, dataset, crop=256, transform='', p_flip=0, p_swap=0, scale_jitter=0, seed=None):25        super().__init__(seed)26        self.dataset = self.with_same_rng(dataset)27        self.transform = instanciate_transforms( transform, rng=self.rng )28        self.crop_size = crop29        self.p_swap = p_swap30        self.p_flip = p_flip31        self.scale_jitter = abs(np.log1p(scale_jitter))32 33    def __len__(self):34        return len(self.dataset)35 36    def __repr__(self):37        fmt_str = f'FastPairLoader({self.dataset},\n'38        short_repr = lambda s: repr(s).strip().replace('\n',', ')[14:-1].replace('    ',' ')39        fmt_str += '    Transform:\t%s\n' % short_repr(self.transform)40        fmt_str +=f'    Crop={self.crop_size}, scale_jitter=x{np.exp(self.scale_jitter):g}, p_swap={self.p_swap:g}'41        return fmt_str42 43    def init_worker(self, tid):44        super().init_worker(tid)45        self.dataset.init_worker(tid)46 47    def set_epoch(self, epoch):48        self.dataset.set_epoch(epoch)49 50    def __getitem__(self, idx):51        self.init_worker(idx) # preserve RNG for this pair52        (img1, img2), gt = self.dataset[idx]53 54        if self.rng.random() < self.p_swap:55            img1, img2 = img2, img156            if 'homography' in gt: gt['homography'] = invh(gt['homography'])57            if 'corres' in gt: gt['corres'] = swap_corres(gt['corres'])58 59        if self.rng.random() < self.p_flip:60            img1, img2, gt = flip_image_pair(img1, img2, gt)61 62        # apply transformations to the second image63        img2 = self.transform(dict(img=img2))64 65        homography, corres = spatial_relationship( img1, img2, gt )66 67        # find a good window68        img1, img2 = map(self._pad_rgb_numpy, (img1, img2['img']))69 70        if not 'debug':71            from tools.viz import show_correspondences72            print(np.median(corres[:,5]))73            show_correspondences(img1, img2, corres, bb=bb)74 75        def windows_from_corres( idx, scale_jitter=1 ):76            c = corres[idx]77            p1, p2, scale = c[0:2], c[2:4], c[6]78            scale *= scale_jitter79 80            # make windows based on scaling81            win1 = window(*p1, self.crop_size, max(1, 1/scale), img1.shape)82            win2 = window(*p2, self.crop_size, max(1, scale/1), img2.shape)83            return win1, win284 85        best = 0, None86        for idx in self.rng.choice(len(corres), size=min(len(corres),5), replace=False):87            # pick a correspondence at random88            win1, win2 = windows_from_corres( idx )89 90            # check how many matches are in the 2 windows91            score = score_windows(is_in(corres[:,0:2],win1), is_in(corres[:,2:4],win2))92            if score > best[0]: best = score, idx93 94        others = {}95        if None in best: # counldn't find a good window96            img1 = img2 = np.zeros((self.crop_size,self.crop_size,3), dtype=np.uint8)97            corres = np.empty((0, 6), dtype=np.float32)98        else:99            # jitter scales100            scale_jitter = np.exp(self.rng.uniform(-self.scale_jitter, self.scale_jitter))101            win1, win2 = windows_from_corres( best[1], scale_jitter )102            # print(win1, win2, img1.shape, img2.shape)103            img1, img2 = imresize(img1[win1], self.crop_size), imresize(img2[win2], self.crop_size)104            trf1, trf2 = wintrf(win1, img1), wintrf(win2, img2)105 106            # fix rotation if necessary107            angle_scores = np.bincount(corres[:,5].astype(int) % 8)108            rot90 = int((((angle_scores.argmax() + 4) % 8) - 4) / 2)109            if rot90: # rectify rotation110                img2, trf = myF.rotate_img_90((img2, np.eye(3)), 90*rot90)111                trf2 = invh(trf) @ trf2112 113            homography = trf2 @ homography @ invh(trf1)114            corres = myF.affmul((trf1,trf2), corres)115 116        f32c = lambda i,**kw: np.require(i, requirements='CWAE', **kw)117        return (f32c(img1), f32c(img2)), dict(homography = f32c(homography, dtype=np.float32), corres=corres, **others)118 119    def _pad_rgb_numpy(self, img):120        if img.mode != 'RGB': 121            img = img.convert('RGB')122        if min(img.size) < self.crop_size:123            w, h = img.size124            result = Image.new('RGB', (max(w,self.crop_size), max(h,self.crop_size)), 0)125            result.paste(img, (0, 0))126            img = result127        return np.asarray(img)128 129 130 131def swap_corres( corres ):  # swap img1 and img2132    res = corres.copy()133    res[:,[0,1,2,3]] = corres[:,[2,3,0,1]]134    if corres.shape[1] > 4: # invert rotation and scale135        scale, rot = myF.decode_scale_rot(corres[:,5])136        res[:,5] = myF.encode_scale_rot(1/scale, -rot)137    return res138 139def flip(img):140    w, h = img.size141    return img.transpose(Image.FLIP_LEFT_RIGHT), np.float32( [[-1,0,w-1],[0,1,0],[0,0,1]] )142 143def flip_image_pair(img1, img2, gt):144    img1, F1 = flip(img1)145    img2, F2 = flip(img2)146    res = {}147    for key, value in gt.items():148        if key == 'homography':149            res['homography'] = F2 @ value @ F1150        elif key == 'aflow':151            assert False, 'flip for aflow: todo'152        elif key == 'corres':153            new_corres = np.c_[applyh(F1,value[:,0:2]), applyh(F2,value[:,2:4])]154            if value.shape[1] == 4: pass155            elif value.shape[1] == 6:156                scale, rot = myF.decode_scale_rot(value[:,5])157                new_code = myF.encode_scale_rot(scale, -rot)158                new_corres = np.c_[new_corres,value[:,4],new_code]159            res['corres'] = new_corres160        else:161            raise ValueError(f"flip_image_pair: bad gt field '{key}'")162    return img1, img2, res163 164 165def spatial_relationship( img1, img2, gt ):166    if 'homography' in gt:167        homography = gt['homography']168        if 'homography' in img2: 169            homography = np.float32(img2['homography']) @ homography170        corres = corres_from_homography(homography, *img1.size)171 172    elif 'corres' in gt:173        homography = np.full((3,3), np.nan, dtype=np.float32)174        corres = gt['corres']175        if 'homography' in img2:176            corres[:,2:4] = applyh(img2['homography'], corres[:,2:4])177        else:178            img2['homography'] = np.eye(3)179        scales = np.sqrt(np.abs(np.linalg.det(jacobianh(img2['homography'], corres[:,0:2]).T)))180 181        if corres.shape[1] == 4:182            scales, rots = scale_rot_from_corres(corres)183            corres = np.c_[corres, np.ones_like(scales), myF.encode_scale_rot(scales,rots*180/np.pi), scales]184        elif corres.shape[1] == 6:185            corres = np.c_[corres, scales * myF.decode_scale_rot(corres[:,5])[0]]186        else:187            assert ValueError(f'bad shape for corres: {corres.shape}')188 189    return homography, corres190 191 192def scale_rot_from_corres( corres, sub=256, nn=16 ):193    # select a subset of relevant correspondences194    sub = np.random.choice(len(corres), size=min(len(corres),sub), replace=False)195    sub = corres[sub]196 197    # for each corres, find the scale change w.r.t. its NNs198    from scipy.spatial.distance import cdist199    nns = cdist(corres, sub, metric='sqeuclidean').argsort(axis=1)[:,:nn]200 201    # affine transform for this set of neighboring correspondences202    pts = sub[nns] # shape = npts x sub x 4203    # [P1,1] @ A = P2  with A = 3x2 matrix204    # A = [P1,1]^-1 @ P2205    P1, P2 = pts[:,:,0:2], pts[:,:,2:4] # each row = list of correspondences206    P1 = np.concatenate((P1,np.ones_like(P1[:,:,:1])),axis=-1)207    A = (np.linalg.pinv(P1) @ P2).transpose(0,2,1)208 209    scale, (angy,angx) = detect_scale_rotation(A.transpose(1,2,0)[:,1::-1])210    rot = np.arctan2(angy, angx)211    return scale.clip(min=0.2, max=5), rot212 213 214def window1(x, size, w):215    l = x - int(0.5 + size / 2)216    r = l + int(0.5 + size)217    if l < 0: l,r = (0, r - l)218    if r > w: l,r = (l + w - r, w)219    if l < 0: l,r = 0,w # larger than width220    return slice(l,r)221 222def window(cx, cy, win_size, scale, img_shape):223    return (window1(int(cy), win_size*scale, img_shape[0]), 224            window1(int(cx), win_size*scale, img_shape[1]))225 226def is_in( pts, window ):227    x, y = pts.T228    sly, slx = window229    return (slx.start <= x) & (x < slx.stop) & (sly.start <= y) & (y < sly.stop)230 231def score_windows( valid1, valid2 ):232    inter = (valid1 & valid2).sum()233    iou1 = inter / (valid1.sum() + 1e-8)234    iou2 = inter / (valid2.sum() + 1e-8)235    return inter * min(iou1, iou2)236 237def imresize( img, max_size, resample=Image.ANTIALIAS):238    if max(img.shape[:2]) > max_size:239        if img.shape[-1] == 2:240            img = np.stack([np.float32(Image.fromarray(img[...,i]).resize((max_size,max_size), resample=resample)) for i in range(2)], axis=-1)241        else:242            img = np.asarray(Image.fromarray(img).resize((max_size,max_size), resample=resample))243    assert img.shape[0] == img.shape[1] == max_size, bb()244    return img245 246def wintrf( window, final_img ):247    wy, wx = window248    H, W = final_img.shape[:2]249    T = np.float32((((wx.stop-wx.start)/W, 0, wx.start),250                    (0, (wy.stop-wy.start)/H, wy.start),251                    (0, 0, 1)) )252    return invh(T)253 254 255def collate_ordered(batch, _use_shared_memory=True):256    pairs, gt = zip(*batch)257    imgs1, imgs2 = zip(*pairs)258    assert len(imgs1) == len(imgs2) == len(gt) and isinstance(gt[0], dict)259    260    # reorder samples (supervised ones first, unsupervised ones last)261    supervised = [i for i,b in enumerate(gt) if np.isfinite(b['homography']).all()]262    unsupervsd = [i for i,b in enumerate(gt) if np.isnan(b['homography']).any()]263    order = supervised + unsupervsd264 265    def collate( tensors, key=None ):266        import torch267        batch = todevice([tensors[i] for i in order], 'cpu')268        if key == 'corres': return batch # cannot concat269        if _use_shared_memory: # shared memory tensor to avoid an extra copy270            numel = sum([x.numel() for x in batch])271            storage = batch[0].storage()._new_shared(numel)272            out = batch[0].new(storage)273        return torch.stack(batch, dim=0, out=out)274 275    return (collate(imgs1), collate(imgs2)), {k:collate([b[k] for b in gt],k) for k in gt[0]}276 277 278if __name__ == '__main__':279    from datasets import *280    from tools.viz import show_random_pairs281 282    db = BalancedCatImagePairs(283                3125, SyntheticImagePairs(RandomWebImages(0,52),distort='RandomTilting(0.5)'),284                4875, SyntheticImagePairs(SfM120k_Images(),distort='RandomTilting(0.5)'),285                8000, SfM120k_Pairs())286 287    db = FastPairLoader(db, 288            crop=256, transform='RandomRotation(20), RandomScale(256,1536,ar=1.3,can_upscale=True), PixelNoise()',289            p_swap=0.5, p_flip=0.5, scale_jitter=0, seed=777)290 291    show_random_pairs(db)292