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

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

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

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

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

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

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

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

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

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

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

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

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

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

Следующий классический шаг в данной ситуации — приравнять градиент к нулю, и выразить из полученного выражения искомый вектор . Записанный определим как или как решение на следующей итерации.

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

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

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

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

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

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

1. Задать параметры
2. Задать
3. Произвести итерацию базового алгоритма
4. Обновить временные переменные
   
5. Добавить сопряженную компоненту к переменной

Теперь о параллельности алгоритма. Цикл для вычисления порога на каждой итерации можно выразить через поэлементные произведения (произведение Адамара) и операции сравнения, которые включены как в 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

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

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

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

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

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

^[1]

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

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

^[2]

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

^[3]

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

^[4]

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

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

• 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

Автор: zedroid

Источник ^[6]