segmentation.py

from __future__ import division
import numpy as np
from _growcut import growcut
from orientation import orientation
from skimage import io
from skimage.morphology.convex_hull import grid_points_inside_poly


def segmentation(image):
    image = np.resize(image, (256, 256))
    m, n = image.shape[0], image.shape[1] 
   
    state = np.zeros((image.shape[0], image.shape[1], 2))

    if orientation(image)=='left':
        v1 = np.array([[47, 92], [182, 92], [182, 217]])
        points1 = grid_points_inside_poly((m,n), v1)
        foreground_pixels = image[points1==True]
        v2 = np.array([[17, 160], [17, 233], [94, 233]])
        points2 = grid_points_inside_poly((m, n), v2)
        background_pixels = image[points2==True]
    else:
        v1 = np.array([[50, 146], [190, 146], [190, 450]])
        points1 = grid_points_inside_poly((m, n), v1)
        foreground_pixels = image[points1==True]
        v2 = np.array([[20, 20], [100, 20], [20, 85]])
        points2 = grid_points_inside_poly((m, n), v2)
        background_pixels = image[points2==True]

    for (r, c) in background_pixels:
        state[r, c] = (0, 1)

    for (r, c) in foreground_pixels:
        state[r, c] = (1, 1)

    out = growcut(image, state, window_size=5, max_iter=200)
    
    return out
image = io.imread("haai1.jpg", plugin='pil') 
out = segmentation(image)