This script demonstrates the use of a convolutional LSTM network.

This network is used to predict the next frame of an artificiallygenerated movie which contains moving squares.

  1. from keras.models import Sequential
  2. from keras.layers.convolutional import Conv3D
  3. from keras.layers.convolutional_recurrent import ConvLSTM2D
  4. from keras.layers.normalization import BatchNormalization
  5. import numpy as np
  6. import pylab as plt
  8. # We create a layer which take as input movies of shape
  9. # (n_frames, width, height, channels) and returns a movie
  10. # of identical shape.
  12. seq = Sequential()
  13. seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
  14.                    input_shape=(None, 40, 40, 1),
  15.                    padding='same', return_sequences=True))
  16. seq.add(BatchNormalization())
  18. seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
  19.                    padding='same', return_sequences=True))
  20. seq.add(BatchNormalization())
  22. seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
  23.                    padding='same', return_sequences=True))
  24. seq.add(BatchNormalization())
  26. seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
  27.                    padding='same', return_sequences=True))
  28. seq.add(BatchNormalization())
  30. seq.add(Conv3D(filters=1, kernel_size=(3, 3, 3),
  31.                activation='sigmoid',
  32.                padding='same', data_format='channels_last'))
  33. seq.compile(loss='binary_crossentropy', optimizer='adadelta')
  36. # Artificial data generation:
  37. # Generate movies with 3 to 7 moving squares inside.
  38. # The squares are of shape 1x1 or 2x2 pixels,
  39. # which move linearly over time.
  40. # For convenience we first create movies with bigger width and height (80x80)
  41. # and at the end we select a 40x40 window.
  43. def generate_movies(n_samples=1200, n_frames=15):
  44.     row = 80
  45.     col = 80
  46.     noisy_movies = np.zeros((n_samples, n_frames, row, col, 1), dtype=np.float)
  47.     shifted_movies = np.zeros((n_samples, n_frames, row, col, 1),
  48.                               dtype=np.float)
  50.     for i in range(n_samples):
  51.         # Add 3 to 7 moving squares
  52.         n = np.random.randint(3, 8)
  54.         for j in range(n):
  55.             # Initial position
  56.             xstart = np.random.randint(20, 60)
  57.             ystart = np.random.randint(20, 60)
  58.             # Direction of motion
  59.             directionx = np.random.randint(0, 3) - 1
  60.             directiony = np.random.randint(0, 3) - 1
  62.             # Size of the square
  63.             w = np.random.randint(2, 4)
  65.             for t in range(n_frames):
  66.                 x_shift = xstart + directionx * t
  67.                 y_shift = ystart + directiony * t
  68.                 noisy_movies[i, t, x_shift - w: x_shift + w,
  69.                              y_shift - w: y_shift + w, 0] += 1
  71.                 # Make it more robust by adding noise.
  72.                 # The idea is that if during inference,
  73.                 # the value of the pixel is not exactly one,
  74.                 # we need to train the network to be robust and still
  75.                 # consider it as a pixel belonging to a square.
  76.                 if np.random.randint(0, 2):
  77.                     noise_f = (-1)**np.random.randint(0, 2)
  78.                     noisy_movies[i, t,
  79.                                  x_shift - w - 1: x_shift + w + 1,
  80.                                  y_shift - w - 1: y_shift + w + 1,
  81.                                  0] += noise_f * 0.1
  83.                 # Shift the ground truth by 1
  84.                 x_shift = xstart + directionx * (t + 1)
  85.                 y_shift = ystart + directiony * (t + 1)
  86.                 shifted_movies[i, t, x_shift - w: x_shift + w,
  87.                                y_shift - w: y_shift + w, 0] += 1
  89.     # Cut to a 40x40 window
  90.     noisy_movies = noisy_movies[::, ::, 20:60, 20:60, ::]
  91.     shifted_movies = shifted_movies[::, ::, 20:60, 20:60, ::]
  92.     noisy_movies[noisy_movies >= 1] = 1
  93.     shifted_movies[shifted_movies >= 1] = 1
  94.     return noisy_movies, shifted_movies
  96. # Train the network
  97. noisy_movies, shifted_movies = generate_movies(n_samples=1200)
  98. seq.fit(noisy_movies[:1000], shifted_movies[:1000], batch_size=10,
  99.         epochs=300, validation_split=0.05)
  101. # Testing the network on one movie
  102. # feed it with the first 7 positions and then
  103. # predict the new positions
  104. which = 1004
  105. track = noisy_movies[which][:7, ::, ::, ::]
  107. for j in range(16):
  108.     new_pos = seq.predict(track[np.newaxis, ::, ::, ::, ::])
  109.     new = new_pos[::, -1, ::, ::, ::]
  110.     track = np.concatenate((track, new), axis=0)
  113. # And then compare the predictions
  114. # to the ground truth
  115. track2 = noisy_movies[which][::, ::, ::, ::]
  116. for i in range(15):
  117.     fig = plt.figure(figsize=(10, 5))
  119.     ax = fig.add_subplot(121)
  121.     if i >= 7:
  122.         ax.text(1, 3, 'Predictions !', fontsize=20, color='w')
  123.     else:
  124.         ax.text(1, 3, 'Initial trajectory', fontsize=20)
  126.     toplot = track[i, ::, ::, 0]
  128.     plt.imshow(toplot)
  129.     ax = fig.add_subplot(122)
  130.     plt.text(1, 3, 'Ground truth', fontsize=20)
  132.     toplot = track2[i, ::, ::, 0]
  133.     if i >= 2:
  134.         toplot = shifted_movies[which][i - 1, ::, ::, 0]
  136.     plt.imshow(toplot)
  137.     plt.savefig('%i_animate.png' % (i + 1))