Dataset generation

First import and install stylegan

In [ ]:
#mount drive for file saving
from google.colab import drive
drive.mount('/content/drive', force_remount=True)

%tensorflow_version 1.x
import tensorflow
print(tensorflow.__version__)
!git clone https://github.com/NVlabs/stylegan.git
!ls /content/stylegan/
import sys
sys.path.insert(0, "/content/stylegan")
import dnnlib

import os
from tqdm import tqdm
import glob
import cv2
from PIL import Image
from IPython.display import clear_output

Now resize (or crop the images)

Ideally the images should be all the same size. Here, I use the PIL image module to sort thru our scrapped images and resize them. Alternatively, you can crop or apply other manipulations to preserve image structure a bit better.

In [ ]:
#from resizeimage import resizeimage
path = os.path.abspath('/content/drive/MyDrive/data/')
filelist = glob.glob(path+"//**//*.jpg", recursive=True)
for i, filename in enumerate(filelist):
  clear_output(wait=True)
  print(i)
  image = Image.open(filename)
  image = image.resize((512,512)).convert('RGB')
  base = os.path.basename(filename)
  try:
    label = int(base.split("_")[1][2:5])
  except:
    label = 0
  print(label)
  image.save('/content/drive/MyDrive/data/resized6/'+base, image.format)

Create the dataset

Now create the dataset using stylegans dataset tool. This is fairly simple. Remember to save the dataset to your drive, as we will need it in other notebooks. This may take up to hour or more depending on dataset size

In [ ]:
! python /content/stylegan/dataset_tool.py create_from_images /content/drive/MyDrive/data/custom-dataset5 /content/drive/MyDrive/data/resized5/