Для создания изображений с помощью GAN я буду использовать Tensorflow.

Генеративно-состязательная сеть (GAN) — это модель машинного обучения (ML), в которой две нейронные сети соревнуются друг с другом, чтобы быть более точными в своих прогнозах.

Как работают GAN?

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

После этого сгенерированные изображения передаются в дискриминатор вместе с фактическими точками данных из исходной концепции. Дискриминатор фильтрует информацию и возвращает вероятность от 0 до 1, чтобы представить подлинность каждого изображения (1 соответствует реальному, а 0 соответствует ложному). Эти значения затем вручную проверяются на успешность и повторяются до тех пор, пока не будет достигнут желаемый результат.

Зачем создавать изображение ЭКГ?

Я создал проект coronarography.ai. В нем на вход подается изображение ЭКГ, а на выходе мы получаем наличие патологии магистральных артерий сердца. Мне стало интересно проверить принципиальную возможность генерации изображений ЭКГ и сравнить полученные изображения с реальными.

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

Используемые библиотеки

import pandas as pd

import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
from PIL import Image
from tensorflow.keras import layers
import time
import tensorflow as tf
from IPython import display
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import matplotlib
from matplotlib.pyplot import figure
from sklearn.preprocessing import MinMaxScaler
import joblib

import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
print(tf.__version__)
%matplotlib inline
matplotlib.rcParams['figure.figsize'] = (18,10)
import moviepy.editor as mpe

Загрузка и подготовка набора данных

Здесь мы загружаем изображения ЭКГ из папки в массив. Преобразовываем в одноканальное (черно-белое) изображение, нормализуем изображение.

data_image = []
for k in os.listdir('../AI_coronarography/DATA_WORK/DATA_WORK/ЭКГ'):
    if k.endswith('.jpg'):
        img = Image.open('../AI_coronarography/DATA_WORK/DATA_WORK/ЭКГ/'+k)
        img = img.convert('L')
        img = img.resize((800, 800))
        data_image += [(np.array(img) - 127.5) / 127.5]

Пример загруженного изображения ЭКГ.

Определите размер партии и смешайте изображения.

train_images = np.array(data_image).reshape(np.array(data_image).shape[0], 800, 800, 1).astype('float32')
BUFFER_SIZE = 100
BATCH_SIZE = 10
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

Мы определяем структуру генератора.

def make_generator_model():
    input_1 = Input(shape=(100, ), name = "Input_image")
    x = Dense(25*25*256, use_bias=False)(input_1)
    x = BatchNormalization()(x)
    x = LeakyReLU()(x)
    x = Reshape((25, 25, 256))(x)
    x = Conv2DTranspose(256, (1, 1), strides=(1, 1), padding='same', use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU()(x)
    x = Conv2DTranspose(128, (3, 3), strides=(2, 2), padding='same', use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU()(x)
    x = Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU()(x)
    x = Conv2DTranspose(32, (7, 7), strides=(2, 2), padding='same', use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU()(x)
    x = Conv2DTranspose(16, (9, 9), strides=(2, 2), padding='same', use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU()(x)
    x = Conv2DTranspose(1, (11, 11), strides=(2, 2), padding='same', use_bias=False, activation='tanh')(x)
    model = Model(inputs=input_1, outputs=x)
    
    return model

Проверим необученный генератор.

generator = make_generator_model()

noise = tf.random.normal([1, 100])
generated_image = generator(noise, training=False)

plt.imshow(generated_image[0, :, :, 0], cmap='gray');

Давайте создадим дискриминатор.

def make_discriminator_model():
    input_1 = Input(shape=(800, 800, 1), name = "Input_image")
    x = Conv2D(16, (11, 11), strides=(2, 2), padding='same')(input_1)
    x = LeakyReLU()(x)
    x = Dropout(0.3)(x)
    x = Conv2D(32, (9, 9), strides=(2, 2), padding='same')(x)
    x = LeakyReLU()(x)
    x = Dropout(0.3)(x)
    x = Conv2D(64, (7, 7), strides=(2, 2), padding='same')(x)
    x = LeakyReLU()(x)
    x = Dropout(0.3)(x)
    x = Conv2D(128, (5, 5), strides=(2, 2), padding='same')(x)
    x = LeakyReLU()(x)
    x = Dropout(0.3)(x)
    x = Conv2D(256, (3, 3), strides=(2, 2), padding='same')(x)
    x = LeakyReLU()(x)
    x = Dropout(0.3)(x)
    x = Conv2D(256, (1, 1), strides=(1, 1), padding='same')(x)
    x = LeakyReLU()(x)
    x = Dropout(0.3)(x)
    x = Flatten()(x)
    x = Dense(1)(x)
    model = Model(inputs=input_1, outputs=x)
    
    return model

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

discriminator = make_discriminator_model()
decision = discriminator(generated_image)
print (decision)
#result 
tf.Tensor([[1.6453338e-05]], shape=(1, 1), dtype=float32)

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

Потеря дискриминатора.
Этот метод определяет, насколько хорошо дискриминатор способен отличать настоящие изображения от подделок. Он сравнивает предсказания дискриминатора для реальных изображений с массивом 1 и предсказания дискриминатора для поддельных (сгенерированных) изображений с массивом 0.

Потери генератора.
Потери генератора определяют, насколько хорошо он смог обмануть дискриминатор. Интуитивно, если генератор работает хорошо, дискриминатор классифицирует поддельные изображения как настоящие (или 1).

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

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

def generator_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

Сохранить контрольные точки

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

Определение цикла обучения

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

EPOCHS = 6000
noise_dim = 100
num_examples_to_generate = 1

seed = tf.random.normal([num_examples_to_generate, noise_dim])

@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
      generated_images = generator(noise, training=True)

      real_output = discriminator(images, training=True)
      fake_output = discriminator(generated_images, training=True)

      gen_loss = generator_loss(fake_output)
      disc_loss = discriminator_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

def train(dataset, epochs):
  for epoch in range(epochs):
    start = time.time()

    for image_batch in dataset:
      train_step(image_batch)

    display.clear_output(wait=True)
    generate_and_save_images(generator,
                             epoch + 1,
                             seed)

    if (epoch + 1) % 500 == 0:
      checkpoint.save(file_prefix = checkpoint_prefix)

    print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

  display.clear_output(wait=True)
  generate_and_save_images(generator,
                           epochs,
                           seed)

Создавать и сохранять изображения

def generate_and_save_images(model, epoch, test_input):
  predictions = model(test_input, training=False)
  fig = plt.figure(figsize=(12, 12))
  plt.imshow(predictions[0, :, :, 0] * 127.5 + 127.5, cmap='gray')
  plt.axis('off')
  plt.savefig('image_per_epoch/image_at_epoch_{:04d}.png'.format(epoch), bbox_inches='tight', pad_inches=0)
  plt.show()

Модель поезда

Вызовите метод train(), определенный выше, для одновременного обучения генератора и дискриминатора. Важно, чтобы генератор и дискриминатор не подавляли друг друга (например, чтобы они обучались с одинаковой скоростью).

В начале обучения сгенерированные изображения выглядят как случайный шум. По мере того, как вы будете учиться, сгенерированные изображения ЭКГ будут выглядеть все более и более реальными.

train(train_dataset, EPOCHS)

Видеообучение нейронной сети, генерирующей изображения ЭКГ.

Сравнение реальных и сгенерированных изображений ЭКГ.