def Raw2Bayer(x, crop_size = cSize, is_rotate = False):
r''' Convert FlatCam raw data to Bayer'''
# Step 1. Convert the Image & rotate
c, b, h, w = x.size()
y = torch.zeros((c, 4, int(h/2), int(w/2)), device = torch.device('cuda'))
if is_rotate: # ---> THIS MODES DOESNOT WORK YET!!! (2019.07.14)
scale = torch.ones(1)
angle = torch.ones(1) * 0.05 * 360 # 0.05 is angle collected from data measurements
center = torch.ones(1, 2)
center[..., 0] = int(h / 4) # x
center[..., 1] = int(w / 4) # y
M = kr.get_rotation_matrix2d(center, angle, scale).cuda()
_, _, h, w = y.size()
y[:, 0, :, : ] = kr.warp_affine(x[:, :, 1::2, 1::2], M, dsize = (h, w))
y[:, 1, :, : ] = kr.warp_affine(x[:, :, 0::2, 1::2], M, dsize = (h, w))
y[:, 2, :, : ] = kr.warp_affine(x[:, :, 1::2, 0::2], M, dsize = (h, w))
y[:, 3, :, : ] = kr.warp_affine(x[:, :, 0::2, 0::2], M, dsize = (h, w))
else:
y[:, 0, :, : ] = x[:, 0, 1::2, 1::2]
y[:, 1, :, : ] = x[:, 0, 0::2, 1::2]
y[:, 2, :, : ] = x[:, 0, 1::2, 0::2]
y[:, 3, :, : ] = x[:, 0, 0::2, 0::2]
# Step 3. Crop the image
start_row = int((y.size()[2] - crop_size[0]) / 2)
end_row = start_row + crop_size[0]
start_col = int((y.size()[3] - crop_size[1])/2)
end_col = start_col + crop_size[1]
return y[:,:, start_row:end_row, start_col:end_col]
def stitch(x, M_matrices, M_rotations, M_flip, label=True):
#Preprocessing: image stitch
data = [] #list to store all the features maps from multi-views
for i in range(6):
#get a batch of *same* view images
img_batch = x[:, i, :, :, :] # torch.stack(x)[:,i,:,:,:] #
img_warp = kornia.warp_perspective(img_batch,
M_matrices[i].unsqueeze(0).repeat(
len(x), 1, 1),
dsize=(219, 306))
img_rotated = kornia.warp_affine(img_warp,
M_rotations[i].unsqueeze(0).repeat(
len(x), 1, 1),
dsize=(219, 306))
data.append(img_rotated)
data = torch.cat(data, dim=0).view(6, len(x), 3, 219, 306)
#max pool feature maps from multi-view:black canvas and ensemble
h, w = 219, 306
#print(h,w)
agg = torch.zeros((x.shape[0], 3, 2 * h,
2 * w)) #[batch_size, 3 ,h, w], twice width/height
if torch.cuda.is_available():
agg = agg.cuda()
#two bases: front and back view
agg[:, :, 0:h, (w - w // 2):(w + w // 2)] = data[1]
agg[:, :, h:, (w - w // 2):(w + w // 2)] = data[4]
#top left
agg[:, :, (0 + 55):(h + 55), (0 + 55):(w + 55)] = torch.max(
data[0], agg[:, :, (0 + 55):(h + 55), (0 + 55):(w + 55)])
#top right
agg[:, :, (0 + 55):(h + 55), (w - 55):(-55)] = torch.max(
data[2], agg[:, :, (0 + 55):(h + 55), (w - 55):(-55)])
#bottom left
agg[:, :, (h - 55):(-55), (0 + 55):(w + 55)] = torch.max(
data[3], agg[:, :, (h - 55):(-55), (0 + 55):(w + 55)])
#bottom right
agg[:, :, (h - 55):(-55),
(w - 55):(-55)] = torch.max(data[5], agg[:, :, (h - 55):(-55),
(w - 55):(-55)])
#center-crop
crop_fn = kornia.augmentation.CenterCrop(size=438)
agg = crop_fn(agg)
#flip 90 degree
agg = kornia.warp_affine(agg,
M_flip.repeat(len(x), 1, 1),
dsize=(438, 438))
#Normalize color
if label:
normalize = K.Normalize(torch.tensor([0.698, 0.718, 0.730]),
torch.tensor([0.322, 0.313, 0.308]))
else:
normalize = K.Normalize(torch.tensor([0.548, 0.597, 0.630]),
torch.tensor([0.339, 0.340, 0.342]))
return normalize(agg)
def kornia_affine(im, parameter, aug_type, data_type='data'):
'''
Get rotation by given angle or scale by given factor
along axis-0 using kornia.
(See https://kornia.readthedocs.io/en/latest/geometry.transform.html)
'''
center = torch.ones(1, 2).cuda()
center[..., 0] = im.shape[1] // 2
center[..., 1] = im.shape[2] // 2
if aug_type == 'rotate':
scale = torch.ones(1).cuda()
angle = parameter * scale
elif aug_type == 'scale':
scale = torch.Tensor([parameter]).cuda()
angle = 0 * scale
# vol_warped = kornia.scale(vol[:, 0, :, :, :], scale, center)
if data_type == 'data':
interpolation = 'bilinear'
elif data_type == 'label':
interpolation = 'nearest'
M = kornia.get_rotation_matrix2d(center, angle, scale)
_, h, w = im.shape
im_warped = kornia.warp_affine(im[None, :, :, :].float(),
M.cuda(),
dsize=(h, w),
flags=interpolation)
# vol_warped = vol_warped[:, None, :, :, :]
return im_warped[0]
def compute_spicy_mayo_rotatedSpeckles_grad_diff_disks(
xd, xp, xl, rotated_data, pa_derotate_matrix, dilatation_matrix,
compute_loss, kernel, mask):
_, t, n, _ = rotated_data.shape
xs = xp + xd
rotated_data.requires_grad = False
kernel.requires_grad = False
L = xl.view(t, n, n)
L.requires_grad = True
L.grad = None
xs.requires_grad = True
#conv_xs_0 = A_pytorch(xs[0],kernel[0])
#conv_xs_1 = A_pytorch(xs[1],kernel[1])
conv_xs_0 = xs[0]
conv_xs_1 = xs[1]
rotated_L = kornia.warp_affine(L.unsqueeze(1).float(),
pa_derotate_matrix,
dsize=(n, n)).squeeze(1)
loss = compute_loss(conv_xs_0 + rotated_L -
rotated_data[0]) + compute_loss(conv_xs_1 + scale_cube(
rotated_L, dilatation_matrix) - rotated_data[1])
loss.backward()
grad_xs = xs.grad
grad_L = L.grad * mask
#grad_d_p = A_adj_pytorch(grad_xs,kernel)*mask
grad_d_p = grad_xs * mask
#print(torch.abs(grad_L).max())
return grad_d_p, grad_d_p, grad_L.view(t, n * n), loss.detach().item()
def __call__(self, clip):
"""
Args:
clip (torch.tensor): Video clip to be rotated. Size is (C, T, H, W)
Returns:
torch.tensor: central cropping of video clip. Size is
(C, T, crop_size, crop_size)
"""
clip = clip.permute(1, 0, 2, 3)
# define the rotation center
center = torch.ones(clip.shape[0], 2)
center[..., 0] = clip.shape[3] / 2 # x
center[..., 1] = clip.shape[2] / 2 # y
# define the scale factor
scale = torch.ones(clip.shape[0])
degree = np.random.randint(-self.degree, self.degree + 2)
degree = torch.ones(clip.shape[0]) * degree
# compute the transformation matrix
M = kornia.get_rotation_matrix2d(center, degree, scale)
# apply the transformation to original image
_, _, h, w = clip.shape
clip = kornia.warp_affine(clip, M, dsize=(h, w))
clip = clip.permute(1, 0, 2, 3)
return clip
def compute_rotatedSpeckles_conv_grad_pytorch(xd, xp, xl, rotated_data,
pa_derotate_matrix, compute_loss,
kernel, mask):
t, n, _ = rotated_data.shape
xs = xd + xp
rotated_data.requires_grad = False
kernel.requires_grad = False
L = xl.view(t, n, n)
L.requires_grad = True
xs.requires_grad = True
conv_xs = A_pytorch(xs, kernel)
#conv_xs.requires_grad = True
rotated_L = kornia.warp_affine(L.unsqueeze(1).float(),
pa_derotate_matrix,
dsize=(n, n)).squeeze(1)
loss = compute_loss(conv_xs + rotated_L - rotated_data)
loss.backward()
grad_xs = xs.grad
grad_L = L.grad
#grad_d_p = A_adj_pytorch(grad_xs,kernel)*mask
grad_d_p = grad_xs * mask
return grad_d_p, grad_d_p, (grad_L * mask).view(t, n *
n), loss.detach().item()
def compute_cube_frame_conv_grad_pytorch(xd, xp, xl, matrix, pa_rotate_matrix,
compute_loss, psf, mask):
t, _ = matrix.shape
xs = xd + xp
n, _ = xs.shape
torch_data = torch.tensor(matrix.reshape(t, n, n), requires_grad=False)
torch_psf = torch.tensor([psf], requires_grad=False)
torch_L = torch.tensor(xl.reshape(t, n, n), requires_grad=True)
torch_xs = torch.tensor(xs, requires_grad=True)
rotated_xs = kornia.warp_affine(torch_xs.expand(t, n,
n).unsqueeze(1).float(),
pa_rotate_matrix,
dsize=(n, n)) #.squeeze(1)
conv_rotated_xs = kornia.filters.filter2D(rotated_xs, torch_psf).squeeze(1)
loss = compute_loss(conv_rotated_xs + torch_L - torch_data)
#rotated_xs = kornia.warp_affine(torch_xs.expand(t,n,n).unsqueeze(1).float(), pa_rotate_matrix, dsize=(n,n)).squeeze(1)
#loss = compute_loss(rotated_xs + torch_L- torch_data)
loss.backward()
torch_grad_xs = torch_xs.grad
torch_grad_L = torch_L.grad
np_grad_xs = torch_grad_xs.detach().numpy()
np_grad_L = torch_grad_L.detach().numpy()
grad_d_p = np_grad_xs * mask
return grad_d_p, grad_d_p, np_grad_L.reshape(t,
n * n), loss.detach().numpy()
def transform(self, x):
""" Transform input with transformation matrices
Input:
x (float torch.Tensor): input data of shape (batch_size, ch, h, w)
Output:
xf_x (float torch.Tensor): transformed data
of shape (batch_size, ch, h, w)
If transformation input is different from the input size defined
in the construction then resize the input before transformation.
Input size is important because transformation matrices depend on
the image size, e.g the center of rotation.
"""
resize_input = False
# Check input size
if x.shape[2:] != self.input_hw:
curr_hw = x.shape[2:]
x = Resize(self.input_hw)(x)
resize_input = True
# Move matrices to the save device as input
self.tf_matrices = self.tf_matrices.to(x.device)
# Transform image
tf_x = kornia.warp_affine(x.float(),
self.tf_matrices,
dsize=self.input_hw)
# Transform back if image has been resized
if resize_input:
tf_x = Resize(curr_hw)(tf_x)
return tf_x
def test_cardinality(self, device, dtype, batch_shape, out_shape):
batch_size, channels, height, width = batch_shape
h_out, w_out = out_shape
aff_ab = torch.eye(2, 3, device=device, dtype=dtype).repeat(batch_size, 1, 1) # Bx2x3
img_b = torch.rand(batch_size, channels, height, width, device=device, dtype=dtype)
img_a = kornia.warp_affine(img_b, aff_ab, (h_out, w_out))
assert img_a.shape == (batch_size, channels, h_out, w_out)
def forward(self, x, M_matrices, M_rotations):
#Preprocessing: image stitch
data = [] #list to store all the features maps from multi-views
for i in range(6):
#get a batch of *same* view images
img_batch = x[:, i, :, :, :] #torch.stack(x)[:,i,:,:,:]
img_warp = kornia.warp_perspective(
img_batch,
M_matrices[i].unsqueeze(0).repeat(len(x), 1, 1),
dsize=(219, 306))
img_rotated = kornia.warp_affine(
img_warp,
M_rotations[i].unsqueeze(0).repeat(len(x), 1, 1),
dsize=(219, 306))
data.append(img_rotated)
data = torch.cat(data, dim=0).view(6, len(x), 3, 219, 306)
#max pool feature maps from multi-view:black canvas and ensemble
h, w = 219, 306
#print(h,w)
agg = torch.zeros((x.shape[0], 3, 2 * h,
2 * w)) #[batch_size, 3 ,h, w], twice width/height
if torch.cuda.is_available():
agg = agg.cuda()
#two bases: front and back view
agg[:, :, 0:h, (w - w // 2):(w + w // 2)] = data[1]
agg[:, :, h:, (w - w // 2):(w + w // 2)] = data[4]
#top left
agg[:, :, (0 + 55):(h + 55), (0 + 55):(w + 55)] = torch.max(
data[0], agg[:, :, (0 + 55):(h + 55), (0 + 55):(w + 55)])
#top right
agg[:, :, (0 + 55):(h + 55), (w - 55):(-55)] = torch.max(
data[2], agg[:, :, (0 + 55):(h + 55), (w - 55):(-55)])
#bottom left
agg[:, :, (h - 55):(-55), (0 + 55):(w + 55)] = torch.max(
data[3], agg[:, :, (h - 55):(-55), (0 + 55):(w + 55)])
#bottom right
agg[:, :, (h - 55):(-55), (w - 55):(-55)] = torch.max(
data[5], agg[:, :, (h - 55):(-55), (w - 55):(-55)])
#center-crop
crop_fn = kornia.augmentation.CenterCrop(size=438)
agg = crop_fn(agg)
###CNN: convolve down
x1 = self.inc(agg)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(
x3
) #shape:[batch_size, 256, 255, 38], scale_factor around 8; pixel shift around 55/8 = 55
###CNN: interpolate up
x = self.up1(x4)
x = self.up2(x)
#x = self.up3(x)
x = self.up_map(x) #last one to match output 800x800
x = self.outc(x)
return x
def crop_chw_torch(image, bbox, out_sz, device='cpu'):
a = (out_sz - 1) / (bbox[2] - bbox[0])
b = (out_sz - 1) / (bbox[3] - bbox[1])
c = -a * bbox[0]
d = -b * bbox[1]
mapping = torch.tensor([[a, 0, c], [0, b, d]], device=device).unsqueeze(0)
return kornia.warp_affine(image, mapping, dsize=(out_sz, out_sz), )
def test_exception(self, device, dtype):
img = torch.rand(1, 2, 3, 4, device=device, dtype=dtype)
aff = torch.eye(2, 3, device=device, dtype=dtype)[None]
size = (4, 5)
with pytest.raises(TypeError):
assert kornia.warp_affine(0., aff, size)
with pytest.raises(TypeError):
assert kornia.warp_affine(img, 0., size)
with pytest.raises(ValueError):
img = torch.rand(2, 3, 4, device=device, dtype=dtype)
assert kornia.warp_affine(img, aff, size)
with pytest.raises(ValueError):
aff = torch.eye(2, 2, device=device, dtype=dtype)[None]
assert kornia.warp_affine(img, aff, size)
def test_translation(self, device, batch_size):
offset = 1.
channels, height, width = 1, 3, 4
aff_ab = torch.eye(2, 3).repeat(batch_size, 1, 1).to(device) # Bx2x3
aff_ab[..., -1] += offset
img_b = torch.arange(float(height * width)).view(
1, channels, height, width).repeat(batch_size, 1, 1, 1).to(device)
img_a = kornia.warp_affine(img_b, aff_ab, (height, width))
assert_allclose(img_b[..., :2, :3], img_a[..., 1:, 1:])
def test_translation(self, batch_size, device, dtype):
offset = 1.
channels, height, width = 1, 3, 4
aff_ab = torch.eye(2, 3, device=device, dtype=dtype).repeat(batch_size, 1, 1) # Bx2x3
aff_ab[..., -1] += offset
img_b = torch.arange(float(height * width), device=device, dtype=dtype).view(
1, channels, height, width).repeat(batch_size, 1, 1, 1)
img_a = kornia.warp_affine(img_b, aff_ab, (height, width), align_corners=True)
assert_allclose(img_b[..., :2, :3], img_a[..., 1:, 1:], rtol=1e-4, atol=1e-4)
def test_rotation_inverse(self, device, dtype):
h, w = 4, 4
img_b = torch.rand(1, 1, h, w, device=device, dtype=dtype)
# create rotation matrix of 90deg (anti-clockwise)
center = torch.tensor([[w - 1, h - 1]], device=device, dtype=dtype) / 2
scale = torch.ones((1, 2), device=device, dtype=dtype)
angle = 90. * torch.ones(1, device=device, dtype=dtype)
aff_ab = kornia.get_rotation_matrix2d(center, angle, scale)
# Same as opencv: cv2.getRotationMatrix2D(((w-1)/2,(h-1)/2), 90., 1.)
# warp the tensor
# Same as opencv: cv2.warpAffine(kornia.tensor_to_image(img_b), aff_ab[0].numpy(), (w, h))
img_a = kornia.warp_affine(img_b, aff_ab, (h, w))
# invert the transform
aff_ba = kornia.convert_affinematrix_to_homography(aff_ab).inverse()[..., :2, :]
img_b_hat = kornia.warp_affine(img_a, aff_ba, (h, w))
assert_allclose(img_b_hat, img_b, atol=1e-3, rtol=1e-3)
def forward(self, img):
# print('img.shape is', img.shape)
bs = img.shape[0]
# print('bs is', bs)
W = img.shape[2]
H = img.shape[3]
tx = ((self.max_tx - self.min_tx) * torch.rand(
(bs, 1)) + self.min_tx).to('cuda')
ty = ((self.max_ty - self.min_ty) * torch.rand(
(bs, 1)) + self.min_ty).to('cuda')
r = ((self.ang - self.ang_neg) * torch.rand(
(bs, 1)) + self.ang_neg).to('cuda')
z = ((self.max_z - self.min_z) * torch.rand(
(bs, 1)) + self.min_z).to('cuda')
hx = ((self.ang - self.ang_neg) * torch.rand(
(bs, 1)) + self.ang_neg).to('cuda')
hy = ((self.ang - self.ang_neg) * torch.rand(
(bs, 1)) + self.ang_neg).to('cuda')
# Transformation Matrix
a = hx - r
b = hy + r
T11 = torch.div(z * torch.cos(a), torch.cos(hx))
T12 = torch.div(z * torch.sin(a), torch.cos(hx))
T13 = torch.div(
W * torch.cos(hx) - W * z * torch.cos(a) +
2 * tx * z * torch.cos(a) - H * z * torch.sin(a) +
2 * ty * z * torch.sin(a), 2 * torch.cos(hx))
T21 = torch.div(z * torch.sin(b), torch.cos(hy))
T22 = torch.div(z * torch.cos(b), torch.cos(hy))
T23 = torch.div(
H * torch.cos(hy) - W * z * torch.cos(b) +
2 * ty * z * torch.cos(b) - W * z * torch.sin(b) +
2 * tx * z * torch.sin(b), 2 * torch.cos(hy))
T = torch.zeros((bs, 2, 3)).to('cuda') #Combined for batch
for i in range(bs):
T[i] = torch.tensor([[T11[i], T12[i], T13[i]],
[T21[i], T22[i], T23[i]]
]) # Transformation Matrix for a batch
Transformed_img = kornia.warp_affine(img, T, dsize=(W, H)).to('cuda')
return Transformed_img
def test_smoke(self, device, dtype):
batch_size, channels, height, width = 1, 2, 3, 4
aff_ab = torch.eye(2, 3, device=device, dtype=dtype)[None] # 1x2x3
img_b = torch.rand(batch_size,
channels,
height,
width,
device=device,
dtype=dtype)
img_a = kornia.warp_affine(img_b, aff_ab, (height, width))
assert img_b.shape == img_a.shape
def make_training_batch(input_tensor, patch, patch_mask):
# determine patch size
H, W = PA_cfg.image_shape[-2:]
PATCH_SIZE = int(np.floor(np.sqrt((H*W*PA_cfg.percentage))))
translate_space = [H-PATCH_SIZE+1, W-PATCH_SIZE+1]
bs = input_tensor.size(0)
training_batch = []
for b in range(bs):
# random translation
u_t = np.random.randint(low=0, high=translate_space[0])
v_t = np.random.randint(low=0, high=translate_space[1])
# random scaling and rotation
scale = np.random.rand() * (PA_cfg.scale_max - PA_cfg.scale_min) + PA_cfg.scale_min
scale = torch.Tensor([scale])
angle = np.random.rand() * (PA_cfg.rotate_max - PA_cfg.rotate_min) + PA_cfg.rotate_min
angle = torch.Tensor([angle])
center = torch.Tensor([u_t+PATCH_SIZE/2, v_t+PATCH_SIZE/2]).unsqueeze(0)
rotation_m = kornia.get_rotation_matrix2d(center, angle, scale)
# warp three tensors
temp_mask = patch_mask.unsqueeze(0)
temp_input = input_tensor[b].unsqueeze(0)
temp_patch = patch.unsqueeze(0)
temp_mask = kornia.translate(temp_mask.float(), translation=torch.Tensor([u_t, v_t]).unsqueeze(0))
temp_patch = kornia.translate(temp_patch.float(), translation=torch.Tensor([u_t, v_t]).unsqueeze(0))
mask_warpped = kornia.warp_affine(temp_mask.float(), rotation_m, temp_mask.size()[-2:])
patch_warpped = kornia.warp_affine(temp_patch.float(), rotation_m, temp_patch.size()[-2:])
# overlay
overlay = temp_input * (1 - mask_warpped) + patch_warpped * mask_warpped
training_batch.append(overlay)
training_batch = torch.cat(training_batch, dim=0)
return training_batch
def inner(image_t):
b, _, h, w = image_t.shape
# kornia takes degrees
alpha = _rads2angle(np.random.choice(angles), units)
angle = torch.ones(b) * alpha
scale = torch.ones(b)
center = torch.ones(b, 2)
center[..., 0] = (image_t.shape[3] - 1) / 2
center[..., 1] = (image_t.shape[2] - 1) / 2
M = kornia.get_rotation_matrix2d(center, angle, scale).to(device)
rotated_image = kornia.warp_affine(image_t.float(), M, dsize=(h, w))
return rotated_image
def forward(self, img: torch.Tensor):
b, c, h, w = img.shape
center = torch.tensor([[w, h]], dtype=torch.float) / 2
transformation_matrix = kornia.get_rotation_matrix2d(
center, self.angle, torch.ones(1))
transformation_matrix = transformation_matrix.expand(b, -1, -1)
transformation_matrix = transformation_matrix.to(img.device)
return kornia.warp_affine(img.float(),
transformation_matrix,
dsize=(h, w),
flags=self.interpolation,
padding_mode=self.padding_mode)
def Rotate(x, alphas, step):
b, c, h, w = x.shape
start, end = alphas
x = CopyN(x, step)
angle = torch.linspace(start,
end, step, device=x.device).unsqueeze(0).repeat(
b, 1).view(b * step)
center = torch.zeros((b * step, 2), device=x.device)
center[:, 0], center[:, 1] = w / 2, h / 2
scale = torch.ones(b * step, device=x.device)
M = kornia.get_rotation_matrix2d(center, angle, scale)
x_hat = kornia.warp_affine(x, M, dsize=(h, w))
return x_hat
def test_translation(self, device, dtype):
offset = 1.
h, w = 3, 4
aff_ab = torch.eye(2, 3, device=device, dtype=dtype)[None]
aff_ab[..., -1] += offset
img_b = torch.arange(float(h * w), device=device, dtype=dtype).view(1, 1, h, w)
expected = torch.zeros_like(img_b)
expected[..., 1:, 1:] = img_b[..., :2, :3]
# Same as opencv: cv2.warpAffine(kornia.tensor_to_image(img_b), aff_ab[0].numpy(), (w, h))
img_a = kornia.warp_affine(img_b, aff_ab, (h, w))
assert_allclose(img_a, expected)
def forward(self, x, params):
angle, scale = params
self.angle.fill_(angle)
self.scale.fill_(scale)
# define the rotation center
self.center[..., 0] = x.shape[3] / 2
self.center[..., 1] = x.shape[2] / 2
M = kornia.get_rotation_matrix2d(self.center, self.angle, self.scale)
return kornia.warp_affine(x,
M,
dsize=(x.shape[2], x.shape[3]),
flags='bilinear',
padding_mode='reflection')
def forward(self, img: torch.tensor):
b, _, h, w = img.shape
# create transformation (rotation)
if not self.same_throughout_batch:
angle = torch.randn(b, device=img.device) * self.angle
else:
angle = torch.randn(1, device=img.device) * self.angle
angle = angle.repeat(b)
center = torch.ones(b, 2, device=img.device)
center[..., 0] = img.shape[3] / 2 # x
center[..., 1] = img.shape[2] / 2 # y
# define the scale factor
scale = torch.ones(b, device=img.device)
M = kornia.get_rotation_matrix2d(center, angle, scale)
img_warped = kornia.warp_affine(img, M, dsize=(h, w))
return img_warped
def augment_batch(batch_tensor):
orig_size = batch_tensor.ndim
batch_size = batch_tensor.shape[0]
if orig_size == 3:
batch_tensor=batch_tensor.unsqueeze(0)
angle = (torch.rand(batch_size)*20)-10
center = torch.ones(batch_size,2)
center[..., 0] = batch_tensor.shape[3] / 2
center[..., 1] = batch_tensor.shape[2] / 2
scale = torch.ones(batch_size)
M = kornia.get_rotation_matrix2d(center, angle, scale)
_, _, h, w = batch_tensor.shape
batch_tensor_warped = kornia.warp_affine(batch_tensor, M, dsize=(h, w))
if orig_size == 3:
batch_tensor_warped=batch_tensor_warped.squeeze(0)
return batch_tensor_warped
def test_performance_speed(device, dtype):
if device.type != 'cuda' or not torch.cuda.is_available():
pytest.skip("Cuda not available in system,")
print("Benchmarking warp_affine")
for input_shape in shapes:
for PS in PSs:
BS = input_shape[0]
inpt = torch.rand(input_shape).to(device)
As = torch.eye(3).unsqueeze(0).repeat(BS, 1, 1)[:, :2, :].to(device)
As += 0.1 * torch.rand(As.size()).to(device)
torch.cuda.synchronize(device)
t = time()
_ = kornia.warp_affine(inpt, As, (PS, PS))
print(f"inp={input_shape}, PS={PS}, dev={device}, {time() - t}, sec")
torch.cuda.synchronize(device)
def compute_cube_frame_grad_pytorch_no_regul(xs, xl, data, pa_rotate_matrix,
compute_loss, mask):
t, n, _ = data.shape
data.requires_grad = False
L = xl.view(t, n, n)
L.requires_grad = True
torch_xs = xs.view(n, n)
torch_xs.requires_grad = True
rotated_xs = kornia.warp_affine(torch_xs.expand(t, n,
n).unsqueeze(1).float(),
pa_rotate_matrix,
dsize=(n, n)).squeeze(1)
loss = compute_loss(rotated_xs + L - data)
loss.backward()
grad_xs = torch_xs.grad
grad_L = L.grad
return grad_xs.detach() * mask, (grad_L.detach()).reshape(t, n * n), loss
def rotate_tensor_along_y_axis(tensor, gamma):
B = tensor.shape[0]
tensor = tensor.to("cpu")
assert tensor.ndim == 6, "Tensors should have 6 dimensions."
tensor = tensor.float()
# B,S,C,D,H,W
__p = lambda x: utils_basic.pack_seqdim(x, B)
__u = lambda x: utils_basic.unpack_seqdim(x, B)
tensor_ = __p(tensor)
tensor_ = tensor_.permute(
0, 1, 3, 2, 4) # Make it BS, C, H, D, W (i.e. BS, C, y, z, x)
BS, C, H, D, W = tensor_.shape
# merge y dimension with channel dimension and rotate with gamma_
tensor_y_reduced = tensor_.reshape(BS, C * H, D, W)
# # gammas will be rotation angles along y axis.
# gammas = torch.arange(10, 360, 10)
# define the rotation center
center = torch.ones(1, 2)
center[..., 0] = tensor_y_reduced.shape[3] / 2 # x
center[..., 1] = tensor_y_reduced.shape[2] / 2 # z
# define the scale factor
scale = torch.ones(1)
gamma_ = torch.ones(1) * gamma
# compute the transformation matrix
M = kornia.get_rotation_matrix2d(center, gamma_, scale)
M = M.repeat(BS, 1, 1)
# apply the transformation to original image
# st()
tensor_y_reduced_warped = kornia.warp_affine(tensor_y_reduced,
M,
dsize=(D, W))
tensor_y_reduced_warped = tensor_y_reduced_warped.reshape(BS, C, H, D, W)
tensor_y_reduced_warped = tensor_y_reduced_warped.permute(0, 1, 3, 2, 4)
tensor_y_reduced_warped = __u(tensor_y_reduced_warped)
return tensor_y_reduced_warped.cuda()
def kornia_translate(im, choice, data_type='data'):
'''
Random translation using kornia translate transform.
Additional function needed for construction of translation Tensor.
Choice has format str(magnitude)+'axis' where axis = x or y.
'''
axis = choice[-1]
transMag = int(choice[:-1])
transVal = torch.zeros(im.shape[0], 2)
if axis == 'x':
transVal[:, 0] = transMag
elif axis == 'y':
transVal[:, 1] = transMag
if data_type == 'data':
interpolation = 'bilinear'
elif data_type == 'label':
interpolation = 'nearest'
M = kornia.geometry.transform.affwarp._compute_translation_matrix(transVal)
_, _, h, w = im.shape
vol_warped = kornia.warp_affine(im.float(),
M.cuda(),
dsize=(h, w),
flags=interpolation)
return vol_warped
def create_synthetic_data_with_disk_planet(empty_data, pa_rotate_matrix,
kernel, add_synthetic_signal):
"""
automatic_load_data(data_name,channel=0,dir='default',RDI=False,quick_look=0,crop=0,center_im=None)
loads ADI datasets automatically
Parameters
----------
data : numpy array
t x n x n the empty (without circumstellar signal) ADI dataset
angles : numpy array
list of angles
psf : numpy array
n x n psf
add_synthetic_signal
tuple containing all the properties of the injected signal
Returns
-------
data : numpy array
t x n x n, ADI dataset (empty_data + circumstellar signal injected)
disk : numpy array
n x n, the injected circumstellar signal (disk+planet)
"""
disk_max_intensity = add_synthetic_signal['disk_max_intensity']
disk_inclination = add_synthetic_signal['disk_inclination']
planets_positions_intensities = tuple(
add_synthetic_signal['planets_positions_intensities'])
if 'xdo' not in add_synthetic_signal:
add_synthetic_signal['xdo'] = 0
if 'ydo' not in add_synthetic_signal:
add_synthetic_signal['ydo'] = 0
if 'pa' not in add_synthetic_signal:
add_synthetic_signal['pa'] = 80
if 'omega' not in add_synthetic_signal:
add_synthetic_signal['omega'] = 45
if 'density_distribution' not in add_synthetic_signal:
add_synthetic_signal['density_distribution'] = {'name': '2PowerLaws'}
if 'phase_function' not in add_synthetic_signal:
add_synthetic_signal['phase_function'] = {
'name': 'HG',
'g': 0.,
'polar': False
}
if 'disk_intensity_thresh' not in add_synthetic_signal:
add_synthetic_signal['disk_intensity_thresh'] = 60
t, n, _ = empty_data.shape
#
# Disk injection
#
my_disk = vip.metrics.scattered_light_disk.ScatteredLightDisk(
nx=n,
ny=n,
xdo=add_synthetic_signal['xdo'],
ydo=add_synthetic_signal['ydo'])
my_disk.set_inclination(add_synthetic_signal['disk_inclination'])
my_disk.set_pa(add_synthetic_signal['pa'])
my_disk.set_omega(add_synthetic_signal['omega'])
my_disk.set_flux_max(100)
my_disk.set_density_distribution(
add_synthetic_signal['density_distribution'])
my_disk.set_phase_function(add_synthetic_signal['phase_function'])
disk = my_disk.compute_scattered_light()
disk = disk * (disk > add_synthetic_signal['disk_intensity_thresh'])
disk = disk / np.max(disk) * disk_max_intensity
#
# Planet injection
#
if planets_positions_intensities:
planet = np.zeros((n, n))
for xx, yy, intensity in planets_positions_intensities:
planet[xx, yy] = intensity
disk += planet
cube_disk = np.zeros((t, n, n))
torch_disk = torch.tensor(disk, requires_grad=False)
rotated_disk = kornia.warp_affine(
torch_disk.expand(t, n, n).unsqueeze(1).float(),
pa_rotate_matrix,
dsize=(n, n)).squeeze(1).detach().numpy()
for k in range(t):
rotated_disk[k, :, :] = mayo_hci.A(rotated_disk[k, :, :], kernel)
data = empty_data + rotated_disk
return data, disk
