【PyTorch】Dataset & DataLorderで画像ペアを扱う

#TL;DR

#画像がこのようなディレクトリに存在するとき
.
└─ my_img_path
   |
   ├── hogehoge
   |     |-img_A.png
   |     |-img_B.png
   |
   ├── fugafuga
   |     |-img_A.png
   |     |-img_B.png
   ...
   |
   └── piyopiyo
         |-img_A.png
         |-img_B.png

from torch.utils.data import DataLoader, Dataset
from PIL import Image
import glob
import os

#データセットの作成
class PairImgs(Dataset):
    # torch.utils.data の Dataset継承
    """
    self.img_paths   画像ペアの入ってるフォルダの一つ上のディレクトリへのパス
    self.imgs_list   画像ペアの入ってるフォルダすべてのList
    self.transform   指定したtransform
    """

    def __init__(self, img_dir, transform):
        self.img_paths = img_dir # 画像ペアを入れたディレクトリの親へのパスを指定する
        self.imgs_list = glob.glob(os.path.join(self.img_paths, "*"))
        self.transform = transform
    
    def __getitem__(self, index):
        # indexで指定したディレクトリ以下のファイルを返す
        
        # 画像をPILとして読み込む
        img_A = Image.open(os.path.join(self.imgs_list[index], "img_A.png"))  #画像の名前は任意に変更すること
        img_B = Image.open(os.path.join(self.imgs_list[index], "img_B.png"))

        if self.transform is not None:
            # 前処理がある場合は行う. 普通はtransforms.ToTensor()でTensorに変換してしまう
            img_A = self.transform(img_A)
            img_B = self.transform(img_B)

        return img_A , img_B
    
    def __len__(self):  # 画像ペアの入ってるフォルダの数 = データセット数を返す
        return len(glob.glob(self.img_paths + "*"))

# ** 利用時 **
# データセットを作成する
data_set = PairImgs("./my_img_path", transform=transforms.ToTensor())

# (必要なら)学習用とテスト用にデータセットを分割する
train_size = int( len(data_set) * 0.8 ) # 教師データのサイズ 全体の80%とする
test_size = n_samples - train_size  # テスト用データのサイズ

train_data, test_data = torch.utils.data.random_split(
        data_set,
        [train_size, test_size ],
        generator=torch.Generator().manual_seed(0)  # 乱数シードの固定
    )

# 作成したデータセットをデータローダ―に読み込ませる
batch_size = 32    # データローダ―のバッチサイズ 任意に変更のこと
# 学習データはシャッフルON テストデータはシャッフルなし
train_loader = DataLoader(train_data, batch_size=batch_size, 
                          shuffle=True, num_workers=2)
test_loader = DataLoader(test_data, batch_size=batch_size,
                         shuffle=False, num_workers=2)

#はじめに
PyTorch Advent Calendar 2020 Advent Calendar 2020の枠が空いてたので急遽魔がさして書きました。

解説

Pytorchで画像ペアを扱おうと考えた。この場合の画像ペアとは2枚の画像に関係がある(たとえば入力データと教師データ)ものであり、その2枚が一緒に取得できる必要があるものである。具体例としてはSRCNNなどのImage-to-Image 変換の際に利用する。
基本的には @mathlive さんのpyTorchのtransforms,Datasets,Dataloaderの説明と自作Datasetの作成と使用 を読んで作った。

Datasetの作成

class PairImgs(Dataset):
    # torch.utils.data の Dataset継承
    """
    self.img_paths   画像ペアの入ってるフォルダの一つ上のディレクトリへのパス
    self.imgs_list   画像ペアの入ってるフォルダすべてのList
    self.transform   指定したtransform
    """

    def __init__(self, img_dir, transform):
        self.img_paths = img_dir # 画像ペアを入れたディレクトリの親へのパスを指定する
        self.imgs_list = glob.glob(os.path.join(self.img_paths, "*"))
        self.transform = transform

--- 略 ---

コンストラクタに引数として**「画像ペアが入っているフォルダの親ディレクトリ」**を渡す。今回は"./my_img_path"となる。
コンストラクタ内の処理としてglob.glob("./my_img_path/*")を実行して今回の対象となるフォルダをリストとして保存する。

※img_pathsはこれ以降利用しないため保存しなくても良いはずだが、__len__をlen(imgs_list)で取得しようとしたらエラーになったため残している。何か書き方を間違えているかもしれない。

--- 略 ---

def __getitem__(self, idx):
        # indexで指定したディレクトリ以下のファイルを返す

        # 画像をPILとして読み込む
        img_A = Image.open(os.path.join(self.imgs_list[idx], "img_A.png"))  #画像の名前は任意に変更すること
        img_B = Image.open(os.path.join(self.imgs_list[idx], "img_B.png"))

        if self.transform is not None:
            # 前処理がある場合は行う. 普通はtransforms.ToTensor()でTensorに変換してしまう
            img_A = self.transform(img_A)
            img_B = self.transform(img_B)

        return img_A , img_B

--- 略 ---

torch.utils.data の Dataset にアクセスする際、配列のインデックス値が引数として送られてくるので、第二引数のidxがそれを受け取る。つまり対象の画像ペアが入っているフォルダはself.img_list[index]となるから、それに画像の名前("img_A.png")を付与してPIL.Image.open()で画像をPILとして読み込む。

--- 略 ---

    def __len__(self):  # 画像ペアの入ってるフォルダの数 = データセット数を返す
        return len(glob.glob(os.path.join(self.img_paths, "*")))

前述の通り、本来ならばlen(imgs_list)で問題ない筈だが、手元の環境ではエラーが出たので再度globを実行してお茶を濁している。

DatasetをDataLoaderに入れて利用

# ** 利用時 **
# データセットを作成する
data_set = PairImgs("./my_img_path", transform=transforms.ToTensor())

--- 略 ---

transformの詳細は参考リンクを参照のこと。

--- 略 ---

# (必要なら)学習用とテスト用にデータセットを分割する
train_size = int( len(data_set) * 0.8 ) # 教師データのサイズ 全体の80%とする
test_size = n_samples - train_size  # テスト用データのサイズ

train_data, test_data = torch.utils.data.random_split(
        data_set,
        [train_size, test_size],
        generator=torch.Generator().manual_seed(0)  # 乱数シードの固定
    )

--- 略 ---

個人的な考えだが、データセットの分割はDataLoaderの作成直前に行った方が楽である。
今回は訓練データとテストデータに分割したが、検証データにも分割したい場合はtorch.utils.data.random_split()の第二引数を [train_size, test_size, valid_size] のような3要素のリストとし、戻り値を3つの変数で受け取れば良い。

--- 略 ---

# 作成したデータセットをデータローダ―に読み込ませる
batch_size = 32    # データローダ―のバッチサイズ 任意に変更のこと
# 学習データはシャッフルON テストデータはシャッフルなし
train_loader = DataLoader(train_data, batch_size=batch_size, 
                          shuffle=True, num_workers=2)
test_loader = DataLoader(test_data, batch_size=batch_size,
                         shuffle=False, num_workers=2)

DataLoaderに上記で分割したDatasetを指定して終わり。
後はfor文などで1バッチずつ取り出しながら学習を進めていけばよい。

#参考リンク
pyTorchのtransforms,Datasets,Dataloaderの説明と自作Datasetの作成と使用