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