190 lines
6.0 KiB
Python
190 lines
6.0 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
|
|
def transformer(U, theta, out_size, **kwargs):
|
|
"""Spatial Transformer Layer
|
|
|
|
Implements a spatial transformer layer as described in [1]_.
|
|
Based on [2]_ and edited by David Dao for Tensorflow.
|
|
|
|
Parameters
|
|
----------
|
|
U : float
|
|
The output of a convolutional net should have the
|
|
shape [num_batch, height, width, num_channels].
|
|
theta: float
|
|
The output of the
|
|
localisation network should be [num_batch, 6].
|
|
out_size: tuple of two ints
|
|
The size of the output of the network (height, width)
|
|
|
|
References
|
|
----------
|
|
.. [1] Spatial Transformer Networks
|
|
Max Jaderberg, Karen Simonyan, Andrew Zisserman, Koray Kavukcuoglu
|
|
Submitted on 5 Jun 2015
|
|
.. [2] https://github.com/skaae/transformer_network/blob/master/transformerlayer.py
|
|
|
|
Notes
|
|
-----
|
|
To initialize the network to the identity transform init
|
|
``theta`` to :
|
|
identity = np.array([[1., 0., 0.],
|
|
[0., 1., 0.]])
|
|
identity = identity.flatten()
|
|
theta = tf.Variable(initial_value=identity)
|
|
|
|
"""
|
|
|
|
def _repeat(x, n_repeats):
|
|
rep = torch.ones(
|
|
[
|
|
n_repeats,
|
|
]
|
|
).unsqueeze(0)
|
|
rep = rep.int()
|
|
x = x.int()
|
|
|
|
x = torch.matmul(x.reshape([-1, 1]), rep)
|
|
return x.reshape([-1])
|
|
|
|
def _interpolate(im, x, y, out_size):
|
|
num_batch, num_channels, height, width = im.size()
|
|
|
|
height_f = height
|
|
width_f = width
|
|
out_height, out_width = out_size[0], out_size[1]
|
|
|
|
zero = 0
|
|
max_y = height - 1
|
|
max_x = width - 1
|
|
|
|
x = (x + 1.0) * (width_f) / 2.0
|
|
y = (y + 1.0) * (height_f) / 2.0
|
|
|
|
# do sampling
|
|
x0 = torch.floor(x).int()
|
|
x1 = x0 + 1
|
|
y0 = torch.floor(y).int()
|
|
y1 = y0 + 1
|
|
|
|
x0 = torch.clamp(x0, zero, max_x)
|
|
x1 = torch.clamp(x1, zero, max_x)
|
|
y0 = torch.clamp(y0, zero, max_y)
|
|
y1 = torch.clamp(y1, zero, max_y)
|
|
dim2 = torch.from_numpy(np.array(width))
|
|
dim1 = torch.from_numpy(np.array(width * height))
|
|
|
|
base = _repeat(torch.arange(0, num_batch) * dim1, out_height * out_width)
|
|
if torch.cuda.is_available():
|
|
dim2 = dim2.cuda()
|
|
dim1 = dim1.cuda()
|
|
y0 = y0.cuda()
|
|
y1 = y1.cuda()
|
|
x0 = x0.cuda()
|
|
x1 = x1.cuda()
|
|
base = base.cuda()
|
|
base_y0 = base + y0 * dim2
|
|
base_y1 = base + y1 * dim2
|
|
idx_a = base_y0 + x0
|
|
idx_b = base_y1 + x0
|
|
idx_c = base_y0 + x1
|
|
idx_d = base_y1 + x1
|
|
|
|
# channels dim
|
|
im = im.permute(0, 2, 3, 1)
|
|
im_flat = im.reshape([-1, num_channels]).float()
|
|
|
|
idx_a = idx_a.unsqueeze(-1).long()
|
|
idx_a = idx_a.expand(out_height * out_width * num_batch, num_channels)
|
|
Ia = torch.gather(im_flat, 0, idx_a)
|
|
|
|
idx_b = idx_b.unsqueeze(-1).long()
|
|
idx_b = idx_b.expand(out_height * out_width * num_batch, num_channels)
|
|
Ib = torch.gather(im_flat, 0, idx_b)
|
|
|
|
idx_c = idx_c.unsqueeze(-1).long()
|
|
idx_c = idx_c.expand(out_height * out_width * num_batch, num_channels)
|
|
Ic = torch.gather(im_flat, 0, idx_c)
|
|
|
|
idx_d = idx_d.unsqueeze(-1).long()
|
|
idx_d = idx_d.expand(out_height * out_width * num_batch, num_channels)
|
|
Id = torch.gather(im_flat, 0, idx_d)
|
|
|
|
x0_f = x0.float()
|
|
x1_f = x1.float()
|
|
y0_f = y0.float()
|
|
y1_f = y1.float()
|
|
|
|
wa = torch.unsqueeze(((x1_f - x) * (y1_f - y)), 1)
|
|
wb = torch.unsqueeze(((x1_f - x) * (y - y0_f)), 1)
|
|
wc = torch.unsqueeze(((x - x0_f) * (y1_f - y)), 1)
|
|
wd = torch.unsqueeze(((x - x0_f) * (y - y0_f)), 1)
|
|
output = wa * Ia + wb * Ib + wc * Ic + wd * Id
|
|
|
|
return output
|
|
|
|
def _meshgrid(height, width):
|
|
x_t = torch.matmul(
|
|
torch.ones([height, 1]),
|
|
torch.transpose(torch.unsqueeze(torch.linspace(-1.0, 1.0, width), 1), 1, 0),
|
|
)
|
|
y_t = torch.matmul(
|
|
torch.unsqueeze(torch.linspace(-1.0, 1.0, height), 1),
|
|
torch.ones([1, width]),
|
|
)
|
|
# x_t = torch.matmul(torch.ones([height, 1]),
|
|
# torch.transpose(torch.unsqueeze(torch.linspace(0.0, width.float(), width), 1), 1, 0))
|
|
# y_t = torch.matmul(torch.unsqueeze(torch.linspace(0.0, height.float(), height), 1),
|
|
# torch.ones([1, width]))
|
|
|
|
x_t_flat = x_t.reshape((1, -1)).float()
|
|
y_t_flat = y_t.reshape((1, -1)).float()
|
|
|
|
ones = torch.ones_like(x_t_flat)
|
|
grid = torch.cat([x_t_flat, y_t_flat, ones], 0)
|
|
if torch.cuda.is_available():
|
|
grid = grid.cuda()
|
|
return grid
|
|
|
|
def _transform(theta, input_dim, out_size):
|
|
num_batch, num_channels, height, width = input_dim.size()
|
|
# Changed
|
|
theta = theta.reshape([-1, 3, 3]).float()
|
|
|
|
out_height, out_width = out_size[0], out_size[1]
|
|
grid = _meshgrid(out_height, out_width)
|
|
grid = grid.unsqueeze(0).reshape([1, -1])
|
|
shape = grid.size()
|
|
grid = grid.expand(num_batch, shape[1])
|
|
grid = grid.reshape([num_batch, 3, -1])
|
|
|
|
T_g = torch.matmul(theta, grid)
|
|
x_s = T_g[:, 0, :]
|
|
y_s = T_g[:, 1, :]
|
|
t_s = T_g[:, 2, :]
|
|
|
|
t_s_flat = t_s.reshape([-1])
|
|
|
|
# smaller
|
|
small = 1e-7
|
|
smallers = 1e-6 * (1.0 - torch.ge(torch.abs(t_s_flat), small).float())
|
|
|
|
t_s_flat = t_s_flat + smallers
|
|
# condition = torch.sum(torch.gt(torch.abs(t_s_flat), small).float())
|
|
# Ty changed
|
|
x_s_flat = x_s.reshape([-1]) / t_s_flat
|
|
y_s_flat = y_s.reshape([-1]) / t_s_flat
|
|
|
|
input_transformed = _interpolate(input_dim, x_s_flat, y_s_flat, out_size)
|
|
|
|
output = input_transformed.reshape(
|
|
[num_batch, out_height, out_width, num_channels]
|
|
)
|
|
output = output.permute(0, 3, 1, 2)
|
|
return output # , condition
|
|
|
|
output = _transform(theta, U, out_size)
|
|
return output # , condition
|