Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import torch | |
| import cv2 | |
| def np_batched_radon(image_batch): | |
| #image_batch: torch tensor, batch_size x 1 x img_size x img_size | |
| # squeeze order #1 and transform to numpy | |
| image_batch = image_batch.squeeze(1).cpu().numpy() | |
| batch_size, img_size = image_batch.shape[:2] | |
| if batch_size > 512: # limit batch size to 512 because cv2.warpAffine fails for batch> 512 | |
| return np.concatenate([np_batched_radon(image_batch[i:i+512]) for i in range(0,batch_size,512)], axis=0) | |
| theta = np.arange(180) | |
| radon_image = np.zeros((image_batch.shape[0], img_size, len(theta)), | |
| dtype='float32') | |
| for i, angle in enumerate(theta): | |
| M = cv2.getRotationMatrix2D(((img_size-1)/2.0,(img_size-1)/2.0),angle,1) | |
| rotated = cv2.warpAffine(np.transpose(image_batch, (1, 2, 0)),M,(img_size,img_size)) | |
| #plt.imshow(rotated[:,:,0]) | |
| #plt.show() | |
| if batch_size == 1: # cv2.warpAffine cancels batch dimension if equal to 1 | |
| rotated = rotated[:,:, np.newaxis] | |
| rotated = np.transpose(rotated, (2, 0, 1)) / 224.0 | |
| #rotated = rotated / np.array(255, dtype='float32') | |
| radon_image[:, :, i] = rotated.sum(axis=1) | |
| #plot the image | |
| # plt.imshow(radon_image[0]) | |
| # plt.show() | |
| return radon_image | |
| def torch_batched_radon(image_batch, neutral_value): | |
| #image_batch: batch_size x 1 x img_size x img_size | |
| #np_batched_radon(image_batch - neutral_value) | |
| image_batch = image_batch - neutral_value # so the 0 value is neutral | |
| batch_size = image_batch.shape[0] | |
| img_size = image_batch.shape[2] | |
| theta = np.arange(180) # we don't need torch here, we will evaluate individual angles below | |
| radon_image = torch.zeros((batch_size, 1, img_size, len(theta)), dtype=torch.float, device=image_batch.device) | |
| for i, angle in enumerate(theta): | |
| #M = cv2.getRotationMatrix2D(((img_size-1)/2.0,(img_size-1)/2.0),angle,1) | |
| #calculate the same rotation matrix but with torch: | |
| M = torch.tensor(cv2.getRotationMatrix2D(((img_size-1)/2.0,(img_size-1)/2.0),angle,1)).to(image_batch.device, dtype=torch.float32) | |
| angle = torch.tensor((angle+90)/180.0*np.pi) | |
| M1 = torch.tensor([[torch.sin(angle), torch.cos(angle), 0], | |
| [torch.cos(angle), -torch.sin(angle), 0]]).to(image_batch.device, dtype=torch.float32) | |
| # we need to add a batch dimension to the rotation matrix | |
| M1 = M1.repeat(batch_size, 1, 1) | |
| grid = torch.nn.functional.affine_grid(M1, image_batch.shape, align_corners=False) | |
| rotated = torch.nn.functional.grid_sample(image_batch, grid, mode='bilinear', padding_mode='zeros', align_corners=False) | |
| rotated = rotated.squeeze(1) | |
| #plt.imshow(rotated[0].cpu().numpy()) | |
| #plt.show() | |
| radon_image[:, 0, :, i] = rotated.sum(axis=1) / 224.0 + neutral_value | |
| #plt.imshow(radon_image[0, 0].cpu().numpy()) | |
| #plt.show() | |
| return radon_image | |