Numpyだけで複数の画像をタイルし1つの画像にまとめる方法

---

「torchvision.utilsのmake_gridやテンソルをタイルして保存するのって便利だよね。でも、いちいちこのためにPyTorchのテンソルに変えるのって面倒だよね」ということで同じことをNumpyでも実装してみました。Numpy配列の扱い方を工夫すればいけます。

コード

make_gridという関数がこれにあたります。「imgs」は複数の画像を格納した4階テンソル（batch, height, width, ch）のshape、「nrow」はタイル後に1行あたり何枚の画像をタイルするか、「padding」はオプションでグリッドの間隔です。

import numpy as np

def make_grid(imgs, nrow, padding=0):
    """Numpy配列の複数枚の画像を、1枚の画像にタイルします

    Arguments:
        imgs {np.ndarray} -- 複数枚の画像からなるテンソル
        nrow {int} -- 1行あたりにタイルする枚数

    Keyword Arguments:
        padding {int} -- グリッドの間隔 (default: {0})

    Returns:
        [np.ndarray] -- 3階テンソル。1枚の画像
    """
    assert imgs.ndim == 4 and nrow > 0
    batch, height, width, ch = imgs.shape
    n = nrow * (batch // nrow + np.sign(batch % nrow))
    ncol = n // nrow
    pad = np.zeros((n - batch, height, width, ch), imgs.dtype)
    x = np.concatenate([imgs, pad], axis=0)
    # border padding if required
    if padding > 0:
        x = np.pad(x, ((0, 0), (0, padding), (0, padding), (0, 0)),
                   "constant", constant_values=(0, 0)) # 下と右だけにpaddingを入れる
        height += padding
        width += padding
    x = x.reshape(ncol, nrow, height, width, ch)
    x = x.transpose([0, 2, 1, 3, 4])  # (ncol, height, nrow, width, ch)
    x = x.reshape(height * ncol, width * nrow, ch)
    if padding > 0:
        x = x[:(height * ncol - padding),:(width * nrow - padding),:] # 右端と下端のpaddingを削除
    return x

解説

例えば、100枚の32×32のカラー画像を、1行あたり12枚のフォームでグリッドを作ることを想定します。このとき、imgsのshapeは(100, 32, 32, 3)、nrow=12となります。

まず、nrowに対応するncolを計算します。入力のnrowが1行あたりの枚数なので、ncolは1列あたりの枚数、つまり行数を表します。小数を使ってもいいですが浮動小数点のバグが怖いので、整数商＋Modのsign関数でやりました。「nrow×ncol-入力の枚数」の不足分の画像を0として埋めます。この0埋め画像をimgsと結合した変数を「x」とします。ここはnp.concatenateを使えばいいです（ディープラーニングやっている人にはおなじみ）

次が少し難しいですが、グリッドを作ったときの画像間の隙間を作ります。発想的にはディープラーニングのZero Paddingと同じです。少し扱い方が難しいですが、np.padで同じことはできます。詳しくは公式ドキュメント見ていろいろ試してみるのが早いですが、「(0, 0), (0, padding), 」の表すものは、テンソルの各軸に対して左右どれだけの要素数を埋めるかを表します。ここでやっているのは画像の右と下にだけZero paddingをすることです。もしここをただのTupleとして指定するとすべての軸（不要なバッチやチャンネルの軸）にPaddingがかかってしまうので、縦と横（axis=1,2）だけ指定します。np.padはmodeにいろんな値が指定できるので、鏡像反転するようにPaddingなどもできます。もし今回のように「mode="constant", constant_values=(0, 0)」だとZero paddingになります。Paddingによって画像の縦と横のサイズが変わるので、忘れないように足しましょう。

次からがポイントで、「n, height, width, ch」というテンソルを、「ncol, nrow, height, width, ch」という5階テンソルにreshapeします。ここでやっているのは1次元に並んでいる画像を、縦横に分解しているということです。画像を横に並べているようなイメージです。

次からが黒魔術で、軸の入れ替えを行います。これはtransposeでできるのですが、「ncol, nrow, height, width, ch」→「ncol, height, nrow, width, ch」となるように入れ替えます。縦は縦、横は横となるように軸を入れ替えているということです。これは後のreshapeのためのもので、軸を入れ替えてreshapeするといい感じに並ぶからです。

最後に軸の結合です。これもreshapeで行います。0・1、2・3の軸を結合し、「ncol, height, nrow, width, ch」→「ncol * height, nrow * width, ch」とします。これは3階テンソルで1枚の画像のshapeとしては適切です。複数の画像をタイルし、グリットとして並べられた1枚の画像ができたということです。ただし、右と下の最後の1回のpaddingが不要なので、スライスして調整します。

例1～RGB画像～

Toy problemとして、1枚目が赤、2枚目が緑、3枚目が青の画像をタイルしてみましょう。

from PIL import Image

def rgb():
    x = np.zeros((3, 64, 64, 3), np.uint8)
    x[0,:,:,0] = 255 # 1枚目はred
    x[1,:,:, 1] = 255 # 2枚目はgreen
    x[2,:,:, 2] = 255 # 3枚目はblue
    stacked = make_grid(x, 2, padding=0)
    with Image.fromarray(stacked) as img:
        img.save("rgb.png")

タイルする方向は、左から右→上から下の順番です。torchvisionのグリッド化と同じです。torchvisionだとテンソル化したり、channels_firstにコンバートしないといけなかったり、Numpy配列をタイルする場合は面倒なんですよね。

Padding=10としてタイルしてみましょう。

隙間がちょっと大きくなりました。

例2～CIFAR-10～

CIFAR-10の訓練画像の最初の100枚をタイルしてみましょう。CIFAR-10のデータの読み込みにはKerasを使っています。

import tensorflow as tf

def cifar():
    (X, _), (_, _) = tf.keras.datasets.cifar10.load_data() # cifar10読み込むためにKerasを使う
    X = X[:100]  # (100, 32, 32, 3)
    stacked = make_grid(X, 12, padding=2)
    with Image.fromarray(stacked) as img:
        img.save("cifar.png")

ほぼtorchvisionですよね。Numpyだけでタイルしたい場合は便利ではないでしょうか。

まとめ

Numpyでもreshapeやtransposeを駆使することで画像のタイルまでできちゃいますよ、という話でした。テンソル使いこなせると楽しいですね。

---

---

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

---