Zenn
Open3

UNetのカスタムデータセットを作る方法を調べる

ディープラーニング
PyTorch
dataset
nabeyangnabeyang

データセットクラスから得た画像とそのマスクをtorch.Tensorから描画可能な状態に戻す方法。

d = BasicDataset(dir_img, dir_mask, 1, mask_suffix='_mask')
it = iter(d)
r = next(it)
img = Image.fromarray((r['image'].to('cpu').detach().numpy() * 255).astype(np.uint8).transpose(1, 2, 0))
mask = Image.fromarray((r['mask'].to('cpu').detach().numpy() * 255).astype(np.uint8)[0])
nabeyangnabeyang

マスクのかかり方チェック

Image.composite(img, Image.new("RGB", mask.size, (0, 0, 0)), mask)