grab_cut.py

import numpy as np
import cv2


class Grab_cut(object):
    suffix = '.jpg'

    def __init__(self, filename=None):
        self.filename = filename
        self.height = None
        self.width = None

    def image_matting(self, image_file, shape, iteration=10):
        points = shape['points']
        xmin, ymin, xmax, ymax = Grab_cut.convertPoints2BndBox(points)
        self.width = xmax - xmin
        self.height = ymax - ymin

        src_img = cv2.imread(image_file)

        mask = np.zeros(src_img.shape[:2], np.uint8)
        bgdModel = np.zeros((1, 65), np.float64)
        fgdModel = np.zeros((1, 65), np.float64)
        rect = (xmin, ymin, self.width, self.height)

        # Grabcut
        cv2.grabCut(src_img, mask, rect, bgdModel, fgdModel,
                    iteration, cv2.GC_INIT_WITH_RECT)

        r_channel, g_channel, b_channel = cv2.split(src_img)
        a_channel = np.where((mask == 2) | (mask == 0), 0, 255).astype('uint8')

        # crop image space
        for row in range(ymin, ymax):
            if sum(r_channel[row, xmin:xmax + 1]) > 0:
                out_ymin = row
                break
        for row in range(ymin, ymax)[::-1]:
            if sum(r_channel[row, xmin:xmax + 1]) > 0:
                out_ymax = row + 1
                break
        for col in range(xmin, xmax):
            if sum(a_channel[ymin:ymax + 1, col]) > 0:
                out_xmin = col
                break
        for col in range(xmin, xmax)[::-1]:
            if sum(a_channel[ymin:ymax + 1, col]) > 0:
                out_xmax = col + 1
                break

        # output image
        img_RGBA = cv2.merge((r_channel[out_ymin:out_ymax, out_xmin:out_xmax],
                              g_channel[out_ymin:out_ymax, out_xmin:out_xmax],
                              b_channel[out_ymin:out_ymax, out_xmin:out_xmax],
                              a_channel[out_ymin:out_ymax, out_xmin:out_xmax]))

        return img_RGBA

    @staticmethod
    def convertPoints2BndBox(points):
        xmin = float('inf')
        ymin = float('inf')
        xmax = float('-inf')
        ymax = float('-inf')
        for p in points:
            x = p[0]
            y = p[1]
            xmin = min(x, xmin)
            ymin = min(y, ymin)
            xmax = max(x, xmax)
            ymax = max(y, ymax)

        # Martin Kersner, 2015/11/12
        # 0-valued coordinates of BB caused an error while
        # training faster-rcnn object detector.
        if xmin < 1:
            xmin = 1

        if ymin < 1:
            ymin = 1

        return (int(xmin), int(ymin), int(xmax), int(ymax))

    @staticmethod
    def resultSave(save_path, image_np):
        cv2.imwrite(save_path, image_np)