天天看點

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

前言

這個降噪的模型來自 Christopher M. Bishop 的 Pattern Recognition And Machine Learning (就是神書 PRML……),問題是如何對一個添加了一定椒鹽噪聲(Salt-and-pepper Noise)(假設噪聲比例不超過 10%)的二值圖(Binary Image)去噪。

原圖 -> 添加 10% 椒鹽噪聲的圖:

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

放在 github 上的可運作完整代碼:https://github.com/joyeecheung/simulated-annealing-denoising

模組化

下文中的數學表示:

  • yi:噪聲圖中的像素
  • xi:原圖中的像素,對應噪聲圖中的 yi

既然噪聲圖是從原圖添加噪聲而來,我們擁有了先驗知識 1:

yi和xi有很強的聯系。

一般圖檔裡,每個像素和與它相鄰的像素值應當較為接近,比如上圖中的黑色筆畫和白色負空間,除了邊緣以外,黑色的像素周圍都是黑色像素,白色像素的周圍都是白色像素(連成一片))。這樣我們就得到了先驗知識 2:

xi和與它相鄰的其他像素也存在較強的聯系

如果我們狠一點,假設原圖像素隻與它的直接相鄰像素有聯系(即具備條件獨立性質),我們就可以得到一個具備局部馬爾可夫性質(Local Markov property)的圖模型:

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

在這樣一個圖模型裡,我們有兩種團(clique):

  1. {xi, yi},即原圖像素與噪聲圖像素對
  2. {xi, xj},其中 xj 表示與 xi 相鄰的像素

這兩種團合并起來,得到的 {xi, yi, xj} 顯然是一個最大團(Maximal Clique),此時我們可以利用它來對這個馬爾可夫随機場進行 factorization,即求得其聯合機率分布關于最大團 xC = {xi, yi, xj} 的函數。

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

其中 Z 為 partition function,是 p(x) 的歸一化常數(normalization constant),求法參見 PRML 8.3.2。因為與我們的實作不相關,這裡不贅述。

ψC(xC) 即所謂的 potential function,為了友善我們通常隻求它的對數形式 E(xC)(按照其實體意義稱為 energy function)

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

關于 factorization 的過程和推導可以參見 PRML 8.3.2,這裡我們需要做的是定義一個 energy function。在降噪的過程中 energy 越低,聯合機率 P(X=x) (降噪過的圖像與原圖一緻的機率)就越大。

因為我們需要處理的是二值圖,首先我們定義 xi ∈ {-1, +1},假設這裡白色為1,黑色為-1。

對于原圖像素與噪聲圖像素構成的團 {xi, yi},我們定義一項 −ηxiyi,其中 η 為一個非負的權重。當兩者同号時,這項的值為−η,異号時為 η。這樣,當噪聲圖像素與原圖像素較為接近時,這一項能夠降低 energy(因為為負)。

對于噪聲圖中相鄰像素構成的團 {xi, xj},我們定義一項 −βxixj,其中 β 為一個非負的權重。這樣,當處理過的圖像裡相鄰的像素較為接近時,這一項能夠降低 energy(因為為負)。

最後,我們再定義一項 hxi,使處理過的圖像整體偏向某一個像素值。

對圖像中的每一個像素,将這三項加起來,就得到我們的 energy function:

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

對應聯合機率

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

顯然 energy 越低,降噪過的圖像與原圖一緻的機率越高。(注意因為我們這裡求的 E 已經對整個矩陣求和,即對應 potential function 的積,是以計算聯合機率分布的時候不需要再求積)

使用 Python 實作這個 energy function 時,我們可以使用一個 closure 來實作一個 function factory,通過傳遞

beta

(β),

eta

(η)和

h

參數,生成對應的 energy function。此外為了友善,我們假設傳入的

x

y

不是一維向量,而是對應圖像的二維矩陣(注意是

np.ndarray

而不是

nd.matrix

,前者的

*

才是array multiplication即逐個元素相乘,後者的

*

是矩陣乘法)。

import numpy as np

def E_generator(beta, eta, h):
    """Generate energy function E.

    Usage: E = E_generator(beta, eta, h)
    Formula:
        E = h * \sum{x_i} - beta * \sum{x_i x_j} - eta * \sum{x_i y_i}
    """
    def E(x, y):
        """Calculate energy for matrices x, y.

        Note: the computation is not localized, so this is quite expensive.
        """
        # sum of products of neighboring paris {xi, yi}
        xxm = np.zeros_like(x)
        xxm[:-1, :] = x[1:, :]  # down
        xxm[1:, :] += x[:-1, :]  # up
        xxm[:, :-1] += x[:, 1:]  # right
        xxm[:, 1:] += x[:, :-1]  # left
        xx = np.sum(xxm * x)
        xy = np.sum(x * y)
        xsum = np.sum(x)
        return h * xsum - beta * xx - eta * xy

    return E      

注意到如果用 xi0 ~ xi3 表示 xi 的四個鄰居,則 xi * xi0 + xi * xi1 + xi * xi2 + xi * xi3 = + xi * (xi1 + ... + xi3),即乘法結合律,是以我們可以先将鄰居相加,再與

x

相乘。

因為邊界元素不一定有四個鄰居,−βxixj這項存在邊界問題,我們需要特别處理,利用 numpy 的 fancy index,寫起來并不困難,即上面代碼中的:

# sum of products of neighboring paris {xi, yi}
xxm = np.empty_like(x)
xxm[:-1, :] = x[1:, :]  # down
xxm[1:, :] += x[:-1, :]  # up
xxm[:, :-1] += x[:, 1:]  # right
xxm[:, 1:] += x[:, :-1]  # left
xx = np.sum(xxm * x)      

即将每個元素的鄰居都都加起來存儲在

x

中,然後用

np.sum(xxm * x)

求得積。

注意這裡生成的

E

每次都要對矩陣中的所有元素進行運算,是以即使有 numpy 加持,開銷依然較大。後面我們會按照需求進行優化。

得到了 energy function 的表示,接下來我們需要做的就是想辦法讓處理過後的圖像具備盡可能低的 energy,進而獲得更高的機率 P(X=x)。

注意如果我們固定 y 作為先驗知識(假設噪聲圖不變),我們所求的機率就變成了 p(x|y),這種模型叫做 Ising Model,在統計實體中有廣泛的應用。這樣我們的問題就成了以 y 為基準,想辦法不斷變動 x,然後找出最接近原圖的 x。

Iterated Conditional Modes

一種最簡單的辦法是我們先将 x 初始化為 y,然後周遊每一個元素,對每個元素分别嘗試 1 和 -1 兩種狀态,選擇能夠得到更高的 energy 的那個,實際上相當于一種貪心的政策。這種方法稱為 Iterated Conditional Modes(ICM),由 Julian Besag 在 1986 年的論文 On the Statistical Analysis of Dirty Pictures 中提出(這篇論文在 80 年代英國數學家所著論文裡引用數排名第一……)。

因為 ICM 的每一步實際上固定住了其他元素,隻變動目前周遊到的那個元素,是以我們可以将 E的計算 localize,隻對受影響的那一小片區域重新計算。我們可以讓 function factory

E_generator

傳回兩個版本的 E,一個是全局的,用于第一次計算 E,一個是局部的,用于計算某個元素兩種狀态下的 E。

def E_generator(beta, eta, h):

    def E(x, y):
        ...  # as shown before

    def is_valid(i, j, shape):
        """Check if coordinate i, j is valid in shape."""
        return i >= 0 and j >= 0 and i < shape[0] and j < shape[1]

    def localized_E(E1, i, j, x, y):
        """Localized version of Energy function E.

        Usage: old_x_ij, new_x_ij, E1, E2 = localized_E(Ecur, i, j, x, y)
        """
        oldval = x[i, j]
        newval = oldval * -1  # flip
        # local computations
        E2 = E1 - (h * oldval) + (h * newval)
        E2 = E2 + (eta * y[i, j] * oldval) - (eta * y[i, j] * newval)
        adjacent = [(0, 1), (0, -1), (1, 0), (-1, 0)]
        neighbors = [x[i + di, j + dj] for di, dj in adjacent
                     if is_valid(i + di, j + dj, x.shape)]
        E2 = E2 + beta * sum(a * oldval for a in neighbors)
        E2 = E2 - beta * sum(a * newval for a in neighbors)
        return oldval, newval, E1, E2

    return E, localized_E      

localized_E

x[i, j]

更改為另一狀态,減去原來狀态在 E 裡對應的那幾項,計算新狀态對應的項,再加上去。

adjacent

neighbors

均是為了過濾掉非法的邊界元素而設。

ICM 的實作如下,為了友善畫圖檢視 energy 變化情況,在每周遊完一行的時候,我們可以記錄目前的時間和 energy:

def ICM(y, E, localized_E):
    """Greedy version of simulated_annealing()."""
    x = np.array(y)
    Ebest = Ecur = E(x, y)  # initial energy
    initial_time = time.time()
    energy_record = [[0.0, ], [Ebest, ]]

    accept, reject = 0, 0
    for idx in np.ndindex(y.shape):  # for each pixel in the matrix
        old, new, E1, E2 = localized_E(Ecur, idx[0], idx[1], x, y)
        if (E2 < Ebest):
            Ecur, x[idx] = E2, new
            Ebest = E2  # update Ebest
        else:
            Ecur, x[idx] = E1, old

    # record time and Ebest of this iteration
    end_time = time.time()
    energy_record[0].append(end_time - initial_time)
    energy_record[1].append(Ebest)

    return x, energy_record      

為了友善将二值圖像素在{0, 255}(黑與白的顔色值)與 {-1, +1}之間轉換,寫一個小函數:

def sign(data, translate):
    """Map a dictionary for the element of data.

    Example:
        To convert every element in data with value 0 to -1, 255 to 1,
        use `signed = sign(data, {0: -1, 255: 1})`
    """
    temp = np.array(data)
    return np.vectorize(lambda x: translate[x])(temp)      

接下來在在 ipython 裡看看結果。設

beta, eta, h = 1e-3, 2.1e-3, 0.0

In [21]: beta, eta, h = 1e-3, 2.1e-3, 0.0
In [22]: im = Image.open('flipped.png')
In [23]: data = sign(im.getdata(), {0: -1, 255: 1})  # convert to {-1, 1}
In [24]: y = data.reshape(im.size[::-1]) # to matrix
In [25]: E, localized_E = E_generator(beta, eta, h)
In [26]: result, energy_record = ICM(y, E, localized_E)
In [27]: result = sign(result, {-1: 0, 1: 255})
In [28]: output_image = Image.fromarray(result).convert('1', dither=Image.NONE)      

檢視去噪結果

output_image.show()      
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

對比一下原圖與噪聲圖,可以看到去掉了相當一部分噪點,但比較密集的噪點形成團塊留了下來(因為我們隻算四連通的相鄰像素,是以很小範圍内的噪點聚集都會形成團塊):

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

檢視 time-energy 關系圖:

In [30]: plt.plot(*energy_record)
Out[30]: [<matplotlib.lines.Line2D at 0xa63ecd0>]

In [31]: plt.xlabel('Time(s)')
Out[31]: <matplotlib.text.Text at 0xa5efb70>

In [32]: plt.ylabel('Energy')
Out[32]: <matplotlib.text.Text at 0xa5f3b70>

In [33]: plt.show()      
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

看看降噪之後與原圖的相似度:

In [60]: original = Image.open('in.png')
In [61]: x = np.reshape(original.getdata(), original.size[::-1])
In [62]: 1 - np.count_nonzero(x - result) / float(x.size)
Out[62]: 0.9621064814814815      

可以看到大約 96% 的像素與原圖一緻。不過光從視覺上看,降噪過後的圖依然有不少明顯的噪點,這是因為 ICM 采取的類似貪心的政策使得它容易陷入局部最優(local optimum)。如果想要得到一個更好的降噪結果,我們顯然要采取一種能夠得到全局最優(global optimum)的政策。

模拟退火

其實這個問題可以看成一個搜尋問題:在所有 x 的兩種狀态組成的狀态空間裡(假設有 n 個像素,那麼狀态空間大小為 2n),找到能使 energy 最低的狀态。由于狀态空間大小呈指數級增長,僅僅是對于一個有 240 × 180 = 43,200 個像素的圖檔來說,這個狀态空間就已經不可能使用暴力搜尋解決了(實際上是 NP-Hard 的),是以我們需要考慮其他的搜尋政策。這裡我們可以嘗試使用模拟退火(Simulated annealing),在有限的時間内找到盡可能好的解。本方法由 Stuart Geman 與 Donald Geman (兄弟喲) 在 1984 年的論文 Stochastic Relaxation, Gibbs Distributions, and the Bayesian Restoration of Images 中提出,并且文中對模拟退火可收斂至全局最優進行了詳細的證明(這篇論文和之前 Julian Besag 的那篇都是 ISI Highly Cited Paper,引用數相當魔性……)。

模拟退火需要一個控制溫度的 schedule,具體怎樣可以自己調,隻需要在疊代次數 k 接近最大疊代次數 kmax 時逼近 0 即可,這裡設定為:

def temperature(k, kmax):
    """Schedule the temperature for simulated annealing."""
    return 1.0 / 500 * (1.0 / k - 1.0 / kmax)      

模拟退火的核心在于當局部變化導緻狀況惡化(這裡為能量變大)時,依據目前溫度 t,設該變化對全局最優有利的機率為 e△E/t,按照這個機率來确定是否保留這個變化,即所謂的 transition probabilities,如下:

def prob(E1, E2, t):
    """Probability transition function for simulated annealing."""
    return 1 if E1 > E2 else np.exp((E1 - E2) / t)      

下面是模拟退火方法的實作,因為模拟退火會進行多次疊代,不斷逼近最優,我們這裡可以将中途結果輸出來看看(這裡順手往參數清單塞了個存中間結果的檔案夾……)

def simulated_annealing(y, kmax, E, localized_E, temp_dir):
    """Simulated annealing process for image denoising.

    Parameters
    ----------
    y: array_like
        The noisy binary image matrix ranging in {-1, 1}.
    kmax: int
        The maximun number of iterations.
    E: function
        Energy function.
    localized_E: function
        Localized version of E.
    temp_dir: path
        Directory to save temporary results.

    Returns
    ----------
    x: array_like
        The denoised binary image matrix ranging in {-1, 1}.
    energy_record:
        [time, Ebest] records for plotting.
    """
    x = np.array(y)
    Ebest = Ecur = E(x, y)  # initial energy
    initial_time = time.time()
    energy_record = [[0.0, ], [Ebest, ]]

    for k in range(1, kmax + 1):  # iterate kmax times
        start_time = time.time()
        t = temperature(k, kmax + 1)
        print "k = %d, Temperature = %.4e" % (k, t)
        accept, reject = 0, 0  # record the acceptance of alteration
        for idx in np.ndindex(y.shape):  # for each pixel in the matrix
            old, new, E1, E2 = localized_E(Ecur, idx[0], idx[1], x, y)
            p, q = prob(E1, E2, t), random()
            if p > q:
                accept += 1
                Ecur, x[idx] = E2, new
                if (E2 < Ebest):
                    Ebest = E2  # update Ebest
            else:
                reject += 1
                Ecur, x[idx] = E1, old

        # record time and Ebest of this iteration
        end_time = time.time()
        energy_record[0].append(end_time - initial_time)
        energy_record[1].append(Ebest)

        print "k = %d, accept = %d, reject = %d" % (k, accept, reject)
        print "k = %d, %.1f seconds" % (k, end_time - start_time)

        # save temporary results
        temp = sign(x, {-1: 0, 1: 255})
        temp_path = os.path.join(temp_dir, 'temp-%d.png' % (k))
        Image.fromarray(temp).convert('1', dither=Image.NONE).save(temp_path)
        print "[Saved]", temp_path

    return x, energy_record      

再看看結果,設

beta, eta, h, kmax = 1e-3, 2.1e-3, 0.0, 15

從上到下,從左到右,k = 1, 5, 10, 15

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

time-energy 關系圖

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

用對 ICM 一樣的方法進行統計,結果是 99.16% 的像素與原圖一緻。

最終結果:原圖->噪聲圖->ICM->模拟退火

用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪
用 Python 通過馬爾可夫随機場(MRF)與 Ising Model 進行二值圖降噪

後話

無論是 ICM 還是模拟退火,都是通過最小化能量來找到原圖的近似。後來 Greig, Porteous 和 Seheult 提出了使用 graph cuts 的模型,将能量值的最小化轉化為最大化解的後驗估計(a posteriori estimate),進而轉為計算機科學裡常見的 max-flow/min-cut 的問題,求解後能夠得到更好的效果(參考 D.M. Greig, B.T. Porteous and A.H. Seheult (1989), Exact maximum a posteriori estimation for binary images, Journal of the Royal Statistical Society Series B, 51, 271–279.)。