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