讀源碼學算法之Colorization

經典論文Colorization using Optimization來自於siggraph2004，論文旨在通過簡單的交互實現灰度圖的全局彩色化，論文短小精悍、思路新穎、實現簡單、效果優良。短短6頁紙的論文，引用1300+。

論文鏈接：https://webee.technion.ac.il/people/anat.levin/papers/colorization-siggraph04.pdf
代碼鏈接：https://github.com/lightalchemist/colorize-image

輸入：原始灰度圖 （下圖第一張圖） + 局部着色（下圖第二張圖）
輸出：全局着色效果（下圖第三張圖）

算法工作在YUV空間，簡單來講，YUV適用於電視信號傳輸，Y表示亮度，UV表示色度，灰度圖的Y通道其實就是本身的灰度值。YUV與RGB顏色空間的相互轉化參見：https://www.jianshu.com/p/cf583040c930。論文核心思想：亮度近似的像素應當具有相近的顏色。

兩個核心能量公式如下所示，如何最下化，其實令大括號內 = 0即可。

$J(U) = \sum\limits_r{(U(r) - \sum\limits_{s\in N(r)}{w_{rs}U(s)})^2}$ (1)

$J(U) = \sum\limits_r{(V(r) - \sum\limits_{s\in N(r)}{w_{rs}V(s)})^2}$ (2)

這裏， $s$是$r$的鄰居，一般限定在一個3*3的鄰域內， $w_{rs}$是像素 $r$ 和 $s$ 根據 Y通道計算的相似度， 並且有 $\sum_{s\in N(r)}{w_{rs}} = 1.0$， $w_{rs}$計算公式如下：

$w_{rs} = exp(-(Y(r)-Y(s))^2 / 2\sigma_r^2)$ (3)

上述公式(1)(2)其實等效爲求解兩個大型線性方程數組，可以寫成$Ax= b$ 的形式，A即是$w$矩陣，b即是根據U/V通道加權的結果。構建線性系統的核心代碼爲：

//Y:Y通道（灰度圖本身的灰度值）, scribbles: 用戶着色的圖，mask：着色區域的mask
//A：係數矩陣，bu，bv關於U/V方程組的右側列矩陣
void setupProblem(const cv::Mat& Y, const cv::Mat& scribbles, const cv::Mat& mask, 
                  Eigen::SparseMatrix<double, Eigen::RowMajor>& A, Eigen::VectorXd& bu,
                  Eigen::VectorXd& bv, double gamma)
{
    typedef Eigen::Triplet<double> TD;

    auto nrows = Y.rows;
    auto ncols = Y.cols;
    auto nPixels = nrows * ncols;
    A.resize(nPixels, nPixels);

    std::vector<TD> coefficients;
    coefficients.reserve(nPixels * 3);
    bu.resize(nPixels);
    bv.resize(nPixels);
    bu.setZero();
    bv.setZero();

    cv::Mat yuvScribbles;
    cv::cvtColor(scribbles, yuvScribbles, cv::COLOR_BGR2YUV);
    yuvScribbles.convertTo(yuvScribbles, CV_64FC3);
    std::vector<cv::Mat> channels;
    cv::split(yuvScribbles, channels);
    cv::Mat& U = channels[1];
    cv::Mat& V = channels[2];

    // TODO: See if we can do this more efficiently using matrix reshape
    std::vector<double> y, u, v;
    std::vector<bool> hasColor;
    to1D(Y, y);
    to1D(U, u);
    to1D(V, v);
    to1D(mask, hasColor);

    const int numNeighbors = 8;
    std::vector<double> weights;
    weights.reserve(numNeighbors);
    std::vector<unsigned long> neighbors;
    neighbors.reserve(numNeighbors);
    for (auto i = 0; i < nrows; ++i) {
        for (auto j = 0; j < ncols; ++j) {
            unsigned long r = i * ncols + j;
            getNeighbours(i, j, nrows, ncols, neighbors);
            getWeights(y, r, neighbors, weights, gamma);
            coefficients.push_back(TD(r, r, 1));
            for (auto k = 0u; k < neighbors.size(); ++k) {
                auto s = neighbors[k];
                auto w = weights[k];
                if (hasColor[s]) {
                    // Move value to RHS of Ax = b
                    bu(r) += w * u[s];
                    bv(r) += w * v[s];
                } else {
                    coefficients.push_back(TD(r, s, -w));
                }
            }
        }
    }

    A.setFromTriplets(coefficients.begin(), coefficients.end());
}

下面給一個直觀的例子：我們給出一個2行3列的圖像, 現在灰度圖上進行了局部着色，如(2)所示，可以抽出着色部分的U通道，接下來我們的任務是求解其他部分的U值，如(3)所示，剩餘的$x_2,x_3,,x_4,x_6$ 是待求解的未知數。V值求解方法一樣，灰度圖的灰度值本身就是Y，最後我們把YUV合起來就得到了最終的彩色圖。

對於$x_1$：因爲這裏已經被着色，那麼必然有： $x_1 = u_1$
對於$x_2$：有5個鄰居，鄰域U的加權平均等於$x_2$本身的U值， 因此 $x_2-w_{21}u_1-w_{23}x_3 -w_{24}x_4 - w_{25}u_5 -w_{26}x_6= 0$
對於$x_3$：有3個鄰居， $x_3-w_{32}x_2-w_{35}u_5 -w_{36}x_6 = 0$
對於$x_4$：有3個鄰居， 因此有： $x_4-w_{41}u_1-w_{42}x_2 -w_{45}u_5= 0$
對於$x_5$：因爲這裏已經被着色，那麼必然有： $x_5 = u_5$
對於$x_6$：有3個鄰居， 因此有： $x_6-w_{62}x_2-w_{63}x_3 -w_{65}u_5= 0$

那麼，上式可以寫成如下的線性系統：(核心就是將上面各式的已知量全部移到右側，未知量全在左側)

$\begin{pmatrix} 1 &0 &0 &0 & 0 & 0 \\ 0 &1 &w_{23} &w_{24} & 0 & w_{26}\\ 0 & w_{32} & 1&0 &0 &w_{36} \\ 0 & w_{42} &0 & 1& 0 & 0\\ 0 & 0& 0& 0 & 1&0 \\ 0 & w_{62}& w_{63}& 0& 0& 1 \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \\ x_3 \\ x_4 \\ x_5 \\ x_6 \end{pmatrix}= \begin{pmatrix} u_1 \\ w_{21}u_1+w_{45}u_5 \\ w_{35}u_5 \\ w_{41}u_1+w_{45}u_5 \\ u_5 \\ w_{65}u_5 \end{pmatrix}$