Created
May 12, 2018 05:39
-
-
Save elistevens/e6064aa6705cf1e97dea5be876fc16ab to your computer and use it in GitHub Desktop.
affine grid for volumetric arrays
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import math | |
| import random | |
| import warnings | |
| import numpy as np | |
| import scipy.ndimage | |
| import torch | |
| from torch.autograd import Function | |
| from torch.autograd.function import once_differentiable | |
| import torch.backends.cudnn as cudnn | |
| import pytest | |
| from util.logconf import logging | |
| log = logging.getLogger(__name__) | |
| # log.setLevel(logging.WARN) | |
| # log.setLevel(logging.INFO) | |
| log.setLevel(logging.DEBUG) | |
| def affine_grid_generator(theta, size): | |
| if theta.data.is_cuda and len(size) == 4: | |
| if not cudnn.enabled: | |
| raise RuntimeError("AffineGridGenerator needs CuDNN for " | |
| "processing CUDA inputs, but CuDNN is not enabled") | |
| if not cudnn.is_acceptable(theta.data): | |
| raise RuntimeError("AffineGridGenerator generator theta not acceptable for CuDNN") | |
| N, C, H, W = size | |
| return torch.cudnn_affine_grid_generator(theta, N, C, H, W) | |
| else: | |
| return AffineGridGenerator.apply(theta, size) | |
| class AffineGridGenerator(Function): | |
| @staticmethod | |
| def _enforce_cudnn(input): | |
| if not cudnn.enabled: | |
| raise RuntimeError("AffineGridGenerator needs CuDNN for " | |
| "processing CUDA inputs, but CuDNN is not enabled") | |
| assert cudnn.is_acceptable(input) | |
| @staticmethod | |
| def forward(ctx, theta, size): | |
| assert type(size) == torch.Size | |
| if len(size) == 5: | |
| N, C, D, H, W = size | |
| ctx.size = size | |
| # if theta.is_cuda: | |
| # AffineGridGenerator._enforce_cudnn(theta) | |
| # assert False | |
| # ctx.is_cuda = False | |
| base_grid = theta.new(N, D, H, W, 4) | |
| w_points = (torch.linspace(-1, 1, W) if W > 1 else torch.Tensor([-1]))#.unsqueeze(0).unsqueeze(0).unsqueeze(0) | |
| h_points = (torch.linspace(-1, 1, H) if H > 1 else torch.Tensor([-1])).unsqueeze(-1)#.unsqueeze(0) | |
| d_points = (torch.linspace(-1, 1, D) if D > 1 else torch.Tensor([-1])).unsqueeze(-1).unsqueeze(-1) | |
| print('w', w_points.size()) | |
| print('h', h_points.size()) | |
| print('d', d_points.size()) | |
| # linear_points = torch.linspace(-1, 1, W) if W > 1 else torch.Tensor([-1]) | |
| # print(size, linear_points.size(), linear_points) | |
| # print(torch.ger(torch.ones(H), linear_points), torch.ger(torch.ones(H), linear_points).size()) | |
| # print(base_grid[:, :, :, :, 0], base_grid[:, :, :, :, 0].size()) | |
| # print() | |
| base_grid[:, :, :, :, 0] = w_points | |
| base_grid[:, :, :, :, 1] = h_points | |
| base_grid[:, :, :, :, 2] = d_points | |
| # base_grid[:, :, :, :, 1] = torch.ger(torch.ger(torch.ones(D), h_points), torch.ones(W)).expand_as(base_grid[:, :, :, :, 1]) | |
| # base_grid[:, :, :, :, 2] = torch.ger(torch.ger(d_points, torch.ones(H)), torch.ones(W)).expand_as(base_grid[:, :, :, :, 2]) | |
| base_grid[:, :, :, :, 3] = 1 | |
| ctx.base_grid = base_grid | |
| grid = torch.bmm(base_grid.view(N, D * H * W, 4), theta.transpose(1, 2)) | |
| grid = grid.view(N, D, H, W, 3) | |
| elif len(size) == 4: | |
| N, C, H, W = size | |
| ctx.size = size | |
| if theta.is_cuda: | |
| AffineGridGenerator._enforce_cudnn(theta) | |
| assert False | |
| ctx.is_cuda = False | |
| base_grid = theta.new(N, H, W, 3) | |
| linear_points = torch.linspace(-1, 1, W) if W > 1 else torch.Tensor([-1]) | |
| base_grid[:, :, :, 0] = torch.ger(torch.ones(H), linear_points).expand_as(base_grid[:, :, :, 0]) | |
| linear_points = torch.linspace(-1, 1, H) if H > 1 else torch.Tensor([-1]) | |
| base_grid[:, :, :, 1] = torch.ger(linear_points, torch.ones(W)).expand_as(base_grid[:, :, :, 1]) | |
| base_grid[:, :, :, 2] = 1 | |
| ctx.base_grid = base_grid | |
| grid = torch.bmm(base_grid.view(N, H * W, 3), theta.transpose(1, 2)) | |
| grid = grid.view(N, H, W, 2) | |
| else: | |
| raise RuntimeError("AffineGridGenerator needs 4d (spatial) or 5d (volumetric) inputs.") | |
| return grid | |
| @staticmethod | |
| @once_differentiable | |
| def backward(ctx, grad_grid): | |
| if len(ctx.size) == 5: | |
| N, C, D, H, W = ctx.size | |
| assert grad_grid.size() == torch.Size([N, D, H, W, 3]) | |
| assert ctx.is_cuda == grad_grid.is_cuda | |
| # if grad_grid.is_cuda: | |
| # AffineGridGenerator._enforce_cudnn(grad_grid) | |
| # assert False | |
| base_grid = ctx.base_grid | |
| grad_theta = torch.bmm( | |
| base_grid.view(N, D * H * W, 4).transpose(1, 2), | |
| grad_grid.view(N, D * H * W, 3)) | |
| grad_theta = grad_theta.transpose(1, 2) | |
| elif len(ctx.size) == 4: | |
| N, C, H, W = ctx.size | |
| assert grad_grid.size() == torch.Size([N, H, W, 2]) | |
| assert ctx.is_cuda == grad_grid.is_cuda | |
| if grad_grid.is_cuda: | |
| AffineGridGenerator._enforce_cudnn(grad_grid) | |
| assert False | |
| base_grid = ctx.base_grid | |
| grad_theta = torch.bmm( | |
| base_grid.view(N, H * W, 3).transpose(1, 2), | |
| grad_grid.view(N, H * W, 2)) | |
| grad_theta = grad_theta.transpose(1, 2) | |
| else: | |
| assert False | |
| return grad_theta, None | |
| if torch.cuda.is_available(): | |
| @pytest.fixture(params=['cpu', 'cuda']) | |
| def device(request): | |
| return request.param | |
| else: | |
| @pytest.fixture(params=['cpu']) | |
| def device(request): | |
| return request.param | |
| @pytest.fixture(params=[0.0, 0.25, 0.125, 'random']) | |
| def angle_rad(request): | |
| if request.param == 'random': | |
| return random.random() * math.pi * 2 | |
| return request.param * math.pi * 2 | |
| @pytest.fixture(params=[torch.nn.functional.affine_grid, affine_grid_generator]) | |
| def affine_func2d(request): | |
| return request.param | |
| @pytest.fixture(params=[[1, 1, 3, 3], [1, 1, 3, 4]]) | |
| def input_size2d(request): | |
| return request.param | |
| @pytest.fixture(params=[[1, 1, 2, 2], [1, 1, 3, 3], [1, 1, 4, 4], [1, 1, 6, 6], ]) | |
| def input_size2dsq(request): | |
| return request.param | |
| @pytest.fixture(params=[[1, 1, 2, 2], [1, 1, 3, 3], [1, 1, 5, 5], [1, 1, 4, 4], [1, 1, 6, 6], ]) | |
| def output_size2dsq(request): | |
| return request.param | |
| @pytest.fixture(params=[[1, 1, 3, 3], [1, 1, 6, 6], [1, 1, 4, 5]]) | |
| def output_size2d(request): | |
| return request.param | |
| def _buildEquivalentTransforms2d(device, input_size, output_size, angle_rad): | |
| input_center = [(x-1)/2 for x in input_size] | |
| output_center = [(x-1)/2 for x in output_size] | |
| s = math.sin(angle_rad) | |
| c = math.cos(angle_rad) | |
| rotation_ary = np.array([ | |
| [c, -s, 0], | |
| [s, c, 0], | |
| [0, 0, 1], | |
| ], dtype=np.float32) | |
| scale_ary = np.array([ | |
| [output_size[2]/input_size[2], 0, 0], | |
| [0, output_size[3]/input_size[3], 0], | |
| [0, 0, 1], | |
| ], dtype=np.float32) | |
| # we use .T because | |
| # https://stackoverflow.com/questions/20161175/how-can-i-use-scipy-ndimage-interpolation-affine-transform-to-rotate-an-image-ab | |
| transform_ary = np.linalg.inv(scale_ary @ rotation_ary)[:2,:2] | |
| # transform_ary = (rotation_ary @ scale_ary)[:2,:2].T | |
| # transform_ary = transform_tensor[0,:,:2].to('cpu').numpy().T | |
| transform_tensor = torch.from_numpy((rotation_ary)).to(device) | |
| # transform_tensor = torch.from_numpy((scale_ary @ rotation_ary)).to(device) | |
| transform_tensor = transform_tensor[:2].unsqueeze(0) | |
| offset = -(transform_ary @ output_center[2:]) + input_center[2:] | |
| print('transform_tensor', transform_tensor.size(), transform_tensor.dtype, transform_tensor.device) | |
| print(transform_tensor) | |
| print('transform_ary', transform_ary.shape, transform_ary.dtype) | |
| print(transform_ary) | |
| print('offset', offset.shape, offset.dtype) | |
| print(offset) | |
| return transform_tensor, transform_ary, offset | |
| def _buildEquivalentTransforms3d(device, input_size, output_size, angle_rad, axis_vector): | |
| input_center = [(x-1)/2 for x in input_size] | |
| output_center = [(x-1)/2 for x in output_size] | |
| s = math.sin(angle_rad) | |
| c = math.cos(angle_rad) | |
| c1 = 1 - c | |
| # axis_vector = torch.rand(3, device='cpu', dtype=torch.float64) | |
| # axis_vector -= 0.5 | |
| # axis_vector /= axis_vector.abs().sum() | |
| l, m, n = axis_vector | |
| transform_tensor = torch.tensor([ | |
| [l*l*c1 + c, m*l*c1 - n*s, n*l*c1 + m*s, 0], | |
| [l*m*c1 + n*s, m*m*c1 + c, n*m*c1 - l*s, 0], | |
| [l*n*c1 - m*s, m*n*c1 + l*s, n*n*c1 + c, 0], | |
| # [0, 0, 0, 1], | |
| ], device=device, dtype=torch.float32).to(device) | |
| # we use .T because | |
| # https://stackoverflow.com/questions/20161175/how-can-i-use-scipy-ndimage-interpolation-affine-transform-to-rotate-an-image-ab | |
| transform_ary = transform_tensor[0,:,:3].to('cpu').numpy().T | |
| # offset = -(transform_ary @ output_center[2:]) + input_center[2:] | |
| offset = -(transform_ary @ input_center[2:]) + output_center[2:] | |
| print('transform_tensor', transform_tensor.size(), transform_tensor.dtype, transform_tensor.device) | |
| print(transform_tensor) | |
| print('transform_ary', transform_ary.shape, transform_ary.dtype) | |
| print(transform_ary) | |
| print('offset', offset.shape, offset.dtype) | |
| print(offset) | |
| return transform_tensor, transform_ary, offset | |
| def test_affine_2d_rotate0(device, affine_func2d): | |
| input_size = [1, 1, 3, 3] | |
| input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| output_size = [1, 1, 3, 3] | |
| angle_rad = 0. | |
| transform_tensor, transform_ary, offset = _buildEquivalentTransforms2d(device, input_size, output_size, angle_rad) | |
| # reference | |
| # https://stackoverflow.com/questions/20161175/how-can-i-use-scipy-ndimage-interpolation-affine-transform-to-rotate-an-image-ab | |
| scipy_ary = scipy.ndimage.affine_transform( | |
| input_ary[0,0], | |
| transform_ary, | |
| offset=offset, | |
| output_shape=output_size[2:], | |
| # output=None, | |
| order=1, | |
| mode='nearest', | |
| # cval=0.0, | |
| prefilter=False) | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| assert np.abs(scipy_ary.mean() - input_ary.mean()) < 1e-6 | |
| assert np.abs(scipy_ary - input_ary).max() < 1e-6 | |
| affine_tensor = affine_func2d( | |
| transform_tensor, | |
| torch.Size(output_size) | |
| ) | |
| print('affine_tensor', affine_tensor.size(), affine_tensor.dtype, affine_tensor.device) | |
| print(affine_tensor) | |
| gridsample_ary = torch.nn.functional.grid_sample( | |
| torch.tensor(input_ary, device=device).to(device), | |
| affine_tensor, | |
| padding_mode='border' | |
| ).to('cpu').numpy() | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('gridsample_ary', gridsample_ary.shape, gridsample_ary.dtype) | |
| print(gridsample_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| assert np.abs(scipy_ary.mean() - gridsample_ary.mean()) < 1e-6 | |
| assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # assert False | |
| def test_affine_2d_rotate90(device, affine_func2d, input_size2dsq, output_size2dsq): | |
| input_size = input_size2dsq | |
| input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| output_size = output_size2dsq | |
| angle_rad = 0.25 * math.pi * 2 | |
| transform_tensor, transform_ary, offset = _buildEquivalentTransforms2d(device, input_size, output_size, angle_rad) | |
| # reference | |
| # https://stackoverflow.com/questions/20161175/how-can-i-use-scipy-ndimage-interpolation-affine-transform-to-rotate-an-image-ab | |
| scipy_ary = scipy.ndimage.affine_transform( | |
| input_ary[0,0], | |
| transform_ary, | |
| offset=offset, | |
| output_shape=output_size[2:], | |
| # output=None, | |
| order=1, | |
| mode='nearest', | |
| # cval=0.0, | |
| prefilter=True) | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| assert np.abs(scipy_ary.mean() - input_ary.mean()) < 1e-6 | |
| assert np.abs(scipy_ary[0,0] - input_ary[0,0,0,-1]).max() < 1e-6 | |
| assert np.abs(scipy_ary[0,-1] - input_ary[0,0,-1,-1]).max() < 1e-6 | |
| assert np.abs(scipy_ary[-1,-1] - input_ary[0,0,-1,0]).max() < 1e-6 | |
| assert np.abs(scipy_ary[-1,0] - input_ary[0,0,0,0]).max() < 1e-6 | |
| # assert np.abs(scipy_ary[1] - input_ary[0,0,:,1]).max() < 1e-6 | |
| # assert np.abs(scipy_ary[-1] - input_ary[0,0,:,0]).max() < 1e-6 | |
| affine_tensor = affine_func2d( | |
| transform_tensor, | |
| torch.Size(output_size) | |
| ) | |
| print('affine_tensor', affine_tensor.size(), affine_tensor.dtype, affine_tensor.device) | |
| print(affine_tensor) | |
| gridsample_ary = torch.nn.functional.grid_sample( | |
| torch.tensor(input_ary, device=device).to(device), | |
| affine_tensor, | |
| padding_mode='border' | |
| ).to('cpu').numpy() | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('gridsample_ary', gridsample_ary.shape, gridsample_ary.dtype) | |
| print(gridsample_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| assert np.abs(scipy_ary.mean() - gridsample_ary.mean()) < 1e-6 | |
| assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # assert False | |
| def test_affine_2d_rotate45(device, affine_func2d): | |
| input_size = [1, 1, 3, 3] | |
| # input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| input_ary = np.array(np.zeros(input_size), dtype=np.float32) | |
| input_ary[0,0,0,:] = 0.5 | |
| input_ary[0,0,2,2] = 1.0 | |
| output_size = [1, 1, 3, 3] | |
| angle_rad = 0.125 * math.pi * 2 | |
| transform_tensor, transform_ary, offset = _buildEquivalentTransforms2d(device, input_size, output_size, angle_rad) | |
| # reference | |
| # https://stackoverflow.com/questions/20161175/how-can-i-use-scipy-ndimage-interpolation-affine-transform-to-rotate-an-image-ab | |
| scipy_ary = scipy.ndimage.affine_transform( | |
| input_ary[0,0], | |
| transform_ary, | |
| offset=offset, | |
| output_shape=output_size[2:], | |
| # output=None, | |
| order=1, | |
| mode='nearest', | |
| # cval=0.0, | |
| prefilter=False) | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| affine_tensor = affine_func2d( | |
| transform_tensor, | |
| torch.Size(output_size) | |
| ) | |
| print('affine_tensor', affine_tensor.size(), affine_tensor.dtype, affine_tensor.device) | |
| print(affine_tensor) | |
| gridsample_ary = torch.nn.functional.grid_sample( | |
| torch.tensor(input_ary, device=device).to(device), | |
| affine_tensor, | |
| padding_mode='border' | |
| ).to('cpu').numpy() | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('gridsample_ary', gridsample_ary.shape, gridsample_ary.dtype) | |
| print(gridsample_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # assert False | |
| def test_affine_2d_rotateRandom(device, affine_func2d, angle_rad, input_size2d, output_size2d): | |
| input_size = input_size2d | |
| input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| output_size = output_size2d | |
| angle_rad = random.random() * math.pi * 2 | |
| transform_tensor, transform_ary, offset = _buildEquivalentTransforms2d(device, input_size, output_size, angle_rad) | |
| # reference | |
| # https://stackoverflow.com/questions/20161175/how-can-i-use-scipy-ndimage-interpolation-affine-transform-to-rotate-an-image-ab | |
| scipy_ary = scipy.ndimage.affine_transform( | |
| input_ary[0,0], | |
| transform_ary, | |
| offset=offset, | |
| output_shape=output_size[2:], | |
| # output=None, | |
| order=1, | |
| mode='nearest', | |
| # cval=0.0, | |
| prefilter=False) | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| affine_tensor = affine_func2d( | |
| transform_tensor, | |
| torch.Size(output_size) | |
| ) | |
| print('affine_tensor', affine_tensor.size(), affine_tensor.dtype, affine_tensor.device) | |
| print(affine_tensor) | |
| gridsample_ary = torch.nn.functional.grid_sample( | |
| torch.tensor(input_ary, device=device).to(device), | |
| affine_tensor, | |
| padding_mode='border' | |
| ).to('cpu').numpy() | |
| print('input_ary', input_ary.shape, input_ary.dtype) | |
| print(input_ary) | |
| print('gridsample_ary', gridsample_ary.shape, gridsample_ary.dtype) | |
| print(gridsample_ary) | |
| print('scipy_ary', scipy_ary.shape, scipy_ary.dtype) | |
| print(scipy_ary) | |
| assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # assert False | |
| # def test_affine_2d_resize(device): | |
| # input_size = [1, 1, 3, 4] | |
| # input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| # output_size = [1, 1, 4, 5] | |
| # | |
| # with warnings.catch_warnings(): | |
| # warnings.simplefilter("ignore") | |
| # # https://docs.scipy.org/doc/scipy-0.16.1/reference/generated/scipy.ndimage.interpolation.zoom.html | |
| # scipy_ary = scipy.ndimage.interpolation.zoom( | |
| # input_ary, | |
| # tuple(o/i for i, o in zip(input_size, output_size)), | |
| # order=1 | |
| # ) | |
| # | |
| # # theta = torch.zeros(1, 2, 3).to(device) | |
| # theta = torch.zeros(1, 2, 3, device=device).to(device) | |
| # print(device) | |
| # print('theta', theta.size(), theta.dtype, theta.device) | |
| # theta[0,0,0] = 1 #output_size[2] / input_size[2] | |
| # theta[0,1,1] = 1 #output_size[3] / input_size[3] | |
| # # theta[2,2] = output_size[4] / input_size[4] | |
| # print('theta', theta.size(), theta.dtype, theta.device) | |
| # | |
| # affine_tensor = affine_grid_generator( | |
| # theta, | |
| # torch.Size(output_size) | |
| # ) | |
| # | |
| # print('theta', theta.size(), theta.dtype, theta.device) | |
| # print(theta) | |
| # print('affine_tensor', affine_tensor.size(), affine_tensor.dtype, theta.device) | |
| # print(affine_tensor[:,:,:,0]) | |
| # print(affine_tensor[:,:,:,1]) | |
| # print('input_ary', input_size, output_size, torch.tensor(input_ary, device=device).device) | |
| # print(input_ary) | |
| # print('scipy_ary', scipy_ary.shape) | |
| # print(scipy_ary) | |
| # | |
| # | |
| # gridsample_ary = torch.nn.functional.grid_sample( | |
| # torch.tensor(input_ary, device=device).to(device), | |
| # affine_tensor, | |
| # ).to('cpu').numpy() | |
| # | |
| # print('gridsample_ary', gridsample_ary.shape) | |
| # print(gridsample_ary) | |
| # | |
| # assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # | |
| # | |
| # def test_affine_2d_rotate(device): | |
| # # input_size = [1, 1, 3, 4] | |
| # input_size = [1, 1, 3, 3] | |
| # input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| # # output_size = [1, 1, 4, 5] | |
| # output_size = [1, 1, 3, 3] | |
| # | |
| # angle_rad = random.random() * math.pi * 2 | |
| # # angle_rad = 0.5 * math.pi * 2 | |
| # angle_deg = angle_rad * 180 / math.pi | |
| # s = math.sin(angle_rad) | |
| # c = math.cos(angle_rad) | |
| # c1 = 1 - c | |
| # # axis_vector = torch.rand(3, device='cpu', dtype=torch.float64) | |
| # # axis_vector -= 0.5 | |
| # # axis_vector /= axis_vector.abs().sum() | |
| # # l, m, n = axis_vector | |
| # | |
| # # l, m, n = 1, 0, 0 | |
| # # | |
| # # theta = torch.tensor([ | |
| # # [l*l*c1 + c, m*l*c1 - n*s, n*l*c1 + m*s, 0], | |
| # # [l*m*c1 + n*s, m*m*c1 + c, n*m*c1 - l*s, 0], | |
| # # [l*n*c1 - m*s, m*n*c1 + l*s, n*n*c1 + c, 0], | |
| # # [0, 0, 0, 1], | |
| # # ], device=device, dtype=torch.float32).to(device) | |
| # | |
| # | |
| # center_matrix = torch.tensor([ | |
| # [1, 0, 0*output_size[2]/2], | |
| # [0, 1, 0*output_size[3]/2], | |
| # [0, 0, 1], | |
| # ], device=device, dtype=torch.float32).to(device) | |
| # rotation_matrix = torch.tensor([ | |
| # [c, -s, 0], | |
| # [s, c, 0], | |
| # [0, 0, 1], | |
| # ], device=device, dtype=torch.float32).to(device) | |
| # shift_matrix = torch.tensor([ | |
| # [1, 0, 0*-output_size[2]/2], | |
| # [0, 1, 0*-output_size[3]/2], | |
| # [0, 0, 1], | |
| # ], device=device, dtype=torch.float32).to(device) | |
| # | |
| # print('rotation_matrix', rotation_matrix.size(), rotation_matrix.dtype, rotation_matrix.device) | |
| # | |
| # theta = (center_matrix @ rotation_matrix @ shift_matrix)[:2].unsqueeze(0) | |
| # | |
| # print('theta', theta.size(), theta.dtype, theta.device) | |
| # | |
| # | |
| # with warnings.catch_warnings(): | |
| # warnings.simplefilter("ignore") | |
| # # https://docs.scipy.org/doc/scipy-0.16.1/reference/generated/scipy.ndimage.interpolation.zoom.html | |
| # scipy_ary = scipy.ndimage.interpolation.zoom( | |
| # input_ary, | |
| # tuple(o/i for i, o in zip(input_size, output_size)), | |
| # order=1 | |
| # ) | |
| # | |
| # scipy_ary = scipy.ndimage.interpolation.rotate( | |
| # scipy_ary, | |
| # angle_deg, | |
| # axes=(-2, -1), | |
| # reshape=False, | |
| # order=1, | |
| # mode='nearest', | |
| # ) | |
| # | |
| # # theta = torch.zeros(1, 2, 3, device=device).to(device) | |
| # # theta[0,0,0] = 1 #output_size[2] / input_size[2] | |
| # # theta[0,1,1] = 1 #output_size[3] / input_size[3] | |
| # # theta[2,2] = output_size[4] / input_size[4] | |
| # | |
| # affine_tensor = affine_grid_generator( | |
| # theta, | |
| # torch.Size(output_size) | |
| # ) | |
| # | |
| # | |
| # gridsample_ary = torch.nn.functional.grid_sample( | |
| # torch.tensor(input_ary, device=device).to(device), | |
| # affine_tensor, | |
| # padding_mode='border', | |
| # ).to('cpu').numpy() | |
| # | |
| # print('theta', angle_deg, angle_rad) | |
| # print(theta) | |
| # print('affine_tensor', affine_tensor.size()) | |
| # print(affine_tensor[:,:,:,0]) | |
| # print(affine_tensor[:,:,:,1]) | |
| # print('input_ary', input_size, output_size) | |
| # print(input_ary) | |
| # print('scipy_ary', scipy_ary.shape) | |
| # print(scipy_ary) | |
| # print('gridsample_ary', gridsample_ary.shape) | |
| # print(gridsample_ary) | |
| # | |
| # assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # | |
| # | |
| # | |
| # def test_affine_3d_rotate(device): | |
| # # input_size = [1, 1, 3, 4, 5] | |
| # input_size = [1, 1, 2, 2, 2] | |
| # input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| # # output_size = [1, 1, 4, 5, 6] | |
| # output_size = [1, 1, 2, 2, 2] | |
| # | |
| # # angle_rad = random.random() * math.pi * 2 | |
| # angle_rad = 0.5 * math.pi * 2 | |
| # angle_deg = angle_rad * 180 / math.pi | |
| # s = math.sin(angle_rad) | |
| # c = math.cos(angle_rad) | |
| # c1 = 1 - c | |
| # | |
| # with warnings.catch_warnings(): | |
| # warnings.simplefilter("ignore") | |
| # # https://docs.scipy.org/doc/scipy-0.16.1/reference/generated/scipy.ndimage.interpolation.zoom.html | |
| # scipy_ary = scipy.ndimage.interpolation.zoom( | |
| # input_ary, | |
| # tuple(o/i for i, o in zip(input_size, output_size)), | |
| # order=1 | |
| # ) | |
| # | |
| # scipy_ary = scipy.ndimage.interpolation.rotate( | |
| # scipy_ary, | |
| # angle_deg, | |
| # axes=(-3, -2), | |
| # reshape=False, | |
| # order=1, | |
| # mode='nearest', | |
| # ) | |
| # | |
| # rotation_matrix = torch.tensor([ | |
| # [c, -s, 0, 0], | |
| # [s, c, 0, 0], | |
| # [0, 0, 1, 0], | |
| # [0, 0, 0, 1], | |
| # ], device=device, dtype=torch.float32).to(device) | |
| # | |
| # theta = (rotation_matrix)[:3].unsqueeze(0) | |
| # | |
| # # theta = torch.zeros(1, 3, 4, device=device).to(device) | |
| # # theta[0,0,0] = 1 #output_size[2] / input_size[2] | |
| # # theta[0,1,1] = 1 #output_size[3] / input_size[3] | |
| # # theta[0,2,2] = 1 #output_size[4] / input_size[4] | |
| # | |
| # affine_tensor = affine_grid_generator( | |
| # theta, | |
| # torch.Size(output_size) | |
| # ) | |
| # | |
| # gridsample_ary = torch.nn.functional.grid_sample( | |
| # torch.tensor(input_ary, device=device).to(device), | |
| # affine_tensor, | |
| # ).to('cpu').numpy() | |
| # | |
| # print('theta') | |
| # print(theta) | |
| # print('affine_tensor', affine_tensor.size()) | |
| # print(affine_tensor[:,:,:,:,0]) | |
| # print(affine_tensor[:,:,:,:,1]) | |
| # print('input_ary', input_size, output_size) | |
| # print(input_ary) | |
| # print('scipy_ary', scipy_ary.shape) | |
| # print(scipy_ary) | |
| # print('gridsample_ary', gridsample_ary.shape) | |
| # print(gridsample_ary) | |
| # | |
| # assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # | |
| # def test_affine_3d_resize(device): | |
| # input_size = [1, 1, 3, 4, 5] | |
| # input_ary = np.array(np.random.random(input_size), dtype=np.float32) | |
| # output_size = [1, 1, 4, 5, 6] | |
| # | |
| # with warnings.catch_warnings(): | |
| # warnings.simplefilter("ignore") | |
| # # https://docs.scipy.org/doc/scipy-0.16.1/reference/generated/scipy.ndimage.interpolation.zoom.html | |
| # scipy_ary = scipy.ndimage.interpolation.zoom( | |
| # input_ary, | |
| # tuple(o/i for i, o in zip(input_size, output_size)), | |
| # order=1 | |
| # ) | |
| # | |
| # theta = torch.zeros(1, 3, 4, device=device).to(device) | |
| # theta[0,0,0] = 1 #output_size[2] / input_size[2] | |
| # theta[0,1,1] = 1 #output_size[3] / input_size[3] | |
| # theta[0,2,2] = 1 #output_size[4] / input_size[4] | |
| # | |
| # affine_tensor = affine_grid_generator( | |
| # theta, | |
| # torch.Size(output_size) | |
| # ) | |
| # | |
| # gridsample_ary = torch.nn.functional.grid_sample( | |
| # torch.tensor(input_ary, device=device).to(device), | |
| # affine_tensor, | |
| # ).to('cpu').numpy() | |
| # | |
| # print('theta') | |
| # print(theta) | |
| # print('affine_tensor', affine_tensor.size()) | |
| # print(affine_tensor[:,:,:,:,0]) | |
| # print(affine_tensor[:,:,:,:,1]) | |
| # print('input_ary', input_size, output_size) | |
| # print(input_ary) | |
| # print('scipy_ary', scipy_ary.shape) | |
| # print(scipy_ary) | |
| # print('gridsample_ary', gridsample_ary.shape) | |
| # print(gridsample_ary) | |
| # | |
| # assert np.abs(scipy_ary - gridsample_ary).max() < 1e-6 | |
| # | |
| # | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment