下一帧视频预测与卷积LSTM

import numpy as np
import matplotlib.pyplot as plt

import keras
from keras import layers

import io
import imageio
from IPython.display import Image, display
from ipywidgets import widgets, Layout, HBox

# 下载和加载数据集。
fpath = keras.utils.get_file(
    "moving_mnist.npy",
    "http://www.cs.toronto.edu/~nitish/unsupervised_video/mnist_test_seq.npy",
)
dataset = np.load(fpath)

# 交换表示帧数和数据样本数的轴。
dataset = np.swapaxes(dataset, 0, 1)
# 我们将选择1000个例子并使用它们。
dataset = dataset[:1000, ...]
# 添加通道维度，因为图像是灰度的。
dataset = np.expand_dims(dataset, axis=-1)

# 使用索引分割为训练和验证集以优化内存。
indexes = np.arange(dataset.shape[0])
np.random.shuffle(indexes)
train_index = indexes[: int(0.9 * dataset.shape[0])]
val_index = indexes[int(0.9 * dataset.shape[0]) :]
train_dataset = dataset[train_index]
val_dataset = dataset[val_index]

# 将数据标准化到0-1范围。
train_dataset = train_dataset / 255
val_dataset = val_dataset / 255

# 我们将定义一个辅助函数来移动帧，其中
# `x`是帧0到n - 1，`y`是帧1到n。
def create_shifted_frames(data):
    x = data[:, 0 : data.shape[1] - 1, :, :]
    y = data[:, 1 : data.shape[1], :, :]
    return x, y

# 将处理函数应用于数据集。
x_train, y_train = create_shifted_frames(train_dataset)
x_val, y_val = create_shifted_frames(val_dataset)

# 检查数据集。
print("训练数据集形状: " + str(x_train.shape) + ", " + str(y_train.shape))
print("验证数据集形状: " + str(x_val.shape) + ", " + str(y_val.shape))

# 构建一个图形以可视化图像。
fig, axes = plt.subplots(4, 5, figsize=(10, 8))

# 绘制一个随机数据示例的每个序列图像。
data_choice = np.random.choice(range(len(train_dataset)), size=1)[0]
for idx, ax in enumerate(axes.flat):
    ax.imshow(np.squeeze(train_dataset[data_choice][idx]), cmap="gray")
    ax.set_title(f"帧 {idx + 1}")
    ax.axis("off")

# 打印信息并显示图形。
print(f"显示示例 {data_choice} 的帧。")
plt.show()

# Construct the input layer with no definite frame size.
inp = layers.Input(shape=(None, *x_train.shape[2:]))

# We will construct 3 `ConvLSTM2D` layers with batch normalization,
# followed by a `Conv3D` layer for the spatiotemporal outputs.
x = layers.ConvLSTM2D(
    filters=64,
    kernel_size=(5, 5),
    padding="same",
    return_sequences=True,
    activation="relu",
)(inp)
x = layers.BatchNormalization()(x)
x = layers.ConvLSTM2D(
    filters=64,
    kernel_size=(3, 3),
    padding="same",
    return_sequences=True,
    activation="relu",
)(x)
x = layers.BatchNormalization()(x)
x = layers.ConvLSTM2D(
    filters=64,
    kernel_size=(1, 1),
    padding="same",
    return_sequences=True,
    activation="relu",
)(x)
x = layers.Conv3D(
    filters=1, kernel_size=(3, 3, 3), activation="sigmoid", padding="same"
)(x)

# Next, we will build the complete model and compile it.
model = keras.models.Model(inp, x)
model.compile(
    loss=keras.losses.binary_crossentropy,
    optimizer=keras.optimizers.Adam(),
)

# 定义一些回调函数以改善训练。
early_stopping = keras.callbacks.EarlyStopping(monitor="val_loss", patience=10)
reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor="val_loss", patience=5)

# 定义可修改的训练超参数。
epochs = 20
batch_size = 5

# 将模型拟合到训练数据上。
model.fit(
    x_train,
    y_train,
    batch_size=batch_size,
    epochs=epochs,
    validation_data=(x_val, y_val),
    callbacks=[early_stopping, reduce_lr],
)

# 从验证数据集中选择一个随机示例。
example = val_dataset[np.random.choice(range(len(val_dataset)), size=1)[0]]

# 从示例中选择前十帧。
frames = example[:10, ...]
original_frames = example[10:, ...]

# 预测一组新的10帧。
for _ in range(10):
    # 提取模型的预测并进行后处理。
    new_prediction = model.predict(np.expand_dims(frames, axis=0))
    new_prediction = np.squeeze(new_prediction, axis=0)
    predicted_frame = np.expand_dims(new_prediction[-1, ...], axis=0)

    # 扩展预测帧的集合。
    frames = np.concatenate((frames, predicted_frame), axis=0)

# 为原始帧和新帧构建一个图形。
fig, axes = plt.subplots(2, 10, figsize=(20, 4))

# 绘制原始帧。
for idx, ax in enumerate(axes[0]):
    ax.imshow(np.squeeze(original_frames[idx]), cmap="gray")
    ax.set_title(f"帧 {idx + 11}")
    ax.axis("off")

# 绘制新帧。
new_frames = frames[10:, ...]
for idx, ax in enumerate(axes[1]):
    ax.imshow(np.squeeze(new_frames[idx]), cmap="gray")
    ax.set_title(f"帧 {idx + 11}")
    ax.axis("off")

# 显示图形。
plt.show()

# 从数据集中选择一些随机示例。
examples = val_dataset[np.random.choice(range(len(val_dataset)), size=5)]

# 遍历示例并预测帧。
predicted_videos = []
for example in examples:
    # 从示例中选择前/后十帧。
    frames = example[:10, ...]
    original_frames = example[10:, ...]
    new_predictions = np.zeros(shape=(10, *frames[0].shape))

    # 预测一组新的 10 帧。
    for i in range(10):
        # 提取模型的预测并进行后处理。
        frames = example[: 10 + i + 1, ...]
        new_prediction = model.predict(np.expand_dims(frames, axis=0))
        new_prediction = np.squeeze(new_prediction, axis=0)
        predicted_frame = np.expand_dims(new_prediction[-1, ...], axis=0)

        # 扩展预测帧的集合。
        new_predictions[i] = predicted_frame

    # 为每一组真实值/预测图像创建和保存 GIF。
    for frame_set in [original_frames, new_predictions]:
        # 从选定的视频帧构建 GIF。
        current_frames = np.squeeze(frame_set)
        current_frames = current_frames[..., np.newaxis] * np.ones(3)
        current_frames = (current_frames * 255).astype(np.uint8)
        current_frames = list(current_frames)

        # 从帧构建 GIF。
        with io.BytesIO() as gif:
            imageio.mimsave(gif, current_frames, "GIF", duration=200)
            predicted_videos.append(gif.getvalue())

# 显示视频。
print(" 真实\t预测")
for i in range(0, len(predicted_videos), 2):
    # 构建并显示一个 `HBox`，其中包含真实值和预测。
    box = HBox(
        [
            widgets.Image(value=predicted_videos[i]),
            widgets.Image(value=predicted_videos[i + 1]),
        ]
    )
    display(box)