Еще один алгоритм для восстановления смазанных изображений

Доброго времени суток. Уже столько сказано о методах деконволюции изображений, кажется добавить больше нечего. Однако всегда найдется алгоритм лучше и новее предыдущих. Не так давно был описан итерационный алгоритм, имеющий линейную скорость сходимости при малых затратах памяти, стабильный и хорошо распараллеливаемый. А через некоторое время он был улучшен еще и до квадратичной сходимости. Встречайте: (Fast) Iterative Shrinkage-Thresholding Algorithm.



Постановка задачи


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

$\mathbf{y}=\mathbf{Ax}+ \mathbf{n}$



Матрица $\mathbf{A}$ является матрицей свертки, и с физической точки зрения показывает взаимосвязь между точками исходных и испорченных данных. В случае с изображениями, эту матрицу можно получить преобразовав ядро свертки. Вектора $\mathbf{x}$ и $\mathbf{y}$ являются исходным и поврежденным изображением в векторизованной форме. Иначе говоря, все столбцы пикселей в изображении конкатенируются одно за другим.

Первая проблема такой постановки задачи: количество выделяемой памяти. Если изображение имеет размеры 640х480, то задача хранит два вектора изображений размерами 307200х1 и матрицу размером 307200х307200 элементов. Матрица свертки разрежена, однако ее размеры даже в разреженной форме будут большими, и займут много памяти. Произведение с матрицей свертки можно заменить на двумерную свертку с ядром, что не требует хранения большой матрицы в памяти, этим мы решим проблему с хранением матрицы $\mathbf{A}$.
Для поиска изображения, максимально близкого к исходному, запишем оптимизируемую функцию в классическом виде второй нормы невязки между поврежденным и искомым изображениями.

$min_{\hat{\mathbf{x}}} \mid\mid \mathbf{y-A\hat{x}} \mid\mid_2$



Отличие от классических методов


Такая постановка задачи является одной из самых популярных, однако мы ее дополним для улучшения производительности. Использованный подход называется «Majorization-maximization», и заключается он в подмене исходной задачи на другую, очень похожую. Новая задача обладает двумя свойствами:

  1. Задача значительно проще
  2. Во всех точках кроме текущей имеет невязку большую, чем исходная задача

Данное действие происходит на каждой итерации алгоритма. Звучит достаточно сложно, гораздо сложнее, чем это выглядит на самом деле. Итоговая функция минимизации для итерации $k$ записана следующим образом:

$\mathbf{J_k(x)}= \mid\mid \mathbf{y-A\hat{x}} \mid\mid_2 + (\mathbf{\hat{x}-\hat{x_k}})^T (\alpha\mathbf{I}- \mathbf{A^TA}) (\mathbf{\hat{x}-\hat{x_k}})$



Матрица правого слагаемого положительно определена. Это достигается тем, что $\alpha>maxeig(\mathbf{A^TA})$. Функция минимизации состоит из суммы двух величин, каждая из которых неотрицательна. При этом, в точке $\mathbf{\hat{x}=\hat{x}_k}$ второе слагаемое равно нулю, благодаря чему выполняется второе условие из списка.

Поиск решения


Поиск решения начнем классическим методом: раскроем скобки, выпишем градиент и приравниванием его к нулю. Важно заметить, это градиент на $k$-той итерации. Внимание много математики!

$\mathbf{J_k(\hat{x})}=( \mathbf{y-A\hat{x}})^T ( \mathbf{y-A\hat{x}} )+ (\mathbf{\hat{x}-\hat{x}_k})^T \cdot(\alpha\mathbf{I}- \mathbf{A^TA})\cdot (\mathbf{\hat{x}-\hat{x}_k})$


$\mathbf{J_k(\hat{x})}=\mathbf{y^Ty}+\mathbf{\hat{x}^TA^TA\hat{x}} -2\mathbf{y^TA\hat{x}} -\\-\mathbf{\hat{x}^T(\alpha\mathbf{I}- \mathbf{A^TA})\hat{x}}+\mathbf{\hat{x}_k^T(\alpha\mathbf{I}- \mathbf{A^TA})\hat{x}_k} -2\mathbf{\hat{x}_k^T(\alpha\mathbf{I}- \mathbf{A^TA})\hat{x}} $



$\frac{\mathbf{J_k(\hat{x})}}{\delta\mathbf{\hat{x}}}=\mathbf{A^TA\hat{x}} -2\mathbf{A^Ty}+\mathbf{2\alpha\mathbf{I}\hat{x}- \mathbf{A^TA}\hat{x}}-2\mathbf{(\alpha\mathbf{I}- \mathbf{A^TA})\hat{x}_k} $



В итоге производная будет иметь следующий вид:

$\delta J_k(x)=2\mathbf{A^Ty}-2(\alpha\mathbf{I-A^TA})\mathbf{\hat{x}_k}+2\alpha\mathbf{\hat{x}}$



Следующий классический шаг в данной ситуации — приравнять градиент к нулю, и выразить из полученного выражения искомый вектор $\mathbf{\hat{x}}$. Записанный $\mathbf{\hat{x}}$ определим как $\mathbf{\hat{x}_{n+1}}$ или как решение на следующей итерации.

$\delta J_k(x)=2\mathbf{A^Ty}-2(\alpha\mathbf{I-A^TA})\mathbf{\hat{x}_k}+2\alpha\mathbf{\hat{x}}=0$


$\mathbf{A^Ty}+\alpha\mathbf{I}\mathbf{\hat{x}_k}-\mathbf{A^TA}\mathbf{\hat{x}_k}=\alpha\mathbf{\hat{x}}$


$\mathbf{\hat{x}_{k+1}}= \mathbf{\hat{x}_k}+\frac{1}{\alpha} \cdot\mathbf{A^T (y-A\hat{x_k})}$



Выписанное в результате выражение называется «Landweber iteration». Это основная формула между итерациями. Используя его, можно гарантировать уменьшение невязки от итерации к итерации с линейной скоростью.

Базовый алгоритм содержит дополнительный шаг, называемый «soft-threshold», и предполагающий разреженность решения. Данный шаг приравнивает нулю все значения искомого вектора по модулю меньшие, чем установленное значение. Это уменьшает влияние шума на результат восстановления.

Как это выглядит


Как видно из решения, у нас есть два параметра, которые мы регулируем на свой вкус. $\lambda$ отвечает за точность приближения, ограничивая сходимость порогом. $\alpha$ отвечает за стабильность работы и скорость сходимости алгоритма.

$\mathbf{\hat{x}_{k+1}}=soft(\mathbf{\hat{x}_k}+\frac{1}{\alpha} \cdot\mathbf{A^T (y-A\hat{x_k})},\lambda/\alpha)$



Модификация алгоритма


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

  1. Задать параметры $\lambda, \alpha $
  2. Задать $t_1=1, t_0=1,\mathbf{y_{temp}=\hat{x}_0=A^Ty}$
  3. Произвести итерацию базового алгоритма $\mathbf{\hat{x}_{k+1}}= soft(\mathbf{y_{temp}}+\frac{1}{\alpha} \cdot\mathbf{A^T (y-Ay_{temp})},\lambda/\alpha)$
  4. Обновить временные переменные
    $t_0=t_1; t_1=0.5+\sqrt{0.25+t_1^2}$
  5. Добавить сопряженную компоненту к переменной $\mathbf{y_{temp}=\hat{x}_k}+\frac{t_0-1}{t_1}\cdot(\mathbf{\hat{x}_k-\hat{x}_{k-1}})$

Теперь о параллельности алгоритма. Цикл для вычисления порога на каждой итерации можно выразить через поэлементные произведения (произведение Адамара) и операции сравнения, которые включены как в OpenCL так и в CUDA, так что их можно легко распараллелить.

Я реализовал алгоритм на Matlab в двух вариантах: для вычисления на CPU и на GPU. В дальнейшем я стал использовать только версию кода для GPU, так как она удобнее. Код представлен ниже.

Код здесь
function [x] = fista_gpu(y,H,Ht,lambda,alpha,Nit)
x=Ht(y);
y_k=x;
t_1=1;
T=lambda/alpha;
for k = 1:Nit
x_old=x;
x=soft_gpu((y_k+(1/alpha)*Ht(y-H(y_k))), T);
t_0=t_1-1;    
t_1=0.5+sqrt(0.25+t_1^2);
y_k=x+(t_0/t_1)*(x-x_old);
end
end
function [y] = soft_gpu(x,T)
eq=ge(abs(x),abs(T));
y=eq.*sign(x).*(abs(x)-abs(T));
end


Давайте спустимся ближе к практике, и протестируем производительность алгоритма в зависимости от размера картинки.

Конфигурация моего ноутбука
Intel core I3-4005U
NVIDIA GeForce 820M

Слева представлен график зависимости работы алгоритма на 100 итераций, в зависимости от количества пикселей для CPU и GPU.

  • Как видно из результатов, алгоритм на GPU практически не зависит от размеров задачи и для действительно больших картинок мы ограничены только памятью GPU. Полсекунды для 2000х2000 (без учета выгрузки из памяти) достойный результат, не находите?
  • Справа представлена зависимость затраченного алгоритмом времени от количества итераций. При увеличении количества итераций, алгоритм на GPU начинает реагировать увеличением времени работы. Вероятнее всего, это связано с моими ошибками написания кода, либо принудительным понижением частоты работы GPU. В большинстве случаев хватает 100 итераций.




На следующих графиках представлено сравнение базового и улучшенного алгоритма, а так же зависимости сходимости алгоритма от $\alpha$ и $\lambda$ параметров.

  • Как видно из графика слева, улучшенный алгоритм сходится значительно быстрее, чем базовый, и после ста итераций практически не показывает увеличения точности. Базовый алгоритм лишь к пятисотой итерации показывает приемлемый результат.
  • По центральному графику видно, в зависимости от $\lambda $ меняется порог сходимости алгоритма, что ограничивает производительность. Это необходимый компромисс, если сигнал сильно зашумлен.
  • По правому графику видно, что при увеличении $\alpha $ алгоритм сходится гораздо медленнее. Аналогично предыдущему случаю, присутствует компромисс между «качеством» матрицы свертки и скоростью сходимости.



Ну и напоследок пример работы алгоритма на FHD картинке,

Исходная картинка


Испорченная и зашумленная картинка


Восстановленная картинка


Как видите результат достаточно хороший, особенно если он занимает всего 15 секунд из которых 1.5 это работа алгоритма и 13.5 выгрузка картинки из памяти GPU.
Весь использованный для алгоритма код выложен в GitHub.

Использованная литература


  • I. Selesnick. (2009) About: Sparse signal restoration.
  • A. Beck and M. Teboulle Fast Gradient-Based Algorithms for Constrained Total Variation Image Denoising and Deblurring Problems IEEE TRANSACTIONS ON IMAGE PROCESSING, VOL. 18, NO. 11, NOVEMBER 2009
  • P. L. Combettes and J.-C. Pesquet Proximal Splitting Methods in Signal Processing

Комментарии 44

    +3
    А слабо exe'шник с gui запилить?
      +3
      «Если бы здесь были и спички, это был бы РАЙ!»
        0
        Чтобы не искать анекдот

        Попал наркоман в ад. Идет он и видит конопляное поле. Решил он курнуть и начал
        рвать коноплю. Подходит черт:


        • Зачем ты ее рвешь? Вон, уже скошенные поля есть! Идет мужик, и точно — конопля
          скошена, в стоги уложена. Взялся было ее сушить, подходит черт:
        • Зачем, дорогой, сушишь? Вон амбары с сушеной коноплей! Идет мужик дальше — амбары полны сушеной конопли! Решил он расфасовать ее по боксам. Подходит черт:
        • Зачем фасуешь? Вон амбары, полные готовых боксов! Идет — и точно, стоят
          амбары, конопля аккуратно расфасованна, боксы в штабеля сложены. Начал было
          забивать. Подходит черт:
        • Слушай, зачем забиваешь? Вон склады с косяками! Идет — склады полны, косяки
          аккуратно забиты, уложены. Взял мужик косяк, глядь — а спичек нет. Подходит
          черт, спрашивает мужик у него:
        • Черт, а черт, дай спички — прикурить!
        • Э, — отвечает черт, — если бы тут были спички, то это был бы РАЙ!

        Отсюда: http://www.narcozona.ru/anek.html

      +2
      В пределе, вот так, из числа Пи, математики восстановят какое нибудь изображение или целую вселенную.
        +2
        13 секунд выгрузка картинки? При типичной скорости(GPU->CPU) более 8 гигабайт в секунду?
        Что-то тут не так…
          0

          А почему это удивляет? Судя по всему, автор на каждой итерации загружает изображение в память видеокарты, считает там что-то простое, затем выгружает обратно.

            0
            [red] = fista_gpu(gpuArray(corrupted_red),obj1_gpu,obj2_gpu,0.000001,3,100);
            red=gather(red);
            

            Первая команда выдает матрицу размещенную в памяти GPU полностью отработав алгоритм и занимает по времени 0.5-1 сек.
            Вторая команда переносит матрицу из памяти GPU в обычную и работает она 4-5 сек. Возможно дело в Matlab, возможно я что-то не так использую, однако утверждение что я загружаю в память GPU картинку на каждой итерации голословно.
              0

              Да, моя ошибка — я ни разу не пользовался ускорением GPU на матлабе, потому неверно понял код.


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

          +2
          Как-то странно проявились тёмные элементы, на испорченной фото ровный мутный белый, а на восстановленной прям пятно:
          image
            0
            Такие артефакты любой алгоритм деконволюции даёт. Причина в том, что при свёртке (смазывании) теряется часть информации из-за недостатка точности (ну и из-за сжатия с потерями, если картинку в джпег пожать после этого). После восстановления изображения эти «обрубленные» значения дают рябь вокруг контрастных участков.
              0

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

            +1

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

              +2

              Вау, картинки в bmp

                +1

                По моему Лену всегда в бмп использовали?

                  –1

                  png или что-то другое без потерь, лишь бы не jpeg

                    +3

                    У меня другой вопрос, почему Лену Сёдерберг не используют в полный рост? :)

                      0
                      Это же риторический вопрос =)
                        0

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

                        +3
                        Используют, но в других целях
                          0

                          Потому что "та самая" Лена — 512*512. Исторически.

                            0

                            Но можно же 512х512 и другой участок взять

                              0

                              Возьмите машину времени, скатайтесь в 1973 и предложите Савчуку взять другой участок.

                                0
                                Ну eldarmusin говорит о более «той самой» версии на отсканированной странице Playboy, где Лена в полный рост…
                                  0

                                  Не-не, "та самая" версия — это не исходная страница Playboy, а тот самый скан 512*512 (вроде размер — это ограничение сканера, который у них в лабе был)

                    +1
                    Всё это, конечно, хорошо на синтетике, но как будет с реальным смазом или расфокусом?
                      0
                      С реальным расфокусом вам надо будет как-то найти матрицу свертки. То есть когда вы к картинке применяете blur, эта матрица вам известна, а когда у вас фото не в фокусе — ну можно попытаться ее угадать, от правильности угадывания конечно будет зависеть результат.
                        0
                        Вы правы, однако это уже совсем другая проблема. Мне не хотелось уходить в экспериментальные моменты, так как в них очень часто результат зависит от конкретного изображения с конкретным алгоритмом.
                          +5

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

                            0

                            А погрешность в ядре свёртки будет всегда, потому что в каждом пикселей ядро будет своё: дефокус, аберрации, неравномерность движения и т.д.

                      0
                      Возможно пропустил, но soft thresholding подразумевает что у нас есть L1 слагаемое в функционале, а в формулах вроде только L2
                      > Добавляя на каждом шаге разницу между результатом двух предыдущих итераций, мы увеличиваем сходимость алгоритма до квадратичной
                      Было бы здорово объяснить почему
                        0
                        Есть еще несколько вопросов. Не сочтите за грубость, просто мне интересны численные оптимизации.
                        В оригинальной постановке задачи (|| y — Ax ||_2) если || x ||_2 = sum_i ((x_i)^2), то вся задача это просто МНК который выливается в решение очень сильно разреженной СЛАУ. Я не большой специалист, но мне кажется что что-нибудь из семейства сопряженных градиентов на GPU должно было показать себя очень хорошо по части производительности.
                        Хотелось бы увидеть обоснование второго «улучшенного» функционала. Если со слагаемым (||alpha * ( x — x_k) ||^2) все более менее ясно (просим чтобы текущее решение было поближе к предыдущему), то с термом -A^t*A все совсем не понятно (откуда он взялся, зачем нужен и т.д.)
                        И да,
                        >Задача значительно проще
                        Вот это вот совсем не очевидно. Было бы здорово рассказать чем проще

                          +1
                          Посмотрел во вторую статью из списка источников, там вроде есть Total variation term (2*lambda*||x||_TV который говорит что градиент по картинке должен быть примерно одинаков, но допускаются разрывы), если его подставить в первую формулу, то все несколько проясняется. Но хотелось бы увидеть все это в данной статье
                            0

                            Такая СЛАУ влёт решается через преобразование Фурье, если А — просто свёртка.

                              0
                              Прикольно, не знал. Честно говоря, я так и не освоил преобразования Фурье и всякие DFT/FFT
                              Собственно еще одна причина думать что с оригинальным функционалом что-то не так
                                0
                                Если Вы имеете ввиду обратную фильтрацию то это не совсем так. Такое решение будет чувствительно к числу обусловленности матрицы свертки. Данный же алгоритм обладает стабильностью за счет дополнительного слагаемого.
                                  0
                                  Если Вы имеете ввиду обратную фильтрацию то это не совсем так. Такое решение будет чувствительно к числу обусловленности матрицы свертки. Данный же алгоритм обладает стабильностью за счет дополнительного слагаемого.

                                  Ну так речь шла именно о СЛАУ. Кстати, дополнительное слагаемое в виде квадратичного стабилизатора вполне себе можно учесть (в приближённом виде), если обращать свёртку как:


                                  IFFT(FFT(I) / (FFT(K) + C)),


                                  где C — константа. Здесь, правда, есть требование, чтобы FFT(K) было вещественным и положительным — не любое ядро размытия подойдёт.


                                  Но, понятное дело, с Total Variation такое уже не прокатит, и нужно применять полноценный итерационный метод минимизации.

                                    0
                                    Я знаю об этом методе, но с точки зрения одномерных задач ЦОС. Вся проблема — это искусство, а не ремесло. Для каждого конкретного ядра свертки и картинки с большой вероятностью пользователю придется самому выбирать C и оценивать качество результата. Здесь же можно программно оценить альфу по ядру свертки, а лямбда слабо влияет на сходимость.
                              0
                              а для размытия и смаза одинаковый алгоритм (я про вообще, а не только ISTA) подходит?
                                0
                                ISTA это вроде как про оптимизацию, а не про смазывание/размытие
                                А так много что используют, к примеру,размытие по гауссу
                                  0
                                  нене, неудачно сократил, в смысле:

                                  а для восстановления после размытия и смаза одинаковый алгоритм (я про вообще, а не только ISTA) подходит?
                                    0
                                    Хороший вопрос =)
                                    Если можно смаз описать сверткой, то можно попробовать ее подсунуть в этот алгоритм да. Иначе — не знаю, надо смотреть статьи.
                                      0
                                      В случае если операция повреждения была линейна, то да, подходит.
                                  +8

                                  Я всё понимаю, интересный алгоритм, исследовательский проект и всё такое, но почему код так ужасно написан? Каждый раз, когда я вижу код на матлабе, то в 99% случаев он написан ужасно. Матлаб к этому располагает, но вы даже отступы не потрудились расставить. Почему?! Отступы в редакторе матлаба автоматические! Вы специально их убрали что-ли? Весь код на гитхабе такой.


                                  Вы же оформили эту статью, почему нельзя оформить код?


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

                                    +2
                                    Тут Вы абсолютно правы. Буду исправляться.

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

                                  Самое читаемое