動画予測のためのベンチマークデータセット「Moving MNIST」を自作する

Moving MNISTが何かの説明は既存の記事に任せるとして
データセット「Moving MNIST」

Moving MNISTは元祖であるMNISTと同様に、数々の動画予測などの深層学習タスクのベンチマークとして用いられています。
それは、タスクとして比較的容易でシンプルであることに加え、その拡張性にも大きな利点があることが理由です。多分。
Moving MNISTは64*64の画像内で数字が動くだけなので、元祖MNISTの手書き文字さえ手元にあれば自作することができます。トロント大学から10000個の動画を含むデータセットが公式に配布されていますが、この程度じゃ過学習しちゃうぜ、さらに1000000000000個の動画が欲しい、と思ったら自作することができます。

.npyファイルをジェネレート

実行前に、必ず元祖MNISTのデータセットをダウンロードしておいてください。

$ wget http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz

または http://yann.lecun.com/exdb/mnist/ から手でダウンロードできます。

トロント大学で配布されているのと同じように、.npy形式でファイルを作成します。

generate_moving_mnist.py

import numpy as np
import gzip
import random
import argparse

parser = argparse.ArgumentParser(description='generate moving mnist dataset')
parser.add_argument('--file_name', type=str, required=True)
parser.add_argument('--source_path', type=str, required=True)
parser.add_argument('--data_len', type=int, required=True)
parser.add_argument('--seq_len', type=int, required=False, default=20)
parser.add_argument('--img_size', type=int, required=False, default=64)
parser.add_argument('--digit_size', type=int, required=False, default=28)
parser.add_argument('--digit_num', type=int, required=False, default=2)
parser.add_argument('--step_len', type=float, required=False, default=0.1)
args = parser.parse_args()

file_name   = args.file_name
source_path = args.source_path
seq_len     = args.seq_len
data_len    = args.data_len
img_size    = args.img_size
digit_size  = args.digit_size
digit_num   = args.digit_num
step_len    = args.step_len

data = np.zeros((seq_len, data_len, img_size, img_size), dtype=np.float32)

with gzip.open(source_path, 'rb') as f:
    mnist = np.frombuffer(f.read(), np.uint8, offset=16)
    mnist = mnist.reshape(-1, digit_size, digit_size)

for i in range(data_len):
    seq = np.zeros((seq_len, img_size, img_size), dtype=np.float32)

    for n in range(digit_num):
        digit_img = mnist[random.randint(0, mnist.shape[0] - 1)]
        canvas_size = img_size - digit_size
        x = random.random()
        y = random.random()
        theta = random.random() * 2 * np.pi
        v_x = np.cos(theta)
        v_y = np.sin(theta)
        seq_x = np.zeros(seq_len)
        seq_y = np.zeros(seq_len)

        for s in range(seq_len):
            x += v_x * step_len
            y += v_y * step_len

            if x <= 0 or 1.0 <= x:
                v_x *= -1
                x = np.clip(x, 0, 1.0)
            if y <= 0 or 1.0 <= y:
                v_y *= -1
                y = np.clip(y, 0, 1.0)

            seq_x[s] = x
            seq_y[s] = y

        seq_x = (canvas_size * seq_x).astype(np.int32)
        seq_y = (canvas_size * seq_y).astype(np.int32)
        for s in range(seq_len):
            left = seq_x[s]
            top  = seq_y[s]
            right  = left + digit_size
            bottom = top  + digit_size
            seq[s, top:bottom, left:right] = np.maximum(seq[s, top:bottom, left:right], digit_img)

    data[:,i,:,:] = seq

data = data[..., np.newaxis]
print(data.shape)
np.save(file_name, data)

内容はそのままで、MNISTからランダムに数字を選び出し、ランダムに配置し、ランダムな方向へ動かすだけです。枠にぶつかると反射しますが、そこは普通に反射します。普通に、というのは、ぶつかった壁と垂直な速度成分は-1倍され、平行な成分は保存されます。
各引数についてはなんとなく察してください。出力するファイル名、MNISTのパス、データの長さはrequiredです。動く数字の数 digit_numなども設定できます。

唯一オリジナルの配布データと異なるのは、チャンネルの次元を追加している点です。　Moving MNISTはグレースケールなのでチャンネルの次元はなくても情報を欠損なく表現できます。しかし、動画予測モデルではチャンネルの次元がないと正常に動作しないことが多いため、ここでは追加しています。
もし必要なければ、下から3行目のdata = data[..., np.newaxis]を削除すればいいと思います。

ex. 標準的な設定のMoving MNISTを1000セット

$　python3 generate_moving_mnist.py --file_name=normal_mmnist.npy --source_path=train-images-idx3-ubyte.gz --data_len=1000

ex. 動く数字が3つ、縦横128*128の大きいMoving MNISTを50セット

$　python3 generate_moving_mnist.py --file_name=3_mmnist.npy --source_path=train-images-idx3-ubyte.gz --data_len=50　--digit_num=3 --img_size=128

作成したファイルをmp4で確認

正常に作成できているか見てみないとわからないので、opencvで動画にして見てみます。

visualize_npy.py

import numpy as np
import cv2
import argparse
import os

parser = argparse.ArgumentParser(description='visualize video prediction datasets')
parser.add_argument('--file_name', type=str, required=True)
parser.add_argument('--source_path', type=str, required=True)
parser.add_argument('--num_samples', type=int, required=True)
args = parser.parse_args()

data = np.load(args.source_path)
if data.ndim == 4:
    data = data[..., np.newaxis]
shape = data.shape
len_seq  = shape[0]
len_data = shape[1]
height   = shape[2]
width    = shape[3]
channel  = shape[4]

print('{} data * {} frames * {} px * {} px * {} channel'.format(len_data, len_seq, height, width, channel))

for i in range(args.num_samples):
    video_name = str(args.file_name) + '_' + str(i+1) + '.mp4'
    fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
    video  = cv2.VideoWriter(video_name, fourcc, 1.00, (width, height))

    for s in range(len_seq):
        if channel == 3:
            frame = data[s,i]
        if channel == 1:
            frame = np.repeat(data[s,i], 3).reshape(height,width,3)
        video.write(frame.astype(np.uint8))
    video.release()

--file_nameには出力するファイルの名前を、--source_pathには先ほど作成したnpyファイルのパスを、--num_samplesには何個のサンプルを出力するかを指定してください。

先ほど作成した動画を見てみます。

ex. 標準的な設定のMoving MNIST

$ python3 visualize_npy.py --file_name=normal_mmnist --source_path=normal_mmnist.npy --num_samples=10

やたら数字がゆっくりですが、動画に出力するときのFPSを1.00に設定しているためです。npyファイルにはフレームの概念はあれどFPSの概念はないので、影響はもちろんありません。
ちなみに、僕の環境だとFPSを2.00や4.00にすると動画が破損します。10.0は正常に出力されます。cv2のこういう謎仕様なんとかしてほしい。

ex. 動く数字が3つ、縦横128*128の大きいMoving MNIST

$ python3 visualize_npy.py --file_name=3_mmnist --source_path=3_mmnist.npy --num_samples=10

枠がでかいというより数字が小さく見えますね。数字が小さいと重なる回数が相対的に少なくなるので、学習は簡単になりそう。
もちろん数字は多いほど難しいと思います。