Selfie StyleGAN2

  1. Quickstart
    1. Generate Images from mp4
    2. Training and using the model
      1. GPU Details

I checked out the repository from github.com/lucidrains/stylegan2-pytorch and trained a basic model to generate my face from a set of ~1000 images derived from a 20 second webcam selfie that I took on my laptop.

Interpolated results of training ~30 epochs are shown below.
The generator and discriminator losses appeared to have converged around epoch 20.

Quickstart

git clone [email protected]:lucidrains/stylegan2-pytorch.git
cd stylegan2-pytorch
pip install stylegan2-pytorch

Generate Images from mp4

The following snippet is modified from this github gist:

'''
Using OpenCV takes a mp4 video and produces a number of images.

Requirements
----
You require OpenCV 3.2 to be installed.

Run
----
Open the main.py and edit the path to the video. Then run:
$ python main.py

Which will produce a folder called data with the images. There will be 2000+ images for example.mp4.
'''
import cv2
import numpy as np
import os

# Playing video from file:
cap = cv2.VideoCapture('example.mp4')

try:
if not os.path.exists('data'):
os.makedirs('data')
except OSError:
print ('Error: Creating directory of data')

currentFrame = 0
while(True):
# Capture frame-by-frame
ret, frame = cap.read()

y, x, _channels = frame.shape
# 512 by 512
y_mid = y // 2
y_min = y_mid - 256
y_max = y_mid + 256

x_mid = x // 2
x_min = x_mid - 256
x_max = x_mid + 256

crop_frame = frame[y_min:y_max, x_min:x_max, :]

# Saves image of the current frame in jpg file
name = './data/frame' + str(currentFrame) + '.jpg'
print ('Creating...' + name)
cv2.imwrite(name, crop_frame)

# To stop duplicate images
currentFrame += 1
# break

# When everything done, release the capture
cap.release()
cv2.destroyAllWindows()

Training and using the model

stylegan2_pytorch --data data/ --name vanity --image-size 128
# using my home server, I did not have enough GPU memory to fit images of size 256 or higher.
stylegan2_pytorch --name vanity --generate-interpolation --interpolation-num-steps 1000

GPU Details

nvidia-smi
Mon Dec 21 10:10:37 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 455.38 Driver Version: 455.38 CUDA Version: 11.1 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 TITAN Xp Off | 00000000:42:00.0 Off | N/A |
| 23% 36C P8 16W / 250W | 0MiB / 12180MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+