def gpu_affine_grid(size, device):
size = ((1, ) + size)
if len(size) == 4:
N, C, H, W = size
grid = torch.FloatTensor(N, H, W, 2).cuda(device)
linear_points = torch.linspace(-1, 1, W) if W > 1 else tensor([-1.])
grid[:, :, :, 0] = torch.ger(torch.ones(H),
linear_points).expand_as(grid[:, :, :, 0])
linear_points = torch.linspace(-1, 1, H) if H > 1 else tensor([-1.])
grid[:, :, :, 1] = torch.ger(linear_points,
torch.ones(W)).expand_as(grid[:, :, :, 1])
return vision.FlowField(size[2:], grid)
elif len(size) == 5:
N, C, D, H, W = size
grid = torch.FloatTensor(N, D, H, W, 2).cuda(device)
linear_points = torch.linspace(-1, 1, W) if W > 1 else tensor([-1.])
grid[:, :, :, :,
0] = torch.ger(torch.ones(H),
linear_points).expand_as(grid[:, :, :, :, 0])
linear_points = torch.linspace(-1, 1, H) if H > 1 else tensor([-1.])
grid[:, :, :, :,
1] = torch.ger(linear_points,
torch.ones(W)).expand_as(grid[:, :, :, :, 1])
return vision.FlowField(size[2:], grid)
def get_image_pt(image, points):
return vision.ImagePoints(vision.FlowField(image.size, points), scale=True)