【Python】 主双対アルゴリズムを用いた Total Variation L1 正則化

最近読んだ論文[^1]の応用例に Total Variation L1 正則化の例があったが，解説が少なかったのでメモ．
[^1]: Thomas Pock and Antonin Chambolle 2012 "Diagonal preconditioning for first order primal-dual algorithms in convex optimization" 2011 International Conference on Computer Vision https://ieeexplore.ieee.org/abstract/document/6126441

## 扱う問題

```math
\min_{u}{\|Du\|_1 + \lambda\|u - f\|_1}
```

ここで $f$ は入力画像をベクトル化したもの，$u$ は決定変数，$D$ は隣接ピクセルの差分を計算する線型作用素，$\lambda$ は正則化パラメータである．

例えば，入力画像 $F$ が $2\times3$ の画像

```math
F = \left(\begin{matrix}
f_1 & f_2 & f_3\\
f_4 & f_5 & f_6
\end{matrix}\right)
```

のとき，

```math
f = (f_1, \ldots, f_6)^\top\hspace{-1pt}\\
Df = \left(\begin{matrix}
f_1-f_2\\ f_2-f_3\\ f_4-f_5\\ f_5-f_6\\ f_1-f_4\\ f_2-f_5\\ f_3-f_6
\end{matrix}\right)\\
D = \left(\begin{matrix}
1 & -1 & 0 & 0 & 0 & 0\\
0 & 1 & -1 & 0 & 0 & 0\\
0 & 0 & 0 & 1 & -1 & 0\\
0 & 0 & 0 & 0 & 1 & -1\\
1 & 0 & 0 & -1 & 0 & 0\\
0 & 1 & 0 & 0 & -1 & 0\\
0 & 0 & 1 & 0 & 0 & -1
\end{matrix}\right)
```

と表される．

### 問題の意味

第 1 項で出力画像を滑らかに，第 2 項で元画像から大きく離れないように調整している．

## L1 ノルムの凸共役関数

L1 ノルム $g(u) = \|u\|_1$ の凸共役関数 $g^*$ は次式で与えられる．

```math
\begin{equation}
\begin{split}
g^*(p) &= \max_{u}{\{\left<u, p\right> - g(u)\}}\\
&= \max_{u}{\{\left<u, p\right> - \|u\|_1\}}\\
&= \max_{u}{\left\{\sum_i(u_i p_i-|u_i|)\right\}}\\
&= \begin{cases}
0 & -1 \leq p \leq 1\\
+\infty & \text{otherwise}
\end{cases}
\end{split}
\end{equation}
```

L1 ノルムは真閉凸関数であるので，

```math
\begin{equation}
\begin{split}
g(u) &= g^{**}(u)\\
&= \max_{p}{\{\left<u, p\right> - g^*(p)\}}\\
&= \max_{-1 \leq p \leq 1}{\left<u, p\right>}
\end{split}
\end{equation}
```

これを用いて，先の最適化問題は，

```math
\min_u{\max_{-1 \leq (p, q) \leq 1}{\left<Du, p\right> + \lambda \left<u-f, q\right>}}
```

とかける．

## 主双対アルゴリズム

近接点法を用いて交互に最小化・最大化問題を解く．

```math
\left\{
\begin{align}
&u^{k+1} = \mathop{\mathrm{arg~min}}\limits_{u}{\left\{\left<Du, p^k\right> + \lambda \left<u-f, q^k\right> + \frac12\|u - u^k\|_X^2\right\}}\\
&\bar{u}^{k+1} = 2u^{k+1} - u^k\\
&(p^{k+1}, q^{k+1}) = \mathop{\mathrm{arg~max}}\limits_{-1 \leq (p, q) \leq 1}{\left\{\left<D\bar{u}^{k+1}, p\right> + \lambda \left<\bar{u}^{k+1}-f, q\right> - \frac12\|(p, q) - (p^k, q^k)\|_Y^2\right\}}
\end{align}
\right.
```

ここで，

```math
K = (D^\top\hspace{-1pt}, \lambda I)^\top\hspace{-1pt}\\
\|x\|_X^2 = \left<x, T^{-1}x\right>\\
T = \mathrm{diag}(\tau)\quad \mathrm{where}\quad \tau_j=\frac{1}{\sum_{i}|K_{i,j}|^{2-\alpha}}\\
\|y\|_Y^2 = \left<y, \Sigma^{-1}y\right>\\
\Sigma = \mathrm{diag}(\sigma)\quad \mathrm{where}\quad \sigma_i=\frac{1}{\sum_{j}|K_{i,j}|^{\alpha}}
```
である．これより，更新式は

```math
\left\{
\begin{align}
&u^{k+1} = u^k - T\left(D^\top\hspace{-1pt}p^k + \lambda q^k\right)\\
&\bar{u}^{k+1} = 2u^{k+1} - u^k\\
&(p^{k+1}, q^{k+1}) = \mathrm{proj}_{[-1, 1]}\left((p^k, q^k) + \Sigma\left(D\bar{u}^{k+1}, \lambda(\bar{u}^{k+1} - f)\right)\right)
\end{align}
\right.
```

となる．

## ソースコード

Python で実装．

- MacOS Mojave : 10.14.3
- Python : 3.7.3
- numpy : 1.16.4

ソースコードや実行した notebook は [Github](https://github.com/Mrrmm510/PrimalDual) に．

```python
from typing import Tuple

import numpy as np
from tqdm import trange


class TotalVariation:
    """
    Total Variation L1 model using the preconditioned primal dual algorithm
    """
    def __init__(self,
            lambd: float = 1.0,
            max_iter: int = 1000,
            coef: np.ndarray = np.array([1, -1]),
            saturation: bool = False,
            extended_output: bool = False):
        """
        Parameters
        ----------
        lambd : float
            A regularization parameter.
        max_iter : int
            The maximum number of iterations.
        coef : np.ndarray
            [1, -1] for total valiation regularization
        saturation : bool
            If True, output will be in a range [0, 1].
        extended_output : bool
            If True, return the value of the objective function of each iteration.
        """
        self.lambd = lambd
        self.max_iter = max_iter
        self.coef = coef
        self.saturation = saturation
        self.extended_output = extended_output

        self.length = len(coef) - 1
        # objective function value
        self.obj = list()

    def _tv(self, u: np.ndarray) -> np.ndarray:
        h, w = u.shape
        ret = np.zeros((2 * h, w))
        for i, c in enumerate(self.coef):
            ret[: h, : w - self.length] += c * u[:, i : w - self.length + i]
            ret[h : 2 * h - self.length] += c * u[i : h - self.length + i]
        return ret

    def _transposed_tv(self, v: np.ndarray) -> np.ndarray:
        h2, w = v.shape
        h = h2 // 2
        ret = np.zeros((h, w))
        for i, c in enumerate(self.coef):
            ret[:, i : w - self.length + i] += c * v[: h, : w - self.length]
            ret[i : h - self.length + i] += c * v[h : 2 * h - self.length]
        return ret

    def _step_size(self, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
        h, w = shape
        abs_coef = np.abs(self.coef)
        abs_sum = np.sum(abs_coef)

        tau = np.full((h, w), self.lambd + 2 * abs_sum)
        for i in range(self.length):
            tau[i] -= np.sum(abs_coef[i + 1 :])
            tau[-(i + 1)] -= np.sum(abs_coef[: -(i + 1)])
            tau[:, i] -= np.sum(abs_coef[i + 1 :])
            tau[:, -(i + 1)] -= np.sum(abs_coef[: -(i + 1)])
        tau = 1. / tau

        sigma = np.zeros((3 * h, w))
        sigma[:2 * h] += abs_sum
        sigma[2 * h:] += self.lambd
        sigma = 1. / sigma
        return tau, sigma

    def transform(self, X: np.ndarray) -> np.ndarray:
        """
        Parameters
        ----------
        X : array, shape = (h, w)
            a 2D image

        Returns
        ----------
        res : array, shape = (h, w)
            a denoised image
        """
        h, w = X.shape
        h2 = h * 2
        tau, sigma = self._step_size((h, w))

        # initialize
        res = np.copy(X)
        dual = np.zeros((3 * h, w))
        dual[: h2] = np.clip(sigma[: h2] * self._tv(res), -1, 1)

        # store objective function value if necessary
        if self.extended_output:
            self.obj.append(np.sum(np.abs(self._tv(res))) + self.lambd * np.sum(np.abs(res - X)))

        # main loop
        for _ in trange(self.max_iter):
            if self.saturation:
                u = np.clip(res - (tau * (self._transposed_tv(dual[: h2]) + self.lambd * dual[h2 :])), 0, 1)
            else:
                u = res - (tau * (self._transposed_tv(dual[: h2]) + self.lambd * dual[h2 :]))
            bar_u = 2 * u - res
            dual[: h2] += sigma[: h2] * self._tv(bar_u)
            dual[h2 :] += sigma[h2 :] * self.lambd * (bar_u - X)
            dual = np.clip(dual, -1, 1)
            res = u

            # store objective function value if necessary
            if self.extended_output:
                self.obj.append(np.sum(np.abs(self._tv(res))) + self.lambd * np.sum(np.abs(res - X)))
        return res
```

実行結果．羽の模様が消えたりしている．
<img width="800" alt="smoothed_insects.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/394475/52f058ff-4f86-762e-2792-083915c6fd3e.png">