Genera volti artificiali con il modello GAN progressivo di CelebA

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza su GitHub Scarica taccuino Vedi il modello del mozzo TF

Questo Colab dimostra l'uso di un modulo Hub TF basato su una rete generativa avversaria (GAN). Il modulo mappa da vettori N-dimensionali, chiamati spazio latente, a immagini RGB.

Vengono forniti due esempi:

  • Mappare dallo spazio latente alle immagini, e
  • Dato un'immagine target, utilizzando gradiente discesa di trovare un vettore latente che genera un'immagine simile all'immagine di destinazione.

Prerequisiti facoltativi

Più modelli

Qui si possono trovare tutti i modelli attualmente ospitati su tfhub.dev che può generare immagini.

Impostare

# Install imageio for creating animations.
pip -q install imageio
pip -q install scikit-image
pip install git+https://github.com/tensorflow/docs

Importazioni e definizioni di funzioni

from absl import logging

import imageio
import PIL.Image
import matplotlib.pyplot as plt
import numpy as np

import tensorflow as tf
tf.random.set_seed(0)

import tensorflow_hub as hub
from tensorflow_docs.vis import embed
import time

try:
  from google.colab import files
except ImportError:
  pass

from IPython import display
from skimage import transform

# We could retrieve this value from module.get_input_shapes() if we didn't know
# beforehand which module we will be using.
latent_dim = 512


# Interpolates between two vectors that are non-zero and don't both lie on a
# line going through origin. First normalizes v2 to have the same norm as v1. 
# Then interpolates between the two vectors on the hypersphere.
def interpolate_hypersphere(v1, v2, num_steps):
  v1_norm = tf.norm(v1)
  v2_norm = tf.norm(v2)
  v2_normalized = v2 * (v1_norm / v2_norm)

  vectors = []
  for step in range(num_steps):
    interpolated = v1 + (v2_normalized - v1) * step / (num_steps - 1)
    interpolated_norm = tf.norm(interpolated)
    interpolated_normalized = interpolated * (v1_norm / interpolated_norm)
    vectors.append(interpolated_normalized)
  return tf.stack(vectors)

# Simple way to display an image.
def display_image(image):
  image = tf.constant(image)
  image = tf.image.convert_image_dtype(image, tf.uint8)
  return PIL.Image.fromarray(image.numpy())

# Given a set of images, show an animation.
def animate(images):
  images = np.array(images)
  converted_images = np.clip(images * 255, 0, 255).astype(np.uint8)
  imageio.mimsave('./animation.gif', converted_images)
  return embed.embed_file('./animation.gif')

logging.set_verbosity(logging.ERROR)

Interpolazione spaziale latente

Vettori casuali

Interpolazione spaziale latente tra due vettori inizializzati casualmente. Useremo un modulo TF Hub progan-128 che contiene un pre-addestrato Progressive GAN.

progan = hub.load("https://tfhub.dev/google/progan-128/1").signatures['default']

def interpolate_between_vectors():
  v1 = tf.random.normal([latent_dim])
  v2 = tf.random.normal([latent_dim])

  # Creates a tensor with 25 steps of interpolation between v1 and v2.
  vectors = interpolate_hypersphere(v1, v2, 50)

  # Uses module to generate images from the latent space.
  interpolated_images = progan(vectors)['default']

  return interpolated_images

interpolated_images = interpolate_between_vectors()
animate(interpolated_images)

gif

Trovare il vettore più vicino nello spazio latente

Correggi un'immagine di destinazione. Ad esempio, usa un'immagine generata dal modulo o carica la tua.

image_from_module_space = True  # @param { isTemplate:true, type:"boolean" }

def get_module_space_image():
  vector = tf.random.normal([1, latent_dim])
  images = progan(vector)['default'][0]
  return images

def upload_image():
  uploaded = files.upload()
  image = imageio.imread(uploaded[list(uploaded.keys())[0]])
  return transform.resize(image, [128, 128])

if image_from_module_space:
  target_image = get_module_space_image()
else:
  target_image = upload_image()

display_image(target_image)

png

Dopo aver definito una funzione di perdita tra l'immagine target e l'immagine generata da una variabile spaziale latente, possiamo utilizzare la discesa del gradiente per trovare valori di variabili che minimizzino la perdita.

tf.random.set_seed(42)
initial_vector = tf.random.normal([1, latent_dim])

display_image(progan(initial_vector)['default'][0])

png

def find_closest_latent_vector(initial_vector, num_optimization_steps,
                               steps_per_image):
  images = []
  losses = []

  vector = tf.Variable(initial_vector)  
  optimizer = tf.optimizers.Adam(learning_rate=0.01)
  loss_fn = tf.losses.MeanAbsoluteError(reduction="sum")

  for step in range(num_optimization_steps):
    if (step % 100)==0:
      print()
    print('.', end='')
    with tf.GradientTape() as tape:
      image = progan(vector.read_value())['default'][0]
      if (step % steps_per_image) == 0:
        images.append(image.numpy())
      target_image_difference = loss_fn(image, target_image[:,:,:3])
      # The latent vectors were sampled from a normal distribution. We can get
      # more realistic images if we regularize the length of the latent vector to 
      # the average length of vector from this distribution.
      regularizer = tf.abs(tf.norm(vector) - np.sqrt(latent_dim))

      loss = target_image_difference + regularizer
      losses.append(loss.numpy())
    grads = tape.gradient(loss, [vector])
    optimizer.apply_gradients(zip(grads, [vector]))

  return images, losses


num_optimization_steps=200
steps_per_image=5
images, loss = find_closest_latent_vector(initial_vector, num_optimization_steps, steps_per_image)

....................................................................................................
....................................................................................................

plt.plot(loss)
plt.ylim([0,max(plt.ylim())])

(0.0, 6696.301751708985)

png

animate(np.stack(images))

gif

Confronta il risultato con l'obiettivo:

display_image(np.concatenate([images[-1], target_image], axis=1))

png

Giocando con l'esempio sopra

Se l'immagine proviene dallo spazio del modulo, la discesa è rapida e converge a un campione ragionevole. Prova scendendo a un'immagine che non è dallo spazio del modulo. La discesa convergerà solo se l'immagine è ragionevolmente vicina allo spazio delle immagini di addestramento.

Come farlo scendere più velocemente e verso un'immagine più realistica? Si può provare:

  • utilizzando una perdita diversa sulla differenza dell'immagine, ad es. quadratica,
  • utilizzando diversi regolarizzatori sul vettore latente,
  • inizializzazione da un vettore casuale in più esecuzioni,
  • eccetera.