def dihedral_mat(x, p=0.5, draw=None):
"Return a random dihedral matrix"
def _def_draw(x):
return torch.randint(0, 8, (x.size(0), ), device=x.device)
idx = _draw_mask(x, _def_draw, draw=draw, p=p).long()
xs = tensor([1, -1, 1, -1, -1, 1, 1, -1], device=x.device).gather(0, idx)
ys = tensor([1, 1, -1, 1, -1, -1, 1, -1], device=x.device).gather(0, idx)
m0 = tensor([1, 1, 1, 0, 1, 0, 0, 0], device=x.device).gather(0, idx)
m1 = tensor([0, 0, 0, 1, 0, 1, 1, 1], device=x.device).gather(0, idx)
return affine_mat(xs * m0, xs * m1, t0(xs), ys * m1, ys * m0,
t0(xs)).float()
mask = mask_tensor(-x.new_ones(x.size(0)), p=p, neutral=1.)
def rotate_mat(x, max_deg=10, p=0.5, draw=None, batch=False):
"Return a random rotation matrix with `max_deg` and `p`"
def _def_draw(x):
return x.new(x.size(0)).uniform_(-max_deg, max_deg)
def _def_draw_b(x):
return x.new_zeros(x.size(0)) + random.uniform(-max_deg, max_deg)
thetas = _draw_mask(
x, _def_draw_b if batch else _def_draw, draw=draw, p=p,
batch=batch) * math.pi / 180
return affine_mat(thetas.cos(), thetas.sin(), t0(thetas), -thetas.sin(),
thetas.cos(), t0(thetas))
def zoom_mat(x, max_zoom=1.1, p=0.5, draw=None, draw_x=None, draw_y=None, batch=False):
"Return a random zoom matrix with `max_zoom` and `p`"
def _def_draw(x): return x.new(x.size(0)).uniform_(1, max_zoom)
def _def_draw_b(x): return x.new_zeros(x.size(0)) + random.uniform(1, max_zoom)
def _def_draw_ctr(x): return x.new(x.size(0)).uniform_(0,1)
def _def_draw_ctr_b(x): return x.new_zeros(x.size(0)) + random.uniform(0,1)
s = 1/_draw_mask(x, _def_draw_b if batch else _def_draw, draw=draw, p=p, neutral=1., batch=batch)
def_draw_c = _def_draw_ctr_b if batch else _def_draw_ctr
col_pct = _draw_mask(x, def_draw_c, draw=draw_x, p=1., batch=batch)
row_pct = _draw_mask(x, def_draw_c, draw=draw_y, p=1., batch=batch)
col_c = (1-s) * (2*col_pct - 1)
row_c = (1-s) * (2*row_pct - 1)
return affine_mat(s, t0(s), col_c,
t0(s), s, row_c)
def dihedral_mat(x, p=0.5, draw=None):
"Return a random dihedral matrix"
def _def_draw():
return random.randint(0, 7)
idx = _draw_mask(x, _def_draw, draw=draw, p=p).long()
xs = tensor([1, -1, 1, -1, -1, 1, 1, -1], device=x.device)[idx]
ys = tensor([1, 1, -1, 1, -1, -1, 1, -1], device=x.device)[idx]
m0 = tensor([1, 1, 1, 0, 1, 0, 0, 0], device=x.device)[idx]
m1 = tensor([0, 0, 0, 1, 0, 1, 1, 1], device=x.device)[idx]
return affine_mat(xs * m0, xs * m1, t0(xs), ys * m1, ys * m0,
t0(xs)).float()
mask = mask_tensor(-x.new_ones(x.size(0)), p=p, neutral=1.)
def zoom_mat(x, max_zoom=1.1, p=0.5, draw=None, draw_x=None, draw_y=None):
"Return a random zoom matrix with `max_zoom` and `p`"
def _def_draw(x):
return x.new(x.size(0)).uniform_(1, max_zoom)
def _def_draw_ctr(x):
return x.new(x.size(0)).uniform_(0, 1)
s = 1 / _draw_mask(x, _def_draw, draw=draw, p=p, neutral=1.)
col_pct = _draw_mask(x, _def_draw_ctr, draw=draw_x, p=1.)
row_pct = _draw_mask(x, _def_draw_ctr, draw=draw_y, p=1.)
col_c = (1 - s) * (2 * col_pct - 1)
row_c = (1 - s) * (2 * row_pct - 1)
return affine_mat(s, t0(s), col_c, t0(s), s, row_c)
def dihedral_mat(x, p=0.5, draw=None, batch=False):
"Return a random dihedral matrix"
def _def_draw(x):
return torch.randint(0, 8, (x.size(0), ), device=x.device)
def _def_draw_b(x):
return random.randint(0, 7) + x.new_zeros((x.size(0), )).long()
idx = _draw_mask(x,
_def_draw_b if batch else _def_draw,
draw=draw,
p=p,
batch=batch).long()
xs = tensor([1, -1, 1, -1, -1, 1, 1, -1], device=x.device).gather(0, idx)
ys = tensor([1, 1, -1, 1, -1, -1, 1, -1], device=x.device).gather(0, idx)
m0 = tensor([1, 1, 1, 0, 1, 0, 0, 0], device=x.device).gather(0, idx)
m1 = tensor([0, 0, 0, 1, 0, 1, 1, 1], device=x.device).gather(0, idx)
return affine_mat(xs * m0, xs * m1, t0(xs), ys * m1, ys * m0,
t0(xs)).float()
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 rotate_mat(x, max_deg=10, p=0.5, draw=None):
"Return a random rotation matrix with `max_deg` and `p`"
def _def_draw(): return random.uniform(-max_deg,max_deg)
thetas = _draw_mask(x, _def_draw, draw=draw, p=p) * math.pi/180
return affine_mat(thetas.cos(), thetas.sin(), t0(thetas),
-thetas.sin(), thetas.cos(), t0(thetas))
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 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 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 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)
return affine_mat(mask, t0(mask), t0(mask),
t0(mask), t1(mask), t0(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)