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 bb6import numpy as np7import torch8 9 10class DatasetWithRng:11    """ Make sure that RNG is distributed properly when torch.dataloader() is used12    """13 14    def __init__(self, seed=None):15        self.seed = seed16        self.rng = np.random.default_rng(seed)17        self._rng_children = set()18 19    def with_same_rng(self, dataset=None):20        if dataset is not None:21            assert isinstance(dataset, DatasetWithRng) and hasattr(dataset, 'rng'), bb()22            self._rng_children.add( dataset )23 24        # update all registered children25        for db in self._rng_children:26            db.rng = self.rng27            db.with_same_rng() # recursive call28        return dataset29 30    def init_worker(self, tid):31        if self.seed is None: 32            self.rng = np.random.default_rng()33        else:34            self.rng = np.random.default_rng(self.seed + tid)35 36 37class WorkerWithRngInit:38    " Dataset inherits from datasets.DatasetWithRng() and has an init_worker() function "39    def __call__(self, tid):40        torch.utils.data.get_worker_info().dataset.init_worker(tid)41 42 43def corres_from_homography(homography, W, H, grid=64):44    s = max(1, min(W, H) // grid) # at least `grid` points in smallest dim45    sx, sy = [slice(s//2, l, s) for l in (W, H)]46    grid1 = np.mgrid[sy, sx][::-1].reshape(2,-1).T # (x1,y1) grid47 48    grid2 = applyh(homography, grid1)49    scale = np.sqrt(np.abs(np.linalg.det(jacobianh(homography, grid1).T)))50 51    corres = np.c_[grid1, grid2, np.ones_like(scale), np.zeros_like(scale), scale]52    return corres53 54 55def invh( H ):56    return np.linalg.inv(H)57 58 59def applyh(H, p, ncol=2, norm=True):60    """ Apply the homography to a list of 2d points in homogeneous coordinates.61 62    H: Homography (...x3x3 matrix/tensor)63    p: numpy/torch/tuple of coordinates. Shape must be (...,2) or (...,3)64    65    Returns an array of projected 2d points.66    """67    if isinstance(H, np.ndarray):68        p = np.asarray(p)69    elif isinstance(H, torch.Tensor):70        p = torch.as_tensor(p, dtype=H.dtype)71 72    if p.shape[-1]+1 == H.shape[-1]:73        H = H.swapaxes(-1,-2) # transpose H74        p = p @ H[...,:-1,:] + H[...,-1:,:]75    else:76        p = H @ p.T77        if p.ndim >= 2: p = p.swapaxes(-1,-2)78 79    if norm: 80        p /= p[...,-1:]81    return p[...,:ncol]82 83 84def jacobianh(H, p):85    """ H is an homography that maps: f_H(x,y) --> (f_1, f_2)86    So the Jacobian J_H evaluated at p=(x,y) is a 2x2 matrix87    Output shape = (2, 2, N) = (f_, xy, N)88 89    Example of derivative:90                  numx    a*X + b*Y + c*Z91        since x = ----- = ---------------92                  denom   u*X + v*Y + w*Z93 94                numx' * denom - denom' * numx   a*denom - u*numx95        dx/dX = ----------------------------- = ----------------96                           denom**2                 denom**297    """98    (a, b, c), (d, e, f), (u, v, w) = H99    numx, numy, denom = applyh(H, p, ncol=3, norm=False).T100 101    #                column x          column x102    J = np.float32(((a*denom - u*numx, b*denom - v*numx),  # row f_1103                    (d*denom - u*numy, e*denom - v*numy))) # row f_2104    return J / np.where(denom, denom*denom, np.nan)105