|
|
|
|
|
|
|
|
|
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 ) |
|
|
|
|
|
for db in self._rng_children: |
|
db.rng = self.rng |
|
db.with_same_rng() |
|
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) |
|
sx, sy = [slice(s//2, l, s) for l in (W, H)] |
|
grid1 = np.mgrid[sy, sx][::-1].reshape(2,-1).T |
|
|
|
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) |
|
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 |
|
|
|
|
|
J = np.float32(((a*denom - u*numx, b*denom - v*numx), |
|
(d*denom - u*numy, e*denom - v*numy))) |
|
return J / np.where(denom, denom*denom, np.nan) |
|
|
