Graph building ...

Author: JeGa

binseg.py

@@ -6,11 +6,15 @@
 from prettytable import PrettyTable
 import progressbar
 import maxflow
-from maxflow.examples_utils import plot_graph_2d
+from examples_utils import plot_graph_2d
 
 
 class GMM:
     def __init__(self, prob):
+        """
+        Arguments:
+            prob: List with mu, sigma and mix.
+        """
         self.mu = prob[0]
         self.sigma = prob[1]
         self.mix = prob[2]
@@ -29,15 +33,29 @@ def prob(self, x):
 
 
 def unaryenergy(fg, bg, img):
+    """
+    Arguments:
+        fg: Foreground GMM
+        bg: Background GMM
+        img: Image array (RGB)
+    """
     logging.info("Calculate unary energy functions.")
     ysize, xsize, _ = img.shape
-    unary = np.empty((ysize, xsize, 2))
+    unary = np.empty((ysize, xsize, 3))
 
     with progressbar.ProgressBar(max_value=xsize, redirect_stdout=True) as progress:
         for x in range(xsize):
             for y in range(ysize):
+                # Background
                 unary[y, x, 0] = -np.log(bg.prob(img[y, x]))
+                # Foreground
                 unary[y, x, 1] = -np.log(fg.prob(img[y, x]))
+
+                # Assign labels
+                if unary[y, x, 0] < unary[y, x, 1]:
+                    unary[y, x, 2] = 0  # Background
+                else:
+                    unary[y, x, 2] = 1  # Foreground
             progress.update(x)
 
     return unary
@@ -55,9 +73,9 @@ def pairwiseenergy(unaries, img):
         for x in range(xsize):
             for y in range(ysize):
                 pass
-                #if unaries[y, x, 3] ==
+                # if unaries[y, x, 3] ==
                 #    np.exp(-l * np.linalg.norm())
-                #else:
+                # else:
                 #    pass
             progress.update(x)
 
@@ -87,22 +105,89 @@ def main():
     img = misc.imread("banana3.png")
     img = np.array(img, dtype=np.float64)
 
-    unaries = unaryenergy(fg, bg, img)
+    # unaries = unaryenergy(fg, bg, img)
+    unaries = np.load("unary.npy")
+    # np.save("unary", unaries)
 
     ysize, xsize, _ = img.shape
 
-    g = maxflow.GraphFloat()
-    nodeids = g.add_grid_nodes((5, 5))
-    structure = np.array([[0, 1, 0],
-                          [1, 0, 1],
-                          [0, 1, 0]])
+    def pairwiseenergy(y1, x1, y2, x2):
+        nonlocal unaries
+        nonlocal img
+
+        l = 0.005
+        w = 110
 
-    g.add_grid_edges(nodeids, structure=structure, symmetric=True)
+        if unaries[y1, x1, 2] != unaries[y2, x2, 2]:
+            delta = 1
+        else:
+            delta = 0
 
-    # Source and sink
-    g.add_grid_tedges(nodeids, 0, 2)
+        # Not same label
+        energy = w * np.exp(-l * np.power(np.linalg.norm(img[y1, x1] - img[y2, x2], ord=2), 2)) * delta
+        #a = (np.linalg.norm(img[y1, x1] - img[y2, x2], ord=2))
+        return energy
 
-    plot_graph_2d(g, nodeids.shape)
+    g = maxflow.GraphFloat()
+    nodeids = g.add_grid_nodes((ysize, xsize))
+
+    for y in range(ysize - 1):
+        for x in range(xsize - 1):
+            e_right = pairwiseenergy(y, x, y, x + 1)
+            if e_right >= 10.0:
+                a = 2
+            g.add_edge(nodeids[y, x], nodeids[y, x + 1], e_right, e_right)
+
+            e_down = pairwiseenergy(y, x, y + 1, x)
+            g.add_edge(nodeids[y, x], nodeids[y + 1, x], e_down, e_down)
+
+            # Source, sink
+            g.add_tedge(nodeids[y, x], unaries[y, x, 0], unaries[y, x, 1])
+
+    for y in range(ysize - 1):
+        e_down = pairwiseenergy(y, xsize - 1, y + 1, xsize - 1)
+        g.add_edge(nodeids[y, xsize - 1], nodeids[y + 1, xsize - 1], e_down, e_down)
+
+        g.add_tedge(nodeids[y, xsize - 1], unaries[y, xsize - 1, 0], unaries[y, xsize - 1, 1])
+
+    for x in range(xsize - 1):
+        e_right = pairwiseenergy(ysize - 1, x, ysize - 1, x + 1)
+
+        g.add_edge(nodeids[ysize - 1, x], nodeids[ysize - 1, x + 1], e_right, e_right)
+        g.add_tedge(nodeids[ysize - 1, x], unaries[ysize - 1, x + 1, 0], unaries[ysize - 1, x, 1])
+    g.add_tedge(nodeids[ysize - 1, xsize - 1], unaries[ysize - 1, xsize - 1, 0], unaries[ysize - 1, xsize - 1, 1])
+
+    # g = maxflow.Graph[float]()
+    # nodeids = g.add_grid_nodes((5, 5))
+    #
+    # # Edges pointing backwards (left, left up and left down) with infinite
+    # # capacity
+    # structure = np.array([[0, 0, 0],
+    #                       [0, 0, 0],
+    #                       [0, 0, 0]
+    #                       ])
+    # g.add_grid_edges(nodeids, structure=structure, symmetric=False)
+    #
+    # # Set a few arbitrary weights
+    # weights = np.array([[100, 110, 120, 130, 140]]).T + np.array([0, 2, 4, 6, 8])
+    #
+    # print(weights)
+    #
+    # structure = np.zeros((3, 3))
+    # structure[1, 2] = 1
+    # g.add_grid_edges(nodeids, structure=structure, weights=weights, symmetric=False)
+
+    # plot_graph_2d(g, nodeids.shape)
+
+    logging.info("Calculate max flow.")
+    g.maxflow()
+
+    for y in range(ysize):
+        for x in range(xsize):
+            if g.get_segment(nodeids[y, x]):
+                img[y, x] = np.array([0, 0, 0])
+            else:
+                img[y, x] = np.array([255, 255, 0])
 
     logging.info("Save image.")
     img = img.astype(np.uint8)