Google Colab で画像をグリッド状に並べて表示する方法

今回は Google Colab で複数の画像をグリッド状に整列して表示する方法をご紹介します。Stable Diffusionで生成した複数の画像を見比べるときなどに便利な方法です。

ライブラリはGoogle Colabに最初からインストール済みのものだけを使います。追加でインストールする必要はありません。

他の外部ライブラリを利用する方法もありますが、Google Colab にはよく用いられるライブラリはほとんどプレインストールされているので、それだけで十分対応できます。

今回３つの方法をご紹介しますが、単に並べて表示したい場合はtorchvisionを使うのが簡単ですが、軸を付けたいときはmatplotlibによる方法、画像処理をともなうときはPillowによる方法というように、適宜使い分けるのがおすすめです。

今回の方法は、Jupyter notebookでもライブラリがインストールしてあれば同じように実行できます。

サンプルに用いる画像

サンプルには並べたときに順番がわかりやすいように、グレーでベタ塗りした画像をサンプルに用います。今回は以下のコードで６枚生成しリストにします。ここで、画像の型は、Stable Diffusionで生成される画像と同じくPIL.Image.Image にします。

サンプルのグレーの画像を６枚生成

from PIL import Image
import numpy as np

# 今回は６枚生成
num_img = 6
colors = np.linspace(32, 224, num_img).astype(int)

# 画像のサイズ
w, h = 128, 128

# 生成した画像のリスト
images = []

for color in colors:
  img = np.zeros((w, h, 3))
  img[:, :, :] = color
  images.append(Image.fromarray(np.uint8(img)))

生成した画像をIPython.displayモジュールを用いて、以下のコードで単純に並べて確認します。

from IPython.display import display
display(*images)

すると以下のように、徐々に薄くなるグレーのサンプル画像が生成できたのを確認できます。次項からこの画像をグリッド状に並び替えます。

torchvisionのmake_grid()関数を用いる方法

torchvisionのmake_grid()関数を用いると、画像を簡単にグリッド状に並べることができます。使い方は、以下のように引数に画像のテンソルを入力し、列数をnrow=に指定します。

from torchvision.transforms import functional, ToPILImage
from torchvision.utils import make_grid

# 列数
ncols = 3

# グリッドの作成
grid = make_grid([functional.pil_to_tensor(img) for img in images], 
                 padding=0, # パディング
                 nrow=ncols, # 列数（ncol ではなく nrow なので注意）
                 pad_value=254 # 画像間の背景色（254: 白、0: 黒）
                 )

# PIL.Imageに変換
ToPILImage()(grid)

今回は列数を3にして、画像どうしを隙間なく配置して（padding=0）表示します。また、画像どうしの隙間があるときの背景色をpad_value=254で「白」にしています。

コードを実行すると以下のように６つの画像が３列のグリッド状に表示できたのを確認できます。

matplotlibのImageGridクラスを用いる方法

Matplotlibのツールキット（mpl_toolkits）には、Matplotlibを便利にしてくれる様々なツールが含まれますが、その中のImageGridクラスを使うと画像をグリッド上に並べることができます。

以下のように列数と行数をnrows_ncols=で指定して、ImageGridのインスタンスを作成します。これをループにかけると、グリッドに配置したグラフ（axes）を１つずつ返すので、imshow()で画像を貼り込みます。

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid

# 列数
ncols = 3

# 行数
nimg = len(images)
nrows = nimg // ncols if nimg % ncols == 0 else nimg // ncols + 1

fig = plt.figure(figsize=(7,7))

# ImageGridインスタンスの作成
grid = ImageGrid(fig, 111, 
                 nrows_ncols=(nrows, ncols),
                 axes_pad=0)

# 画像を貼り込む
for ax, img in zip(grid, images):
    ax.imshow(img)

# 軸を消す
for ax in grid.axes_all:
    ax.axis('off')

plt.show()

コードを実行するとtorchvisionを用いたときと同じように６つの画像が３列のグリッド状に表示できたのを確認できます。

軸を表示するには？

上のコードでは、ImageGridのインスタンスを作成するときにaxes_pad=0で画像どうしの隙間をなくし、ax.axis('off')で軸を消しています。これを軸を表示するように変更するには、以下のようにaxes_pad=0.2で隙間の値を入力し、軸を消すコードを削除します。

コードを実行すると以下のように軸が表示されたのを確認できます。

Pillowで１枚ずつ画像を貼り付ける方法

Stable Diffusion のサンプルコードで用いられている方法です。Pythonでの画像処理の定番ライブラリであるPillow （パッケージ名はPIL）を利用します。

最初に「列数x行数」分の白紙の画像をPIL.Image.new()で作成しておき、そこにpaste()メソッドで１枚ずつ左上から順に画像を貼り付けます。

from PIL import Image

def image_grid(imgs, rows, cols):
  """グリッド画像を作成する関数"""

  # 貼り付け用画像の作成
  w, h = imgs[0].size
  grid = Image.new('RGB', size=(cols*w, rows*h), color='white')
  grid_w, grid_h = grid.size
  
  # 画像を１枚ずつ貼り付ける
  for i, img in enumerate(imgs):
    grid.paste(img, box=(i % cols*w, i // cols*h))
  
  return grid

# 列数
ncols = 3

# 行数
nimg = len(images)
nrows = nimg // ncols if nimg % ncols == 0 else nimg // ncols + 1

# グリッド画像の作成
image_grid(images, nrows, ncols)

コードを実行すると上記のtorchvisionやmatplotlibの場合と同じく、以下のように６つの画像が３列のグリッド状に表示できたのを確認できます。