A tensorflow/keras implementation of StyleGAN to generate images of new Pokemon.
The model has been trained on dataset that includes 819 pokémon.
You can download dataset from this kaggle link.
I have used the following versions for code work:
- python==3.8.8
- tensorflow==2.4.1
- tensorflow-gpu==2.4.1
- numpy==1.19.1
- h5py==2.10.0
There are several difficulties in pokemon generation using GAN :
- The difficulty of GAN training is well known; changing a hyperparameter can greatly change the results.
- The dataset size is too small! 819 different pokemon images are not enough. For this reason, I applied data augmentation on the data; these are the transformations applied :
img_transf = tf.keras.Sequential([
tf.keras.layers.experimental.preprocessing.RandomContrast(factor=(0.05, 0.15)),
image_aug.RandomBrightness(brightness_delta=(-0.15, 0.15)),
image_aug.PowerLawTransform(gamma=(0.8,1.2)),
image_aug.RandomSaturation(sat=(0, 2)),
image_aug.RandomHue(hue=(0, 0.15)),
tf.keras.layers.experimental.preprocessing.RandomFlip("horizontal"),
tf.keras.layers.experimental.preprocessing.RandomTranslation(height_factor=(-0.10, 0.10), width_factor=(-0.10, 0.10)),
tf.keras.layers.experimental.preprocessing.RandomZoom(height_factor=(-0.10, 0.10), width_factor=(-0.10, 0.10)),
tf.keras.layers.experimental.preprocessing.RandomRotation(factor=(-0.10, 0.10))])
- StyleGAN training is very expensive! I trained the model starting from a 4x4 resolution up to the final resolution of 256x256. The model was trained for 8 days using a Tesla V100 32GB SXM2.
To get better results you need to use higher resolutions and train for longer time.
These are some examples of new pokémon generated by the model :
New Generated Pokémon |
---|
You can see hundreds of new pokemon here.
I repeat again it : to get better results (better details in pokemon) is necessary to train for more time.
This code implementation is inspired by the unofficial keras implementation of styleGAN.