【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
```
```python:

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 + "*"))
```
```python:
# ** 利用時 **
# データセットを作成する
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](https://qiita.com/advent-calendar/2020/pytorch)の枠が空いてたので急遽魔がさして書きました。

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

## Datasetの作成
```python:
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)`で取得しようとしたらエラーになったため残している。何か書き方を間違えているかもしれない。

```python:
--- 略 ---

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として読み込む。

```python:
--- 略 ---

    def __len__(self):  # 画像ペアの入ってるフォルダの数 = データセット数を返す
        return len(glob.glob(os.path.join(self.img_paths, "*")))
```
前述の通り、本来ならば`len(imgs_list)`で問題ない筈だが、手元の環境ではエラーが出たので再度`glob`を実行してお茶を濁している。

## DatasetをDataLoaderに入れて利用
```python:
# ** 利用時 **
# データセットを作成する
data_set = PairImgs("./my_img_path", transform=transforms.ToTensor())

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

```python:
--- 略 ---

# (必要なら)学習用とテスト用にデータセットを分割する
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つの変数で受け取れば良い。


```python:
--- 略 ---

# 作成したデータセットをデータローダ―に読み込ませる
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の作成と使用](https://qiita.com/mathlive/items/2a512831878b8018db02)