def construct_gcgraph(self, lam): vertex_count = self.cols*self.rows edge_count = 2*(4*vertex_count - 3*(self.rows + self.cols) + 2) self.graph = GCGraph(vertex_count, edge_count) for y in range(self.rows): for x in range(self.cols): vertex_index = self.graph.add_vertex() # add-node and return its index color = self.img[y, x] if self._mask[y, x] == self._GC_PR_BGD or self._mask[y, x] == self._GC_PR_FGD: fromSource = -np.log(self.BGD_GMM.prob_pixel_GMM(color)) toSink = -np.log(self.FGD_GMM.prob_pixel_GMM(color)) elif self._mask[y, x] == self._GC_BGD: fromSource = 0 toSink = lam else: fromSource = lam toSink = 0 self.graph.add_term_weights(vertex_index, fromSource, toSink) if x > 0: w = self.left_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-1, w, w) if x > 0 and y > 0: w = self.upleft_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-self.cols-1, w, w) if y > 0: w = self.up_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-self.cols, w, w) if x < self.cols - 1 and y > 0: w = self.upright_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-self.cols+1, w, w)
def construct_gcgraph(self, lam):
'''Construct a GCGraph with the Gibbs Energy'''
'''The vertexs of the graph are the pixels, and the edges are constructed by two parts,
the first part of which are the edges that connect each vertex with Sink Point(the background) and the Source Point(the foreground),
and the weight of which is the first term in Gibbs Energy;
the second part of the edges are those that connect each vertex with its neighbourhoods,
and the weight of which is the second term in Gibbs Energy.'''
vertex_count = self.cols * self.rows
edge_count = 2 * (4 * vertex_count - 3 * (self.rows + self.cols) + 2)
self.graph = GCGraph(vertex_count, edge_count)
for y in range(self.rows):
for x in range(self.cols):
vertex_index = self.graph.add_vertex(
) # add-node and return its index
color = self.img[y, x]
# '''Set t-weights: Calculate the weight of each vertex with Sink node and Source node'''
if self._mask[y, x] == self._GC_PR_BGD or self._mask[
y, x] == self._GC_PR_FGD:
# For each vertex, calculate the first term of G.E. as it be the BGD or FGD, and set them respectively as weight to t/s.
fromSource = -np.log(self.BGD_GMM.prob_pixel_GMM(color))
toSink = -np.log(self.FGD_GMM.prob_pixel_GMM(color))
# print(np.exp(-fromSource), np.exp(-toSink))
# print(fromSource)
elif self._mask[y, x] == self._GC_BGD:
# For the vertexs that are Background pixels, t-weight with Source = 0, with Sink = lam
fromSource = 0
toSink = lam
else:
# GC_FGD
fromSource = lam
toSink = 0
# print(fromSource, toSink)
self.graph.add_term_weights(vertex_index, fromSource, toSink)
'''Set n-weights and n-link, Calculate the weights between two neighbour vertexs, which is also the second term in Gibbs Energy(the smooth term)'''
if x > 0:
w = self.left_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index - 1, w, w)
if x > 0 and y > 0:
w = self.upleft_weight[y, x]
self.graph.add_edges(vertex_index,
vertex_index - self.cols - 1, w, w)
if y > 0:
w = self.up_weight[y, x]
self.graph.add_edges(vertex_index,
vertex_index - self.cols, w, w)
if x < self.cols - 1 and y > 0:
w = self.upright_weight[y, x]
self.graph.add_edges(vertex_index,
vertex_index - self.cols + 1, w, w)
def construct_gcgraph(self, lam):
'''Construct a GCGraph with the Gibbs Energy'''
'''The vertexs of the graph are the pixels, and the edges are constructed by two parts,
the first part of which are the edges that connect each vertex with Sink Point(the background) and the Source Point(the foreground),
and the weight of which is the first term in Gibbs Energy;
the second part of the edges are those that connect each vertex with its neighbourhoods,
and the weight of which is the second term in Gibbs Energy.'''
print('Constructing grabcut graph...')
vertex_count = self.cols * self.rows * self.depth
edge_count = 2 * (9 * self.rows * self.cols * self.depth - 5 *
(self.rows * self.cols + self.rows * self.depth +
self.cols * self.depth) + 2 *
(self.rows + self.cols + self.depth)
) #有向图的边数，每两个顶点之间有正反两条边
self.graph = GCGraph(vertex_count, edge_count)
[
self._construct_gcgraph(lam, z, y, x) for z in range(self.depth)
for y in range(self.rows) for x in range(self.cols)
]
def construct_gcgraph(self, lam): '''Construct a GCGraph with the Gibbs Energy''' '''The vertexs of the graph are the pixels, and the edges are constructed by two parts, the first part of which are the edges that connect each vertex with Sink Point(the background) and the Source Point(the foreground), and the weight of which is the first term in Gibbs Energy; the second part of the edges are those that connect each vertex with its neighbourhoods, and the weight of which is the second term in Gibbs Energy.''' vertex_count = self.cols*self.rows edge_count = 2*(4*vertex_count - 3*(self.rows + self.cols) + 2) self.graph = GCGraph(vertex_count, edge_count) for y in range(self.rows): for x in range(self.cols): vertex_index = self.graph.add_vertex() # add-node and return its index color = self.img[y, x] # '''Set t-weights: Calculate the weight of each vertex with Sink node and Source node''' if self._mask[y, x] == self._GC_PR_BGD or self._mask[y, x] == self._GC_PR_FGD: # For each vertex, calculate the first term of G.E. as it be the BGD or FGD, and set them respectively as weight to t/s. fromSource = -np.log(self.BGD_GMM.prob_pixel_GMM(color)) toSink = -np.log(self.FGD_GMM.prob_pixel_GMM(color)) # print(np.exp(-fromSource), np.exp(-toSink)) # print(fromSource) elif self._mask[y, x] == self._GC_BGD: # For the vertexs that are Background pixels, t-weight with Source = 0, with Sink = lam fromSource = 0 toSink = lam else: # GC_FGD fromSource = lam toSink = 0 # print(fromSource, toSink) self.graph.add_term_weights(vertex_index, fromSource, toSink) '''Set n-weights and n-link, Calculate the weights between two neighbour vertexs, which is also the second term in Gibbs Energy(the smooth term)''' if x > 0: w = self.left_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-1, w, w) if x > 0 and y > 0: w = self.upleft_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-self.cols-1, w, w) if y > 0: w = self.up_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-self.cols, w, w) if x < self.cols - 1 and y > 0: w = self.upright_weight[y, x] self.graph.add_edges(vertex_index, vertex_index-self.cols+1, w, w)
class GCClient:
def __init__(self, img, k):
self.k = k # The number of components in each GMM model
self.img = np.asarray(img, dtype = np.float32)
self.img2 = img
self.rows, self.cols = get_size(img)
self.gamma = 50
self.lam = 9*self.gamma
self.beta = 0
self._BLUE = [255,0,0] # rectangle color
self._RED = [0,0,255] # PR BG
self._GREEN = [0,255,0] # PR FG
self._BLACK = [0,0,0] # sure BG
self._WHITE = [255,255,255] # sure FG
self._DRAW_BG = {'color':self._BLACK, 'val':0}
self._DRAW_FG = {'color':self._WHITE, 'val':1}
self._DRAW_PR_FG = {'color':self._GREEN, 'val':3}
self._DRAW_PR_BG = {'color':self._RED, 'val':2}
self._rect = [0, 0, 1, 1]
self._drawing = False # flag for drawing curves
self._rectangle = False # flag for drawing rect
self._rect_over = False # flag to check if rect drawn
self._thickness = 5 # brush thickness
self._GC_BGD = 0 #{'color' : BLACK, 'val' : 0}
self._GC_FGD = 1 #{'color' : WHITE, 'val' : 1}
self._GC_PR_BGD = 2 #{'color' : GREEN, 'val' : 3}
self._GC_PR_FGD = 3 #{'color' : RED, 'val' : 2}
self.calc_beta()
self.calc_nearby_weight()
self._DRAW_VAL = None
self._mask = np.zeros([self.rows, self.cols], dtype = np.uint8) # Init the mask
self._mask[:, :] = self._GC_BGD
def calc_beta(self):
beta = 0
self._left_diff = self.img[:, 1:] - self.img[:, :-1] # Left-difference
self._upleft_diff = self.img[1:, 1:] - self.img[:-1, :-1] # Up-Left difference
self._up_diff = self.img[1:, :] - self.img[:-1, :] # Up-difference
self._upright_diff = self.img[1:, :-1] - self.img[:-1, 1:] # Up-Right difference
beta = (self._left_diff*self._left_diff).sum() + (self._upleft_diff*self._upleft_diff).sum() \
+ (self._up_diff*self._up_diff).sum() + (self._upright_diff*self._upright_diff).sum() # According to the formula
self.beta = 1/(2*beta/(4*self.cols*self.rows - 3*self.cols - 3*self.rows + 2))
def calc_nearby_weight(self):
self.left_weight = np.zeros([self.rows, self.cols])
self.upleft_weight = np.zeros([self.rows, self.cols])
self.up_weight = np.zeros([self.rows, self.cols])
self.upright_weight = np.zeros([self.rows, self.cols])
for y in range(self.rows):
for x in range(self.cols):
color = self.img[y, x]
if x >= 1:
diff = color - self.img[y, x-1]
self.left_weight[y, x] = self.gamma*np.exp(-self.beta*(diff*diff).sum())
if x >= 1 and y >= 1:
diff = color - self.img[y-1, x-1]
self.upleft_weight[y, x] = self.gamma/np.sqrt(2) * np.exp(-self.beta*(diff*diff).sum())
if y >= 1:
diff = color - self.img[y-1, x]
self.up_weight[y, x] = self.gamma*np.exp(-self.beta*(diff*diff).sum())
if x+1 < self.cols and y >= 1:
diff = color - self.img[y-1, x+1]
self.upright_weight[y, x] = self.gamma/np.sqrt(2)*np.exp(-self.beta*(diff*diff).sum())
def init_mask(self, event, x, y, flags, param):
if event == cv2.EVENT_RBUTTONDOWN:
self._rectangle = True
self._ix,self._iy = x,y
elif event == cv2.EVENT_MOUSEMOVE:
if self._rectangle == True:
self.img = self.img2.copy()
cv2.rectangle(self.img,(self._ix,self._iy),(x,y),self._BLUE,2)
self._rect = [min(self._ix,x),min(self._iy,y),abs(self._ix-x),abs(self._iy-y)]
self.rect_or_mask = 0
elif event == cv2.EVENT_RBUTTONUP:
self._rectangle = False
self._rect_over = True
cv2.rectangle(self.img,(self._ix,self._iy),(x,y),self._BLUE,2)
self._rect = [min(self._ix,x),min(self._iy,y),abs(self._ix-x),abs(self._iy-y)]
self.rect_or_mask = 0
self._mask[self._rect[1]+self._thickness:self._rect[1]+self._rect[3]-self._thickness, self._rect[0]+self._thickness:self._rect[0]+self._rect[2]-self._thickness] = self._GC_PR_FGD
if event == cv2.EVENT_LBUTTONDOWN:
if self._rect_over == False:
print("Continue")
else:
self._drawing == True
cv2.circle(self.img, (x, y), self._thickness, self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness, self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_MOUSEMOVE:
if self._drawing == True:
cv2.circle(self.img, (x, y), self._thickness, self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness, self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_LBUTTONUP:
if self._drawing == True:
self._drawing = False
cv2.circle(self.img, (x, y), self._thickness, self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness, self._DRAW_VAL['val'], -1)
def init_with_kmeans(self):
print(self.cols*self.rows)
print(len(list(np.where(self._mask == 0))[1]))
max_iter = 2 # Max-iteration count for Kmeans
self._bgd = np.where(np.logical_or(self._mask == self._GC_BGD, self._mask == self._GC_PR_BGD)) # Find the places where pixels in the mask MAY belong to BGD.
self._fgd = np.where(np.logical_or(self._mask == self._GC_FGD, self._mask == self._GC_PR_FGD)) # Find the places where pixels in the mask MAY belong to FGD.
self._BGDpixels = self.img[self._bgd]
self._FGDpixels = self.img[self._fgd]
KMB = kmeans(self._BGDpixels, dim = 3, n = 5, max_iter = max_iter) # The Background Model by kmeans
KMF = kmeans(self._FGDpixels, dim = 3, n = 5, max_iter = max_iter) # The Foreground Model by kmeans
KMB.run()
KMF.run()
self._BGD_by_components = KMB.output()
self._FGD_by_components = KMF.output()
self.BGD_GMM = GMM() # The GMM Model for BGD
self.FGD_GMM = GMM() # The GMM Model for FGD
for ci in range(5):
for pixel in self._BGD_by_components[ci]:
self.BGD_GMM.add_pixel(pixel, ci)
for pixel in self._FGD_by_components[ci]:
self.FGD_GMM.add_pixel(pixel, ci)
self.BGD_GMM.learning()
self.FGD_GMM.learning()
def assign_GMM_components(self):
self.components_index = np.zeros([self.rows, self.cols], dtype = np.uint)
for y in range(self.rows):
for x in range(self.cols):
pixel = self.img[y, x]
self.components_index[y, x] = self.BGD_GMM.most_likely_pixel_component(pixel) if (self._mask[y, x] \
== self._GC_BGD or self._mask[y, x] == self._GC_PR_BGD) else self.FGD_GMM.most_likely_pixel_component(pixel)
def _assign_GMM_components(self):
self.components_index = np.zeros([self.rows, self.cols], dtype = np.uint)
self.components_index[self._bgd] = [i[0] for i in self.BGD_GMM.vec_pix_comp(self.img[self._bgd])]
self.components_index[self._fgd] = [i[0] for i in self.FGD_GMM.vec_pix_comp(self.img[self._fgd])]
def learn_GMM_parameters(self):
for ci in range(5):
bgd_ci = np.where(np.logical_and(self.components_index == ci, np.logical_or(self._mask == self._GC_BGD, self._mask == self._GC_PR_BGD)))
fgd_ci = np.where(np.logical_and(self.components_index == ci, np.logical_or(self._mask == self._GC_FGD, self._mask == self._GC_PR_FGD)))
for pixel in self.img[bgd_ci]:
self.BGD_GMM.add_pixel(pixel, ci)
for pixel in self.img[fgd_ci]:
self.FGD_GMM.add_pixel(pixel, ci)
self.BGD_GMM.learning()
self.FGD_GMM.learning()
def construct_gcgraph(self, lam):
vertex_count = self.cols*self.rows
edge_count = 2*(4*vertex_count - 3*(self.rows + self.cols) + 2)
self.graph = GCGraph(vertex_count, edge_count)
for y in range(self.rows):
for x in range(self.cols):
vertex_index = self.graph.add_vertex() # add-node and return its index
color = self.img[y, x]
if self._mask[y, x] == self._GC_PR_BGD or self._mask[y, x] == self._GC_PR_FGD:
fromSource = -np.log(self.BGD_GMM.prob_pixel_GMM(color))
toSink = -np.log(self.FGD_GMM.prob_pixel_GMM(color))
elif self._mask[y, x] == self._GC_BGD:
fromSource = 0
toSink = lam
else:
fromSource = lam
toSink = 0
self.graph.add_term_weights(vertex_index, fromSource, toSink)
if x > 0:
w = self.left_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-1, w, w)
if x > 0 and y > 0:
w = self.upleft_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-self.cols-1, w, w)
if y > 0:
w = self.up_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-self.cols, w, w)
if x < self.cols - 1 and y > 0:
w = self.upright_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-self.cols+1, w, w)
def estimate_segmentation(self):
a = self.graph.max_flow()
for y in range(self.rows):
for x in range(self.cols):
if self._mask[y, x] == self._GC_PR_BGD or self._mask[y, x] == self._GC_PR_FGD:
if self.graph.insource_segment(y*self.cols+x): # Vertex Index
self._mask[y, x] = self._GC_PR_FGD
else:
self._mask[y, x] = self._GC_PR_BGD
def iter(self, n):
for i in range(n):
self.assign_GMM_components()
self.learn_GMM_parameters()
self.construct_gcgraph(self.lam)
self.estimate_segmentation()
def run(self):
self.init_with_kmeans()
self.iter(1)
def _smoothing(self):
for y in range(1, self.rows-2):
for x in range(1, self.cols-2):
a = self._mask[x-1, y]
b = self._mask[x+1, y]
c = self._mask[x, y-1]
d = self._mask[x, y+1]
if a==b==3 or a==c==3 or a==d==3 or b==c==3 or b==d==3 or c==d==3:
self._mask[x, y] = 3
def show(self, output):
FGD = np.where((self._mask == 1) + (self._mask == 3), 255, 0).astype('uint8')
output = cv2.bitwise_and(self.img2, self.img2, mask = FGD)
return output
class GCClient:
'''The engine of grabcut'''
def __init__(self, img, k):
self.k = k # The number of components in each GMM model
self.img3d = np.asarray(img, dtype=np.float32)
self.img3d_2 = self.img3d.copy()
self.depth, self.rows, self.cols = img.shape
print('img.shape:', img.shape)
self.gamma = 50
self.lam = 28 * self.gamma
self.beta = 0
self._BLACK = 0 # sure BG
self._GRAY1 = 80 # PR BG
self._GRAY2 = 160 # PR FG
self._WHITE = 255 # sure FG
self._GC_BGD = 0
self._GC_FGD = 1
self._GC_PR_BGD = 2
self._GC_PR_FGD = 3
self._DRAW_BG = {'color': self._BLACK, 'val': self._GC_BGD}
self._DRAW_FG = {'color': self._WHITE, 'val': self._GC_FGD}
self._DRAW_PR_BG = {'color': self._GRAY1, 'val': self._GC_PR_BGD}
self._DRAW_PR_FG = {'color': self._GRAY2, 'val': self._GC_PR_FGD}
# setting up flags
self._rect1 = [0, 0, 1, 1, -1] #[x1,x2, y1,y2, z1] 矩形1左上角的坐标以及宽、高、深度坐标
self._rect2 = [0, 0, 1, 1, -1] #[x1,x2, y1,y2, z2] 矩形2左上角的坐标以及宽、高、深度坐标
self._cube = [0, 0, 1, 1, 0, 1] #[x1,x2, y1,y2, z1,z2]前景存在的立方体的坐标
self._drawing = False # flag for drawing curves
self._rectangle1 = False # flag for drawing rectangle1
self._rectangle2 = False # flag for drawing rectangle2
self._rect1_over = False # flag to check if rect drawn
self._rect2_over = False # flag to check if rect drawn
self._rect_or_mask = 0 # flag for selecting rect or mask mode
self._thickness = 2 # brush thickness
self._DRAW_VAL = None #color of brush
self._mask = np.zeros([self.depth, self.rows, self.cols],
dtype=np.uint8) # Init the mask
self._mask[:, :, :] = self._GC_BGD
self._mask3d = self._mask.astype('float32')
self.calc_nearby_weight()
def calc_nearby_weight(self):
'''STEP 1:'''
'''Calculate Beta -- The Exp Term of Smooth Parameter in Gibbs Energy'''
'''beta = 1/(2*average(sqrt(||pixel[i] - pixel[j]||)))'''
'''Beta is used to adjust the difference of two nearby pixels in high or low contrast rate'''
'''STEP 2:'''
'''Calculate the weight of the edge of each pixel with its nearby pixel, as each pixel is regarded
as a vertex of the graph'''
'''The weight of each direction is saved in a image the same size of the original image'''
'''weight = gamma * 1/distance<m,n> * exp(-beta*(diff*diff))'''
#diff1 = self.img3d[:-1,:-1,:-1] - self.img3d[1:,1:,1:]
diff2 = self.img3d[:-1, :-1, :] - self.img3d[1:, 1:, :]
#diff3 = self.img3d[:-1,:-1,1:] - self.img3d[1:,1:,:-1]
diff4 = self.img3d[:-1, :, :-1] - self.img3d[1:, :, 1:]
diff5 = self.img3d[:-1, :, :] - self.img3d[1:, :, :]
diff6 = self.img3d[:-1, :, 1:] - self.img3d[1:, :, :-1]
#diff7 = self.img3d[:-1,1:,:-1] - self.img3d[1:,:-1,1:]
diff8 = self.img3d[:-1, 1:, :] - self.img3d[1:, :-1, :]
#diff9 = self.img3d[:-1,1:,1:] - self.img3d[1:,:-1,:-1]
diff10 = self.img3d[:, :-1, :-1] - self.img3d[:, 1:, 1:]
diff11 = self.img3d[:, :-1, :] - self.img3d[:, 1:, :]
diff12 = self.img3d[:, :-1, 1:] - self.img3d[:, 1:, :-1]
diff13 = self.img3d[:, :, :-1] - self.img3d[:, :, 1:]
beta = (diff2*diff2).sum() + (diff4*diff4).sum() + (diff5*diff5).sum() \
+ (diff6*diff6).sum() + (diff8*diff8).sum() + (diff10*diff10).sum() \
+ (diff11*diff11).sum() + (diff12*diff12).sum() + (diff13*diff13).sum()
self.beta = (9 * self.rows * self.cols * self.depth -
5 * self.rows * self.cols - 5 * self.rows * self.depth -
5 * self.cols * self.depth + 2 * self.rows +
2 * self.cols + 2 * self.depth) / (2 * beta)
print('self.beta:', self.beta)
# Use the formula to calculate the weight
self.weight2 = self.gamma * 0.707 * np.exp(-self.beta *
(diff2 * diff2))
self.weight4 = self.gamma * 0.707 * np.exp(-self.beta *
(diff4 * diff4))
self.weight5 = self.gamma * np.exp(-self.beta * (diff5 * diff5))
self.weight6 = self.gamma * 0.707 * np.exp(-self.beta *
(diff6 * diff6))
self.weight8 = self.gamma * 0.707 * np.exp(-self.beta *
(diff8 * diff8))
self.weight10 = self.gamma * 0.707 * np.exp(-self.beta *
(diff10 * diff10))
self.weight11 = self.gamma * np.exp(-self.beta * (diff11 * diff11))
self.weight12 = self.gamma * 0.707 * np.exp(-self.beta *
(diff12 * diff12))
self.weight13 = self.gamma * np.exp(-self.beta * (diff13 * diff13))
'''The following function is derived from the sample of opencv sources'''
def init_mask(self, event, x, y, flags, param):
'''Init the mask with interactive movements'''
'''Notice: the elements in the mask should be within the follows:
"GC_BGD":The pixel belongs to background;
"GC_FGD":The pixel belongs to foreground;
"GC_PR_BGD":The pixel MAY belongs to background;
"GC_PR_FGD":The pixel MAY belongs to foreground;'''
# Draw Rectangle1
if self._rect1_over == False:
if event == cv2.EVENT_LBUTTONDOWN:
self._rectangle1 = True
self._ix, self._iy = x, y
elif event == cv2.EVENT_LBUTTONUP:
if self._rectangle1 == True:
self._rect1_over = True
self._rectangle1 == False
cv2.rectangle(self.img3d_2[index1], (self._ix, self._iy),
(x, y), self._WHITE, self._thickness)
self._rect1 = [
min(self._ix, x),
max(self._ix, x),
min(self._iy, y),
max(self._iy, y), index1
]
elif self._rect2_over == False:
if event == cv2.EVENT_LBUTTONDOWN:
if index1 == self._rect1[4]:
print('Please change the depth to draw anather rectangle!')
else:
self._rectangle2 = True
self._ix, self._iy = x, y
elif event == cv2.EVENT_LBUTTONUP:
if self._rectangle2 == True:
self._rect2_over = True
cv2.rectangle(self.img3d_2[index1], (self._ix, self._iy),
(x, y), self._WHITE, self._thickness)
self._rect2 = [
min(self._ix, x),
max(self._ix, x),
min(self._iy, y),
max(self._iy, y), index1
]
self._cube = [
min(self._rect1[0], self._rect2[0]),
max(self._rect1[1], self._rect2[1]),
min(self._rect1[2], self._rect2[2]),
max(self._rect1[3], self._rect2[3]),
min(self._rect1[4], self._rect2[4]),
max(self._rect1[4], self._rect2[4])
]
print('self._cube:', self._cube)
self._mask[self._cube[4]:self._cube[5] + 1,
self._cube[2]:self._cube[3],
self._cube[0]:self._cube[1]] = self._GC_PR_FGD
self._mask3d = self._mask.astype('float32')
# Notice : The x and y axis in CV2 are inversed to those in numpy.
elif self._DRAW_VAL:
if event == cv2.EVENT_LBUTTONDOWN:
self._drawing = True
cv2.circle(self.img3d_2[index1], (x, y), self._thickness,
self._DRAW_VAL['color'], -1)
cv2.circle(self._mask3d[index1], (x, y), self._thickness,
self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_MOUSEMOVE:
if self._drawing == True:
cv2.circle(self.img3d_2[index1], (x, y), self._thickness,
self._DRAW_VAL['color'], -1)
cv2.circle(self._mask3d[index1], (x, y), self._thickness,
self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_LBUTTONUP:
if self._drawing == True:
self._drawing = False
cv2.circle(self.img3d_2[index1], (x, y), self._thickness,
self._DRAW_VAL['color'], -1)
cv2.circle(self._mask3d[index1], (x, y), self._thickness,
self._DRAW_VAL['val'], -1)
def init_with_kmeans(self):
'''Initialise the BGDGMM and FGDGMM, which are respectively background-model and foreground-model,
using kmeans algorithm'''
print('init with k-means processing...')
self._mask = self._mask3d.astype('uint8')
max_iter = 3 # Max-iteration count for Kmeans
'''In the following two indexings, the np.logical_or is needed in place of or'''
self._bgd = np.where(
np.logical_or(self._mask == self._GC_BGD,
self._mask == self._GC_PR_BGD)
) # Find the places where pixels in the mask MAY belong to BGD.
self._fgd = np.where(
np.logical_or(self._mask == self._GC_FGD,
self._mask == self._GC_PR_FGD)
) # Find the places where pixels in the mask MAY belong to FGD.
self._BGDpixels = self.img3d[self._bgd]
self._FGDpixels = self.img3d[self._fgd]
KMB = kmeans(self._BGDpixels, n=self.k,
max_iter=max_iter) # The Background Model by kmeans
KMF = kmeans(self._FGDpixels, n=self.k,
max_iter=max_iter) # The Foreground Model by kmeans
KMB.run()
KMF.run()
self._BGD_by_components = KMB.output()
self._FGD_by_components = KMF.output()
self.BGD_GMM = GMM() # The GMM Model for BGD
self.FGD_GMM = GMM() # The GMM Model for FGD
self.BGD_GMM.learning(self._BGD_by_components)
self.FGD_GMM.learning(self._FGD_by_components)
print('BGD_GMM:', '\nweights:\n', list(self.BGD_GMM.weights),
'\nmeans:\n', list(self.BGD_GMM.means), '\nvars:\n',
list(self.BGD_GMM.vars), '\n')
print('FGD_GMM:', '\nweights:\n', list(self.FGD_GMM.weights),
'\nmeans:\n', list(self.FGD_GMM.means), '\nvars:\n',
list(self.FGD_GMM.vars), '\n')
'''The first step of the iteration in the paper: Assign components of GMMs to pixels,
(the kn in the paper), which is saved in self.components_index'''
def assign_GMM_components(self):
print('Refreshing GMM components...')
self._mask = self._mask3d.astype('uint8')
self.components_index = np.asarray([
self.BGD_GMM.most_likely_pixel_component(pixel) if
(mask == self._GC_BGD) or (mask == self._GC_PR_BGD) else
self.FGD_GMM.most_likely_pixel_component(pixel)
for mat1, mat2 in zip(self.img3d, self._mask)
for row1, row2 in zip(mat1, mat2)
for pixel, mask in zip(row1, row2)
],
dtype='uint8').reshape(
self.depth, self.rows,
self.cols)
'''The second step in the iteration: Learn the parameters from GMM models'''
def learn_GMM_parameters(self):
'''Calculate the parameters of each component of GMM'''
print('Learning GMM parameters...')
self._BGD_by_components = []
self._FGD_by_components = []
for ci in range(self.k):
# The places where the pixel belongs to the ci_th model and background model.
bgd_ci = np.where(
np.logical_and(
self.components_index == ci,
np.logical_or(self._mask == self._GC_BGD,
self._mask == self._GC_PR_BGD)))
self._BGD_by_components.append(self.img3d[bgd_ci])
# The places where the pixel belongs to the ci_th model and foreground model.
fgd_ci = np.where(
np.logical_and(
self.components_index == ci,
np.logical_or(self._mask == self._GC_FGD,
self._mask == self._GC_PR_FGD)))
self._FGD_by_components.append(self.img3d[fgd_ci])
self.BGD_GMM.learning(self._BGD_by_components)
self.FGD_GMM.learning(self._FGD_by_components)
print('BGD_GMM:', '\nweights:\n', list(self.BGD_GMM.weights),
'\nmeans:\n', list(self.BGD_GMM.means), '\nvars:\n',
list(self.BGD_GMM.vars), '\n')
print('FGD_GMM:', '\nweights:\n', list(self.FGD_GMM.weights),
'\nmeans:\n', list(self.FGD_GMM.means), '\nvars:\n',
list(self.FGD_GMM.vars), '\n')
def _construct_gcgraph(self, lam, z, y, x):
'''Set t-weights: Calculate the weight of each vertex with Sink node and Source node'''
vertex_index = self.graph.add_vertex() # add-node and return its index
color = self.img3d[z, y, x]
if self._mask[z, y, x] == self._GC_PR_BGD or self._mask[
z, y, x] == self._GC_PR_FGD:
# For each vertex, calculate the first term of G.E. as it be the BGD or FGD, and set them respectively as weight to t/s.
fromSource = -np.log(self.BGD_GMM.prob_pixel_GMM(color))
toSink = -np.log(self.FGD_GMM.prob_pixel_GMM(color))
elif self._mask[z, y, x] == self._GC_BGD:
# For the vertexs that are Background pixels, t-weight with Source = 0, with Sink = lam
fromSource = 0
toSink = lam
else:
fromSource = lam
toSink = 0
self.graph.add_term_weights(vertex_index, fromSource, toSink)
'''Set n-weights and n-link, Calculate the weights between two neighbour vertexs, which is also the second term in Gibbs Energy(the smooth term)'''
if z > 0:
w = self.weight5[z - 1, y, x]
self.graph.add_edges(vertex_index,
vertex_index - (self.rows * self.cols), w, w)
if y > 0:
w = self.weight2[z - 1, y - 1, x]
self.graph.add_edges(
vertex_index,
vertex_index - (self.rows * self.cols + self.cols), w, w)
if x > 0:
w = self.weight4[z - 1, y, x - 1]
self.graph.add_edges(
vertex_index, vertex_index - (self.rows * self.cols + 1),
w, w)
if x < self.cols - 1:
w = self.weight6[z - 1, y, x]
self.graph.add_edges(
vertex_index, vertex_index - (self.rows * self.cols - 1),
w, w)
if y < self.rows - 1:
w = self.weight8[z - 1, y, x]
self.graph.add_edges(
vertex_index,
vertex_index - (self.rows * self.cols - self.cols), w, w)
if y > 0:
w = self.weight11[z, y - 1, x]
self.graph.add_edges(vertex_index, vertex_index - self.cols, w, w)
if x > 0:
w = self.weight10[z, y - 1, x - 1]
self.graph.add_edges(vertex_index,
vertex_index - (self.cols + 1), w, w)
if x < self.cols - 1:
w = self.weight12[z, y - 1, x]
self.graph.add_edges(vertex_index,
vertex_index - (self.cols - 1), w, w)
if x > 0:
w = self.weight13[z, y, x - 1]
self.graph.add_edges(vertex_index, vertex_index - 1, w, w)
def construct_gcgraph(self, lam):
'''Construct a GCGraph with the Gibbs Energy'''
'''The vertexs of the graph are the pixels, and the edges are constructed by two parts,
the first part of which are the edges that connect each vertex with Sink Point(the background) and the Source Point(the foreground),
and the weight of which is the first term in Gibbs Energy;
the second part of the edges are those that connect each vertex with its neighbourhoods,
and the weight of which is the second term in Gibbs Energy.'''
print('Constructing grabcut graph...')
vertex_count = self.cols * self.rows * self.depth
edge_count = 2 * (9 * self.rows * self.cols * self.depth - 5 *
(self.rows * self.cols + self.rows * self.depth +
self.cols * self.depth) + 2 *
(self.rows + self.cols + self.depth)
) #有向图的边数，每两个顶点之间有正反两条边
self.graph = GCGraph(vertex_count, edge_count)
[
self._construct_gcgraph(lam, z, y, x) for z in range(self.depth)
for y in range(self.rows) for x in range(self.cols)
]
def estimate_segmentation(self):
print('Estimating segmentation...')
a = self.graph.max_flow()
self._mask = np.asarray([
self._GC_PR_FGD if
self.graph.insource_segment(self.rows * self.cols * z +
self.cols * y + x) else self._GC_PR_BGD
for z in range(self.depth) for y in range(self.rows)
for x in range(self.cols)
]).reshape(self.depth, self.rows, self.cols)
def iter(self, n):
for i in range(n):
self.assign_GMM_components()
self.learn_GMM_parameters()
self.construct_gcgraph(self.lam)
self.estimate_segmentation()
def run(self):
self.init_with_kmeans()
self.construct_gcgraph(self.lam)
self.estimate_segmentation()
class GCClient:
'''The engine of grabcut'''
def __init__(self, img, k):
self.k = k # The number of components in each GMM model
self.img = np.asarray(img, dtype = np.float32)
self.img2 = img
self.rows, self.cols = get_size(img)
self.gamma = 50
self.lam = 9*self.gamma
self.beta = 0
self._BLUE = [255,0,0] # rectangle color
self._RED = [0,0,255] # PR BG
self._GREEN = [0,255,0] # PR FG
self._BLACK = [0,0,0] # sure BG
self._WHITE = [255,255,255] # sure FG
self._DRAW_BG = {'color':self._BLACK, 'val':0}
self._DRAW_FG = {'color':self._WHITE, 'val':1}
self._DRAW_PR_FG = {'color':self._GREEN, 'val':3}
self._DRAW_PR_BG = {'color':self._RED, 'val':2}
# setting up flags
self._rect = [0, 0, 1, 1]
self._drawing = False # flag for drawing curves
self._rectangle = False # flag for drawing rect
self._rect_over = False # flag to check if rect drawn
# se;f._rect_or_mask = 100 # flag for selecting rect or mask mode
# self._value = DRAW_FG # drawing initialized to FG
self._thickness = 3 # brush thickness
self._GC_BGD = 0 #{'color' : BLACK, 'val' : 0}
self._GC_FGD = 1 #{'color' : WHITE, 'val' : 1}
self._GC_PR_BGD = 2 #{'color' : GREEN, 'val' : 3}
self._GC_PR_FGD = 3 #{'color' : RED, 'val' : 2}
self.calc_beta()
self.calc_nearby_weight()
self._DRAW_VAL = None
self._mask = np.zeros([self.rows, self.cols], dtype = np.uint8) # Init the mask
self._mask[:, :] = self._GC_BGD
def calc_beta(self):
'''Calculate Beta -- The Exp Term of Smooth Parameter in Gibbs Energy'''
'''beta = 1/(2*average(sqrt(||pixel[i] - pixel[j]||)))'''
'''Beta is used to adjust the difference of two nearby pixels in high or low contrast rate'''
beta = 0
self._left_diff = self.img[:, 1:] - self.img[:, :-1] # Left-difference
self._upleft_diff = self.img[1:, 1:] - self.img[:-1, :-1] # Up-Left difference
self._up_diff = self.img[1:, :] - self.img[:-1, :] # Up-difference
self._upright_diff = self.img[1:, :-1] - self.img[:-1, 1:] # Up-Right difference
beta = (self._left_diff*self._left_diff).sum() + (self._upleft_diff*self._upleft_diff).sum() \
+ (self._up_diff*self._up_diff).sum() + (self._upright_diff*self._upright_diff).sum() # According to the formula
self.beta = 1/(2*beta/(4*self.cols*self.rows - 3*self.cols - 3*self.rows + 2))
# print(self.beta) # According to the paper
@timeit
def calc_nearby_weight(self):
'''Calculate the weight of the edge of each pixel with its nearby pixel, as each pixel is regarded
as a vertex of the graph'''
'''The weight of each direction is saved in a image the same size of the original image'''
'''weight = gamma*exp(-beta*(diff*diff))'''
self.left_weight = np.zeros([self.rows, self.cols])
self.upleft_weight = np.zeros([self.rows, self.cols])
self.up_weight = np.zeros([self.rows, self.cols])
self.upright_weight = np.zeros([self.rows, self.cols])
# Use the formula to calculate the weight
# for y in range(self.rows):
# for x in range(self.cols):
# color = self.img[y, x]
# if x >= 1:
# diff = color - self.img[y, x-1]
# diff.shape = (1, 3)
# self.left_weight[y, x] = self.gamma*np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# if x >= 1 and y >= 1:
# diff = color - self.img[y-1, x-1]
# diff.shape = (1, 3)
# self.upleft_weight[y, x] = self.gamma/np.sqrt(2) * np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# if y >= 1:
# diff = color - self.img[y-1, x]
# diff.shape = (1, 3)
# self.up_weight[y, x] = self.gamma*np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# if x+1 < self.cols and y >= 1:
# diff = color - self.img[y-1, x+1]
# diff.shape = (1, 3)
# self.upright_weight[y, x] = self.gamma/np.sqrt(2)*np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# Use the formula to calculate the weight
for y in range(self.rows):
for x in range(self.cols):
color = self.img[y, x]
if x >= 1:
diff = color - self.img[y, x-1]
# print(np.exp(-self.beta*(diff*diff).sum()))
self.left_weight[y, x] = self.gamma*np.exp(-self.beta*(diff*diff).sum())
if x >= 1 and y >= 1:
diff = color - self.img[y-1, x-1]
self.upleft_weight[y, x] = self.gamma/np.sqrt(2) * np.exp(-self.beta*(diff*diff).sum())
if y >= 1:
diff = color - self.img[y-1, x]
self.up_weight[y, x] = self.gamma*np.exp(-self.beta*(diff*diff).sum())
if x+1 < self.cols and y >= 1:
diff = color - self.img[y-1, x+1]
self.upright_weight[y, x] = self.gamma/np.sqrt(2)*np.exp(-self.beta*(diff*diff).sum())
# self.left_weight[:, 1:] = self.gamma*np.exp(-self.beta*((self._left_diff*self._left_diff).sum()))
# self.upleft_weight[1:, 1:] = self.gamma*np.exp(-self.beta*((self._upleft_diff*self._upleft_diff).sum()))
# self.up_weight[1:, :] = self.gamma*np.exp(-self.beta*(self._up_diff*self._up_diff).sum())
# self.upright_weight = self.gamma*np.exp(-self.beta*(self._upright_diff*self._upright_diff).sum())
'''The following function is derived from the sample of opencv sources'''
def init_mask(self, event, x, y, flags, param):
'''Init the mask with interactive movements'''
'''Notice: the elements in the mask should be within the follows:
"GC_BGD":The pixel belongs to background;
"GC_FGD":The pixel belongs to foreground;
"GC_PR_BGD":The pixel MAY belongs to background;
"GC_PR_FGD":The pixel MAY belongs to foreground;'''
# Draw Rectangle
if event == cv2.EVENT_RBUTTONDOWN:
self._rectangle = True
self._ix,self._iy = x,y
elif event == cv2.EVENT_MOUSEMOVE:
if self._rectangle == True:
self.img = self.img2.copy()
cv2.rectangle(self.img,(self._ix,self._iy),(x,y),self._BLUE,2)
self._rect = [min(self._ix,x),min(self._iy,y),abs(self._ix-x),abs(self._iy-y)]
self.rect_or_mask = 0
elif event == cv2.EVENT_RBUTTONUP:
self._rectangle = False
self._rect_over = True
cv2.rectangle(self.img,(self._ix,self._iy),(x,y),self._BLUE,2)
self._rect = [min(self._ix,x),min(self._iy,y),abs(self._ix-x),abs(self._iy-y)]
# print(" Now press the key 'n' a few times until no further change \n")
self.rect_or_mask = 0
self._mask[self._rect[1]+self._thickness:self._rect[1]+self._rect[3]-self._thickness, self._rect[0]+self._thickness:self._rect[0]+self._rect[2]-self._thickness] = self._GC_PR_FGD
# Notice : The x and y axis in CV2 are inversed to those in numpy.
if event == cv2.EVENT_LBUTTONDOWN:
if self._rect_over == False:
print('Draw a rectangle on the image first.')
else:
self._drawing == True
cv2.circle(self.img, (x, y), self._thickness, self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness, self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_MOUSEMOVE:
if self._drawing == True:
cv2.circle(self.img, (x, y), self._thickness, self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness, self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_LBUTTONUP:
if self._drawing == True:
self._drawing = False
cv2.circle(self.img, (x, y), self._thickness, self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness, self._DRAW_VAL['val'], -1)
# @timeit
def init_with_kmeans(self):
print(self.cols*self.rows)
print(len(list(np.where(self._mask == 0))[1]))
'''Initialise the BGDGMM and FGDGMM, which are respectively background-model and foreground-model,
using kmeans algorithm'''
max_iter = 2 # Max-iteration count for Kmeans
'''In the following two indexings, the np.logical_or is needed in place of or'''
self._bgd = np.where(np.logical_or(self._mask == self._GC_BGD, self._mask == self._GC_PR_BGD)) # Find the places where pixels in the mask MAY belong to BGD.
self._fgd = np.where(np.logical_or(self._mask == self._GC_FGD, self._mask == self._GC_PR_FGD)) # Find the places where pixels in the mask MAY belong to FGD.
self._BGDpixels = self.img[self._bgd]
self._FGDpixels = self.img[self._fgd]
KMB = kmeans(self._BGDpixels, dim = 3, n = self.k, max_iter = max_iter) # The Background Model by kmeans
KMF = kmeans(self._FGDpixels, dim = 3, n = self.k, max_iter = max_iter) # The Foreground Model by kmeans
KMB.run()
KMF.run()
# self._BGD_types = KMB.output()
# self._FGD_types = KMF.output()
# print(self._BGD_types)
self._BGD_by_components = KMB.output()
self._FGD_by_components = KMF.output()
self.BGD_GMM = GMM() # The GMM Model for BGD
self.FGD_GMM = GMM() # The GMM Model for FGD
'''Add the pixels by components to GMM'''
for ci in range(self.k):
# print(len(self._BGD_by_components[ci]))
# print(self._BGD_by_components[ci])
for pixel in self._BGD_by_components[ci]:
# pixel = np.asarray([j for j in pixel], dtype = np.float32)
self.BGD_GMM.add_pixel(pixel, ci)
for pixel in self._FGD_by_components[ci]:
self.FGD_GMM.add_pixel(pixel, ci)
# for ci in range(self.k):
# bgd_index = np.where(self._BGD_types == ci)
# fgd_index = np.where(self._FGD_types == ci)
# for pixel in self.img[bgd_index]:
# self.BGD_GMM.add_pixel(pixel, ci)
# for pixel in self.img[fgd_index]:
# self.FGD_GMM.add_pixel(pixel, ci)
self.BGD_GMM.learning()
self.FGD_GMM.learning()
'''The first step of the iteration in the paper: Assign components of GMMs to pixels,
(the kn in the paper), which is saved in self.components_index'''
# @timeit
def assign_GMM_components(self):
self.components_index = np.zeros([self.rows, self.cols], dtype = np.uint)
# self.components_index[self._bgd] = [i[0] for i in self.BGD_GMM.vec_pix_comp(self.img[self._bgd])]
# self.components_index[self._fgd] = [i[0] for i in self.FGD_GMM.vec_pix_comp(self.img[self._fgd])]
for y in range(self.rows):
for x in range(self.cols):
pixel = self.img[y, x]
self.components_index[y, x] = self.BGD_GMM.most_likely_pixel_component(pixel) if (self._mask[y, x] \
== self._GC_BGD or self._mask[y, x] == self._GC_PR_BGD) else self.FGD_GMM.most_likely_pixel_component(pixel)
# @timeit
def _assign_GMM_components(self):
self.components_index = np.zeros([self.rows, self.cols], dtype = np.uint)
self.components_index[self._bgd] = [i[0] for i in self.BGD_GMM.vec_pix_comp(self.img[self._bgd])]
self.components_index[self._fgd] = [i[0] for i in self.FGD_GMM.vec_pix_comp(self.img[self._fgd])]
'''The second step in the iteration: Learn the parameters from GMM models'''
# @timeit
def learn_GMM_parameters(self):
for ci in range(self.k):
# The places where the pixel belongs to the ci_th model and background model.
bgd_ci = np.where(np.logical_and(self.components_index == ci, np.logical_or(self._mask == self._GC_BGD, self._mask == self._GC_PR_BGD)))
fgd_ci = np.where(np.logical_and(self.components_index == ci, np.logical_or(self._mask == self._GC_FGD, self._mask == self._GC_PR_FGD)))
for pixel in self.img[bgd_ci]:
self.BGD_GMM.add_pixel(pixel, ci)
for pixel in self.img[fgd_ci]:
self.FGD_GMM.add_pixel(pixel, ci)
self.BGD_GMM.learning()
self.FGD_GMM.learning()
# @timeit
def construct_gcgraph(self, lam):
'''Construct a GCGraph with the Gibbs Energy'''
'''The vertexs of the graph are the pixels, and the edges are constructed by two parts,
the first part of which are the edges that connect each vertex with Sink Point(the background) and the Source Point(the foreground),
and the weight of which is the first term in Gibbs Energy;
the second part of the edges are those that connect each vertex with its neighbourhoods,
and the weight of which is the second term in Gibbs Energy.'''
vertex_count = self.cols*self.rows
edge_count = 2*(4*vertex_count - 3*(self.rows + self.cols) + 2)
self.graph = GCGraph(vertex_count, edge_count)
for y in range(self.rows):
for x in range(self.cols):
vertex_index = self.graph.add_vertex() # add-node and return its index
color = self.img[y, x]
# '''Set t-weights: Calculate the weight of each vertex with Sink node and Source node'''
if self._mask[y, x] == self._GC_PR_BGD or self._mask[y, x] == self._GC_PR_FGD:
# For each vertex, calculate the first term of G.E. as it be the BGD or FGD, and set them respectively as weight to t/s.
fromSource = -np.log(self.BGD_GMM.prob_pixel_GMM(color))
toSink = -np.log(self.FGD_GMM.prob_pixel_GMM(color))
# print(np.exp(-fromSource), np.exp(-toSink))
# print(fromSource)
elif self._mask[y, x] == self._GC_BGD:
# For the vertexs that are Background pixels, t-weight with Source = 0, with Sink = lam
fromSource = 0
toSink = lam
else:
# GC_FGD
fromSource = lam
toSink = 0
# print(fromSource, toSink)
self.graph.add_term_weights(vertex_index, fromSource, toSink)
'''Set n-weights and n-link, Calculate the weights between two neighbour vertexs, which is also the second term in Gibbs Energy(the smooth term)'''
if x > 0:
w = self.left_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-1, w, w)
if x > 0 and y > 0:
w = self.upleft_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-self.cols-1, w, w)
if y > 0:
w = self.up_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-self.cols, w, w)
if x < self.cols - 1 and y > 0:
w = self.upright_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index-self.cols+1, w, w)
# @timeit
def estimate_segmentation(self):
a = self.graph.max_flow()
for y in range(self.rows):
for x in range(self.cols):
if self._mask[y, x] == self._GC_PR_BGD or self._mask[y, x] == self._GC_PR_FGD:
if self.graph.insource_segment(y*self.cols+x): # Vertex Index
self._mask[y, x] = self._GC_PR_FGD
else:
# print(y, x)
self._mask[y, x] = self._GC_PR_BGD
def iter(self, n):
for i in range(n):
self.assign_GMM_components()
self.learn_GMM_parameters()
self.construct_gcgraph(self.lam)
self.estimate_segmentation()
# self._smoothing()
def run(self):
self.init_with_kmeans()
self.iter(1)
def _smoothing(self):
for y in range(1, self.rows-2):
for x in range(1, self.cols-2):
# if self._mask[x-1, y] == self._mask[x+1, y] == self._mask[x, y-1] == self._mask[x, y+1]:
a = self._mask[x-1, y]
b = self._mask[x+1, y]
c = self._mask[x, y-1]
d = self._mask[x, y+1]
if a==b==3 or a==c==3 or a==d==3 or b==c==3 or b==d==3 or c==d==3:
self._mask[x, y] = 3
def show(self, output):
#
# FGD = np.where(np.logical_and(np.logical_or(self._mask == 1, self._mask == 3), self._mask0 == 3))
# FGD = np.where(np.logical_or(self._mask==1, self._mask==3))
FGD = np.where((self._mask == 1) + (self._mask == 3), 255, 0).astype('uint8')
# output[FGD] = self.img[FGD]
# output = output.astype(np.uint8)
output = cv2.bitwise_and(self.img2, self.img2, mask = FGD)
# print('Press N to continue')
return output
class GCClient:
'''The engine of grabcut'''
def __init__(self, img, k):
self.k = k # The number of components in each GMM model
self.img = np.asarray(img, dtype=np.float32)
self.img2 = img
self.rows, self.cols = get_size(img) # 이미지 사이즈
self.gamma = 50
self.lam = 9 * self.gamma
self.beta = 0
self._BLUE = [255, 0, 0] # rectangle color
self._RED = [0, 0, 255] # PR BG
self._GREEN = [0, 255, 0] # PR FG
self._BLACK = [0, 0, 0] # sure BG
self._WHITE = [255, 255, 255] # sure FG
self._DRAW_BG = {'color': self._BLACK, 'val': 0}
self._DRAW_FG = {'color': self._WHITE, 'val': 1}
self._DRAW_PR_FG = {'color': self._GREEN, 'val': 3}
self._DRAW_PR_BG = {'color': self._RED, 'val': 2}
# setting up flags
self._rect = [0, 0, 1, 1]
self._drawing = False # flag for drawing curves
self._rectangle = False # flag for drawing rect
self._rect_over = False # flag to check if rect drawn
# se;f._rect_or_mask = 100 # flag for selecting rect or mask mode
# self._value = DRAW_FG # drawing initialized to FG
self._thickness = 3 # brush thickness
self._GC_BGD = 0 #{'color' : BLACK, 'val' : 0}
self._GC_FGD = 1 #{'color' : WHITE, 'val' : 1}
self._GC_PR_BGD = 2 #{'color' : GREEN, 'val' : 3}
self._GC_PR_FGD = 3 #{'color' : RED, 'val' : 2}
self.calc_beta()
self.calc_nearby_weight()
self._DRAW_VAL = None
self._mask = np.zeros([self.rows, self.cols],
dtype=np.uint8) # Init the mask
self._mask[:, :] = self._GC_BGD
flag = False
## jaejin add
self._rectangle = False
self._rect_over = True
self._rect = [0, 0, self.cols, self.rows] # 전체선택.
self.rect_or_mask = 0
self._mask[self._rect[1] + self._thickness:self._rect[1] +
self._rect[3] - self._thickness,
self._rect[0] + self._thickness:self._rect[0] +
self._rect[2] - self._thickness] = self._GC_PR_FGD
# --- ---- - --- - - --- - - -- end 시작하자마자 선택됨
## _mask 바꾸면 될듯함
def calc_beta(self):
'''Calculate Beta -- The Exp Term of Smooth Parameter in Gibbs Energy'''
'''beta = 1/(2*average(sqrt(||pixel[i] - pixel[j]||)))'''
'''Beta is used to adjust the difference of two nearby pixels in high or low contrast rate'''
beta = 0
self._left_diff = self.img[:, 1:] - self.img[:, :-1] # Left-difference
self._upleft_diff = self.img[
1:, 1:] - self.img[:-1, :-1] # Up-Left difference
self._up_diff = self.img[1:, :] - self.img[:-1, :] # Up-difference
self._upright_diff = self.img[
1:, :-1] - self.img[:-1, 1:] # Up-Right difference
beta = (self._left_diff*self._left_diff).sum() + (self._upleft_diff*self._upleft_diff).sum() \
+ (self._up_diff*self._up_diff).sum() + (self._upright_diff*self._upright_diff).sum() # According to the formula
self.beta = 1 / (
2 * beta /
(4 * self.cols * self.rows - 3 * self.cols - 3 * self.rows + 2))
# print(self.beta) # According to the paper
@timeit
def calc_nearby_weight(self):
'''Calculate the weight of the edge of each pixel with its nearby pixel, as each pixel is regarded
as a vertex of the graph'''
'''The weight of each direction is saved in a image the same size of the original image'''
'''weight = gamma*exp(-beta*(diff*diff))'''
self.left_weight = np.zeros([self.rows, self.cols])
self.upleft_weight = np.zeros([self.rows, self.cols])
self.up_weight = np.zeros([self.rows, self.cols])
self.upright_weight = np.zeros([self.rows, self.cols])
# Use the formula to calculate the weight
# for y in range(self.rows):
# for x in range(self.cols):
# color = self.img[y, x]
# if x >= 1:
# diff = color - self.img[y, x-1]
# diff.shape = (1, 3)
# self.left_weight[y, x] = self.gamma*np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# if x >= 1 and y >= 1:
# diff = color - self.img[y-1, x-1]
# diff.shape = (1, 3)
# self.upleft_weight[y, x] = self.gamma/np.sqrt(2) * np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# if y >= 1:
# diff = color - self.img[y-1, x]
# diff.shape = (1, 3)
# self.up_weight[y, x] = self.gamma*np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# if x+1 < self.cols and y >= 1:
# diff = color - self.img[y-1, x+1]
# diff.shape = (1, 3)
# self.upright_weight[y, x] = self.gamma/np.sqrt(2)*np.exp(-self.beta*np.dot(diff, np.transpose(diff)))
# Use the formula to calculate the weight
for y in range(self.rows):
for x in range(self.cols):
color = self.img[y, x]
if x >= 1:
diff = color - self.img[y, x - 1]
# print(np.exp(-self.beta*(diff*diff).sum()))
self.left_weight[y, x] = self.gamma * np.exp(
-self.beta * (diff * diff).sum())
if x >= 1 and y >= 1:
diff = color - self.img[y - 1, x - 1]
self.upleft_weight[y,
x] = self.gamma / np.sqrt(2) * np.exp(
-self.beta * (diff * diff).sum())
if y >= 1:
diff = color - self.img[y - 1, x]
self.up_weight[y, x] = self.gamma * np.exp(
-self.beta * (diff * diff).sum())
if x + 1 < self.cols and y >= 1:
diff = color - self.img[y - 1, x + 1]
self.upright_weight[y,
x] = self.gamma / np.sqrt(2) * np.exp(
-self.beta * (diff * diff).sum())
# self.left_weight[:, 1:] = self.gamma*np.exp(-self.beta*((self._left_diff*self._left_diff).sum()))
# self.upleft_weight[1:, 1:] = self.gamma*np.exp(-self.beta*((self._upleft_diff*self._upleft_diff).sum()))
# self.up_weight[1:, :] = self.gamma*np.exp(-self.beta*(self._up_diff*self._up_diff).sum())
# self.upright_weight = self.gamma*np.exp(-self.beta*(self._upright_diff*self._upright_diff).sum())
'''The following function is derived from the sample of opencv sources'''
def init_mask(self, event, x, y, flags, param):
'''Init the mask with interactive movements'''
'''Notice: the elements in the mask should be within the follows:
"GC_BGD":The pixel belongs to background;
"GC_FGD":The pixel belongs to foreground;
"GC_PR_BGD":The pixel MAY belongs to background;
"GC_PR_FGD":The pixel MAY belongs to foreground;'''
# Draw Rectangle
if event == cv2.EVENT_RBUTTONDOWN: # 오른쪽 마우스 눌렀을때
self._rectangle = True
self._ix, self._iy = x, y
elif event == cv2.EVENT_MOUSEMOVE:
if self._rectangle == True:
self.img = self.img2.copy()
cv2.rectangle(
self.img, (self._ix, self._iy), (x, y), self._BLUE, 2
) # 사각형 그리기 cv2.rectangle(img, start, end, color, thickness) 이미지 시작좌표 종료좌표 color 선의두께
self._rect = [
min(self._ix, x),
min(self._iy, y),
abs(self._ix - x),
abs(self._iy - y)
]
self.rect_or_mask = 0
elif event == cv2.EVENT_RBUTTONUP:
self._rectangle = False
self._rect_over = True
cv2.rectangle(self.img, (self._ix, self._iy), (x, y), self._BLUE,
2)
self._rect = [
min(self._ix, x),
min(self._iy, y),
abs(self._ix - x),
abs(self._iy - y)
] # x와 ix 작은값 y와 iy 작은값 ix-x 의 절대값 , iy-y의 절대값
# add jaejin
self._rect = [0, 0, self.cols, self.rows] # 전체선택.
# ---
# print(" Now press the key 'n' a few times until no further change \n")
self.rect_or_mask = 0
self._mask[self._rect[1] + self._thickness:self._rect[1] +
self._rect[3] - self._thickness,
self._rect[0] + self._thickness:self._rect[0] +
self._rect[2] - self._thickness] = self._GC_PR_FGD
## _mask 바꾸면 될듯함
# Notice : The x and y axis in CV2 are inversed to those in numpy.
if event == cv2.EVENT_LBUTTONDOWN: # 왼쪽마우스 눌렸을때
if self._rect_over == False:
print('Draw a rectangle on the image first.')
else:
self._drawing == True
cv2.circle(self.img, (x, y), self._thickness,
self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness,
self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_MOUSEMOVE:
if self._drawing == True:
cv2.circle(self.img, (x, y), self._thickness,
self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness,
self._DRAW_VAL['val'], -1)
elif event == cv2.EVENT_LBUTTONUP:
if self._drawing == True:
self._drawing = False
cv2.circle(self.img, (x, y), self._thickness,
self._DRAW_VAL['color'], -1)
cv2.circle(self._mask, (x, y), self._thickness,
self._DRAW_VAL['val'], -1)
# @timeit
def init_with_kmeans(self):
print(self.cols * self.rows)
print(len(list(np.where(self._mask == 0))[1]))
'''Initialise the BGDGMM and FGDGMM, which are respectively background-model and foreground-model,
using kmeans algorithm'''
max_iter = 2 # Max-iteration count for Kmeans
'''In the following two indexings, the np.logical_or is needed in place of or'''
self._bgd = np.where(
np.logical_or(self._mask == self._GC_BGD,
self._mask == self._GC_PR_BGD)
) # Find the places where pixels in the mask MAY belong to BGD.
self._fgd = np.where(
np.logical_or(self._mask == self._GC_FGD,
self._mask == self._GC_PR_FGD)
) # Find the places where pixels in the mask MAY belong to FGD.
self._BGDpixels = self.img[self._bgd]
self._FGDpixels = self.img[self._fgd]
KMB = kmeans(self._BGDpixels, dim=3, n=self.k,
max_iter=max_iter) # The Background Model by kmeans
KMF = kmeans(self._FGDpixels, dim=3, n=self.k,
max_iter=max_iter) # The Foreground Model by kmeans
KMB.run()
KMF.run()
# self._BGD_types = KMB.output()
# self._FGD_types = KMF.output()
# print(self._BGD_types)
self._BGD_by_components = KMB.output()
self._FGD_by_components = KMF.output()
self.BGD_GMM = GMM() # The GMM Model for BGD
self.FGD_GMM = GMM() # The GMM Model for FGD
'''Add the pixels by components to GMM'''
for ci in range(self.k):
# print(len(self._BGD_by_components[ci]))
# print(self._BGD_by_components[ci])
for pixel in self._BGD_by_components[ci]:
# pixel = np.asarray([j for j in pixel], dtype = np.float32)
self.BGD_GMM.add_pixel(pixel, ci)
for pixel in self._FGD_by_components[ci]:
self.FGD_GMM.add_pixel(pixel, ci)
# for ci in range(self.k):
# bgd_index = np.where(self._BGD_types == ci)
# fgd_index = np.where(self._FGD_types == ci)
# for pixel in self.img[bgd_index]:
# self.BGD_GMM.add_pixel(pixel, ci)
# for pixel in self.img[fgd_index]:
# self.FGD_GMM.add_pixel(pixel, ci)
self.BGD_GMM.learning()
self.FGD_GMM.learning()
'''The first step of the iteration in the paper: Assign components of GMMs to pixels,
(the kn in the paper), which is saved in self.components_index'''
# @timeit
def assign_GMM_components(self):
self.components_index = np.zeros([self.rows, self.cols], dtype=np.uint)
# self.components_index[self._bgd] = [i[0] for i in self.BGD_GMM.vec_pix_comp(self.img[self._bgd])]
# self.components_index[self._fgd] = [i[0] for i in self.FGD_GMM.vec_pix_comp(self.img[self._fgd])]
for y in range(self.rows):
for x in range(self.cols):
pixel = self.img[y, x]
self.components_index[y, x] = self.BGD_GMM.most_likely_pixel_component(pixel) if (self._mask[y, x] \
== self._GC_BGD or self._mask[y, x] == self._GC_PR_BGD) else self.FGD_GMM.most_likely_pixel_component(pixel)
# @timeit
def _assign_GMM_components(self):
self.components_index = np.zeros([self.rows, self.cols], dtype=np.uint)
self.components_index[self._bgd] = [
i[0] for i in self.BGD_GMM.vec_pix_comp(self.img[self._bgd])
]
self.components_index[self._fgd] = [
i[0] for i in self.FGD_GMM.vec_pix_comp(self.img[self._fgd])
]
'''The second step in the iteration: Learn the parameters from GMM models'''
# @timeit
def learn_GMM_parameters(self):
for ci in range(self.k):
# The places where the pixel belongs to the ci_th model and background model.
bgd_ci = np.where(
np.logical_and(
self.components_index == ci,
np.logical_or(self._mask == self._GC_BGD,
self._mask == self._GC_PR_BGD)))
fgd_ci = np.where(
np.logical_and(
self.components_index == ci,
np.logical_or(self._mask == self._GC_FGD,
self._mask == self._GC_PR_FGD)))
for pixel in self.img[bgd_ci]:
self.BGD_GMM.add_pixel(pixel, ci)
for pixel in self.img[fgd_ci]:
self.FGD_GMM.add_pixel(pixel, ci)
self.BGD_GMM.learning()
self.FGD_GMM.learning()
# @timeit
def construct_gcgraph(self, lam):
'''Construct a GCGraph with the Gibbs Energy'''
'''The vertexs of the graph are the pixels, and the edges are constructed by two parts,
the first part of which are the edges that connect each vertex with Sink Point(the background) and the Source Point(the foreground),
and the weight of which is the first term in Gibbs Energy;
the second part of the edges are those that connect each vertex with its neighbourhoods,
and the weight of which is the second term in Gibbs Energy.'''
vertex_count = self.cols * self.rows
edge_count = 2 * (4 * vertex_count - 3 * (self.rows + self.cols) + 2)
self.graph = GCGraph(vertex_count, edge_count)
for y in range(self.rows):
for x in range(self.cols):
vertex_index = self.graph.add_vertex(
) # add-node and return its index
color = self.img[y, x]
# '''Set t-weights: Calculate the weight of each vertex with Sink node and Source node'''
if self._mask[y, x] == self._GC_PR_BGD or self._mask[
y, x] == self._GC_PR_FGD:
# For each vertex, calculate the first term of G.E. as it be the BGD or FGD, and set them respectively as weight to t/s.
fromSource = -np.log(self.BGD_GMM.prob_pixel_GMM(color))
toSink = -np.log(self.FGD_GMM.prob_pixel_GMM(color))
# print(np.exp(-fromSource), np.exp(-toSink))
# print(fromSource)
elif self._mask[y, x] == self._GC_BGD:
# For the vertexs that are Background pixels, t-weight with Source = 0, with Sink = lam
fromSource = 0
toSink = lam
else:
# GC_FGD
fromSource = lam
toSink = 0
# print(fromSource, toSink)
self.graph.add_term_weights(vertex_index, fromSource, toSink)
'''Set n-weights and n-link, Calculate the weights between two neighbour vertexs, which is also the second term in Gibbs Energy(the smooth term)'''
if x > 0:
w = self.left_weight[y, x]
self.graph.add_edges(vertex_index, vertex_index - 1, w, w)
if x > 0 and y > 0:
w = self.upleft_weight[y, x]
self.graph.add_edges(vertex_index,
vertex_index - self.cols - 1, w, w)
if y > 0:
w = self.up_weight[y, x]
self.graph.add_edges(vertex_index,
vertex_index - self.cols, w, w)
if x < self.cols - 1 and y > 0:
w = self.upright_weight[y, x]
self.graph.add_edges(vertex_index,
vertex_index - self.cols + 1, w, w)
# @timeit
def estimate_segmentation(self):
a = self.graph.max_flow()
for y in range(self.rows):
for x in range(self.cols):
if self._mask[y, x] == self._GC_PR_BGD or self._mask[
y, x] == self._GC_PR_FGD:
if self.graph.insource_segment(y * self.cols +
x): # Vertex Index
self._mask[y, x] = self._GC_PR_FGD
else:
# print(y, x)
self._mask[y, x] = self._GC_PR_BGD
def iter(self, n):
for i in range(n):
self.assign_GMM_components()
self.learn_GMM_parameters()
self.construct_gcgraph(self.lam)
self.estimate_segmentation()
# self._smoothing()
def run(self):
self.init_with_kmeans()
self.iter(1)
def _smoothing(self):
for y in range(1, self.rows - 2):
for x in range(1, self.cols - 2):
# if self._mask[x-1, y] == self._mask[x+1, y] == self._mask[x, y-1] == self._mask[x, y+1]:
a = self._mask[x - 1, y]
b = self._mask[x + 1, y]
c = self._mask[x, y - 1]
d = self._mask[x, y + 1]
if a == b == 3 or a == c == 3 or a == d == 3 or b == c == 3 or b == d == 3 or c == d == 3:
self._mask[x, y] = 3
def show(self, output):
#
# FGD = np.where(np.logical_and(np.logical_or(self._mask == 1, self._mask == 3), self._mask0 == 3))
# FGD = np.where(np.logical_or(self._mask==1, self._mask==3))
FGD = np.where((self._mask == 1) + (self._mask == 3), 255,
0).astype('uint8')
# output[FGD] = self.img[FGD]
# output = output.astype(np.uint8)
output = cv2.bitwise_and(self.img2, self.img2, mask=FGD)
# print('Press N to continue')
return output