Внедрение CycleGAN для преобразования изображения в изображение

Преобразование изображений из одного стиля в другой

Знаете ли вы, что вы можете сделать снимок лета и преобразовать его так, чтобы он выглядел как снимок, сделанный зимой? Для этого типа преобразования у нас есть набор различных алгоритмов, и CycleGAN — один из них, используемый для преобразования изображения в изображение.

Преобразование изображения в изображение — это класс задач компьютерного зрения и глубокого обучения, целью которого является изучение сопоставления между входным и выходным изображениями с использованием обучающего набора пар изображений.

Что такое CycleGAN?

Cycle Generative Adversarial Networks(CycleGAN) – это хорошо известный передовой алгоритм глубокого обучения. Он был разработан для решения проблемы перехода от изображения к изображению. CycleGAN является самым популярным алгоритмом и относится к набору алгоритмов, называемых генеративными моделями, и эти алгоритмы относятся к области техники обучения без учителя. Эта модель была впервые описана Джун-Ян Чжу в 2017 году.

Как мы знаем, для решения проблемы преобразования изображения в изображение алгоритм должен изучить сопоставление между входным изображением и выходным изображением с помощью обучающего набора парных изображений. Однако алгоритмы CycleGAN изучают сопоставление между входным и выходным изображением, не требуя набора парных входных и выходных изображений.

Архитектура CycleGAN

CycleGAN использует архитектуру модели Generative Adversarial Network (GAN). Он имеет две нейронные сети, называемые gгенератор и дискриминатор.

Работа генератора состоит в том, чтобы произвести выборку из желаемого дистрибутива, что означает создание реалистичных изображений, подобных ожидаемому. Работа дискриминатора состоит в том, чтобы отличать реальные изображения от поддельных, созданных генератором.

Обе сети (генератор и дискриминатор)обучаются одновременно, пока не наступает этап, когда генератор создает настолько реалистичные изображения, что дискриминатор не может определить, какое изображение реальное, а какое искусственно создано генератором. Этот этап называется равновесием. Как только достигается стадия равновесия, обучение модели автоматически прекращается.

Ниже у нас есть два сопоставления: одно G: X → Y, а второе F: Y → X. Генератор (G) пытается преобразовать (X) в выходные данные → (Y), и он проходит через Dy, чтобы проверить, являются ли изображения реальными или поддельными в соответствии с доменом (Y). Точно так же генератор (F) пытается преобразовать (Y) в выходные данные → (X), и он проходит через Dx, чтобы проверить, неотличимы ли изображения от домена (X).

Генератор (F и G) состоит из трех секций:

1. Кодировщик — на самом первом этапе входное изображение передается в кодировщик. В этом разделе функция изображения извлекается из входного изображения с помощью сверток и сжимает представление изображения. Кодировщик имеет три свертки, которые уменьшают представление на 1/4 фактического размера изображения. Если входное изображение имеет размер 256, 256, 3, при вводе в кодировщик выход кодировщика будет 64, 64, 256.
2. Transformer — для этого раздела ввод будет поступать от кодировщика после применения функции активации к выходу кодировщика. В зависимости от размера входа преобразователь содержит 6/9 остаточных блоков.
3. Декодер — выход преобразователя является входом для декодера. Он использует два блока деконволюции дробных шагов, чтобы увеличить размер представления до исходного размера.

Функции потерь

Потеря согласованности цикла:как и в CycleGAN, нет парных данных (входное изображение и соответствующее выходное изображение) для обучения модели, поэтому мы не можем предсказать, что вход X и целевая пара Y имеет смысл во время обучения. Чтобы гарантировать, что модель изучает правильное сопоставление, была введена потеря согласованности цикла.

Генератор (G) берет входное изображение (изображение летнего сезона) из домена A и преобразует его в выходное изображение (изображение зимнего сезона) домена B, снова выходное изображение переведены обратно в домен A с помощью генератора (F). Таким образом, во время полного цикла разница между этими двумя изображениями называется потерей согласованности цикла, как показано на изображении ниже.

Потери со стороны противника:потери, рассчитанные как для наших отображений генераторов, так и для дискриминаторов, называются потерями со стороны противника и рассчитываются по приведенной ниже формуле.

Таким образом, общая потеря, рассчитанная как сумма потерь противника и потери согласованности цикла, равна окончательной формуле, приведенной ниже.

Выполнение

Давайте посмотрим, как реализовать CycleGAN шаг за шагом. Во-первых, нам нужно установить необходимые библиотеки.

pip install tensorflow
pip install matplotlib

Импортируйте указанные ниже библиотеки.

import tensorflow as tf
import os
import numpy as np
import matplotlib.pyplot as plt
import tensorflow_datasets as tfds
import tensorflow_addons as tfa
from tensorflow import keras
from tensorflow.keras import layers

Следующим шагом является подготовка набора данных. Используя наборы данных TensorFlow, загрузите набор данных «лошадь-зебра» и определите размер изображения, размер буфера и размер пакета.

 
dataset, _ = tfds.load("cycle_gan/horse2zebra", with_info=True, as_supervised=True)
train_horse, train_zebra = dataset["trainA"], dataset["trainB"]
test_horse, test_zebra = dataset["testA"], dataset["testB"]
buffer_size = 256
batch_size = 1
# Define image size.
org_image_size=(286,286)
# Size to be used during training.
input_image_size = (256, 256, 3)
# Gamma initializer for instance normalization.
gamma_init = keras.initializers.RandomNormal(mean=0.0,stddev=0.02)
# Weights initializer for the layer.
kernel_init = keras.initializers.RandomNormal(mean=0.0,stddev=0.02)

Определены функции для нормализации изображения, а также для обработки и обучения изображения и тестовых изображений по отдельности.

def normalize_image(image):
    image = tf.cast(image, dtype=tf.float32)
    return (image/127.5) - 1.0
def preprocess_train_image(image, label):
    image = tf.image.random_flip_left_right(image)
    image = tf.image.resize(image, [*org_image_size])
    image = tf.image.random_crop(image, size=[*input_image_size])
    # Normalize the pixel values in the range [-1, 1]
    image = normalize_image(image)
    return image
def preprocess_test_image(image, label):
    # used for resizing and normalization for the test images.
    image = tf.image.resize(image, [input_image_size[0], input_image_size[1]])
    image = normalize_image(image)
    return image

Далее мы обработаем обучающие данные и тестовые данные, используя приведенный ниже код Python для лошади (домен A) и зебры (домен B).

train_horse = (
    train_horse.map(preprocess_train_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)
train_zebra = (
    train_zebra.map(preprocess_train_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)
# Test data processing 
test_horse = (
    test_horse.map(preprocess_test_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)
test_zebra = (
    test_zebra.map(preprocess_test_image, num_parallel_calls=autotune)
    .cache()
    .shuffle(buffer_size)
    .batch(batch_size)
)

Теперь мы определим функцию cycleGAN_resnet_generator()для создания модели длягенератора. Код Python для этого приведен ниже.

def cycleGAN_resnet_generator(
    filters=64,
    num_downsampling_blocks=2,
    num_residual_blocks=9,
    num_upsample_blocks=2,
    gamma_initializer=gamma_init,
    name=None,
):
    img_input = layers.Input(shape=input_image_size, name=name + "_img_input")
    x = Reflection_Padding2D(padding=(3, 3))(img_input)
    x = layers.Conv2D(filters, (7, 7), kernel_initializer=kernel_init, use_bias=False)(
        x
    )
    x = tfa.layers.InstanceNormalization(gamma_initializer=gamma_initializer)(x)
    x = layers.Activation("relu")(x)
# Downsampling
    for _ in range(num_downsampling_blocks):
        filters *= 2
        x = downsample1(x, filters=filters, activation=layers.Activation("relu"))
# Residual blocks
    for _ in range(num_residual_blocks):
        x = residual_block1(x, activation=layers.Activation("relu"))
# Upsampling
    for _ in range(num_upsample_blocks):
        filters //= 2
        x = upsample1(x, filters, activation=layers.Activation("relu"))
# Final block
    x = Reflection_Padding2D(padding=(3, 3))(x)
    x = layers.Conv2D(3, (7, 7), padding="valid")(x)
    x = layers.Activation("tanh")(x)
model = keras.models.Model(img_input, x, name=name)
return model

Точно так же функция cycleGAN_distributor() определена для создания модели дискриминатора.

def cycleGAN_discriminator(
    filters=64,kernel_initializer=kernel_init,num_downsampling=3, name=None
):
img_input = layers.Input(shape=input_image_size, name=name + "_img_input")
    x = layers.Conv2D(
        filters,
        (4,4),
        strides=(2, 2),
        padding="same",
        kernel_initializer=kernel_initializer,
    )(img_input)
    x = layers.LeakyReLU(0.2)(x)
num_filters = filters
    for num_downsample_block in range(3):
        num_filters *= 2
        if num_downsample_block < 2:
            x = downsample1(
                x,
                filters=num_filters,
                activation=layers.LeakyReLU(0.2),
                kernel_size=(4, 4),
                strides=(2, 2),
            )
        else:
            x = downsample1(
                x,
                filters=num_filters,
                activation=layers.LeakyReLU(0.2),
                kernel_size=(4, 4),
                strides=(1, 1),
            )
x = layers.Conv2D(
        1, (4, 4), strides=(1, 1), padding="same", kernel_initializer=kernel_initializer
    )(x)
model = keras.models.Model(inputs=img_input, outputs=x, name=name)
return model

Как и в CycleGAN, существуют два генератора (G и F) и два дискриминатора (X и Y), поэтому отдельные значения назначаются в соответствии с определением функции, написанным на предыдущем шаге.

# Get the generators
Gen_G = get_resnet_generator(name = "generator_G")
Gen_F = get_resnet_generator(name = "generator_F")
# Get the discriminators
Disc_X = get_discriminator(name = "discriminator_X")
Disc_Y = get_discriminator(name = "discriminator_Y")

Построить модель

Теперь наша цель — построить модель CycleGAN. На этом этапе мы рассчитаем различия и потери между изображениями (входными и выходными) для генераторов и дискриминаторов. Для этого мы последовательно выполним следующие шаги.

1. Пропустите входные изображения через генераторы и получите сгенерированные изображения.
2. Передайте сгенерированные изображения обратно в генераторы, чтобы убедиться, что мы можем предсказать исходное изображение из сгенерированного изображения.
3. Используя генераторы, выполните сопоставление идентичности реальных изображений.
4. Передайте сгенерированные изображения из шага 1 в соответствующие дискриминаторы.
5. Найдите общие потери генераторов = (противодействие + цикл + идентичность).
6. Найдите потери дискриминаторов.
7. Обновите веса генераторов.
8. Обновите веса дискриминаторов.
9. Верните потери в словарь.

class cycleGan(keras.Model):
    def __init__(
        self,
        generator_G,
        generator_F,
        discriminator_X,
        discriminator_Y,
        lambda_cycle=10.0,
        lambda_identity=0.5,
    ):
        super(CycleGan, self).__init__()
        self.gen_G = generator_G
        self.gen_F = generator_F
        self.disc_X = discriminator_X
        self.disc_Y = discriminator_Y
        self.lambda_cycle = lambda_cycle
        self.lambda_identity = lambda_identity
def compile(
        self,
        gen_G_optimizer,
        gen_F_optimizer,
        disc_X_optimizer,
        disc_Y_optimizer,
        gen_loss_fn,
        disc_loss_fn,
    ):
        super(CycleGan, self).compile()
        self.gen_G_optimizer = gen_G_optimizer
        self.gen_F_optimizer = gen_F_optimizer
        self.disc_X_optimizer = disc_X_optimizer
        self.disc_Y_optimizer = disc_Y_optimizer
        self.generator_loss_function = gen_loss_fn
        self.discriminator_loss_function = disc_loss_fn
        self.cycle_loss_fn = keras.losses.MeanAbsoluteError()
        self.identity_loss_fn = keras.losses.MeanAbsoluteError()
def train_step(self, batch_data):
             real_x, real_y = batch_data
with tf.GradientTape(persistent=True) as tape:
            # Horse to fake new zebra
            fake_y = self.gen_G(real_x, training=True)
            # Zebra to fake new horse -> y2x
            fake_x = self.gen_F(real_y, training=True)
# Cycle (Horse to new fake zebra to fake horse): x -> y -> x
            cycled_x = self.gen_F(fake_y, training=True)
            # Cycle (Zebra to fake horse to fake zebra) y -> x -> y
            cycled_y = self.gen_G(fake_x, training=True)
# Identity mapping
            same_x = self.gen_F(real_x, training=True)
            same_y = self.gen_G(real_y, training=True)
# Discriminator output
            Disc_real_x = self.disc_X(real_x, training=True)
            Disc_fake_x = self.disc_X(fake_x, training=True)
disc_real_y = self.disc_Y(real_y, training=True)
            Disc_fake_y = self.Disc_Y(fake_y, training=True)
# Generator adverserial loss
            gen_G_loss = self.generator_loss_function(Disc_fake_y)
            gen_F_loss = self.generator_loss_function(Disc_fake_x)
# Generator cycle loss
            cycle_loss_G = self.cycle_loss_fn(real_y, cycled_y) * self.lambda_cycle
            cycle_loss_F = self.cycle_loss_fn(real_x, cycled_x) * self.lambda_cycle
# Generator identity loss
            id_loss_G = (
                self.identity_loss_fn(real_y, same_y)
                * self.lambda_cycle
                * self.lambda_identity
            )
            id_loss_F = (
                self.identity_loss_fn(real_x, same_x)
                * self.lambda_cycle
                * self.lambda_identity
            )
# Total generator loss
            total_loss_G = gen_G_loss + cycle_loss_G + id_loss_G
            total_loss_F = gen_F_loss + cycle_loss_F + id_loss_F
# Discriminator loss
            Disc_X_loss = self.discriminator_loss_function(Disc_real_x, Disc_fake_x)
            Disc_Y_loss = self.discriminator_loss_function(Disc_real_y, Disc_fake_y)
# find the gradients for the generators
        grads_G = tape.gradient(total_loss_G, self.gen_G.trainable_variables)
        grads_F = tape.gradient(total_loss_F, self.gen_F.trainable_variables)
# Find the gradients for the discriminators
        disc_X_grads = tape.gradient(disc_X_loss, self.disc_X.trainable_variables)
        disc_Y_grads = tape.gradient(disc_Y_loss, self.disc_Y.trainable_variables)
# Update the weights of the generators
        self.gen_G_optimizer.apply_gradients(
            zip(grads_G, self.gen_G.trainable_variables)
        )
        self.gen_F_optimizer.apply_gradients(
            zip(grads_F, self.gen_F.trainable_variables)
        )
# Update the weights of the discriminators
        self.disc_X_optimizer.apply_gradients(
            zip(disc_X_grads, self.disc_X.trainable_variables)
        )
        self.disc_Y_optimizer.apply_gradients(
            zip(disc_Y_grads, self.disc_Y.trainable_variables)
        )
return {
            "G_loss"  : total_loss_G,
            "F_loss"  : total_loss_F,
            "D_X_loss": disc_X_loss,
            "D_Y_loss": disc_Y_loss,
        }

Определены функции для расчета потерь от генераторов и дискриминаторов с помощью keras.losses.MeanSquaredError() .Мы компилируем модель после генерации, и код Python написан ниже.

# Loss function for evaluating  loss
adv_loss_fn = keras.losses.MeanSquaredError()
# Here we are defining the loss function for the generators
def generator_loss_function(fake):
    fake_loss = adv_loss_fn(tf.ones_like(fake), fake)
    return fake_loss
# Define the loss function for the discriminators
def discriminator_loss_function(real, fake):
    real_loss = adv_loss_fn(tf.ones_like(real), real)
    fake_loss = adv_loss_fn(tf.zeros_like(fake), fake)
    return (real_loss + fake_loss) * 0.5

# Model Creation steps
cycle_gan_model = CycleGan(
                           generator_G=gen_G, 
                           generator_F=gen_F, 
                           discriminator_X=disc_X, 
                           discriminator_Y=disc_Y
                          )

# Next, will Compile the model
cycle_gan_model.compile(
gen_G_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
gen_F_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
disc_X_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
disc_Y_optimizer=keras.optimizers.Adam(learning_rate=2e-4, beta_1=0.5),
    gen_loss_fn=generator_loss_function,
    disc_loss_fn=discriminator_loss_function,
)
# Callbacks
plotter = GANMonitor()
checkpoint_filepath = "./model_checkpoints/cyclegan_checkpoints.{epoch:03d}"
model_checkpoint_callback = keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath
)
# Here we will train the model for just one epoch.
 
cycle_gan_model.fit(
    tf.data.Dataset.zip((train_horses,train_zebras)),
    epochs =1,
    callbacks=[plotter, model_checkpoint_callback],
   )

Прогноз — вывод

Это последний этап, на котором мы найдем выходные данные в соответствии с моделью, созданной путем предоставления входных данных.

weight_files = "./saved_checkpoints/cyclegan_checkpoints.080"
cycle_gan_model.load_weights(weight_files).expect_partial()


_, ax = plt.subplots(4,2, figsize=(15,26))
for i, img in enumerate(test_horses.take(4)):
    prediction = cycle_gan_model.gen_G(img, training=False)[0].numpy()
    prediction = (prediction * 127.5 + 127.5).astype(np.uint8)
    img = (img[0] * 127.5+127.5).numpy().astype(np.uint8)

    ax[i,0].imshow(img)
    ax[i,1].imshow(prediction)
    ax[i,0].set_title("Input image")
    ax[i,1].set_title("Translated image")
    ax[i,0].axis("off")
    ax[i,1].axis("off")

    prediction = keras.preprocessing.image.array_to_img(prediction)
    prediction.save("predicted_img_{i}.png".format(i=i))
    
    plt.tight_layout()
    plt.show()

Из-за нехватки места я показываю только две пары ввода-вывода, но вы можете легко понять конвертированные изображения лошади в изображения зебры.

Вывод

Мы рассмотрели CycleGAN и его архитектуру. В настоящее время концепции модели CycleGAN используются во многих приложениях — одним из них является FaceApp.

Итак, куда идти отсюда? Мы можем больше поэкспериментировать с условной частью модели, чтобы посмотреть, сможем ли мы добиться таких вещей, как выполнение нескольких задач одновременно, и мы увидим, как генератор будет вести себя с другим условным входным изображением. Есть много возможностей для экспериментов и улучшений, над которыми я буду работать и призываю вас делать то же самое!

Примечание редактора. Heartbeat — это интернет-издание и сообщество, созданное участниками и посвященное предоставлению лучших образовательных ресурсов для специалистов по науке о данных, машинному обучению и глубокому обучению. Мы стремимся поддерживать и вдохновлять разработчиков и инженеров из всех слоев общества.

Независимая от редакции, Heartbeat спонсируется и публикуется Comet, платформой MLOps, которая позволяет специалистам по данным и командам машинного обучения отслеживать, сравнивать, объяснять и оптимизировать свои эксперименты. Мы платим нашим авторам и не продаем рекламу.

Если вы хотите внести свой вклад, перейдите к нашему призыву к участию. Вы также можете подписаться на получение наших еженедельных информационных бюллетеней (Еженедельник глубокого обучения и Информационный бюллетень Comet), присоединиться к нам в Slack и следить за Comet в Twitter и LinkedIn для получения ресурсов, событий и многое другое, что поможет вам быстрее создавать лучшие модели машинного обучения.