【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つのマスク画像の対をロードすることなど出来そう