I've been trying to add custom augmentation to a model. However, while library functions such as tf.image.random_flip_left_right or tf.image.random_brightness actually work and create random variations for every image in every sample, my custom functions always have the same effect for ALL images in ALL batches.
For instance, these are my functions for "random" gaussian and motion blurs:
@tf.function
def gaussian(image, n_channels, dtype):
size = np.random.randint(2, 6)
sigma = np.random.uniform(0.5, 4.0)
x = tf.range(-size // 2 + 1, size // 2 + 1, dtype=dtype)
g = tf.math.exp(-(tf.pow(x, 2) / (2 * tf.pow(tf.cast(sigma, dtype), 2))))
g_norm2d = tf.pow(tf.reduce_sum(g), 2)
k = tf.tensordot(g, g, axes=0) / g_norm2d
k = tf.expand_dims(k, axis=-1)
k = tf.expand_dims(tf.tile(k, (1, 1, n_channels)), axis=-1)
return tf.nn.depthwise_conv2d(image[None], k, [1,1,1,1], 'SAME')[0]
@tf.function
def motion(image, n_channels, dtype):
size = np.random.randint(2, 11)
k = np.zeros((size, size))
k[int((size-1)/2), :] = np.ones(size)
k = k / size
k = tf.convert_to_tensor(k, dtype=dtype)
k = tf.expand_dims(k, axis=-1)
k = tf.expand_dims(tf.tile(k, (1, 1, n_channels)), axis=-1)
return tf.nn.depthwise_conv2d(image[None], k, [1,1,1,1], 'SAME')[0]
As I said, my images get random brightness, flip, etc. But they ALL get the same motion and Gaussian blur, which is not what I wanted.
EDIT: To clarify, I call all augmentation functions sequentially in an augment(image, label) function, which is called via dataset.map(augment). This works just fine for the other augmentations.
Okay, so I tried MANY things but what made it work was using only TF functions... which forced me to change the code a bit and make things more complicated but hey that's TF for you.
The code looks now like this:
@tf.function
def gaussian(image, n_channels):
size = tf.random.uniform([1], 2, 6, dtype=tf.int32)[0]
sigma = tf.random.uniform([1], 0, 4.0, dtype=tf.float32)[0]
x = tf.range(-size // 2 + 1, size // 2 + 1, dtype=tf.float32)
g = tf.math.exp(-(tf.pow(x, 2) / (2 * tf.pow(tf.cast(sigma, dtype=tf.float32), 2))))
g_norm2d = tf.pow(tf.reduce_sum(g), 2)
k = tf.tensordot(g, g, axes=0) / g_norm2d
k = tf.expand_dims(k, axis=-1)
k = tf.expand_dims(tf.tile(k, (1, 1, n_channels)), axis=-1)
return tf.nn.depthwise_conv2d(image[None], k, [1,1,1,1], 'SAME')[0]
@tf.function
def motion(image, n_channels):
size = tf.random.uniform([1], 2, 11, dtype=tf.int32)[0]
a = tf.zeros(((size-1)/2, size))
b = tf.reshape(tf.ones((size)), [1,size])
k = tf.concat([a,b,a], 0)
k = k / tf.cast(size, tf.float32)
k = tf.expand_dims(k, axis=-1)
k = tf.expand_dims(tf.tile(k, (1, 1, n_channels)), axis=-1)
return tf.nn.depthwise_conv2d(image[None], k, [1,1,1,1], 'SAME')[0]