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)