Single Image Super Resolution
Введение
Super Resolution — задача в области компьютерного зрения, направленная на восстановление изображения или видео высокого разрешения из версий низкого разрешения. Существует много различных методов ее решения — это и классические, и обучаемые методы. Кроме того, сама задача может формулироваться по-разному. Например, для спутниковых снимков необходимо решать many-to-one задачу, или задачу отображения множества изображений низкого разрешения в одно или несколько изображений высокого разрешения, а для видео — many-to-many задачу.
В этом обзоре мы рассмотрим one-to-one задачу, так называемую Single Image Super Resolution (далее SISR). Изучим популярные метрики оценки качества, используемые функции потерь и подходы к проектированию моделей, а также разберем самого популярного представителя этой задачи — семейство моделей SRGAN.
План статьи следующий:
- метрики качества;
- функции потерь;
- подходы к проектированию;
- семейство моделей SRGAN;
- итоги.
Давайте начинать 😉
Метрики качества
В оценке качества изображений есть много аспектов, таких как резкость, контраст и отсутствие шума. Поэтому справедливая оценка — сложная задача. Все методы можно разделить на субъективные и объективные. Субъективные методы основаны на естественном восприятии изображений человеком, в то время как объективные методы представляют собой количественные подходы к оценке качества изображения. В этой главе мы рассмотрим популярные в применении методы для задачи SISR.
PSNR
Peak Signal-to-Noise Ratio (PSNR) — логарифмическая величина в децибелах, которая измеряет отношение максимально возможной силы сигнала к уровню шума.
Более высокое значение PSNR указывает на более высокое качество восстановленного изображения.
Формула для вычисления PSNR между исходным и восстановленным изображениями выглядит следующим образом:
PSNR фокусируется на разнице, выявленной на уровне пикселей, из-за чего слабо коррелирует с визуальным восприятием. Например, небольшой шифт в пикселях может привести к большой деградации метрики, хотя на глаз шифт едва заметен.
SSIM
Structural Similarity Index (SSIM) — метрика, которая оценивает степень сходства двух изображений, учитывая их яркость, контраст и структуру.
Она лучше отражает визуальное восприятие человека и точнее оценивает структурные характеристики изображений в сравнении, например, с PSNR. Значения SSIM лежат в отрезке от 0 до 1, где 1 означает идеальное сходство, а 0 — полное несоответствие.
SSIM оценивает яркость как среднее значение пикселей изображения, а контраст как стандартное отклонение:
\( \mu_x = \frac{1}{N} \sum_i x_i \\ \sigma_x = \frac{1}{N-1} \sum_i (y_i — \mu(x))^2 \)
Для их сравнения авторы SSIM вводят функцию сходства \( S \):
\( S(x, y, c) = \frac{2 * x * y + c}{x^2 + y^2 + c} \)
Где \( x \) и \( y \) — сравниваемые величины, а \( c \) — константа для избежания нестабильностей.
Для сравнения же структур исходного и восстановленного изображений авторы предлагают использовать ковариацию:
\( \sigma_{y, y^*} = \frac{1}{N-1} \sum_i (y — \mu(y)) * (y^* — \mu(y^*)) \)
Окончательная формула для SSIM выглядит следующим образом:
\( SSIM(y, y^*) = S(\mu_y, \mu_{y^*}, c_1) * S(\sigma_y, \sigma_{y^*}, c_2) * \frac{\sigma_{y, y^*} +c_3}{\sigma_y * \sigma_{y^*} + c_3} \)
LPQ
Learning-Based Perceptual Quality (LPQ) — семейство подходов, в основе которых лежат попытки использовать deep learning для оценки качества.
Например, при первых попытках исследователи обучали простые сверточные сети предсказывать оценки фотографиям, которые были размечены как средняя оценка людей от 1 (плохое качество) до 5 (отличное качество). Минус такого подхода — сложность в разметке большой выборки фотографий для обучения сети. Одно из возможных решений проблемы — подходы, основанные на ранжировании фотографий и их деградированных версий (деградация осуществляется за счет добавления шума и других техник). Примерами таких метрик являются dipIQ и RankIQA.
Сегодня самым популярным и эффективным считается подход, основанный на сравнении признаков, извлеченных из предобученных сетей. Ярким представителем здесь является LPIPS, где авторы показывают, что их метрика лучше коррелирует с человеческими суждениями в сравнении с PSNR или SSIM. Однако качество этого метода зависит от сети, из которой извлекаются признаки.
Функции потерь
Цель SISR — создать модель, которая преобразует изображение с низким разрешением в изображение с высоким разрешением. Следовательно, возникает вопрос: как обучать такую модель?
Pixel loss
Как видно из названия, эта функция потерь измеряет разницу между исходным и восстановленным изображениями в виде разницы пикселей. Самые популярные функции — это L1 и L2:
\( {L}_{L1}(y, y^*) = \frac{1}{N}\sum_i |y_i — y_i^*| \\{L}_{L2}(y, y^*) = \frac{1}{N}\sum_i (y_i — y_i^*)^2 \)
На рисунке ниже представлено сравнение функции потерь MSE (L2) между исходным изображением и его различными версиями деградации. Мы видим, как она слабо коррелирует с визуальным качеством, что приводит к субоптимальным результатам. Например, сети, обученные с такими функциями потерь, склонны выдавать размытые изображения.
Content loss
Качество визуального восприятия, как упоминалось ранее, имеет важное значение для оценки качества моделей. Поэтому была предложена функция потерь, которая измеряет расстояние между изображениями в пространстве признаков предобученной сети, что позволяет сравнивать картинки не с точки зрения индивидуальных пикселей, а со стороны контента и структуры.
Более формально: пусть \( \phi \) — сверточная нейронная сеть, например, VGG, а \( \phi^l \) — признаки, извлеченные после \( l \)-го слоя. Тогда функцию потерь можно записать следующим образом:
\( {L}_{Content}(y, y^*, l) = \sum_l||\phi^l(y) — \phi^l(y^*)||_2 \)
Первое упоминание этой функции потерь датируется мартом 2016-го года: авторы статьи Perceptual Losses for Real-Time Style Transfer and Super-Resolution предложили использовать VGG-16, предобученную на задачу классификации на датасете ImageNet для извлечения признаков. В качестве \( l \) они взяли слой \( \text{relu3-3} \), но в последующих статьях не ограничивались одним слоем.
В 2018-ом году исследователи из UC Berkeley и Adobe Research предложили функцию потерь LPIPS, где авторы сделали два изменения:
- Собрали свой датасет BAPPS — Berkeley Adobe Perceptual Patch Similarity, куда положили 20к патчей размером 64×64 из рандомных изображений, применили к ним разные виды деградаций (от простых шумов до NN based) и разметили в триплет-режиме, какая из двух картинок с разной силой деградации ближе к исходной.
- Обучили дополнительные слои поверх извлеченных признаков из VGG решать задачу ранжирования с помощью размеченных триплетов.
Следующее развитие LPIPS — метрика DreamSim. Здесь авторы сделали следующее:
- Вместо датасета BAPPS собрали NIGHTS — Novel Image Generations with Human-Tested Similarities. В нем теперь не рандомные патчи размером 64х64, а 60 тысяч полноценных изображений размером 224х224, сгенерированных через Stable Diffusion. Деградации для формирования триплетов генерируются за счет аугментаций промптов Stable Diffusion. Авторы делают упор на то, что такие деградации позволяют обращать внимание обучаемых слоев не только на текстуру изображений, но и на контекст.
- Вместо VGG-16 извлекли признаки из DINO, CLIP и OpenCLIP.
Adversarial loss
Суть этой функции потерь заключается в использовании двух различных нейронных сетей: генератора и дискриминатора. Задача генератора — генерировать изображения, максимально похожие на данные из реального набора, и пытаться обмануть дискриминатор. В свою очередь, дискриминатор стремится отличить реальные изображения от тех, что были сгенерированы генератором. В контексте задачи SISR генератор выступает в роли Super Resolution модели, которая создает изображения с высоким разрешением из изображений с низким разрешением.
Адверсариал функция потерь основана на кросс-энтропии:
\( {L}_G^{CE} = \mathbb{E}_{x\sim\mathbb{P}_{LR}(x)} [\log(1 — D_\theta[G_\phi(x)])] \\ {L}_D^{CE} = -\mathbb{E}_{y\sim\mathbb{P}_{HR}(y)} [\log D_\theta(y)] -\mathbb{E}_{x\sim\mathbb{P}_{LR}(x)} [\log(1 — D_\theta[G_\phi(x)])] \)
Где \( G_\phi \) и \( D_\theta \) — генератор и дискриминатор, \( x \) — изображение низкого разрешения, \( y \) — изображение высокого разрешения.
Альтернатива кросс-энтропии — метод наименьших квадратов (Least Square), который приводит к более стабильному обучению:
\( {L}_G^{CE} = \mathbb{E}_{x\sim\mathbb{P}_{LR}(x)} (D_\theta[G_\phi(x)]-1)^2 \\ {L}_D^{CE} = \mathbb{E}_{y\sim\mathbb{P}_{HR}(y)} (D_\theta(y)-1)^2 +\mathbb{E}_{x\sim\mathbb{P}_{LR}(x)} (D_\theta[G_\phi(x)])^2 \)
Существуют более продвинутые вариации, такие как Вассерштейн лосс с градиентным штрафом (WGAN-GP), решающие проблемы сходимости и стабильности. Несмотря на большее разнообразие вариаций адверсариал лоссов, часто в задачах Super Resolution ограничиваются кросс-энтропией или методом наименьших квадратов.
Подходы к проектированию
В этом разделе обсудим идеи и тенденции проектирования архитектур в задаче SISR.
В зависимости от типа архитектур различаются потребление памяти и эффективность обучения. Это критически важно при решении задачи на высоких разрешениях.
В основном выделяют три типа архитектур:
- Pre-Upsampling — изображение низкого разрешения интерполируется вверх, а затем применяется последовательность сверточных слоев, выполняющих роль ‘рефайнера’. Такой подход хоть и более требователен к потреблению памяти, зато сходится быстрее и обучается проще, чем остальные типы архитектур.
- Post-Upsampling — последовательность сверточных слоев применяется к изображению низкого разрешения, а затем интерполируется вверх. В отличие от предыдущего подхода, этот метод менее требователен к потреблению памяти и позволяет использовать более глубокие сети и больший размер батч сайза при обучении.
- Progressive-Upsampling — комбинация последовательных маленьких Post-Upsampling сетей. С одной стороны, это менее ресурсозатратный подход в сравнении с Pre-Upsampling, но более эффективный, чем Post-Upsampling.
Семейство моделей SRGAN
Пожалуй, Real-ESRGAN — самая известная публичная модель для решения SISR. Если вам нужно завести бейзлайн для ваших данных — советуем начать с нее. Но для начала нужно разобраться с ее устройством, ведь ей предшествуют две другие версии — SRGAN и ESRGAN. Давайте рассмотрим эволюцию и устройство каждой из этих моделей.
SRGAN
Впервые использовать Content loss и Adversarial loss в задаче SISR предложили авторы статьи SRGAN в 2017-ом году. До этого все работы в основном были сосредоточены на минимизации pixel loss, что приводило к размытым реконструкциям и не позволяло добиваться высокочастотных деталей. Кроме того, они были первыми, кто предложил архитектуру для решения задачи x4 SISR. Давайте рассмотрим эту архитектуру и пайплайн обучения 🙂
В основе архитектуры генератора заложены B (B=16) residual блоки, каждый из которых состоит из двух последовательностей свертки k3n64s1 и батч-нормализации, между которыми стоит нелинейность ParametricReLU. На входе генератора — входная свертка k9n64s1 и нелинейность ParametricReLU, а на выходе — две последовательности сверток k3n256s1, PixelShuffle x2 слоя и ParametricReLU, за которыми следует финальная свертка k9n3s1. PixelShuffle x2 представляет собой необучаемую операцию для повышения разрешения тензора. Например, если размерность тензора (B, C*r*r, H, W), то после применения PixelShuffle x2 она станет (B, C, H*r, W*r). Такая архитектура относится к типу Post-Upsampling, где роль интерполяции вверх выполняют два слоя PixelShuffle x2.
Дискриминатор же состоит из 8 сверток с ядром размером 3х3. Четные свертки имеют страйд 2 для уменьшения пространственной размерности тензора в два раза, а нечетные — страйд 1. Каждая пара сверток увеличивает количество каналов в два раза. Между свертками находятся батч-нормализации и нелинейность LeakyReLU с \( \alpha=0.2 \). На выходе дискриминатора — два линейных слоя с нелинейностью и сигмоидой для предсказания вероятности.
В качестве датасета для обучения использовали рандомные патчи размером 96х96 из 350к картинок датасета ImageNet, а для получения изображений низкого разрешения — bicubic интерполяцию x4 вниз.
Во время обучения авторы предлагают использовать комбинацию Perceptual loss (признаки извлекаются из VGG-19) и Adversarial loss:
\( {L} = {L}_{Perceptual} + 10^{-3} * {L}_{Adversarial} \)
На рисунке ниже представлены результаты сравнения SRGAN и его предшественников SRResNet, обученных с помощью Pixel loss. Стоит отметить: SRResNet страдает размытыми результатами, что вполне закономерно, если учитывать его функцию потерь. SRGAN же вносит намного больше деталей, и его результаты выглядят визуально более приятными.
ESRGAN
В этой статье авторы решили улучшить визуальное качество предыдущего подхода, изучив три главных компонента — архитектуру, adversarial loss и perceptual loss.
Они предложили два изменения в архитектуре ESRGAN:
- убрали все батч-нормализации для улучшения качества генератора. Эмпирически обнаружили, что батч-нормализация приводит к нежелательными артефактам (что подтверждали и последующие многие статьи по GAN’ам). На рисунке 12 мы видим, как генератор с батч-нормализациями учится нестабильно, поэтому нужно тщательно подбирать итерацию обучения, где артефактов меньше.
- оставили верхнеуровнево архитектуру такой же, но добавили новые базовые блоки — Reidual in Residual Dense Block (RRDB). Поскольку увеличение слоев и добавление residual коннекшенов улучшает качество генератора, авторы предложили блок с большим количеством слоев и коннекшенов.
Вдохновившись идеями Relativistic GAN, авторы предложили предсказывать дискриминатором не вероятность реалистичности изображения, а степень его реалистичности по отношению к другому изображению. Если конкретнее: заменили \( D(x) = \sigma(C(x)) \) на \( D(x, y) = \sigma(C(x) — \mathbb{E} [C(y)]) \), где \( C(y) \) — ответы дискриминатора до слоя сигмоиды, а \( \mathbb{E}[C(y)] \) — среднее ответов по батчу \( y \). Тогда adversarial лосс становится следующим:
\( {L}_G = -\mathbb{E}_{y\sim\mathbb{P}_{HR}(y)} [\log(1 — D_\theta[y, G_\phi(x)])] + \mathbb{E}_{x\sim\mathbb{P}_{LR}(x)} [\log(1 — D_\theta[G_\phi(x), y])] \\ {L}_D = -\mathbb{E}_{y\sim\mathbb{P}_{HR}(y)} [\log D_\theta(y, x)] -\mathbb{E}_{x\sim\mathbb{P}_{LR}(x)} [\log(1 — D_\theta[G_\phi(x), y])] \)
Такая модификация дискриминатора помогает генератору выучивать более четкие края и детализированные текстуры.
Последнее изменение связано с извлечением признаков из VGG-19 до активаций. В оригинале было предложено извлекать признаки после активаций ReLU, но авторы ESRGAN обнаружили — такие признаки более разряженные, следовательно, несут меньше информации о входной картинке. Это может оказаться критичным при обучении генератора. Разряженность признаков до и после активаций ReLU показана на рисунке 16.
REAL-ESRGAN
Прежде, чем мы обсудим последнюю работу этой серии, давайте остановимся на важном моменте в обучении SISR. Мы уже упомянули, что при обучении изображения на входе в модель подвергаются bicubic интерполяции, но что происходит на инференсе? Здесь мы подаем в модель изображения из реальной жизни, которые подверглись неизвестным нам искажениям. Например, фотографируя на камеру телефона, мы случайными движениями руки создаем motion блюр, а сама камера накладывает sensor noise. Это приводит к трейн-тест шифту, и на инференсе мы наблюдаем качество на порядок хуже, чем на трейне. Большое количество статей посвящено обучению модели, робастной к изображениям из реального мира. REAL-ESRGAN — не исключение. Авторы задались этим же вопросом.
Вместо bicubic интерполяции авторы последовательно используют следующие операции дважды:
- применение изотропного или анизотропного фильтра Гаусса (в простонародье — блюр);
- интерполяция вниз рандомным методом: bicubic, bilinear, area;
- прибавление цветного или черно-белого шума Гаусса или Пуассона;
- добавление JPEG компрессии.
Еще одно важное изменение — архитектура дискриминатора. Поскольку генератору нужно генерализовать свои способности на более сложные деградации, архитектура дискриминатора тоже должна стать более сложной. Зачем? Чтобы градиенты, текущие от дискриминатора, помогали генератору справляться с реалистичными локальными текстурами, а не только с реалистичным глобальным стилем фотографии. Для этого мы используем Unet-подобный дискриминатор со спектральными нормализациями, а adversarial loss применяем к каждому пикселю в выходной одноканальной карте дискриминатора. Эмпирически было найдено, что спектральные нормализации помогают стабилизировать обучение.
Итоги
Задача Super Resolution очень обширна. С каждым годом появляется все больше и больше работ, разобраться в которых становится все сложнее 🙂. В нашем обзоре мы рассмотрели задачу SISR, применяемые в ней метрики и функции потерь, возможные подходы к проектированию архитектур, а также разобрали семейство моделей SRGAN. Стоит отметить: из-за большого ажиотажа вокруг диффузионных моделей в последние годы появилось немало работ по их применению к задачам SISR. Однако недавняя статья GigaGan показала — в задачах SISR диффузионные модели не лучше GAN’ов. Здесь ключевой рецепт для их решения — чистый датасет.