Repository: salaee/pegbis
Branch: master
Commit: d409712e19f2
Files: 5
Total size: 9.5 KB

Directory structure:
gitextract_8ifmyllb/

├── README.md
├── disjoint_set.py
├── filter.py
├── main.py
└── segment_graph.py

================================================
FILE CONTENTS
================================================

================================================
FILE: README.md
================================================
# PEGBIS (Python Efficient Graph-Based Image Segmentation)
Python implementation of "Efficient Graph-Based Image Segmentation" paper written by P. Felzenszwalb, D. Huttenlocher. 
The paper is available: http://cs.brown.edu/~pff/papers/seg-ijcv.pdf <br>
C++ implementation is written by the author and is available on:
http://cs.brown.edu/~pff/segment/ <br>
The C++ implementation is much more faster than python implementation (obviously). 
<br>
<br>
### Results
parameters: (Sigma=0.5, K=300, Min=50) <br>
![alt text](https://github.com/salaee/egbis/blob/master/results/results_1.png)
<br>
parameters: (Sigma=0.5, K=300, Min=50) <br>
![alt text](https://github.com/salaee/egbis/blob/master/results/results_2.png)
<br>
parameters: (Sigma=0.5, K=1000, Min=50) <br>
![alt text](https://github.com/salaee/egbis/blob/master/results/results_3.png)
<br>
parameters: (Sigma=0.8, K=500, Min=10) <br>
![alt text](https://github.com/salaee/egbis/blob/master/results/results_4.png)
<br>
parameters: (Sigma=0.5, K=500, Min=50) <br>
![alt text](https://github.com/salaee/egbis/blob/master/results/results_5.png)
<br>
<br>
### Requirements
Python 3.5<br>

##### Libraries used: 
scipy and matplotlib



================================================
FILE: disjoint_set.py
================================================
import numpy as np


# disjoint-set forests using union-by-rank and path compression (sort of).
class universe:
    def __init__(self, n_elements):
        self.num = n_elements
        self.elts = np.empty(shape=(n_elements, 3), dtype=int)
        for i in range(n_elements):
            self.elts[i, 0] = 0  # rank
            self.elts[i, 1] = 1  # size
            self.elts[i, 2] = i  # p

    def size(self, x):
        return self.elts[x, 1]

    def num_sets(self):
        return self.num

    def find(self, x):
        y = int(x)
        while y != self.elts[y, 2]:
            y = self.elts[y, 2]
        self.elts[x, 2] = y
        return y

    def join(self, x, y):
        # x = int(x)
        # y = int(y)
        if self.elts[x, 0] > self.elts[y, 0]:
            self.elts[y, 2] = x
            self.elts[x, 1] += self.elts[y, 1]
            self.elts[y, 1] = self.elts[x, 1]
        else:
            self.elts[x, 2] = y
            self.elts[y, 1] += self.elts[x, 1]
            self.elts[x, 1] = self.elts[y, 1]
            if self.elts[x, 0] == self.elts[y, 0]:
                self.elts[y, 0] += 1
        self.num -= 1


================================================
FILE: filter.py
================================================
import numpy as np
import math
np.seterr(over='ignore')


# some constants
WIDTH = 4.0


# convolve image with gaussian filter
def smooth(src, sigma):
    mask = make_fgauss(sigma)
    mask = normalize(mask)
    dst = convolve_even(src, mask)
    return dst


# gaussian filter
def make_fgauss(sigma):
    sigma = max(sigma, 0.01)
    length = int(math.ceil(sigma * WIDTH)) + 1
    mask = np.zeros(shape=(length, length), dtype=float)
    for i in range(length):
        for j in range(length):
            mask[i, j] = math.exp(-0.5 * (math.pow(i / sigma, 2) + math.pow(j / sigma, 2)))
    return mask

# normalize mask so it integrates to one
def normalize(mask):
    sum = 4 * np.sum(np.absolute(mask)) - 3 * abs(mask[0]) - \
          2 * np.sum(np.absolute(mask[0, :])) - 2 * np.sum(np.absolute(mask[:, 0]))
    return np.divide(mask, sum)


# convolve src with mask.  output is flipped!
def convolve_even(src, mask):
    output = np.zeros(shape=src.shape, dtype=float)
    height, width = src.shape
    length = len(mask)

    for y in range(height):
        for x in range(width):
            sum = float(mask[0, 0] * src[y, x])
            for i in range(0, length):
                for j in range(0, length):
                    if i != 0 and j != 0:
                        sum += mask[i, j] * (src[max(y - j, 0), max(x - i, 0)] + src[max(y - j, 0), min(x + i, width - 1)] + \
                                             src[min(y + j, height - 1), min(x + i, width - 1)] + src[min(y + j, height - 1), max(x - i, 0)])
            output[y, x] = sum
    return output


================================================
FILE: main.py
================================================
from skimage import io
import matplotlib.pyplot as plt
from filter import *
from segment_graph import *
import time


# --------------------------------------------------------------------------------
# Segment an image:
# Returns a color image representing the segmentation.
#
# Inputs:
#           in_image: image to segment.
#           sigma: to smooth the image.
#           k: constant for threshold function.
#           min_size: minimum component size (enforced by post-processing stage).
#
# Returns:
#           num_ccs: number of connected components in the segmentation.
# --------------------------------------------------------------------------------
def segment(in_image, sigma, k, min_size):
    start_time = time.time()
    height, width, band = in_image.shape
    print("Height:  " + str(height))
    print("Width:   " + str(width))
    smooth_red_band = smooth(in_image[:, :, 0], sigma)
    smooth_green_band = smooth(in_image[:, :, 1], sigma)
    smooth_blue_band = smooth(in_image[:, :, 2], sigma)

    # build graph
    edges_size = width * height * 4
    edges = np.zeros(shape=(edges_size, 3), dtype=object)
    num = 0
    for y in range(height):
        for x in range(width):
            if x < width - 1:
                edges[num, 0] = int(y * width + x)
                edges[num, 1] = int(y * width + (x + 1))
                edges[num, 2] = diff(smooth_red_band, smooth_green_band, smooth_blue_band, x, y, x + 1, y)
                num += 1
            if y < height - 1:
                edges[num, 0] = int(y * width + x)
                edges[num, 1] = int((y + 1) * width + x)
                edges[num, 2] = diff(smooth_red_band, smooth_green_band, smooth_blue_band, x, y, x, y + 1)
                num += 1

            if (x < width - 1) and (y < height - 1):
                edges[num, 0] = int(y * width + x)
                edges[num, 1] = int((y + 1) * width + (x + 1))
                edges[num, 2] = diff(smooth_red_band, smooth_green_band, smooth_blue_band, x, y, x + 1, y + 1)
                num += 1

            if (x < width - 1) and (y > 0):
                edges[num, 0] = int(y * width + x)
                edges[num, 1] = int((y - 1) * width + (x + 1))
                edges[num, 2] = diff(smooth_red_band, smooth_green_band, smooth_blue_band, x, y, x + 1, y - 1)
                num += 1
    # Segment
    u = segment_graph(width * height, num, edges, k)

    # post process small components
    for i in range(num):
        a = u.find(edges[i, 0])
        b = u.find(edges[i, 1])
        if (a != b) and ((u.size(a) < min_size) or (u.size(b) < min_size)):
            u.join(a, b)

    num_cc = u.num_sets()
    output = np.zeros(shape=(height, width, 3))

    # pick random colors for each component
    colors = np.zeros(shape=(height * width, 3))
    for i in range(height * width):
        colors[i, :] = random_rgb()

    for y in range(height):
        for x in range(width):
            comp = u.find(y * width + x)
            output[y, x, :] = colors[comp, :]

    elapsed_time = time.time() - start_time
    print(
        "Execution time: " + str(int(elapsed_time / 60)) + " minute(s) and " + str(
            int(elapsed_time % 60)) + " seconds")

    # displaying the result
    fig = plt.figure()
    a = fig.add_subplot(1, 2, 1)
    plt.imshow(in_image)
    a.set_title('Original Image')
    a = fig.add_subplot(1, 2, 2)
    output = output.astype(int)
    plt.imshow(output)
    a.set_title('Segmented Image')
    plt.show()


if __name__ == "__main__":
    sigma = 0.5
    k = 500
    min = 50
    input_path = "data/paris.jpg"

    # Loading the image
    input_image = io.imread(input_path)
    print("Loading is done.")
    print("processing...")
    segment(input_image, sigma, k, min)


================================================
FILE: segment_graph.py
================================================
from disjoint_set import *
import math
import numpy as np
import random


# ---------------------------------------------------------
# Segment a graph:
# Returns a disjoint-set forest representing the segmentation.
#
# Inputs:
#           num_vertices: number of vertices in graph.
#           num_edges: number of edges in graph
#           edges: array of edges.
#           c: constant for threshold function.
#
# Output:
#           a disjoint-set forest representing the segmentation.
# ------------------------------------------------------------
def segment_graph(num_vertices, num_edges, edges, c):
    # sort edges by weight (3rd column)
    edges[0:num_edges, :] = edges[edges[0:num_edges, 2].argsort()]
    # make a disjoint-set forest
    u = universe(num_vertices)
    # init thresholds
    threshold = np.zeros(shape=num_vertices, dtype=float)
    for i in range(num_vertices):
        threshold[i] = get_threshold(1, c)

    # for each edge, in non-decreasing weight order...
    for i in range(num_edges):
        pedge = edges[i, :]

        # components connected by this edge
        a = u.find(pedge[0])
        b = u.find(pedge[1])
        if a != b:
            if pedge[2] <= min(threshold[a], threshold[b]):
                u.join(a, b)
                a = u.find(a)
                b = u.find(b)
                threshold[a] = pedge[2] + get_threshold(u.size(a), c)
                threshold[b] = threshold[a]

    return u


def get_threshold(size, c):
    return c / size


# returns square of a number
def square(value):
    return value * value


# randomly creates RGB
def random_rgb():
    rgb = np.zeros(3, dtype=int)
    rgb[0] = random.randint(0, 255)
    rgb[1] = random.randint(0, 255)
    rgb[2] = random.randint(0, 255)
    return rgb


# dissimilarity measure between pixels
def diff(red_band, green_band, blue_band, x1, y1, x2, y2):
    result = math.sqrt(
        square(red_band[y1, x1] - red_band[y2, x2]) + square(green_band[y1, x1] - green_band[y2, x2]) + square(
            blue_band[y1, x1] - blue_band[y2, x2]))
    return result