Разработка нейросети, делающей любое изображение более красивым

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

Что такое абстрактная красота

Вы когда нибудь задумывались, почему нам красиво. Казалось бы какую пользу для нашего выживания может давать любование облаками, графиками математических формул, листвой растений или красотой человеческих построек. А польза определённо должна быть, так как удовольствие от созерцания прекрасного нам дают центры потребностей, а это довольно сложные конструкции мозга, для формирования которых в процессе эволюции необходимо подкрепление. Те, кто не любуются красотой, должны таки умирать почаще любующихся. И некоторое время назад учёные в этом вопросе наконец-то разобрались, вот ссылки на научную статью и видео-лекции от Евгении Тимоновой:
https://naked-science.ru/article/nakedscience/what-is-beauty
https://www.youtube.com/watch?v=6kzud41yEGc
Если попытаться максимально обобщить прочитанное то красиво то, что предсказуемо. Точнее предсказуемое будет просто успокаивать, а вот удовольствие наши центры потребностей в виде выброса дофамина дают в тот момент, когда мозг научился предсказывать что-то новое. Например, верхнюю часть листика по нижней, или середину трели соловья по её началу. А предсказывать это нам нужно, чтобы обращать внимание на непредсказуемое. Вот заметила часть нашего мозга, что в данном месте должно быть продолжение листика, а там что-то другое, обратил человек на это внимание, а там клык гепарда. А в результате осталась кошечка без 1.5 кг вкуснейших мозговых калорий и в муках умерла от голода.

Идея нейросети

Для начала попробуем сделать картинку более предсказуемой.
Чтобы определить предсказуемость вначале сделаем и обучим нейросеть, которая предсказывает по окружению(128,128,3) его центр (32,32,3). Предсказуемость определим как попиксельную разницу между оригинальной картинкой и картинкой составленной из предсказанных центров.
Однако, если опираться только на предсказуемость, то ожидаемо всегда будет получаться "квадрат Малевича", хоть и не всегда чёрный. Поэтому картинка должна в какой-то степени стремится походить и на оригинал.
Функцию потерь определим как k1*MAE(X,Y)+k2*MAE(Y,Z) где:
X - батч оригинальных картинок
Y - батч улучшенных картинок
Z - батч из картинок, полученных из Y последовательным предсказанием центров
k1,k2 - коэффициенты, определяющие соотношение похожести на оригинал и предсказуемости.

Реализация

Так как мы должны использовать другую нейросеть в функции ошибок, то использовать стандартный цикл обучения model.fit() не получится, придётся писать его самому.
Вот наиболее значимая часть цикла опроса:

# Перебираем батчи, запуская шаг тренировки, иногда сохраняем веса
def train(self, generator, model_name, save_period):
    for i in range (0, len(generator)):
      x = generator[i]
      self.train_step(x)
      if i%save_period == 0:
        self.improver_model.save_weights(model_name)

# Шаг тренировки        
def train_step(self, x):
  with tf.GradientTape() as improver_tape:  
    y = self.improver_model(x, training=True) # инференс модели
    pr_loss = self.predicting_loss(x, y) # выполняем функцию ошибки
    self.aver_loss += pr_loss
  # Получаем градиенты   
  improve_gradients = improver_tape.gradient(pr_loss, self.improver_model.trainable_variables)
  # Вычитаем градиенты из весов
  self.optimizer.apply_gradients(zip(improve_gradients, self.improver_model.trainable_variables))

# Функция ошибки
def predicting_loss(self, x, y):
  # Последовательным предсказыванием центров создаём изображение
  full_pred = self._create_full_pred(y, self.PIC_SIZE)
  my_mse = tf.keras.losses.MeanAbsoluteError()
  return 0.5*my_mse(y,full_pred) + 0.5*my_mse(x,y)

Это ноутбук с моделью, предсказывающей центр:
https://colab.research.google.com/drive/1aZPTblFUTxHs6B5zL31iok-mOUirapKI?usp=sharing

Далее ноутбуки находятся на google, но обучались на Яндексе, так как для двух моделей на бесплатной версии не хватает памяти, а платная недоступна. Но на Яндекс нельзя открывать ссылку для всех. Из за этого же я пока остановился, так как на Яндекс вычисления стоят очень дорого. После копирования на google, результаты в ноутбуках иногда вместо оригинальных изображений видны как синие квадраты.

Это ноутбук с моделью, улучшающей изображения(Основная):
https://drive.google.com/file/d/1KPIMvKegTXBo0Aq_25GxT9z1e5iOoTxV/view?usp=sharing

Результат сейчас примерно такой:

Потом у меня возникла мысль,что будет, если я попробую максимизировать разницу между изображением полученным с помощью хорошо и плохо обученного предсказателя центра.
Функция потерь будет MAE(Y,Zп)/MAE(Y,Zх) где:
X - батч оригинальных картинок
Y - батч улучшенных картинок
Zп - батч из картинок, полученных из Y последовательным предсказанием центров(плохим предсказателем)
Zх - батч из картинок, полученных из Y последовательным предсказанием центров(хорошим предсказателем)
Как видно здесь, я вообще не использовал исходное изображение в функции потерь.

Ноутбук с моделью: https://drive.google.com/file/d/1PXKosb6Nvd8ka5dJ9kYNYFwNaa0EAJsn/view?usp=sharing
Результат:

Далее я попробовал сделать более сложную Основную модель(до этого сама модель была очень простой), но падение loss на ней происходило очень медленно и мне уже стало жалко тратить на это деньги. Со временем, я как-нибудь постараюсь довести эту идею до более интересных результатов. Ну и жду Ваших предложений. Если у кого есть свободные мощности для обучения, можете поэкспериментировать и отправить мне результаты. Буду очень благодарен.
Вот ссылка на ноутбук со сложной моделью:
https://drive.google.com/file/d/1Kfa2-BJ4QgUnggQWvXgzeCMvPllW6eXS/view?usp=sharing

Если кому интересно как использовать разработанные нейросети в приложениях на с++, в том числе на мобильных устройствах, можете прочитать мою статью: https://habr.com/ru/articles/719888/