graphcut - networkx

networkxを用いたgraphcut

guraphcutの記事が見当たらなかったので書いてみました。
data項は大津の閾値判定法から前景らしさ、背景らしさを決定し
smooth項は上下左右の画素値の差を元に決定しました。

graphcut.py

import networkx as nx
from scipy.misc import imread
import numpy as np
import math
import matplotlib.pyplot as plt
import cv2

lumda = 1000
k = 0.001


# 大津の手法
def threshold_otsu(gray, min_value=0, max_value=255):
    # ヒストグラムの算出
    hist = [np.sum(gray == i) for i in range(256)]

    s_max = (0, -10)

    for th in range(256):

        # クラス1とクラス2の画素数を計算
        n1 = sum(hist[:th])
        n2 = sum(hist[th:])

        # クラス1とクラス2の画素値の平均を計算
        if n1 == 0:
            mu1 = 0
        else:
            mu1 = sum([i * hist[i] for i in range(0, th)]) / n1
        if n2 == 0:
            mu2 = 0
        else:
            mu2 = sum([i * hist[i] for i in range(th, 256)]) / n2

        # クラス間分散の分子を計算
        s = n1 * n2 * (mu1 - mu2) ** 2

        # クラス間分散の分子が最大のとき、クラス間分散の分子と閾値を記録
        if s > s_max[1]:
            s_max = (th, s)

    # クラス間分散が最大のときの閾値を取得
    t = s_max[0]
    print(t)

    return t


def add_tedges(g, t, img):
    row = img.shape[0]
    column = img.shape[1]
    for i in range(row):
        for j in range(column):
            p_v = img[i, j]
            if img[i, j] < t:
                s_c = p_v / t * 125
                t_c = 255 - s_c

            else:
                s_c = (p_v - t) / (255 - t) * 125 + 125
                t_c = 255 - s_c
            temp = i * 1392 + j
            g.add_edge('s', temp, capacity=s_c)
            g.add_edge(temp, 't', capacity=t_c)

    gro = imread('%05d.png' % 200)
    x, y = np.where(gro > 200)
    for i in range(x.shape[0]):
        if img[x[i], y[i]] >= 100:
            temp = x[i] * column + y[i]
            g.add_edge('s', temp, capacity=1000)
    x, y = np.where(img == 0)
    for i in range(x.shape[0]):
        temp = x[i] * column + y[i]
        g.add_edge(temp, 't', capacity=1000)
    return g


def column_components(g, img):
    for i in range(img.shape[0] - 1):
        for j in range(img.shape[1]):
            capacity = lumda * math.exp((-k) * (img[i, j] - img[i + 1, j]) ** 2)
            row_in = img.shape[1] * i
            g.add_edge(j + row_in, img.shape[1] + j + row_in, capacity=capacity)
            g.add_edge(img.shape[1] + j + row_in, j + row_in, capacity=capacity)
    return g


def row_components(g, img):
    for i in range(img.shape[0]):
        for j in range(img.shape[1] - 1):
            capacity = lumda * math.exp((-k) * (img[i, j] - img[i, j + 1]) ** 2)
            row_in = img.shape[1] * i
            g.add_edge(j + row_in, 1 + j + row_in, capacity=capacity)
            g.add_edge(1 + j + row_in, j + row_in, capacity=capacity)
    return g


def create_graph(img, t):
    g = nx.DiGraph()
    N = img.shape[0] * img.shape[1]
    g.add_nodes_from(range(N))
    g = column_components(g, img)
    g = row_components(g, img)
    g.add_nodes_from(['s', 't'])
    g = add_tedges(g, t, img)
    return g


if __name__ == '__main__':
    img = imread("result200.tif")
    # img = np.asarray([[1, 255, 3], [4, 5, 255], [0, 8, 255]])

    # 大津の閾値決定法
    t = threshold_otsu(img)

    row = img.shape[0]
    column = img.shape[1]
    g = create_graph(img, t)
    cut_value, partition = nx.minimum_cut(g, 's', 't')
    s, t = partition
    s = list(s)
    t = list(t)
    print(s)
    result = np.zeros((row, column))
    for i in s:
        if i != 's':
            if i == 0:
                result[0, 0] = 1
            else:
                row_n = i // column
                column_n = i % column
                result[row_n, column_n] = 1
    m = np.zeros((1040, 1392, 3))
    m[:, :, 0] = result * 255
    m = m.astype('uint8')
    print(m.shape)
    img = cv2.imread("result200.tif")
    dst = cv2.addWeighted(img, 0.5, m, 0.5, 0)
    plt.imshow(dst), plt.show()