Improved code structure. Todo: Pairwise energy.

Author: JeGa

binseg.py

@@ -6,7 +6,7 @@
 from prettytable import PrettyTable
 import progressbar
 import maxflow
-from examples_utils import plot_graph_2d
+import os
 
 
 class GMM:
@@ -38,6 +38,12 @@ def unaryenergy(fg, bg, img):
         fg: Foreground GMM
         bg: Background GMM
         img: Image array (RGB)
+
+    Returns:
+        Array A (ysize, xsize, 3) with
+            A[y, x, 0] = energy background
+            A[y, x, 1] = energy foreground
+            A[y, x, 2] = label
     """
     logging.info("Calculate unary energy functions.")
     ysize, xsize, _ = img.shape
@@ -61,27 +67,6 @@ def unaryenergy(fg, bg, img):
     return unary
 
 
-def pairwiseenergy(unaries, img):
-    logging.info("Calculate pairwise energy functions.")
-    ysize, xsize, _ = img.shape
-    pairwise = np.array()
-
-    l = 0.5
-    w = 1.0
-
-    with progressbar.ProgressBar(max_value=xsize, redirect_stdout=True) as progress:
-        for x in range(xsize):
-            for y in range(ysize):
-                pass
-                # if unaries[y, x, 3] ==
-                #    np.exp(-l * np.linalg.norm())
-                # else:
-                #    pass
-            progress.update(x)
-
-    return pairwise
-
-
 def readprobfile(filename):
     file = np.load(filename)
     mu = file["arr_0"]
@@ -91,106 +76,153 @@ def readprobfile(filename):
     return [mu, sigma, mix]
 
 
-def main():
-    logging.basicConfig(level=logging.INFO)
-
+def generateunaries(img):
     logging.info("Read GMM for unaries.")
     probf = readprobfile("prob_foreground.npz")
     probb = readprobfile("prob_background.npz")
 
     fg = GMM(probf)
     bg = GMM(probb)
 
-    logging.info("Read image.")
-    img = misc.imread("banana3.png")
-    img = np.array(img, dtype=np.float64)
+    logging.info("Generate unaries.")
+    unaries = unaryenergy(fg, bg, img)
+    np.save("unary", unaries)
 
-    # unaries = unaryenergy(fg, bg, img)
-    unaries = np.load("unary.npy")
-    # np.save("unary", unaries)
 
-    ysize, xsize, _ = img.shape
+class Nodegrid:
+    def __init__(self, ysize, xsize):
+        self.g = maxflow.GraphFloat()
+        self.nodeids = self.g.add_grid_nodes((ysize, xsize))
 
-    def pairwiseenergy(y1, x1, y2, x2):
-        nonlocal unaries
-        nonlocal img
+        self.ysize = ysize
+        self.xsize = xsize
+
+    def loop(self, edgecallback, nodecallback):
+        """
+        Loops over the grid of nodes. Two callback functions are required:
+
+        :param edgecallback: Called for every edge with node id (i, j)
+            as parameter.
+        :param nodecallback: called for every node with node id as parameter.
+        """
+        logging.info("Iterate through graph.")
+
+        for y in range(self.ysize - 1):
+            for x in range(self.xsize - 1):
+                # Right edge
+                edgecallback(self.nodeids[y, x],
+                             self.nodeids[y, x + 1], self.g)
+
+                # Down edge
+                edgecallback(self.nodeids[y, x],
+                             self.nodeids[y + 1, x], self.g)
+
+                # Node
+                nodecallback(self.nodeids[y, x], y, x, self.g)
+
+        # Last column
+        for y in range(self.ysize - 1):
+            # Down edge
+            edgecallback(self.nodeids[y, self.xsize - 1],
+                         self.nodeids[y + 1, self.xsize - 1], self.g)
+
+            # Node
+            nodecallback(self.nodeids[y, self.xsize - 1], y, x, self.g)
+
+        # Last row
+        for x in range(self.xsize - 1):
+            # Right edge
+            edgecallback(self.nodeids[self.ysize - 1, x],
+                         self.nodeids[self.ysize - 1, x + 1], self.g)
+
+            # Node
+            nodecallback(self.nodeids[self.ysize - 1, x], y, x, self.g)
+
+        # Last node
+        nodecallback(self.nodeids[self.ysize - 1, self.xsize - 1], y, x, self.g)
 
-        l = 0.005
-        w = 110
+    def loopnodes(self, callback):
+        for y in range(self.ysize):
+            for x in range(self.xsize):
+                callback(self.nodeids[y, x], y, x, self.g)
 
-        if unaries[y1, x1, 2] != unaries[y2, x2, 2]:
-            delta = 1
+    def maxflow(self):
+        logging.info("Calculate max flow.")
+        self.g.maxflow()
+
+
+class Binseg:
+    def __init__(self, img, unaries):
+        self.img = img
+        self.unaries = unaries
+
+        self.nodegrid = Nodegrid(img.shape[0], img.shape[1])
+
+        self.l = 0.005
+        self.w = 100
+
+    def edge(self, nodeid_i, nodeid_j, graph):
+        """
+        Callback for pairwise energy.
+        """
+        pass
+
+    def node_assign(self, nodeid_i, y, x, graph):
+        """
+        Callback for assigning unary energy.
+        """
+        graph.add_tedge(nodeid_i, self.unaries[y, x, 1], self.unaries[y, x, 0])
+
+    def node_segment(self, nodeid_i, y, x, graph):
+        """
+        Callback for segmentation.
+        """
+        if graph.get_segment(nodeid_i) == 0:
+            self.img[y, x] = np.array([0, 0, 0])
         else:
-            delta = 0
-
-        # 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
-
-    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])
+            self.img[y, x] = np.array([255, 255, 0])
+
+    def pairwiseenergy(y1, x1, y2, x2):
+        pass
+        # if unaries[y1, x1, 2] != unaries[y2, x2, 2]:
+        #     delta = 1
+        # else:
+        #     delta = 0
+        #
+        # # 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
+
+    def segment(self):
+        self.nodegrid.loop(self.edge, self.node_assign)
+
+        self.nodegrid.maxflow()
+
+        self.nodegrid.loopnodes(self.node_segment)
+
+    def getimg(self):
+        return self.img
+
+
+def main():
+    logging.basicConfig(level=logging.INFO)
+
+    logging.info("Read image.")
+    img = misc.imread("banana3.png")
+    img = np.array(img, dtype=np.float64)
+
+    if not os.path.exists("unary.npy"):
+        generateunaries(img)
+
+    logging.info("Load unaries.")
+    unaries = np.load("unary.npy")
+
+    binseg = Binseg(img, unaries)
+    binseg.segment()
 
     logging.info("Save image.")
-    img = img.astype(np.uint8)
+    img = binseg.getimg().astype(np.uint8)
 
     plt.imshow(img)
     plt.show()