seam carve

import numpy as np
from skimage.io import imread, imsave


def get_grad(img_lum, mask):
    grad_x = np.zeros_like(img_lum)
    grad_y = np.zeros_like(img_lum)
    grad_y[1:-1,:] = img_lum[2:,:] - img_lum[:-2,:]
    grad_x[:,1:-1] = img_lum[:,2:] - img_lum[:,:-2]
    grad_y[0,:]    = img_lum[1, :] - img_lum[0, :]
    grad_x[:,0]    = img_lum[:, 1] - img_lum[:, 0]
    grad_y[-1,:]   = img_lum[-1,:] - img_lum[-2,:]
    grad_x[:,-1]   = img_lum[:,-1] - img_lum[:,-2]
    #normalize to -1, 1
    #grad_x = grad_x / 255.0
    #grad_y = grad_y / 255.0
    grad = np.sqrt(grad_x*grad_x + grad_y*grad_y)
    #to 0, 1
    #grad = grad / np.sqrt(2.0)
    #imsave("grad.png", grad.astype(np.uint8))
    if (mask is not None):
        grad += mask*(256*grad.shape[0]*grad.shape[1])
    return grad

def gen_shifted(row):
    shifted = np.zeros((*row.shape, 3))
    shifted[:-1, 2] = row[1:]
    shifted[ -1, 2] = np.inf
    shifted[  :, 1] = row[:]
    shifted[ 1:, 0] = row[:-1]
    shifted[  0, 0] = np.inf
    return shifted

def seam_carve(img, mode, mask = None):
    mode1, mode2 = tuple(mode.split(' '))
    is_vertical = mode1 == 'vertical'
    is_shrinking = mode2 == 'shrink'
    if (is_vertical):
        img = np.swapaxes(img, 0, 1)
        if (mask is not None):
            mask = mask.T

    img_lum = 0.299*img[:,:,0] + \
              0.587*img[:,:,1] + 0.114*img[:,:,2]

    grad = get_grad(img_lum, mask)
    cum_sum = np.copy(grad)
    ind = np.zeros_like(grad, dtype=np.int32)

    for i in range(1, cum_sum.shape[0]):
        shifted = gen_shifted(cum_sum[i-1])
        ind[i] = np.argmin(shifted, axis=1)
        tmp = np.arange(ind[i].shape[0])
        cum_sum[i,:] = cum_sum[i,:] + shifted[tmp, ind[i]]

    seam_mask = np.zeros_like(grad, dtype=np.int32)
    cur_col = np.argmin(cum_sum[-1])
    cur_row = cum_sum.shape[0] - 1
    x_list = [0]*cum_sum.shape[0]
    while cur_row >= 0:
        x_list[cur_row] = cur_col
        seam_mask[cur_row, cur_col] = 1
        chg = -1 if ind[cur_row, cur_col] == 0 else \
              0 if ind[cur_row, cur_col] == 1 else 1
        cur_row -= 1
        cur_col += chg

    t_seam_mask = np.logical_not(seam_mask).astype(np.bool_)
    t_seam_mask = np.stack([t_seam_mask]*3, axis=2)

    if (is_shrinking):
        res_shape = (img.shape[0], img.shape[1]-1, img.shape[2])
        img_res = np.zeros(res_shape, dtype=np.int32)
        for i in range(res_shape[0]):
            img_res[i, 0:x_list[i],:] = img[i, 0:x_list[i],:]
            img_res[i, x_list[i]:, :] = img[i, x_list[i]+1:,:]
        #img_res = img[t_seam_mask].reshape(res_shape)
        out_mask = None
        if (mask is not None):
            mask_shape = (mask.shape[0], mask.shape[1]-1)
            out_mask = mask[t_seam_mask[:,:,0]].\
                       reshape(mask_shape)
            if (is_vertical):
                out_mask = out_mask.T
    else:
        res_shape = (img.shape[0], img.shape[1]+1, img.shape[2])
        img_res = np.zeros(res_shape, dtype=np.int32)
        for i in range(res_shape[0]):
            img_res[i, 0:x_list[i],  :] = img[i, 0:x_list[i], :]
            img_res[i, x_list[i]+1:, :] = img[i, x_list[i]:,:]
            x_ind = min(img.shape[1]-1, x_list[i]+1)
            img_res[i, x_list[i], :] = (img[i, x_list[i],:].
                                        astype(np.int32) + \
                                        img[i, x_ind,:])//2
        out_mask = None
        if (mask is not None):
            mask_shape = (mask.shape[0], mask.shape[1]+1)
            out_mask = np.zeros(mask_shape, dtype=np.int32)
            for i in range(res_shape[0]):
                out_mask[i, 0:x_list[i]] = mask[i, 0:x_list[i]]
                out_mask[i, x_list[i]+1:] = mask[i, x_list[i]:]
                x_ind  = max(x_list[i]-1, 0)
                x_ind2 = min(out_mask.shape[1], x_list[i]+2)
                out_mask[i,x_ind:x_ind2] = 1
            if (is_vertical):
                out_mask = out_mask.T

    if (is_vertical):
        return np.swapaxes(img_res, 0, 1), out_mask, seam_mask.T
    return img_res, out_mask, seam_mask

def convert_img_to_mask(img):
        return ((img[:, :, 0] != 0) * -1 + (img[:, :, 1] != 0)).astype('int8')

if __name__ == "__main__":
    img = imread("seyana.png")
    # img = img[:,:,0:3]
    # mask = imread("Mask2.png")
    mask = np.zeros_like(img)
    mask = convert_img_to_mask(mask)

    for i in range(60):
        img, mask, _ = seam_carve(img, "vertical shrink", mask)
    imsave("out.png", img)