Commit 060d732b authored by Jean Ibarz's avatar Jean Ibarz
Browse files

Fixed a bug in CoordConv layer implementation, where coordinates were not...

Fixed a bug in CoordConv layer implementation, where coordinates were not scaled correctly along each dimensions.
parent 90e7cad1
......@@ -24,7 +24,7 @@ from tensorflow.python.layers import base
import tensorflow as tf
class AddCoords(base.Layer):
class AddCoords(tf.keras.layers.Layer):
"""Add coords to a tensor"""
def __init__(self, x_dim=64, y_dim=64, with_r=False, skiptile=False):
......@@ -40,34 +40,38 @@ class AddCoords(base.Layer):
In the first case, first tile the input_tensor to be (batch, x_dim, y_dim, c)
In the second case, skiptile, just concat
"""
x_dim = self.x_dim
y_dim = self.y_dim
if not self.skiptile:
input_tensor = tf.tile(input_tensor, [1, self.x_dim, self.y_dim, 1]) # (batch, 64, 64, 2)
input_tensor = tf.tile(input_tensor, [1, x_dim, y_dim, 1]) # (batch, x, y, 2)
input_tensor = tf.cast(input_tensor, 'float32')
batch_size_tensor = tf.shape(input_tensor)[0] # get batch size
xx_ones = tf.ones([batch_size_tensor, self.x_dim],
dtype=tf.int32) # e.g. (batch, 64)
xx_ones = tf.expand_dims(xx_ones, -1) # e.g. (batch, 64, 1)
xx_range = tf.tile(tf.expand_dims(tf.range(self.y_dim), 0),
[batch_size_tensor, 1]) # e.g. (batch, 64)
xx_range = tf.expand_dims(xx_range, 1) # e.g. (batch, 1, 64)
xx_ones = tf.ones([batch_size_tensor, y_dim],
dtype=tf.int32) # e.g. (batch, y)
xx_ones = tf.expand_dims(xx_ones, 1) # e.g. (batch, 1, y)
xx_range = tf.tile(tf.expand_dims(tf.range(x_dim), 0),
[batch_size_tensor, 1]) # e.g. (batch, x)
xx_range = tf.expand_dims(xx_range, -1) # e.g. (batch, x, 1)
xx_channel = tf.matmul(xx_ones, xx_range) # e.g. (batch, 64, 64)
xx_channel = tf.expand_dims(xx_channel, -1) # e.g. (batch, 64, 64, 1)
xx_channel = tf.matmul(xx_range, xx_ones) # e.g. (batch, x, y)
xx_channel = tf.expand_dims(xx_channel, -1) # e.g. (batch, x, y, 1)
yy_ones = tf.ones([batch_size_tensor, self.y_dim],
dtype=tf.int32) # e.g. (batch, 64)
yy_ones = tf.expand_dims(yy_ones, 1) # e.g. (batch, 1, 64)
yy_range = tf.tile(tf.expand_dims(tf.range(self.x_dim), 0),
[batch_size_tensor, 1]) # (batch, 64)
yy_range = tf.expand_dims(yy_range, -1) # e.g. (batch, 64, 1)
yy_ones = tf.ones([batch_size_tensor, x_dim],
dtype=tf.int32) # e.g. (batch, x)
yy_ones = tf.expand_dims(yy_ones, -1) # e.g. (batch, x, 1)
yy_range = tf.tile(tf.expand_dims(tf.range(y_dim), 0),
[batch_size_tensor, 1]) # (batch, y)
yy_range = tf.expand_dims(yy_range, 1) # e.g. (batch, 1, y)
yy_channel = tf.matmul(yy_range, yy_ones) # e.g. (batch, 64, 64)
yy_channel = tf.expand_dims(yy_channel, -1) # e.g. (batch, 64, 64, 1)
yy_channel = tf.matmul(yy_ones, yy_range) # e.g. (batch, x, y)
yy_channel = tf.expand_dims(yy_channel, -1) # e.g. (batch, x, y, 1)
xx_channel = tf.cast(xx_channel, 'float32') / (self.x_dim - 1)
yy_channel = tf.cast(yy_channel, 'float32') / (self.y_dim - 1)
xx_channel = tf.cast(xx_channel, 'float32') / (x_dim - 1)
yy_channel = tf.cast(yy_channel, 'float32') / (y_dim - 1)
xx_channel = xx_channel * 2 - 1 # [-1,1]
yy_channel = yy_channel * 2 - 1
......@@ -84,7 +88,7 @@ class AddCoords(base.Layer):
return ret
class CoordConv(base.Layer):
class CoordConv(tf.keras.layers.Layer):
"""CoordConv layer as in the paper."""
def __init__(self, x_dim, y_dim, with_r, *args, **kwargs):
......@@ -93,7 +97,7 @@ class CoordConv(base.Layer):
y_dim=y_dim,
with_r=with_r,
skiptile=True)
self.conv = tf.layers.Conv2D(*args, **kwargs)
self.conv = tf.keras.layers.Conv2D(*args, **kwargs)
def call(self, input_tensor):
ret = self.addcoords(input_tensor)
......
......@@ -126,4 +126,20 @@ def test_gtf_layer2():
x_right = gtf_layer(x[:, :, 1, :])
x_band = tf.stack([x_left, x_right], axis=-2)
np_r = x_band.numpy()
print('b')
\ No newline at end of file
def test_add_coords():
from core.coordconv import AddCoords
batch_size = 3
x_dim = 14
y_dim = 6
input = np.zeros(shape=(batch_size, x_dim, y_dim, 1))
input_with_coords = AddCoords(x_dim=x_dim, y_dim=y_dim, with_r=False, skiptile=True)(input).numpy()
assert input_with_coords.shape == (batch_size, x_dim, y_dim, 1 + 2)
assert np.allclose(a=input_with_coords[0, 0, :, 1], b=-1.0 * np.ones(y_dim, ))
assert np.allclose(a=input_with_coords[0, -1, :, 1], b=np.ones(y_dim, ))
assert np.allclose(a=input_with_coords[0, :, 0, 2], b=-1.0 * np.ones(x_dim, ))
assert np.allclose(a=input_with_coords[0, :, -1, 2], b=np.ones(x_dim, ))
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment