基于像素的传统图像修补算法实现

图像修补算法可以用来修复图像中的瑕疵，划痕等，或者移除不需要的内容，比如水印或其他物体。传统的图像修补算法有基于像素和基于区域的两种分类，本文介绍基于像素的传统图像修补算法的实现。论文可以在这里找到An image inpainting technique based on the fast marching method。基于像素的图像修补算法主要步骤包括两个，一是如何在要修不的图像区域中确定下一个要修补的像素，即要通过一种方法找到下一个待修补的像素，其周围有最多的已知像素，保证最好的修补效果。二是如何用周围的已知像素来估算要修补的像素。

找到要修补的下一个像素的基本思想是从需要进行修补区域的边界开始，由边界到中心逐渐填充待修补区域中的所有像素。越靠近边界的像素自然周围有越多的已知像素，而越向修补区域中心去，周围已知像素就越少。当然修补区域的形状是不确定的，所谓的中心是远离边界，向修补区域内部。完成这一步骤的就是所谓的FMM(Fast Marching Method)算法。记待修补的图像区域为Ω，∂Ω为Ω的边界。为修补Ω中的每一个像素，从∂Ω的初始位置∂Ω_i开始，向Ω内部前进。下一个带修补像素的选择，以diatance的大小顺序，由小到大选择。这里distance指的是distance map的distance，所有已知像素的distance为0。如果我们得到一张包含待修补区域的distance map，distance的大小代表了待修补像素距离已知像素的远近。以distance的升序排列，即可以得到修补像素的顺序。

通常要计算distance map的计算量较大，也有一些近似的快速算法。而FMM算法提供了一种即时计算distance的方法。FMM维护一个所谓的narrow band的像素集合，它就是边界∂Ω的所有像素。对每一个像素，保存它的distance值T，它的像素值I，还有一个flag值f。f可以是以下三种

BAND：属于narrow band的像素
KNOWN：在∂Ω外部的像素，已知像素。T，I已知
INSIDE：在∂Ω内部的像素，像素待修补。T，I未知

FMM初始时，对于所有位于边界∂Ω以及其外部区域的像素赋值T为0，而其内部待修补像素的T值为某一个大数（实际使用1e6）。f值按照上述定义给所有图像像素赋值。所有BAND像素放入一个以T值升序排列的有序集合，就是所谓的narrow band。因为对这个集合主要涉及到插入删除操作，实际中我使用的链表。

运行时按照下面的伪代码。需要注意的是，不知道什么原因，论文中有一些我认为明显错漏的地方，特别是在伪代码中。所以我这里的伪代码和原文中的有所不同。

while (NarrowBand not empty)
{
    extract P(i,j) = head(NarrowBand);
    f(i,j) = KNOWN;
    for (k,l) in (i,j-1), (i-1,j), (i+1,j), (i,j+1)
    {
        if (f(k,l) == INSIDE)
        {
            T(k,l) = min(solve(k-1,l,k,l-1),
                solve(k+1,l,k,l-1),
                solve(k-1,l,k,l+1),
                solve(k+1,l,k,l+1));
            Inpaint(k,l);
            f(k,l) = BAND;
            Insert (k,l) in NarrowBand;
        }
    }
}

原文中这段伪代码最大的问题是对f(k,l) != KNOWN的判断。不是KNOWN的话，就是BAND或者INSIDE，那么就可能会重复把已经是BAND的像素重新加入到NarrowBand了。所以这里直接改成判断f(k,l)是否是INSIDE，只有INSIDE的像素才需要更新T并修补。FMM首先从NarrowBand中取出一个有最小T值的BAND像素，修改其flag为KNOWN。然后检查其临近的上下左右四联通像素，如果是INSIDE，则依次做

更新其distance值T。
实际修补其像素值。
修改其flag为BAND，加入NarrowBand。

这几步做完，实际上把边界∂Ω向边界Ω内部前进了一步。注意这里我是先做了更新T值，是因为在实际inpaint像素值时是有可能需要用到它本身T值的（后面会讲到）。而在原文中，更新distance值T是放在inpaint之后，那么inpaint时就无法用了。这也是我无法理解，认为原文这里是错误的原因。

这里更新(k,l)的T值需要其上下左右四联通像素的T值，解如下方程

其中，，对y也类似。如果两个max值都不为0的话，设其中一个max值中的已知量，比如T(i-1,j)为T₁，另一个max值中的已知量比如T(i,j+1)为T₂，上述方程实际是一个一元二次方程，很容易求得其两个解为。而如果要求解同时大于T₁和T₂，则只有一个解，并且abs(T₁-T₂)<1。所以这里也不需要像原文中的伪代码来。解得四个值，取最小值。比较简单，就不写伪代码了，直接看后面solve()函数的源代码。

这样FMM算法的主要流程就完成了，重复以上过程直到NarrowBand为空，就是从边界一直深入到Ω内部，到所有像素都修补完毕。

剩下的问题就是如何修补单个像素。待修补像素由其邻域中所有已知像素的加权和得到。

p为要inpaint的像素，为p的大小为ε的领域。领域中单个像素q对p的贡献为，为q点的梯度。

w(p,q)为权重，是三个分量的乘积。

其中dir(p,q)是方向分量，其保证越接近distance变化的法线方向（即T的梯度方向）的像素对inpaint像素的贡献越大。这也是前面FMM中为什么要把对T的更新放在实际inpaint之前的原因。虽然计算T的中央差分可以不用它本身的T值，但是如果其他相邻的像素有一个是INSIDE的话，那么就需要用到它本身的值了。dst(p,q)也很容易理解，是地理距离分量，q距离p越近，对p的影响越大。lev(p,q)是T等值集合分量，让越接近通过p的轮廓等值线的像素对inpaint贡献更大，也就是q距离已知像素和p距离已知像素的差值越小越好，即两者值越接近越好。在实际应用中，d₀和T₀都使用1。这三个分量结合，共同决定领域内每个像素像素对修补值的影响。

下面给出Inpaint的代码。PriorityHeap是一个以T值排序的像素集合，内部使用双向链表，方便快速插入删除。初始化narrow band，用同样的narrow band对三通道数据做Inpaint。可以对三通道各用一个线程加快处理速度。

enum {
    KNOWN,
    BAND,
    INSIDE,
};

typedef struct PaintType {
    int flag;
    float T;
} PaintType;

typedef struct HeapElement {
    float T;
    CPoint point;
    HeapElement(float T, CPoint& p) {
        this->T = T;
        point = p;
    }
} HeapElement;

class PriorityHeap
{
protected:
    list<HeapElement> m_band;
public:
    PriorityHeap() {}
    PriorityHeap(PriorityHeap& ph)
    {
        m_band = ph.m_band;
    }
    PriorityHeap(PriorityHeap&& ph) noexcept
    {
        m_band = move(ph.m_band);
    }
    void Push(float T, CPoint& p)
    {
        auto rit = m_band.rbegin();
        for (; rit != m_band.rend(); ++rit)
            if (rit->T <= T)
                break;
        m_band.emplace(rit.base(), T, p);
    }
    void Push(float T, CPoint&& p)
    {
        auto rit = m_band.rbegin();
        for (; rit != m_band.rend(); ++rit)
            if (rit->T <= T)
                break;
        m_band.emplace(rit.base(), T, p);
    }
    BOOL Pop(CPoint* pp)
    {
        if (m_band.size() == 0) return FALSE;
        *pp = m_band.front().point;
        m_band.pop_front();
        return TRUE;
    }
    size_t size()
    {
        return m_band.size();
    }
};

BOOL Inpaint(int radius)
{
    CRect* pRect = nullptr;
    CRect Mask;
    // sMask包含要修补的区域
    unique_ptr<BYTE[]> sMask;
    if (pDibMask)
    {
        GetMaskCircumRect(Mask);
        GetMask(sMask);
        pRect = &Mask;
    }
    else
        return FALSE;

    long i, j;
    long Width = GetWidth();
    long Height = GetHeight();

    long Hstart, Hend, Wstart, Wend;
    long len;

    pRect->NormalizeRect();
    if (pRect->top < 0) pRect->top = 0;
    if (pRect->left < 0) pRect->left = 0;
    if (pRect->bottom > Height) pRect->bottom = Height;
    if (pRect->right > Width) pRect->right = Width;
    Hstart = pRect->top;
    Hend = pRect->bottom;
    Wstart = pRect->left;
    Wend = pRect->right;
    ASSERT(Hstart >= 0 && Hstart <= Hend && Hend <= Height);
    ASSERT(Wstart >= 0 && Wstart <= Wend && Wend <= Width);
    
    unique_ptr<BYTE[]> pRed, pGreen, pBlue;
    if (GetRGB(pRed, pGreen, pBlue, &len) == FALSE) return FALSE;

   // initialize narrow band
    unique_ptr<PaintType[]> sPT(new PaintType[len]);
    PriorityHeap narrow_band;
    BYTE* pm = sMask.get();
    PaintType* ppt = sPT.get();
    for (i = 0; i < Hstart - 1; ++i)        // to (Hstart-1) because 1 pixel extension to pRect needs to be checked
    {
        for (j = 0; j < Width; ++j)
        {
            // known
            ppt->flag = KNOWN;
            ppt->T = 0;
            ++ppt;
        }
        pm += Width;
    }
    for (; i < min(Hend + 1, Height); ++i)
    {
        for (j = 0; j < Wstart - 1; ++j)
        {
            // known
            ppt->flag = KNOWN;
            ppt->T = 0;
            ++ppt;
            ++pm;
        }
        for (; j < min(Wend + 1, Width); ++j)
        {
            if (*pm == 0) {
                // check 4 connected pixels
                if (i - 1 >= 0 && *(pm - Width)) {
                    // band
                    ppt->flag = BAND;
                    ppt->T = 0;
                    narrow_band.Push(0, CPoint(j, i));
                    ++ppt;
                    ++pm;
                    continue;
                }
                else if (j - 1 >= 0 && *(pm - 1)) {
                    // band
                    ppt->flag = BAND;
                    ppt->T = 0;
                    narrow_band.Push(0, CPoint(j, i));
                    ++ppt;
                    ++pm;
                    continue;
                }
                else if (j + 1 < Width && *(pm + 1)) {
                    // band
                    ppt->flag = BAND;
                    ppt->T = 0;
                    narrow_band.Push(0, CPoint(j, i));
                    ++ppt;
                    ++pm;
                    continue;
                }
                else if (i + 1 < Height && *(pm + Width)) {
                    // band
                    ppt->flag = BAND;
                    ppt->T = 0;
                    narrow_band.Push(0, CPoint(j, i));
                    ++ppt;
                    ++pm;
                    continue;
                }
                else {
                    // known
                    ppt->flag = KNOWN;
                    ppt->T = 0;
                    ++ppt;
                    ++pm;
                }
            }
            else {
                // inside
                ppt->flag = INSIDE;
                ppt->T = 1e6f;
                ++ppt;
                ++pm;
            }
        }
        for (; j < Width; ++j)
        {
            // known
            ppt->flag = KNOWN;
            ppt->T = 0;
            ++ppt;
            ++pm;
        }
    }
    for (; i < Height; ++i)
    {
        for (j = 0; j < Width; ++j)
        {
            // known
            ppt->flag = KNOWN;
            ppt->T = 0;
            ++ppt;
        }
    }

    PriorityHeap narrow_band2(narrow_band), narrow_band3(narrow_band);
    unique_ptr<PaintType[]> sPT2, sPT3;
    sPT2.reset(new PaintType[len]);
    memcpy(sPT2.get(), sPT.get(), sizeof(PaintType) * len);
    sPT3.reset(new PaintType[len]);
    memcpy(sPT3.get(), sPT.get(), sizeof(PaintType) * len);

#if 0
    Inpaint(pRed.get(), sPT.get(), narrow_band, radius);
    Inpaint(pGreen.get(), sPT2.get(), narrow_band2, radius);
    Inpaint(pBlue.get(), sPT3.get(), narrow_band3, radius);
#else
    auto InpaintFunc = bind((void(*)(BYTE*, PaintType*, PriorityHeap&, const int)) & Inpaint, placeholders::_1, placeholders::_2, placeholders::_3, placeholders::_4);
    auto thread2 = thread(InpaintFunc, pGreen.get(), sPT2.get(), narrow_band2, radius);
    auto thread3 = thread(InpaintFunc, pBlue.get(), sPT3.get(), narrow_band3, radius);

    Inpaint(pRed.get(), sPT.get(), narrow_band, radius);
    thread2.join();
    thread3.join();
#endif
    PutRGB(nullptr, pRed, pGreen, pBlue);
    return TRUE;
}

static float solve(const PaintType* pPT, const long offset1, const long offset2)
{
    float sol = 1e6;
    float T1, T2;
    T1 = pPT[offset1].T;
    T2 = pPT[offset2].T;
    if (pPT[offset1].flag != INSIDE)
    {
        if (pPT[offset2].flag != INSIDE) {
            if (abs(T1 - T2) >= 1.0f)
                sol = 1 + min(T1, T2);
            else
                sol = (T1 + T2 + sqrtf(2 - (T1 - T2) * (T1 - T2))) / 2;
        }
        else
            sol = 1 + T1;
    }
    else if (pPT[offset2].flag != INSIDE)
        sol = 1 + T2;
    return sol;
}

void Inpaint(BYTE* pData, PaintType* pPT, PriorityHeap& narrow_band, const int radius)
{
    long Width = GetWidth();
    long Height = GetHeight();
    long len = Width * Height;
    unique_ptr<float[]> pfData(new float[len]);        // 保存中间数据，提高精度
    for (long k = 0; k < len; ++k)
        pfData[k] = pData[k];
    long offset;
    CPoint point;
    while (narrow_band.Pop(&point)) {
        offset = point.y * Width + point.x;
        pPT[offset].flag = KNOWN;
        for (int ii = 0; ii < 4; ++ii)
        {
            long off_i;
            CPoint point_i(point);
            switch (ii) {
            case 0:
                if (--point_i.y < 0) continue; off_i = offset - Width; break;
            case 1:
                if (--point_i.x < 0) continue; off_i = offset - 1; break;
            case 2:
                if (++point_i.x >= Width) continue; off_i = offset + 1; break;
            case 3:
                if (++point_i.y >= Height) continue; off_i = offset + Width; break;
            }
            int& flag = pPT[off_i].flag;
            if (flag == INSIDE)
            {
                bool off_ivalid0 = true, off_ivalid1 = true, off_ivalid2 = true, off_ivalid3 = true;
                long off_i0 = off_i - Width;
                if (point_i.y - 1 < 0) off_ivalid0 = false;
                long off_i1 = off_i - 1;
                if (point_i.x - 1 < 0) off_ivalid1 = false;
                long off_i2 = off_i + 1;
                if (point_i.x + 1 >= Width) off_ivalid2 = false;
                long off_i3 = off_i + Width;
                if (point_i.y + 1 >= Height) off_ivalid3 = false;
                // calculate distance map T
                {
                    float minimum = 1e6;

                    if (off_ivalid1 && off_ivalid0)
                        minimum = min(minimum, solve(pPT, off_i1, off_i0));
                    if (off_ivalid2 && off_ivalid0)
                        minimum = min(minimum, solve(pPT, off_i2, off_i0));
                    if (off_ivalid1 && off_ivalid3)
                        minimum = min(minimum, solve(pPT, off_i1, off_i3));
                    if (off_ivalid2 && off_ivalid3)
                        minimum = min(minimum, solve(pPT, off_i2, off_i3));

                    pPT[off_i].T = minimum;
                }
                // inpaint pixel
                {
                    long i, j;
                    long hstart = point_i.y - radius;
                    long hend = point_i.y + radius;
                    long wstart = point_i.x - radius;
                    long wend = point_i.x + radius;
                    if (hstart < 0) hstart = 0;
                    if (hend > Height - 1) hend = Height - 1;
                    if (wstart < 0) wstart = 0;
                    if (wend > Width - 1) wend = Width - 1;
                    // gradT of the pixel being inpainted
                    float gradTx = 0, gradTy = 0;
                    {
                        if (off_ivalid2 && pPT[off_i2].flag != INSIDE) {
                            if (off_ivalid1 && pPT[off_i1].flag != INSIDE)
                                gradTx = (pPT[off_i2].T - pPT[off_i1].T) * 0.5f;
                            else
                                gradTx = pPT[off_i2].T - pPT[off_i].T;
                        }
                        else {
                            if (off_ivalid1 && pPT[off_i1].flag != INSIDE)
                                gradTx = pPT[off_i].T - pPT[off_i1].T;
                        }
                        if (off_ivalid3 && pPT[off_i3].flag != INSIDE) {
                            if (off_ivalid0 && pPT[off_i0].flag != INSIDE)
                                gradTy = (pPT[off_i3].T - pPT[off_i0].T) * 0.5f;
                            else
                                gradTy = pPT[off_i3].T - pPT[off_i].T;
                        }
                        else {
                            if (off_ivalid0 && pPT[off_i0].flag != INSIDE)
                                gradTy = pPT[off_i].T - pPT[off_i0].T;
                        }
                    }
                    float s_weight = 0, sum = 0;
                    for (i = hstart; i <= hend; ++i)
                    {
                        long row_offset = i * Width;
                        float r_y = float(point_i.y - i);
                        for (j = wstart; j <= wend; ++j)
                        {
                            long coff = row_offset + j;
                            if (coff == off_i) continue;        // skip the pixel being inpainted
                            float r_x = float(point_i.x - j);
                            if (pPT[coff].flag != INSIDE)
                            //if (pPT[coff].flag == KNOWN)
                            {
                                float dir = abs(gradTx * r_x + gradTy * r_y);
                                if (dir < 1e-6f) dir = 1e-6f;
                                float dst = 1.0f / (r_y * r_y + r_x * r_x);
                                float lev = 1.0f / (1 + abs(pPT[coff].T - pPT[off_i].T));
                                float w = dir * dst * lev;
                                float gradIx = 0, gradIy = 0;
                                if (j + 1 < Width && pPT[coff + 1].flag != INSIDE) {
                                    if (j - 1 >= 0 && pPT[coff - 1].flag != INSIDE)
                                        gradIx = (pfData[coff + 1] - pfData[coff - 1]) * 0.5f;
                                    else
                                        gradIx = float(pfData[coff + 1] - pfData[coff]);
                                }
                                else {
                                    if (j - 1 >= 0 && pPT[coff - 1].flag != INSIDE)
                                        gradIx = float(pfData[coff] - pfData[coff - 1]);
                                }
                                if (i + 1 < Height && pPT[coff + Width].flag != INSIDE) {
                                    if (i - 1 >= 0 && pPT[coff - Width].flag != INSIDE)
                                        gradIy = (pfData[coff + Width] - pfData[coff - Width]) * 0.5f;
                                    else
                                        gradIy = float(pfData[coff + Width] - pfData[coff]);
                                }
                                else {
                                    if (i - 1 >= 0 && pPT[coff - Width].flag != INSIDE)
                                        gradIy = float(pfData[coff] - pfData[coff - Width]);
                                }
                                float kx, ky;
                                kx = gradIx * r_x;
                                ky = gradIy * r_y;
                                sum += w * (pfData[coff] + 0.35f * (kx + ky));
                                s_weight += w;
                            }
                        }
                    }
                    pfData[off_i] = sum / s_weight;
                    CLAMP0255(pfData[off_i]);
                    pData[off_i] = BYTE(roundf(pfData[off_i]));
                }
                flag = BAND;
                // push into narrow band
                narrow_band.Push(pPT[off_i].T, point_i);
            }
        }
    }
}