def find_coeffs(p1, p2):
"Find coefficients for warp tfm from `p1` to `p2`"
m = []
p = p1[:, 0, 0]
#The equations we'll need to solve.
for i in range(p1.shape[1]):
m.append(
stack([
p2[:, i, 0], p2[:, i, 1],
t1(p),
t0(p),
t0(p),
t0(p), -p1[:, i, 0] * p2[:, i, 0], -p1[:, i, 0] * p2[:, i, 1]
]))
m.append(
stack([
t0(p),
t0(p),
t0(p), p2[:, i, 0], p2[:, i, 1],
t1(p), -p1[:, i, 1] * p2[:, i, 0], -p1[:, i, 1] * p2[:, i, 1]
]))
#The 8 scalars we seek are solution of AX = B
A = stack(m).permute(2, 0, 1)
B = p1.view(p1.shape[0], 8, 1)
return torch.solve(B, A)[0]
def flip_affine(x, p=0.5):
"Flip as an affine transform"
mask = -2 * x.new_empty(x.size(0)).bernoulli_(p) + 1
return stack([
stack([mask, t0(mask), t0(mask)], dim=1),
stack([t0(mask), t1(mask), t0(mask)], dim=1),
stack([t0(mask), t0(mask), t1(mask)], dim=1)
],
dim=1)
def flip_mat(x, p=0.5, draw=None, batch=False):
"Return a random flip matrix"
def _def_draw(x): return x.new_ones(x.size(0))
mask = x.new_ones(x.size(0)) - 2*_draw_mask(x, _def_draw, draw=draw, p=p, batch=batch)
#mask = mask_tensor(-x.new_ones(x.size(0)), p=p, neutral=1.)
return affine_mat(mask, t0(mask), t0(mask),
t0(mask), t1(mask), t0(mask))
def apply_perspective(coords, coeffs):
"Apply perspective tranfom on `coords` with `coeffs`"
sz = coords.shape
coords = coords.view(sz[0], -1, 2)
coeffs = torch.cat([coeffs, t1(coeffs[:,:1])], dim=1).view(coeffs.shape[0], 3,3)
coords = coords @ coeffs[...,:2].transpose(1,2) + coeffs[...,2].unsqueeze(1)
coords.div_(coords[...,2].unsqueeze(-1))
return coords[...,:2].view(*sz)
def affine_mat(*ms):
"Restructure length-6 vector `ms` into an affine matrix with 0,0,1 in the last line"
return stack([
stack([ms[0], ms[1], ms[2]], dim=1),
stack([ms[3], ms[4], ms[5]], dim=1),
stack([t0(ms[0]), t0(ms[0]), t1(ms[0])], dim=1)
],
dim=1)
def randomize(self, x):
idx = mask_tensor(torch.randint(0, 8, (x.size(0),), device=x.device), p=self.p)
xs = 1 - 2*(idx & 1)
ys = 1 - (idx & 2)
m0,m1 = (idx<4).long(),(idx>3).long()
self.mat = stack([stack([xs*m0, xs*m1, t0(xs)], dim=1),
stack([ys*m1, ys*m0, t0(xs)], dim=1),
stack([t0(xs), t0(xs), t1(xs)], dim=1)], dim=1).float()
def flip_mat(x, p=0.5):
"Return a random flip matrix"
mask = mask_tensor(-x.new_ones(x.size(0)), p=p, neutral=1.)
return affine_mat(mask, t0(mask), t0(mask),
t0(mask), t1(mask), t0(mask))
def flip_affine(x, p=0.5):
"Flip as an affine transform"
mask = -2 * x.new_empty(x.size(0)).bernoulli_(p) + 1
return affine_mat(mask, t0(mask), t0(mask), t0(mask), t1(mask), t0(mask),
t0(mask), t0(mask), t1(mask))
def randomize(self, x):
mask = -2*x.new_empty(x.size(0)).bernoulli_(self.p)+1
self.mat = stack([stack([mask, t0(mask), t0(mask)], dim=1),
stack([t0(mask), t1(mask), t0(mask)], dim=1),
stack([t0(mask), t0(mask), t1(mask)], dim=1)], dim=1)
