Как с помощью нейросети стилизовать изображение под работу известного художника: разбираемся с нейронным переносом стиля

В этой статье рассказывается о том, как использовать deep learning для стилизации изображения по заданному образцу. Это возможно благодаря нейронному переносу стиля (англ. neural style transfer). Эта техника описана в статье Leon A. Gatys, A Neural Algorithm of Artistic Style.

Нейронная передача стиля — это процесс оптимизации, который работает с 3 изображениями: картинкой содержания, картинкой стиля (например произведением художника) и входной картинкой. Если «смешать» их, то получится входная картинка, подогнанная по композиции под картинку содержания в образе копируемого стиля.

Для примера возьмём фотографию черепахи и гравюру Кацусики Хокусая «Большая волна в Канагаве»:

И что бы вышло, если бы художник решил стилизовать фотографию черепахи под свою гравюру? У него получилось бы что-то подобное:

Принцип передачи стиля заключается в определении двух функций расстояния. Одна из них описывает, насколько друг от друга отличаются содержания двух изображений (Lcontent). Вторая функция описывает разницу между двумя стилями изображений (Lstyle). Получив три изображения (желаемый стиль, желаемый контент и входное изображение), сеть пытается преобразовать входное изображение так, чтобы минимизировать его расстояние Lcontent с изображением контента и расстояние Lstyle с изображением стиля.

О чём статья

Статья освещает следующие аспекты:

• моментальное исполнение (англ. Eager Execution) — использование библиотеки TensorFlow, которая позволяет выполнять операции незамедлительно, без построения графов. Тут можно узнать больше о моментальном исполнении, а увидеть в действии можно тут;
• работа с functional API для определения модели — вы будете использовать подмножество моделей, чтобы получить доступ к важным промежуточным функциям активации с помощью functional API;
• использование карт признаков подготовленной модели;
• создание собственных циклов обучения — вы научитесь минимизировать заданные потери входных параметров.

Выполняя перенос стиля, вы проделаете следующие шаги:

1. Визуализация данных.
2. Базовая предварительная обработка/подготовка данных.
3. Настройка функций потери.
4. Создание модели.
5. Оптимизация функции потери.

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

• https://tproger.ru/translations/6-step-for-building-machine-learning-projects/
• https://tproger.ru/translations/math-for-ai-linear-algebra/
• https://tproger.ru/video/machine-learning-2014/?autoplay=1
• https://tproger.ru/digest/learning-neuroweb-all-for-begin/
• https://tproger.ru/experts/required-ml-skills/

Код

Вы сможете найти полные исходники здесь. Если вы хотите детально разобрать примеры из этой статьи, то можно перейти на Colab.

Реализация

Начать стоит с включения моментального исполнения. Это позволит вам работать с техникой переноса стиля наиболее эффективным и понятным образом.

			tf.enable_eager_execution()
print("Eager execution: {}".format(tf.executing_eagerly()))
 
# Изображения контента и стиля, которые будут использованы: 
plt.figure(figsize=(10,10))

content = load_img(content_path).astype('uint8')
style = load_img(style_path)

plt.subplot(1, 2, 1)
imshow(content, 'Content Image')

plt.subplot(1, 2, 2)
imshow(style, 'Style Image')
plt.show()
		

Определите представления содержания и стиля

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

Почему именно промежуточные слои?

Вы можете задаться вопросом: почему эти промежуточные выводы дают возможность определить стиль и контент изображения? Чтобы сеть могла классифицировать изображение (чему она уже была обучена), она должна понимать это изображение. Это включает в себя построение из группы пикселей сложных представлений объектов на изображении. Отчасти это объясняет, почему свёрточные нейронные сети могут хорошо обобщать: они способны заметить постоянство и определить особенности, характерные для какого-либо класса (чтобы отличить, например, кота от собаки), не обращая внимания на фоновый шум. Таким образом, где-то между подачей изображения на вход и выводом результата классификации этого изображения, стоит модель, которая находит признаки во входных данных. Соответственно, обращаясь к этой самой промежуточной точке (т. е. слоям), можно без труда получить представление стиля и содержания изображения.

Вот как выглядит работа с промежуточными слоями сети:

			# Слой контента, в который помещается карта объектов
content_layers = ['block5_conv2'] 

# Со слоем стиля немного по-другому
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1', 
                'block4_conv1', 
		     'block5_conv1'
               ]

num_content_layers = len(content_layers)
num_style_layers = len(style_layers)
		

Модель

Сначала нужно загрузить VGG19 и подать тензор на вход модели. Это даст возможность получать карты признаков, а впоследствии — представления стиля и контента.

Плюсом VGG19 является её относительная простота (по сравнению с ResNet, Inception и им подобным). Поэтому карты признаков будут лучше подходить для переноса стиля.

Чтобы получить доступ к промежуточным слоям, соответствующим картам признаков стиля и контента, нужно получить характерные выходные данные, используя Keras functional API для определения модели с требуемыми выходными функциями активации.

Благодаря functional API определение модели сводится к банальному определению входных и выходных данных:

			model = Model(inputs, outputs).
def get_model():
  """ Создание модели с доступом к промежуточным слоям 
  
  Эта функция будет подгружать модель VGG19 и давать доступ к промежуточным слоям.
  В дальнейшем эти слои будут использоваться для создания собственной модели для изображения.
  Возвращает данные с промежуточных слоёв VGG19 модели.
 
  """
  # Тут подгружается модель (weights=’imagenet’)
  vgg = tf.keras.applications.vgg19.VGG19(include_top=False, weights='imagenet')
  vgg.trainable = False
  # Получение соответствующих слоёв стиля и контента 
  style_outputs = [vgg.get_layer(name).output for name in style_layers]
  content_outputs = [vgg.get_layer(name).output for name in content_layers]
  model_outputs = style_outputs + content_outputs
  # Построение модели
  return models.Model(vgg.input, model_outputs)
		

В приведённом выше коде подгружается подготовленная сеть классификации изображений. После этого нужно взять необходимые слои, про которые говорилось ранее. Затем нужно определить модель. Это можно сделать, настроив входы для изображения и выходы для слоёв стиля и контента. Таким образом вы сможете создать модель, которая на входе принимает изображение, а на выходе выдаёт промежуточные слои для стиля и контента.

Определение и создание функций потерь (расстояний Lcontent и Lstyle)

Функция потерь для контента

Определить функцию потерь для содержимого на самом деле довольно просто. Нужно передать сети два изображения: изображение желаемого стиля и базовое. После этого вы получите промежуточные слои вашей модели. И единственное, что остаётся, это рассчитать Евклидово расстояние между двумя промежуточными представлениями этих изображений.

Если быть точным, то функция потерь описывает расстояние содержимого (Lcontent) между входным изображением x и изображением контентаp.Пусть тогда Cₙₙ будет предварительно обученной глубокой свёртываемой нейронной сетью. Опять же в этом случае будет использоваться VGG19.

Допустим, X— это любое изображение, тогда Cₙₙ(x) — это сеть, на вход которой подаётся X. Пусть тогда Fˡᵢⱼ(x) ∈ Cₙₙ(x)иPˡᵢⱼ(x) ∈ Cₙₙ(x) описывает соответствующие промежуточные представления объектов сети, принимающей X и P. Тогда Lcontent можно будет рассчитать по следующей формуле:

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

Реализовать это довольно просто. Как и в прошлом случае, на вход нужно подать карту признаков со слоя L сети со входом X, входное изображение и P — изображение контента. На выходе получится расстояние Lcontent.

			def get_content_loss(base_content, target):
  return tf.reduce_mean(tf.square(base_content - target))
		

Функция потерь для стиля

Расчёт функции потерь для стиля немного сложнее, но базируется на том же принципе. В этот раз на вход сети нужно подавать входное изображение и картинку стиля. Но теперь, вместо того чтобы сравнивать «сырые» данные с выходов базового и стиля изображения, нужно сравнить матрицы Грама этих двух выходов.

С математической точки зрения этот процесс заключается в описании функции потерь для стиля главного изображения (X) и изображения стиля (A) и расстояния между представлениями (матрица Грама) стиля этих двух картинок.

Представление стиля картинки можно описать как корреляцию между различными ответами фильтра матрицы Gˡ, где Gˡᵢⱼ — это внутреннее произведение между векторизированной картой признаков i и j в слое L.

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

где Gˡᵢⱼ и Aˡᵢⱼ — это соответствующие представления на слое L входного изображения X и изображения стиля A. Nl описывает количество карт объектов, каждая из которых имеет размер Ml = высота * ширина. Исходя из этого, функция потерь всех слоёв будет такой:

где взвешивается влияние потери каждого слоя от какого-либо фактора wl. В этом случае все слои «взвешиваются» одинаково:

А вот, собственно, и реализация:

			def gram_matrix(input_tensor):
  # Сначала идёт канал изображения
  channels = int(input_tensor.shape[-1])
  a = tf.reshape(input_tensor, [-1, channels])
  n = tf.shape(a)[0]
  gram = tf.matmul(a, a, transpose_a=True)
  return gram / tf.cast(n, tf.float32)
 
def get_style_loss(base_style, gram_target):
  """Принимает два изображения измерений h, w, c"""
  # высота, ширина и количество фильтров в каждом слое
  height, width, channels = base_style.get_shape().as_list()
  gram_style = gram_matrix(base_style)
  
  return tf.reduce_mean(tf.square(gram_style - gram_target))
		

Градиентный спуск

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

Чтобы минимизировать потери при переносе стиля, понадобится оптимизатор Adam. Для минимизации нужно многократно обновлять выходное изображение: не стоит как-либо изменять веса в сети. Вместо этого можно тренировать вход изображения. Чтобы это сделать, нужно понять, каким образом рассчитываются потери и градиенты. Используя Adam, можно понять функциональность autograd/gradient tape в собственных циклах обучения.

Расчёт потери и градиентов

Нужно создать всего лишь одну маленькую функцию, которая будет подгружать изображения стиля и контента, а потом передавать их сети. В будущем это даст представления признаков стиля и контента в модели:

			def get_feature_representations(model, content_path, style_path):
  """Функция, которая рассчитывает признаки стиля и контента
 
  Эта функция будет просто предварительно подгружать и обрабатывать содержимое и стиль. 
  Затем эти представления пройдут через сеть, чтобы получить промежуточные слои.
  
  Аргументы:
    model: Используемая модель.
    content_path: Путь к изображению содержимого.
    style_path: Путь к изображению стиля.
    
  Возвращает:
    Признаки стиля и контента. 
  """
  # Подгрузка изображений
  content_image = load_and_process_img(content_path)
  style_image = load_and_process_img(style_path)
  
  # Одновременная обработка признаков стиля и контента
  stack_images = np.concatenate([style_image, content_image], axis=0)
  model_outputs = model(stack_images)
  
  # Получение представлений признаков 
  style_features = [style_layer[0] for style_layer in model_outputs[:num_style_layers]]
  content_features = [content_layer[1] for content_layer in model_outputs[num_style_layers:]]
  return style_features, content_features
		

Для расчёта градиента тут используется tf.GradientTape. Этот способ даёт преимущество использования автоматического дифференцирования, доступного благодаря трассировке последующих вычислений градиента. В этом случае во время прямого прохода операции кэшируются. Это даст возможность рассчитать потери градиента на обратном проходе.

			def compute_loss(model, loss_weights, init_image, gram_style_features, content_features):
  """Эта функция рассчитывает полную потерю.
  
  Аргументы:
    model: Модель с нужными промежуточными слоями.
    loss_weights: Вес каждого компонента для каждой функции потерь. 
      (вес для стиля, для контента и общий).
    init_image: Первичное изображение. Это то изображение, которое в процессе оптимизации будет обновляться.
    gram_style_features: Предварительные вычисления матрицы Грама соответствующих слоёв.
    content_features: Предварительные вычисления нужных слоёв контента.
      
  Возвращает:
    Общие потери, потери для стиля, контента и вариационные потери
  """
  style_weight, content_weight, total_variation_weight = loss_weights
  
  # Прогонка изображение через модель. Это даст представления контента и стиля.
  # Из-за использования мгновенного выполнения, эта модель вызывается как и любая другая функция.
  model_outputs = model(init_image)
  
  style_output_features = model_outputs[:num_style_layers]
  content_output_features = model_outputs[num_style_layers:]
  
  style_score = 0
  content_score = 0

  # Суммирует потерю стиля со всех слоёв
  # Тут одинаково взвешиваются потери каждого слоя.
  weight_per_style_layer = 1.0 / float(num_style_layers)
  for target_style, comb_style in zip(gram_style_features, style_output_features):
    style_score += weight_per_style_layer * get_style_loss(comb_style[0], target_style)
    
  # Суммирование потерь контента со всех слоёв
  weight_per_content_layer = 1.0 / float(num_content_layers)
  for target_content, comb_content in zip(content_features, content_output_features):
    content_score += weight_per_content_layer* get_content_loss(comb_content[0], target_content)
  
  style_score *= style_weight
  content_score *= content_weight
  total_variation_score = total_variation_weight * total_variation_loss(init_image)

  # Получение суммарной потери
  loss = style_score + content_score + total_variation_score 
  return loss, style_score, content_score, total_variation_score
		

В итоге расчёт градиента сводится к этому:

			def compute_grads(cfg):
  with tf.GradientTape() as tape: 
    all_loss = compute_loss(**cfg)
  # Расчёт градиента изображения
  total_loss = all_loss[0]
  return tape.gradient(total_loss, cfg['init_image']), all_loss
		

Запуск процесса переноса стиля

Вот так выглядит фактический запуск сети:

			def run_style_transfer(content_path, 
                       style_path,
                       num_iterations=1000,
                       content_weight=1e3, 
                       style_weight = 1e-2): 
  display_num = 100
  # В этом случае не нужно обучать каждый слой модели. Поэтому параметр trainability нужно выставить в false.
  model = get_model() 
  for layer in model.layers:
    layer.trainable = False
  
  # Получение представлений признаков стиля и контента (из промежуточных слоёв)
  style_features, content_features = get_feature_representations(model, content_path, style_path)
  gram_style_features = [gram_matrix(style_feature) for style_feature in style_features]
  
  # Загрузка изначального изображения
  init_image = load_and_process_img(content_path)
  init_image = tfe.Variable(init_image, dtype=tf.float32)

  # Создание оптимизатора
  opt = tf.train.AdamOptimizer(learning_rate=10.0)

  # Отображение промежуточных изображений
  iter_count = 1
  
  # Сохранение лучшего результата
  best_loss, best_img = float('inf'), None
  
  # Создание конфигурации 
  loss_weights = (style_weight, content_weight)
  cfg = {
      'model': model,
      'loss_weights': loss_weights,
      'init_image': init_image,
      'gram_style_features': gram_style_features,
      'content_features': content_features
  }
    
  # Отображение
  plt.figure(figsize=(15, 15))
  num_rows = (num_iterations / display_num) // 5
  start_time = time.time()
  global_start = time.time()
  
  norm_means = np.array(1)
  min_vals = -norm_means
  max_vals = 255 - norm_means   
  for i in range(num_iterations):
    grads, all_loss = compute_grads(cfg)
    loss, style_score, content_score = all_loss
    # grads, _ = tf.clip_by_global_norm(grads, 5.0)
    opt.apply_gradients([(grads, init_image)])
    clipped = tf.clip_by_value(init_image, min_vals, max_vals)
    init_image.assign(clipped)
    end_time = time.time() 
    
    if loss < best_loss:
      # Обновление лучшей потери и изображения 
      best_loss = loss
      best_img = init_image.numpy()

    if i % display_num == 0:
      print('Iteration: {}'.format(i))        
      print('Total loss: {:.4e}, ' 
            'style loss: {:.4e}, '
            'content loss: {:.4e}, '
            'time: {:.4f}s'.format(loss, style_score, content_score, time.time() - start_time))
      start_time = time.time()
      
      # Отображение промежуточных изображений
      if iter_count > num_rows * 5: continue 
      plt.subplot(num_rows, 5, iter_count)
      # Используйте метод .numpy(), чтобы получить конкретный numpy-массив
      plot_img = init_image.numpy()
      plot_img = deprocess_img(plot_img)
      plt.imshow(plot_img)
      plt.title('Iteration {}'.format(i + 1))

      iter_count += 1
  print('Total time: {:.4f}s'.format(time.time() - global_start))
      
  return best_img, best_loss
		

На этом всё!

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

			best, best_loss = run_style_transfer(content_path, 
                                     style_path,
                                     verbose=True,
                                     show_intermediates=True)
		

Вот ещё крутые примеры работы сети:

Ключевые моменты

В этой статье были разобраны следующие этапы:

• создание нескольких различных функций потерь и использование обратного распространения для входного изображения;
• для этого использовалась предварительно обученная модель и изученные карты признаков для описания содержимого на изображении;
• функциями потерь в основном являлись вычисления расстояний различных представлений;
• всё это выполнялось благодаря собственной модели и моментальным исполнениям;
• построение модели осуществлялось благодаря Functional API;
• моментальное исполнение позволило динамически работать с тензорами, используя естественный поток управления Python;
• управление тензорами велось напрямую, а это в свою очередь облегчило отладку и работу в целом.