Remove blank rows and columns of an array inside a keras Sequential model-CodePudding

I have a keras model, which takes a 10x10x1 array as input. For example:

array = np.array([
[[0],[0],[0],[0],[0],[0],[0],[0],[0],[0]],
[[0],[0],[0],[0],[0],[0],[0],[0],[0],[0]],
[[0],[0],[0],[0],[0],[0],[0],[0],[0],[0]],
[[0],[0],[0],[0],[0],[0],[0],[0],[0],[0]],
[[0],[0],[0],[0],[1],[1],[0],[0],[0],[0]],
[[0],[0],[0],[0],[1],[0],[0],[0],[0],[0]],
[[0],[0],[0],[0],[1],[1],[0],[0],[0],[0]],
[[0],[0],[0],[0],[0],[1],[0],[0],[0],[0]],
[[0],[0],[0],[0],[1],[1],[0],[0],[0],[0]],
[[0],[0],[0],[0],[0],[0],[0],[0],[0],[0]]
])

This array represents the number '5'

What I want to do is, to remove all the 'blank' (pixel value of 0) rows and columns.
The resultant array would have a shape of 5x2x1,

array = np.array([
[[1],[1]],
[[1],[0]],
[[1],[1]],
[[0],[1]],
[[1],[1]]
])

And it will look like this:

I want to do this in a keras model, as a layer, so it has to be done with tensorflow.keras.backend.
It can be as a Lambda layer, like this:

def remove_blank(x):
   ...
   return x_

model.add(Lambda(remove_blank))

Also, this function should be able to do this with RGB images too, meaning a 3D array of shape (HEIGHT, WIDTH, 3).

An example of RGB image would be:

array = np.array([
[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[1,1,1],[1,1,1],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[1,1,1],[0,0,0],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[1,1,1],[1,1,1],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[1,1,1],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[1,1,1],[1,1,1],[0,0,0]],
[[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]]
], dtype='float')

The function should remove the rows and columns which contain only black pixels and transform this image to this:

array = np.array([
[[1,1,1],[1,1,1]],
[[1,1,1],[0,0,0]],
[[1,1,1],[1,1,1]],
[[0,0,0],[1,1,1]],
[[1,1,1],[1,1,1]]
], dtype='float')

CodePudding user response：

You could try something like this, but note that iterating over a tensor in graph mode will not work, so rather use tf.while_loop for a batch_size > 1. If batch_size = 1, you do not need any loop:

import tensorflow as tf
import matplotlib.pyplot as plt

five = tf.constant([
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]]
], dtype=tf.float32)

eight = tf.constant([
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [0], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [0], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]]
], dtype=tf.float32)

def remove_blank(t):
  ragged = []
  for i in tf.range(tf.shape(t)[0]):
    tensor = tf.expand_dims(t[i, :, :, 0], axis=0)
    x = tf.where(tf.not_equal(tensor, 0.0))
    _min = tf.reduce_min(x, axis=0)
    _max = tf.reduce_max(x, axis=0)
    col_range = tf.range(_min[1], _max[1]   1)
    row_range = tf.range(_min[2], _max[2]   1)

    i = tf.repeat(col_range, repeats=tf.shape(row_range)[0])
    j = tf.tile(row_range, multiples=[tf.shape(col_range)[0]])
    b = tf.zeros(tf.shape(i)[0], dtype=tf.int64)

    z = tf.gather_nd(tensor, tf.transpose(tf.stack([b, i, j])))
    z = tf.reshape(z, (tf.shape(tf.unique(x[:, 1])[0])[0], tf.shape(tf.unique(x[:, 2])[0])[0], tf.shape(t)[-1]))
    ragged.append(z)
  return tf.ragged.stack(ragged)

x = tf.stack([five, eight])

model = tf.keras.Sequential()
model.add(tf.keras.layers.Lambda(remove_blank))
output = model(x)

f = plt.figure(figsize=(10,3))
ax1 = f.add_subplot(121)
ax2 = f.add_subplot(122)
ax1.imshow(five[:, :, 0], cmap='gray')
ax2.imshow(eight[:, :, 0], cmap='gray')

f1 = plt.figure(figsize=(10,3))
ax1 = f1.add_subplot(121)
ax2 = f1.add_subplot(122)
ax1.imshow(output[0].to_tensor()[:, :, 0], cmap='gray')
ax2.imshow(output[1].to_tensor()[:, :, 0], cmap='gray')

You will have to resize the images to a uniform size when training your model. I will leave it to you to figure out how to extend this method for RGB images ;)

Update 1:

Here is an example of the method integrated into a Keras model:

import tensorflow as tf
import matplotlib.pyplot as plt

five = tf.constant([
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]]
], dtype=tf.float32)

eight = tf.constant([
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [0], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [0], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [1], [1], [1], [0], [0], [0], [0]],
    [[0], [0], [0], [0], [0], [0], [0], [0], [0], [0]]
], dtype=tf.float32)

def remove_blank(t):
  def _remove_blank(i0, v, t):
    tensor = tf.expand_dims(t[i0, :, :, 0], axis=0)
    x = tf.where(tf.not_equal(tensor, 0.0))
    _min = tf.reduce_min(x, axis=0)
    _max = tf.reduce_max(x, axis=0)
    col_range = tf.range(_min[1], _max[1]   1)
    row_range = tf.range(_min[2], _max[2]   1)

    i = tf.repeat(col_range, repeats=tf.shape(row_range)[0])
    j = tf.tile(row_range, multiples=[tf.shape(col_range)[0]])
    b = tf.zeros(tf.shape(i)[0], dtype=tf.int64)

    z = tf.gather_nd(tensor, tf.transpose(tf.stack([b, i, j])))
    z = tf.reshape(z, (tf.shape(tf.unique(x[:, 1])[0])[0], tf.shape(tf.unique(x[:, 2])[0])[0], tf.shape(t)[-1]))
    z = tf.expand_dims(tf.image.resize(z, [5, 2]), axis=0)

    h = tf.range(tf.shape(z)[1], dtype=tf.int32)
    ij = tf.expand_dims(tf.stack([tf.repeat([i0], repeats=tf.shape(h)[0]), h], axis=-1), axis=0)
    v = tf.tensor_scatter_nd_update(v, ij, z)

    return tf.add(i0, 1), v, t

  i0 = tf.constant(0, dtype=tf.int32)
  v = tf.zeros((tf.shape(t)[0], 5, 2, 1))
  while_condition = lambda i0, v, t: tf.less(i0, tf.shape(t)[0])
  _, v, _ = tf.while_loop(while_condition, _remove_blank, loop_vars=(i0, v, t))
  return v

batch_size = 2
inputs = tf.keras.layers.Input((10, 10, 1))
outputs = tf.keras.layers.Lambda(remove_blank)(inputs)
model = tf.keras.Model(inputs, outputs)

y = tf.random.normal((2, 5, 2, 1))
model.compile(optimizer='adam', loss='mse')
model.fit(x, y, batch_size=2, epochs=4)

Epoch 1/4
1/1 [==============================] - 1s 860ms/step - loss: 1.1193
Epoch 2/4
1/1 [==============================] - 0s 17ms/step - loss: 1.1193
Epoch 3/4
1/1 [==============================] - 0s 16ms/step - loss: 1.1193
Epoch 4/4
1/1 [==============================] - 0s 16ms/step - loss: 1.1193
<keras.callbacks.History at 0x7f11b0350510>