如何为 PNG 的图像序列预测或下一帧预测创建 TensorFlow 数据集？

Question

我正在尝试在 Google Colab 上使用 Tensorflow 制作图像序列预测模型。基本上，模型应该预测给定图像序列的下一帧。

但是在将模型拟合到训练数据时出现以下错误：

ValueError: Input 0 of layer conv_lst_m2d_1 is incompatible with the layer: expected ndim=5, found ndim=4. Full shape received: (None, 900, 900, 3)

我按照本教程进行下一帧预测

https://keras.io/examples/vision/conv_lstm/

并尝试使用此代码加载我的图像

https://www.tensorflow.org/tutorials/load_data/images

我的图片是 PNG (900x900px, rgb) 并像这样存储，文件名是以毫秒为单位的 unix 时间戳：

raw_images/
    1626008400000.png
    1626008700000.png
    1626009000000.png
    1626009300000.png
    1626009600000.png
    1626009900000.png
    ...

我的代码（最低工作版本）：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
IMAGE_FOLDER = "raw_images/"

# Load the dataset
dataset = keras.preprocessing.image_dataset_from_directory(
    IMAGE_FOLDER,
    None,
    None,
    None,
    "rgb",
    32,
    (900, 900),
    False,
    None,
    None,
    None,
    "bilinear",
    False,
    True
)

dataset = dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)


model = tf.keras.Sequential([
  tf.keras.layers.experimental.preprocessing.Rescaling(1./255),
  tf.keras.layers.ConvLSTM2D(
    filters=64,
    kernel_size=(5, 5),
    padding="same",
    return_sequences=True,
    activation="relu",
),
  tf.keras.layers.BatchNormalization(),
  tf.keras.layers.ConvLSTM2D(
    filters=64,
    kernel_size=(3, 3),
    padding="same",
    return_sequences=True,
    activation="relu",
),
  tf.keras.layers.BatchNormalization(),
  tf.keras.layers.ConvLSTM2D(
    filters=64,
    kernel_size=(1, 1),
    padding="same",
    return_sequences=True,
    activation="relu",
),
  tf.keras.layers.Conv3D(
    filters=1, kernel_size=(3, 3, 3), activation="sigmoid", padding="same"
)
])


model.compile(
    loss=keras.losses.binary_crossentropy, 
    optimizer=keras.optimizers.Adam(),
)


# Define some callbacks to improve training.
early_stopping = keras.callbacks.EarlyStopping(monitor="val_loss", patience=10)
reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor="val_loss", patience=5)

# Define modifiable training hyperparameters.
epochs = 20
batch_size = 32

# Fit the model to the training data.
model.fit(
    dataset,
    None,
    batch_size=batch_size,
    epochs=epochs
)

model.save("model")

我怀疑我需要以某种方式将时间戳添加到数据集，但我找不到任何方法。

Answer 1

这很容易。 LSTM 需要多一维。所以你只需要添加一个tf.keras.layers.Reshape(target_shape) 来将输入到该 ConvLSTM2D 层扩展一维。这应该可以解决问题。