!git clone https://github.com/Puzer/stylegan-encoder.git
!mkdir -p data
!wget -nc -O data/starr.jpg https://www.dropbox.com/s/oyr35cz55lry5my/starr.jpg?dl=1
import sys
sys.path.append('stylegan-encoder')
import os
import bz2
import dlib
import PIL.Image
import numpy as np
from keras.utils import get_file
import matplotlib.pyplot as plt
from google.colab import files
import pickle
import config
import dnnlib
import dnnlib.tflib as tflib
from tqdm.autonotebook import tqdm
import tensorflow as tf
from keras.models import Model
from keras.applications.vgg16 import VGG16, preprocess_input
from keras.preprocessing import image
import keras.backend as K
from keras.applications.vgg16 import VGG16, preprocess_input
from scipy.optimize import fmin_l_bfgs_b
import scipy
#@title >> SKRYTÁ IMPLEMENTÁCIA RUTÍN << { display-mode: "form" }
def unpack_bz2(src_path):
data = bz2.BZ2File(src_path).read()
dst_path = src_path[:-4]
with open(dst_path, 'wb') as fp:
fp.write(data)
return dst_path
def image_align(img, face_landmarks, output_size=1024, transform_size=4096, enable_padding=True):
# Align function from FFHQ dataset pre-processing step
# https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py
lm = np.array(face_landmarks)
lm_chin = lm[0 : 17] # left-right
lm_eyebrow_left = lm[17 : 22] # left-right
lm_eyebrow_right = lm[22 : 27] # left-right
lm_nose = lm[27 : 31] # top-down
lm_nostrils = lm[31 : 36] # top-down
lm_eye_left = lm[36 : 42] # left-clockwise
lm_eye_right = lm[42 : 48] # left-clockwise
lm_mouth_outer = lm[48 : 60] # left-clockwise
lm_mouth_inner = lm[60 : 68] # left-clockwise
# Calculate auxiliary vectors.
eye_left = np.mean(lm_eye_left, axis=0)
eye_right = np.mean(lm_eye_right, axis=0)
eye_avg = (eye_left + eye_right) * 0.5
eye_to_eye = eye_right - eye_left
mouth_left = lm_mouth_outer[0]
mouth_right = lm_mouth_outer[6]
mouth_avg = (mouth_left + mouth_right) * 0.5
eye_to_mouth = mouth_avg - eye_avg
# Choose oriented crop rectangle.
x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
x /= np.hypot(*x)
x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
y = np.flipud(x) * [-1, 1]
c = eye_avg + eye_to_mouth * 0.1
quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
qsize = np.hypot(*x) * 2
# Shrink.
shrink = int(np.floor(qsize / output_size * 0.5))
if shrink > 1:
rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))
dst_img = img.resize(rsize, PIL.Image.ANTIALIAS)
quad /= shrink
qsize /= shrink
else:
dst_img = img.copy()
# Crop.
border = max(int(np.rint(qsize * 0.1)), 3)
crop = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))),
int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1]))))
crop = (max(crop[0] - border, 0),
max(crop[1] - border, 0),
min(crop[2] + border, dst_img.size[0]),
min(crop[3] + border, dst_img.size[1]))
if crop[2] - crop[0] < dst_img.size[0] or crop[3] - crop[1] < dst_img.size[1]:
dst_img = dst_img.crop(crop)
quad -= crop[0:2]
# Pad.
pad = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1]))))
pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - dst_img.size[0] + border, 0), max(pad[3] - dst_img.size[1] + border, 0))
if enable_padding and max(pad) > border - 4:
pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
dst_img = np.pad(np.float32(dst_img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
h, w, _ = dst_img.shape
y, x, _ = np.ogrid[:h, :w, :1]
mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w-1-x) / pad[2]), 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h-1-y) / pad[3]))
blur = qsize * 0.02
dst_img += (scipy.ndimage.gaussian_filter(dst_img, [blur, blur, 0]) - dst_img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
dst_img += (np.median(dst_img, axis=(0,1)) - dst_img) * np.clip(mask, 0.0, 1.0)
dst_img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(dst_img), 0, 255)), 'RGB')
quad += pad[:2]
# Transform.
dst_img = dst_img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR)
if output_size < transform_size:
dst_img = dst_img.resize((output_size, output_size), PIL.Image.ANTIALIAS)
return dst_img
class LandmarksDetector:
def __init__(self, predictor_model_path):
"""
:param predictor_model_path: path to shape_predictor_68_face_landmarks.dat file
"""
self.detector = dlib.get_frontal_face_detector() # cnn_face_detection_model_v1 also can be used
self.shape_predictor = dlib.shape_predictor(predictor_model_path)
def get_landmarks(self, img):
dets = self.detector(img, 1)
for detection in dets:
face_landmarks = [(item.x, item.y) for item in self.shape_predictor(img, detection).parts()]
yield face_landmarks
def convert_images_loss(images):
drange=[-1,1]
nchw_to_nhwc=True
shrink=1
images = tf.cast(images, tf.float32)
if shrink > 1:
ksize = [1, 1, shrink, shrink]
images = tf.nn.avg_pool(images, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW")
if nchw_to_nhwc:
images = tf.transpose(images, [0, 2, 3, 1])
scale = 255 / (drange[1] - drange[0])
images = images * scale + (0.5 - drange[0] * scale)
return images
def convert_images_gen(images):
images = tf.saturate_cast(images, tf.uint8)
return images
class Evaluator(object):
def __init__(self, aligned_img, loss_grad_func, latent_shape):
self.loss_value = None
self.grads_values = None
self.latent_shape = latent_shape
aligned_img = np.asarray(aligned_img)
if len(aligned_img.shape) == 4:
self.aligned_img = aligned_img
elif len(aligned_img.shape) == 3:
self.aligned_img = np.expand_dims(aligned_img, 0)
else:
raise RuntimeError("Unsupported image shape '{}'.".format(aligned_img.shape))
self.loss_grad_func = loss_grad_func
self.eval_iter = 0
def loss(self, latent):
assert self.loss_value is None
latent = latent.reshape(self.latent_shape)
outs = self.loss_grad_func([self.aligned_img, latent])
self.loss_value = outs[0]
self.grad_values = np.array(outs[1:]).flatten().astype('float64')
# clip the gradients
self.grad_values = np.maximum(np.minimum(self.grad_values, 1.0), -1.0)
self.eval_iter += 1
print("eval {}, loss {}".format(self.eval_iter, self.loss_value))
return self.loss_value
def grads(self, x):
assert self.loss_value is not None
grad_values = np.copy(self.grad_values)
self.loss_value = None
self.grad_values = None
return grad_values
def move_and_show(dlatent, direction, coeffs):
fig,ax = plt.subplots(1, len(coeffs), figsize=(12, 10), dpi=80)
dlatent = dlatent.reshape(dlatent_shape)
for i, coeff in enumerate(coeffs):
new_latent_vector = dlatent.copy()
new_latent_vector[:8] = (dlatent + coeff*direction)[:8]
ax[i].imshow(gen_func([new_latent_vector])[0][0])
ax[i].set_title('Coeff: %0.1f' % coeff)
[x.axis('off') for x in ax]
plt.show()
def blend_and_show(dlatent1, dlatent2, coeffs):
fig,ax = plt.subplots(1, len(coeffs), figsize=(12, 10), dpi=80)
dlatent1 = dlatent1.reshape(dlatent_shape)
for i, coeff in enumerate(coeffs):
new_latent_vector = coeff * dlatent1 + (1-coeff) * dlatent2
ax[i].imshow(gen_func([new_latent_vector])[0][0])
ax[i].set_title('Coeff: %0.1f' % coeff)
[x.axis('off') for x in ax]
plt.show()
Generovanie ľudských tvári pomocou StyleGAN¶
V notebook-u predvedieme, ako sa dajú generovať obrázky ľudských tvárí pomocou metódy StyleGAN z článku "A Style-Based Generator Architecture for Generative Adversarial Networks" od firmy NVIDIA. Oficiálna implementácia metódy sa dá nájsť v GitHub repozitári. Okrem toho používame latentné vektory a časť kódu z ďalšieho GitHub repozitára.
Stiahnutie kódu a načítanie modelu¶
Na začiatku si stiahneme repozitár s oficiálnou implementáciou.
Definícia parametrov a načítanie modelov¶
Definujeme niekoľko základných parametrov.
URL_FFHQ = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ'
LANDMARKS_MODEL_URL = 'http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2'
batch_size = 1
feature_img_size = 256
aligned_img_size = 1024
vgg_layer = 9
Načítame predtrénované modely – GAN siete na generovanie obrázkov.
tflib.init_tf()
with dnnlib.util.open_url(URL_FFHQ, cache_dir=config.cache_dir) as f:
generator_network, discriminator_network, Gs_network = pickle.load(f)
Modely, ktorými sa zarovnajú tváre.
landmarks_model_path = unpack_bz2(
get_file('shape_predictor_68_face_landmarks.dat.bz2',
LANDMARKS_MODEL_URL, cache_subdir='temp')
)
landmarks_detector = LandmarksDetector(landmarks_model_path)
Časti modelu, ktoré sa pri generovaní tvárí používajú, si oddelíme a vytvoríme funkcie, ktoré ich umožnia použiť. Časť, ktorá mapuje pôvodný latentný vektor na štýlový latentný vektor:
tf_map_in = Gs_network.components.mapping.input_templates[0]
tf_map_out = Gs_network.components.mapping.output_templates[0]
map_func = K.function([tf_map_in], [tf_map_out])
Časť, ktorá zo štýlového latentného vektora generuje obrázky:
tf_dlatents = Gs_network.components.synthesis.input_templates[0]
tf_output = Gs_network.components.synthesis.output_templates[0]
tf_loss_img_out = convert_images_loss(tf_output)
tf_img_out = convert_images_gen(tf_loss_img_out)
gen_func = K.function([tf_dlatents], [tf_img_out])
Tvar latentného a štýlového latentného vektora si uložíme do pomocných premenných.
latent_shape = (1,) + K.int_shape(tf_map_in)[1:]
dlatent_shape = (1,) + K.int_shape(tf_dlatents)[1:]
Generovanie náhodných tvárí¶
Ďalej si vygenerujeme náhodnú tvár. Začneme vygenerovaním latentného vektora. Jeho prvky budeme brať z normálneho rozdelenia a jeho tvar bude podľa premennej latent_shape.
latents = np.random.randn(*latent_shape)
Štýlový latentný vektor získame tak, že na ten pôvodný aplikujeme funkciu map_func, ktorú sme vyššie definovali.
dlatents = map_func([latents])[0]
Ďalej môžeme dlatents použiť ako vstup pre funkciu gen_func, ktorá vygeneruje samotný obrázok.
img = gen_func([dlatents])[0][0]
Potom nezostáva už nič iné než si obrázok vizualizovať, prípadne uložiť do súboru.
plt.imshow(img)
plt.axis('off')
Generovanie ďalších obrázkov si môžeme vyskúšať tu – stačí zakaždým vygenerovať iný latentný vektor.
fig, axes = plt.subplots(3, 4, figsize=(10, 10))
for row in axes:
for ax in row:
latents = np.random.randn(*latent_shape)
dlatents = map_func([latents])[0]
img = gen_func([dlatents])[0][0]
ax.imshow(img)
ax.axis('off')
Štýlový latentný vektor existujúcej tváre¶
GAN sieť by sa dala použiť aj na zaujímavé manipulácie s existujúcimi tvárami. Potrebovali by sme však najprv poznať latentné vektory, ktoré im zodpovedajú. StyleGAN však funguje len jednosmerne: generuje tváre z latentných vektorov, nie naopak.
Vieme však využiť obdobný princíp, aký používame pri generovaní predobrazov a protivníckych príkladov. Neurónová sieť je diferencovateľná a optimalizáciou vieme nájsť latentný vektor, ktorý bude generovať tvár čo najpodobnejšiu cieľovej tvári.
Podobnosť tvárí¶
Podobnosť tvárí nebudeme merať v zmysle pixelovej vzdialenosti, pretože tá dobre nevyjadruje ich skutočnú podobnosť. Namiesto toho si necháme obrázok tváre najprv predspracovať sieťou predtrénovanou na dátovej množine ImageNet. Porovnávať budeme až extrahované príznaky.
Vytvoríme si teda jednotlivé tenzory a načítame predtrénovaný model.
vgg16 = VGG16(include_top=False, input_shape=(feature_img_size, feature_img_size, 3))
perceptual_model = Model(vgg16.input, vgg16.layers[vgg_layer].output)
tf_img_ref = K.placeholder((1, aligned_img_size, aligned_img_size, 3))
tf_out_resized = preprocess_input(tf.image.resize_images(tf_loss_img_out,
(feature_img_size, feature_img_size), method=1))
tf_out_features = perceptual_model(tf_out_resized)
tf_ref_resized = preprocess_input(tf.image.resize_images(tf_img_ref,
(feature_img_size, feature_img_size), method=1))
tf_ref_features = perceptual_model(tf_ref_resized)
Zadefinujeme si chybovú funkciu – ako strednú kvadratickú chybu medzi príznakmi generovaného obrázku a pôvodného obrázku. Chybu si môžeme preškálovať tak, aby čísla boli z rozumnejšieho rozsahu:
loss = tf.losses.mean_squared_error(tf_out_features, tf_ref_features) / 85000
Ďalej určíme gradient chybovej funkcie voči štýlovému latentnému vektoru, aby sme ho neskôr mohli použiť na minimalizáciu chybovej funkcie.
tf_grads = K.gradients(loss, tf_dlatents)
loss_grad_func = K.function([tf_img_ref, tf_dlatents], [loss] + tf_grads)
Minimalizácia chybovej funkcie¶
Načítame obrázok a vykonáme minimalizáciu chybovej funkcie, aby sme našli jeho latentný vektor.
face_img_path = "data/starr.jpg"
# face_img_path = list(files.upload())[0]
face_img = PIL.Image.open(face_img_path)
plt.imshow(face_img)
plt.axis('off')
Ďalej si tvár mierne predspracujeme. Extrahujeme kľúčové body a obrázok tváre podľa nich zarovnáme na určité preddefinované pozície.
face_landmarks = next(landmarks_detector.get_landmarks(np.asarray(face_img)))
aligned_img = image_align(face_img, face_landmarks, output_size=aligned_img_size)
plt.imshow(aligned_img)
plt.axis('off')
Začneme optimalizáciu z nulového latentného vektora pomocou metódy LBFGS:
evaluator = Evaluator(aligned_img, loss_grad_func, dlatent_shape)
latent_init = np.zeros(latent_shape)
dlatent = map_func([latent_init])[0]
dlatent, min_val, info = fmin_l_bfgs_b(evaluator.loss, dlatent.flatten(),
fprime=evaluator.grads, maxfun=400, disp=1)
Generovanie obrázka¶
Po minimalizácii chybovej funkcie získame štýlový latentný vektor, ktorý približne zodpovedá pôvodnému obrázku. Keď na základe neho vygenerujeme novú tvár, mala by sa podobať pôvodnej tvári.
img = gen_func([dlatent.reshape((dlatent_shape))])[0][0]
plt.imshow(img)
plt.axis('off')
Manipulácia latentného vektora¶
Latentný vektor fotografie je možné ďalej modifikovať. Podobne ako pri rôznych iných typoch GAN a embedovaní, niektoré aritmetické operácie s vektormi nesú sémantický význam. Je napríklad možné identifikovať vektor, ktorý približne zodpovedá úsmevu, veku, pohlaviu a pod. Niekoľko takýchto vektorov si načítame:
smile_direction = np.load('stylegan-encoder/ffhq_dataset/latent_directions/smile.npy')
gender_direction = np.load('stylegan-encoder/ffhq_dataset/latent_directions/gender.npy')
age_direction = np.load('stylegan-encoder/ffhq_dataset/latent_directions/age.npy')
Aplikácia úsmevu:
move_and_show(dlatent.reshape((dlatent_shape)), smile_direction, [-1, 0, 1])
Zmena pohlavia:
move_and_show(dlatent.reshape((dlatent_shape)), gender_direction, [-1.5, 0, 2])
Zmena veku:
move_and_show(dlatent.reshape((dlatent_shape)), age_direction, [-2, 0, 1.5])
Zmiešavanie štýlov¶
Alternatívne vieme zmiešať štýly z viacerých obrázkov.
latent2 = np.random.RandomState(1855).randn(*latent_shape)
dlatent2 = map_func([latent2])[0]
img2 = gen_func([dlatent2])[0][0]
plt.imshow(img2)
plt.axis('off')
style_range = [0, 1, 2, 3, 4, 5, 6]
dlatent3 = dlatent2.copy()
dlatent3[:, style_range] = dlatent.reshape(dlatent_shape)[:, style_range]
img2 = gen_func([dlatent3])[0][0]
plt.imshow(img2)
plt.axis('off')
blend_and_show(dlatent.reshape(dlatent_shape), dlatent2, [0, 0.25, 0.5, 0.75, 1])
Aritmetika¶
def find_mean_dlatent(seeds):
mean_dlatent = np.zeros(dlatent_shape)
for s in seeds:
h_latent = np.random.RandomState(s).randn(*latent_shape)
h_dlatent = map_func([h_latent])[0]
mean_dlatent += h_dlatent / len(seeds)
return mean_dlatent
female_long_hair = [517, 519, 521, 523, 525, 528, 529, 538, 539, 540, 618, 642, 655]
female_short_hair = [537, 546, 561, 597, 599, 602, 610, 616, 627, 637, 652]
# male_long_hair = [629]
# male_short_hair = [535, 536, 549, 558, 559, 560]
long_hair_dlatent = find_mean_dlatent(female_long_hair)
short_hair_dlatent = find_mean_dlatent(female_short_hair)
dlatent2 = dlatent.reshape(dlatent_shape) + 0.5 * (short_hair_dlatent - long_hair_dlatent)
img2 = gen_func([dlatent2])[0][0]
plt.imshow(img2)
plt.axis('off')
