Posted at

graphcut - networkx

More than 1 year has passed since last update.


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()