StyleGAN2-adaでValueError: not enough image data

以下のリポジトリのStyleGAN2-ada でグレースケール画像を学習した重みを使って generate.py を実行した時に表題のエラー

GitHub - NVlabs/stylegan2-ada-pytorch: StyleGAN2-ADA - Official PyTorch implementation

StyleGAN2-ADA - Official PyTorch implementation. Contribute to NVlabs/stylegan2-ada-pytorch development by creating an a...

generate.py の最後の方にPILを使った画像の保存でつまづいている様子。

PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')

saveする時にRGBを指定しているので、グレースケールの画像ではチャンネル数が合わずにエラーが出たと予想させる。

生成画像がグレースケールかカラーかによって場合分けすることで、グレースケールの画像を生成することが出来た。

グレースケールをpillowで保存するときはチャンネルを表すベクトルは不要なので、np.squeezeを入れて(width, height, 1)のndarrayを(width, height)の2次元配列に圧縮し忘れないようにする。

最後にRGBをLに変えて完成。

if img.shape[-1] == 3:
    PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
else:
    PIL.Image.fromarray(np.squeeze(img[0][:-2].cpu().numpy()), 'L').save(f'{outdir}/seed{seed:04d}.png')

参考

Image Module

The Image module provides a class with the same name which is used to represent a PIL image. The module also provides a ...