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

К счастью, я знаю, что архитектуры U-net обычно хорошо подходят для такого рода задач. Идея U-net довольно близка к автоэнкодерам: у вас есть часть нейронной сети для извлечения признаков, за которой следует часть реконструкции. А поскольку слои свертки/деконволюции используются на всем протяжении сети, выходные пиксели не являются независимыми друг от друга.

Вот как будет выглядеть нейронная сеть:

Сам проект будет состоять из 3 основных частей, которые будут реализованы:

  • Создайте конвейер данных для обучения и проверки
  • Спроектировать и обучить нейронную сеть
  • Результаты теста визуально

Итак, начнем.

Конвейер данных:

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

Для собственного удобства я организовал поток данных «на лету»: загружается одно большое изображение, разбивается на кучу изображений размером 64x64 (или 128x128, или сколько угодно) и передается в нейросеть. Этот процесс будет повторяться для всех изображений, которые у меня есть. Конвейер «на лету» упрощает эксперименты: разные размеры изображений, разные шаги для разбиения, разные уровни шума и так далее.

Вот как этот пайплайн будет выглядеть в Python:

Шум также добавляется на лету: определенный процент пикселей будет заполнен случайными «темными пикселями» размером от 1x1 до 2x2. Определенно не 100% натуральный шум, но достаточно хороший шум для игрушек на 1 вечер.

Теперь, когда у меня есть генератор 2-в-1, совместимый с наборами данных TensorFlow, я могу легко настроить его и использовать для реального обучения:

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

Архитектура нейронной сети:

Как я уже говорил выше, нейронная сеть для такого рода задач — это, как правило, архитектура U-net. Что-то, что можно было бы назвать «Сверточный автоэнкодер» — здесь у меня есть структура с двумя ветвями: ветвь X для больших карт объектов, ветвь Y для меньших карт функций. В конце концов они объединяются и используются для окончательной реконструкции изображения.

Как видите, некоторые аргументы ссылаются на словарь params, я использовал его для настройки гиперпараметров. Вот пример параметров:

model = build_model({'lr': 0.001,
                     'conv_activation': 'swish',
                     'last_activation': 'tanh',
                     'out_activation': 'sigmoid',
                     'loss': 'mean_squared_error'})

Теперь, когда все биты установлены — давайте нажмем кнопку «Выполнить», и примерно через 4714 секунд (используя RTX 2080 Ti) вы увидите, что эпоха закончилась.

76200/76200 — 4714s 62ms/step — loss: 5.0034e-04 — val_loss: 2.3591e-04

Проверка модели:

Для быстрых экспериментов с обученной моделью я создал минималистичный блокнот Jupyter (доступен в репозитории GitHub, ссылка внизу). Jupyter очень хорош для таких вещей: можно просто отредактировать 1 строку в 1 ячейке и повторно запустить только эту конкретную ячейку — и это именно то, что мне нужно здесь для «визуальной проверки».

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

Я думаю, что можно с уверенностью сказать, что результаты не идеальны (поскольку мы можем видеть артефакты даже при уровне шума 10%), но с другой стороны — эй, это сеть на 100 тысяч параметров, способная очистить 256 участков изображения за 10–15 миллисекунд! Таким образом, применив некоторую дополнительную настройку, мы, вероятно, могли бы уменьшить его до 30-50 тысяч параметров и 3-7 миллисекунд.

Спасибо за чтение. Не стесняйтесь обращаться ко мне, если у вас есть какие-либо вопросы :)

Ссылка на репозиторий с исходным кодом проекта: https://github.com/raver119/denoiser