PyTorchでImageFolderを使ってみる

PyTorchで画像認識などの学習を行うときに、お試しでtorchvisionのdatasetsを使用することがよくあります。特にMNISTの手書き文字の画像はよく利用されていて、練習にとても便利です。

datasetsを使用した場合は、手書き文字が収録されたバイナリデータをPyTorchのテンソルに取り込んでいるので、PNGやGIFなどの画像ファイルを直接扱っているわけではありません。

しかしながら、自己が作成しようとしているモデルの学習に使用する画像データは、必ずしもdatasetsから得られるわけではなく、独自に収集した画像データを使用することがほとんどです。

そこで、torchvisionのImageFolderを使用して、イメージ画像データをテンソル取り込む方法について解説したいとおもいます。

まずは、大量の画像ファイルが手元にないのでMNISTの０～９の手書き文字の画像ファイルをdatasetsから作成してみることにします。その後でImageFolderを使って作成した画像ファイルをテンソルに取り込み、データローダを作成していきます。

なぜこの解説を書くことにしたかと言うと、ディープラーニングでモデル学習を行うためには、データローダを作成することが最も重要なタスクだからです。そもそもデータが準備できなければ、モデル学習はできないのです。

NISTの手書きデータセットを画像ファイルにする

今回のトライアルでは、以下のモジュールをインポートします。

import os
from PIL import Image
import torch
import torch.utils.data as data
import torchvision
import torchvision.datasets as datasets
import torchvision.transforms as transforms

from tqdm import tqdm

import matplotlib.pyplot as plt
%matplotlib inline

MNISTデータセットは、torchvision.datasetsをインポートすることで利用することが可能になります。データを取り込むには以下のようにします。

mnist_data = datasets.MNIST(root='./', train=True, transform=None, download=True)
print(mnist_data)

Dataset MNIST
  Number of datapoints: 60000
  Root location: ./
  Split: Train

上記の場合、root=’./’としているので、カレントディレクトリにデータがダウンロードされます。この訓練用データセットには、6万枚の手書きデータが収録されています。
ダウンロードされたディレクトリは、以下のような構造になっています。

./MNIST
├── processed
│   ├── test.pt
│   └── training.pt
├── raw
├── t10k-images-idx3-ubyte
├── t10k-images-idx3-ubyte.gz
├── t10k-labels-idx1-ubyte
├── t10k-labels-idx1-ubyte.gz
├── train-images-idx3-ubyte
├── train-images-idx3-ubyte.gz
├── train-labels-idx1-ubyte
└── train-labels-idx1-ubyte.gz

2 directories, 10 files

変数mnist_dataに取り込まれたデータは、画像データとラベルが一対になった配列構造になっています。例えば取り込んだデータの先頭の画像データにアクセスするには、mnist[0][0]とします。また、ラベルにアクセスするには、mnist[0][1]とすることで、画像データのラベルを見ることができます。

６万枚目の画像データをmatplotlibで可視化するとこうなります。

plt.imshow(mnist_data[59999][0], cmap = 'gray')

８が表示されました。

データセットをイメージ画像で保存する

データセットから取り込んだ０から８の手書き文字画像データをPNG画像として保存してみます。

データセットは画像データとラベルが一対になっているので、数字毎の画像数とラベルを活用してPNGファイルの画像名を設定することとします。作成した画像ファイルは数字毎に分別したいので、以下のようなディレクトリの階層にして保存していきます。

./MNIST_PNG
├── 0
├── 1
├── 2
├── 3
├── 4
├── 5
├── 6
├── 7
├── 8
└── 9

10 directories

画像ファイルを生成する関数を作成します。forループで１データごと画像ファイルを生成します。

def makeMnistPng(image_dsets):
  for idx in tqdm(range(10)):
    print("Making image file for index {}".format(idx))
    num_img = 0
    dir_path = './MNIST_PNG/' + str(idx)
    if not os.path.exists(dir_path):
      os.makedirs(dir_path)
      for image, label in image_dsets:
        if label == idx:
          filename = dir_path +'/' + 'mnist_'+ str(idx) + '_' + str(num_img) + '.png'
          if not os.path.exists(filename):
            image = image.resize((64, 64))
            image.save(filename)
          num_img += 1

  print('Success to make MNIST PNG image files. index={}'.format(idx))  

関数を実行すると、画像を作成する処理が実行されます。tqdmをビッグループに仕込んであるので、処理の進捗が表示されるようになっています。

makeMnistPng(mnist_data)

処理が完了すると、MNIST_PNGディレクトリの下に数字毎のディレクトリが作成され、そのディレクトリ中には、64×64ピクセルの画像ファイルが格納されています。

ImageFolderをつかってみる

次に、ImageFolderを使用して、先ほど作成したMNISTの手書き文字画像ファイルを取り込んでみます。最終的にバッチ分割されたデータローダーを作成します。

まずは、画像をテンソル化した後に、イメージ画像のデータ変換（標準化など）を行うクラスを定義します。

class ImageTransform():
  def __init__(self, mean, std):
    self.data_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
    ])

  def __call__(self, img):
    return self.data_transform(img)

次にImageFolderを使用して画像データの取り込みを行います。イメージ画像の格納状態は以下のようになっています。各数字の画像ファイルを格納しているディレクトリ名を、イメージ画像のラベルとして使用します。

./MNIST_PNG
├── 0
│   ├── mnist_0_0.png
│   ├── mnist_0_1.png
│   ├── ・・・
│
├── 2
│   ├── mnist_2_0.png
│   ├── mnist_2_1.png
│   ├── ・・・
・・・

10 directories, 60000 files

torchvision.datasets.ImageFolderは、イメージ画像ファイルを格納したディレクトリと画像変換設定を与えるだけなので、イメージ画像ファイルの整理さえできていれば、データ取込みがとても簡単にできてしまいます。PyTorchを使うならImageFolderを利用することで開発効率が格段にアップします。

以下のコードでイメージ画像の取り込みを行っています。

#画像データをImageFolderを使って取込みする
mean = (0.5,)
std = (0.5,)
images = ImageFolder( "./MNIST_PNG", transform = ImageTransform(mean, std))

データローダーを作成する

最後にImageFolderで取り込んだイメージ画像のデータを使用して、データローダーを作成します。データローダーは、torch.utils.data.DataLoaderにImageFolderで取り込んだイメージ画像データのオブジェクト変数（ここではimages）、1バッチ当たりのデータ数とシャッフル方法を与えて作成します。

#1バッチに含む画像の枚数を指定する
batch_size = 64

#ImageFolderで取り込んだ画像からデータローダーを作成する
train_dataloader = data.DataLoader(images, batch_size = batch_size, shuffle = True)

せっかくなので、バッチ化されたデータローダーからイメージ画像を可視化してみます。データローダーから画像とラベルを取り出すには、iterとnextを組み合わせて使用します。

#データローダから、画像とラベルのテンソルを取り出す。
imgs, labels = iter(train_dataloader).next()

#バッチから取り出した画像の大きさを確認する
print("image shape ==>;",imgs[0].shape)

image shape ==> torch.Size([3, 64, 64])ちゃんと取り出せたようですので、可視化してみます。

#バッチから取り出した画像のイメージとラベルを表示する
pic = transforms.ToPILImage(mode='RGB')(imgs[0])
plt.imshow(pic)
print("Label is ",labels[0].numpy())

Label is 5

データローダーはshuffle = Trueで作成されているので、何が出るかは毎回異なりますが、今回の場合は、５が表示されました。

これで、ImageFolderを使用したデータローダーの作成は終わりです。モデル学習の基本事項となる学習データの準備について、ImageFolderを使って実践してみました。

データローダーの作成についても、PyTorchはKerasと比べて簡単に実装できるので、効率よく開発できるところがとてもよいと実感しました。