Added alpha expansion.

Author: JeGa

binseg.py

@@ -3,11 +3,11 @@
 import scipy.stats
 import logging
 import matplotlib.pyplot as plt
-from prettytable import PrettyTable
 import progressbar
 import maxflow
 import os
-
+import click
+import random
 
 class GMM:
     def __init__(self, prob):
@@ -40,14 +40,13 @@ def unaryenergy(fg, bg, img):
         img: Image array (RGB)
 
     Returns:
-        Array A (ysize, xsize, 3) with
+        Array A (ysize, xsize, 2) 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
-    unary = np.empty((ysize, xsize, 3))
+    unary = np.empty((ysize, xsize, 2))
 
     with progressbar.ProgressBar(max_value=xsize, redirect_stdout=True) as progress:
         for x in range(xsize):
@@ -56,12 +55,6 @@ def unaryenergy(fg, bg, img):
                 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
@@ -155,6 +148,7 @@ def loop(self, edgecallback, nodecallback):
         nodecallback(self.getNode(self.ysize - 1, self.xsize - 1), self.g)
 
     def loopnodes(self, callback):
+        logging.info("Iterate through nodes.")
         for y in range(self.ysize):
             for x in range(self.xsize):
                 callback(self.getNode(y, x), self.g)
@@ -193,7 +187,7 @@ def edge(self, node_i, node_j, graph):
         C = self.pairwiseenergy(1, 0, xi, xj)
         D = self.pairwiseenergy(1, 1, xi, xj)
 
-        #energy = self.pairwiseenergy(self.unaries[i[0], i[1], 2],
+        # energy = self.pairwiseenergy(self.unaries[i[0], i[1], 2],
         #                             self.unaries[j[0], j[1], 2],
         #                             xi, xj)
 
@@ -250,9 +244,156 @@ def getimg(self):
         return self.img
 
 
-def main():
-    logging.basicConfig(level=logging.INFO)
+class BinsegAlphaexp:
+    def __init__(self, img, unaries, numlabel):
+        self.img = img
+        self.unaries = unaries
+
+        # Available labels.
+        self.label = range(numlabel)
+
+        # Initial labeling. All = 0
+        self.y = np.empty((img.shape[0], img.shape[1]))
+
+        self.l = 0.5
+        self.w = 3.5
+
+        # Current alpha
+        self.alpha = 0
+
+    def constructgraph(self):
+        nodegrid = Nodegrid(self.img.shape[0], self.img.shape[1])
+        return nodegrid
+
+    def edge(self, node_i, node_j, graph):
+        """
+        Callback for pairwise energy.
+        """
+
+        # Pixel coordinates.
+        i = [node_i.y, node_i.x]
+        j = [node_j.y, node_j.x]
+
+        # Current label.
+        i_label = self.y[i[0], i[1]]
+        j_label = self.y[j[0], j[1]]
+
+        # Pixel values
+        xi = self.img[i[0], i[1]]
+        xj = self.img[j[0], j[1]]
+
+        # Only for nodes that are not alpha.
+        if i_label == self.alpha:
+            return
+
+        sourceenergy = self.pairwiseenergy(i_label, j_label, xi, xj)
+        graph.add_tedge(node_i.nodeid, sourceenergy, 0)
+
+        if j_label == self.alpha:
+            return
+
+        energy = self.pairwiseenergy(self.alpha, j_label, xi, xj)
+        energy += self.pairwiseenergy(i_label, self.alpha, xi, xj)
+        energy -= self.pairwiseenergy(i_label, j_label, xi, xj)
+
+        graph.add_edge(node_i.nodeid, node_j.nodeid, energy, energy)
+
+        # graph.add_tedge(node_i.nodeid, C, A)
+        # graph.add_tedge(node_j.nodeid, D, C)
+
+    def node_assign(self, node_i, graph):
+        """
+        Callback for assigning unary energy.
+        """
+
+        # Pixel
+        y = node_i.y
+        x = node_i.x
+
+        # Label of pixel
+        label = self.y[y, x]
+
+        # Just for nodes that are not alpha.
+        if label == self.alpha:
+            return
+
+        # Get unary for assigned label.
+        source = self.unaries[y, x, label]
+
+        # Get unary for alpha.
+        sink = self.unaries[y, x, self.alpha]
+
+        graph.add_tedge(node_i.nodeid, source, sink)
+
+    def node_segment(self, node_i, graph):
+        """
+        Callback for segmentation.
+        """
+
+        # Pixel
+        y = node_i.y
+        x = node_i.x
+
+        # Label of pixel
+        label = self.y[y, x]
+
+        # Just for nodes that are not alpha.
+        if label == self.alpha:
+            return
+
+        if graph.get_segment(node_i.nodeid) == 0:  # Change to alpha
+            self.y[y, x] = self.alpha
+
+    def pairwiseenergy(self, y1, y2, x1, x2):
+        """
+        Returns pairwise energy between node i and node j using the Potts model.
+
+        :param y1: Label of i node.
+        :param y2: Label of j node.
+        :param x1: Pixel value at node i.
+        :param x2: Pixel value at node j.
+        :return: Pairwise energy.
+        """
+        if y1 == y2:
+            return 0.0
 
+        # Not same label
+        # np.sum(np.power(x1 - x2, 2), 0)
+        energy = self.w * np.exp(-self.l * np.power(np.linalg.norm(x1 - x2, 2), 2))
+        return energy
+
+    def segment(self):
+        for i in range(1):
+            # For each label: Change current label to alpha?
+            for alpha in self.label:
+                self.alpha = alpha
+
+                logging.info("Alpha: " + str(alpha))
+
+                # Get graph with all nodes.
+                nodegrid = self.constructgraph()
+
+                nodegrid.loop(self.edge, self.node_assign)
+
+                nodegrid.maxflow()
+
+                # Sets label to alpha if it should change.
+                nodegrid.loopnodes(self.node_segment)
+
+        # Assign color.
+        colors = []
+        for i in self.label:
+            colors.append([random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)])
+
+        for y in range(self.img.shape[0]):
+            for x in range(self.img.shape[1]):
+                self.img[y, x] = colors[int(self.y[y, x])]
+
+    def getimg(self):
+        return self.img
+
+
+def binseg():
     logging.info("Read image.")
     img = misc.imread("banana3.png")
     img = np.array(img, dtype=np.float64) / 255
@@ -274,5 +415,38 @@ def main():
     plt.imsave("banana_out", img)
 
 
+def alphaexp():
+    logging.info("Read image.")
+    img = misc.imread("banana3.png")
+    img = np.array(img, dtype=np.float64) / 255
+
+    if not os.path.exists("unary.npy"):
+        generateunaries(img)
+
+    logging.info("Load unaries.")
+    unaries = np.load("unary.npy")
+
+    binseg = BinsegAlphaexp(img, unaries, 2)
+    binseg.segment()
+
+    logging.info("Save image.")
+    img = binseg.getimg().astype(np.uint8)
+
+    plt.imshow(img)
+    plt.show()
+    plt.imsave("banana_out", img)
+
+
+@click.command()
+@click.option('--usealphaexp', is_flag=True)
+def main(usealphaexp):
+    logging.basicConfig(level=logging.INFO)
+
+    if usealphaexp:
+        alphaexp()
+    else:
+        binseg()
+
+
 if __name__ == '__main__':
     main()