【PyTorch】DataLoaderで(images, labels)だけでなく(images, labels, paths)をロードできるようにする

まずglob.glob()を用いて, すべての画像データのパスを保持したリストを作る.

def make_datapath_list(phase="train"):
     rootpath = "./data/"
     target_path = os.path.join(rootpath, phase,  "**", "*.jpg") # 最初の**はクラスのディレクトリ
     path_list = []
     # globを利用してサブディレクトリまでファイルパスを格納
     for path in glob.glob(target_path):
          path_list.append(path)
     return path_list
# 動作確認
train_list = make_datapath_list(phase="train")

次に前処理クラスを作っておく.
前処理クラスはImage.open("path")で読み込んだ1画像配列を受け取り, 変形などをしtorch.tensorを返すクラス

class ImageTransform():
     def __init__(self, resize, mean, std):
          self.data_transform = transforms.Composee([transforms.ToTensor(), transforms.Normalize(mean, std)])
     def __call__(self, img):
          return self.data_transform(img)

次にDatasetクラスを作成する. この際に, __getitem__()メソッドで返す値を3要素のタプルにする.

class Dataset(data.Dataset):
     def __init__(slef, file_list, transform=None):
          self.file_list = file_lsit
          self.file_transform = transform # 前処理クラスのインスタンス
     
     def __len__(self):
          return len(slef.file_list)
     
     def __getitem__(self, index):
           img_path = self.file_list[index] # index番目のpath
           img = Image.open(img_path) # index番目の画像ロード
           img_transformed = slef.transform(img) # 前処理クラスでtensorに変換
           label = img_path.split("/")[2] # クラスを表すディレクトリの文字列
           label = convert(label) # クラスを表すディレクトリ文字列から数値に変更
           return img_transformed, label, img_path

train_dataset = Dataset(file_list=train_list = make_datapath_list(phase="train"), transform=ImageTransform(size, mean, std))

データローダを作成

batch_size = 32
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 動作確認: イテレータを作って初めのやつを表示する
tmp = iter(train_dataloader)
inputs, labels, paths = next(tmp)
print(inputs.size()) # torch.Size([32, 3, h, w])
print(labels) # tensor([label0, ...., label31])
print(paths) # [file_path0, ..., file_path31]

尚、動作の保証なし

Datasetクラスの__getitem__を工夫すれば, 良いことがわかった.
例えば, 2つの画像の組をロードすることや, 1つの画像とそれをセグメンテーションモデルにかけたもう1つのマスク画像の対をロードすることなど出来そう