# --------------------------------------------------------
# ipynb by Trenton Tabor
# Modified from proj2_starter.py by Yufei Ye (https://github.com/JudyYe)
# Convert from MATLAB code https://inst.eecs.berkeley.edu/~cs194-26/fa18/hw/proj3/gradient_starter.zip
# --------------------------------------------------------
from __future__ import print_function

import numpy as np
import cv2
import imageio
import matplotlib.pyplot as plt
from skimage.transform import SimilarityTransform, warp
import skimage.io as skio
import scipy.sparse as scisparse
import skimage.color as skolor

import matplotlib as mpl

%matplotlib inline

def toy_recon(image):
    # extemely expensive identify function for 2D matrices

    # set up a coordinate to index lut
    imh, imw = image.shape
    im2var = np.arange(imh * imw).reshape((imh, imw)).astype(int) # sort of like range or 0:(imh*imw) reshaped into image shape
    
    # pre-allocate linalg
    countConstraints = 2 * imh * imw + 1 - imh - imw
    #A=np.zeros([countConstraints, imh * imw])
    A = scisparse.lil_matrix((countConstraints, imh * imw))
    b=np.zeros([countConstraints, 1])

    
    e = 0
    # add constraints for x finite difference terms
    for x in range(imw-1):
        for y in range(imh):
            A[e, im2var[y, x + 1]] = 1
            A[e, im2var[y, x]] = -1
            b[e] = int(image[y, x + 1]) - int(image[y, x])
            e += 1
    # add constraints for y finite difference terms
    for x in range(imw):
        for y in range(imh-1):
            A[e, im2var[y+1, x]] = 1
            A[e, im2var[y, x]] = -1
            b[e] = int(image[y+1, x]) - int(image[y, x])
            e += 1
    # add constraints for corner
    A[e, im2var[0, 0]] = 1
    b[e] = image[0, 0]
    # solve, reshape to matrix, demote data type
    #image = np.linalg.lstsq(A,b)[0].reshape((imh, imw)).astype('uint8')
    image = scisparse.linalg.lsqr(A,b)[0].reshape((imh, imw)).astype('uint8')
    
    return image

image = imageio.imread('toy_problem.png')
image_hat = toy_recon(image)

skio.imshow(image)
skio.show()
skio.imshow(image_hat)
skio.show()

def poisson_blend(fg, mask, bg):
    """
    Poisson Blending.
    :param fg: (H, W, C) source texture / foreground object
    :param mask: (H, W)
    :param bg: (H, W, C) target image / background
    :return: (H, W, C)
    """


    # set up a coordinate to index lut
    imh, imw, _ = mask.shape
    pixelCount = imh * imw
    im2var = np.arange(pixelCount).reshape((imh, imw)).astype(int) # sort of like range or 0:(imh*imw) reshaped into image shape
    
    # make some handy transformations of the mask
    maskCount = np.sum(mask[:,:,1])
    maskLinear = mask[:,:,1].reshape([imh*imw])
    listMask = np.argwhere(maskLinear)
    
    # initialize image to store answer in
    fgConstructed = np.zeros(fg.shape)
    
    # can't have more constraints than this, we will reduce constraints later. 
    # With sparse representation, this size doesn't matter much
    constMax = 4 * pixelCount - imh - imw
    for channel in range(3):
        # initialize sparse matrix, dense vector, and count for constraints
        A = scisparse.lil_matrix((constMax, maskCount))
        b = np.zeros([constMax, 1])
        countConstraints = 0
        
        # loop over the mask as indices and subscripts
        for indA, indImage in enumerate(listMask):
            y,x = np.unravel_index(indImage,mask[:,:,1].shape)

            if mask[y,x,1]:
                # Check if pixel below is also in mask, then add internal gradient constraint
                if (y+1 < imh) and mask[y+1,x,1]:
                    indimLocal = np.ravel_multi_index([y+1,x],mask[:,:,1].shape)[0]
                    indALocal = np.argwhere(listMask==indimLocal)[0][0]
                    A[countConstraints,indALocal] = -1
                    A[countConstraints,indA] = 1
                    blocal =  (fg[y,x,channel]) - fg[y+1,x,channel]
                    b[countConstraints,0] = blocal
                    countConstraints += 1

                # Check if pixel to right is also in mask, then add internal gradient constraint
                if (x+1 < imw) and mask[y,x+1,1]:
                    indimLocal = np.ravel_multi_index([y,x+1],mask[:,:,1].shape)[0]
                    indALocal = np.argwhere(listMask==indimLocal)[0][0]
                    A[countConstraints,indALocal] = -1
                    A[countConstraints,indA] = 1
                    b[countConstraints,0] =  (fg[y,x,channel]) - (fg[y,x+1,channel])
                    countConstraints += 1
                    
                # Check if pixel to right is outside mask, then add edge gradient constraint
                if (x+1 < imw) and not mask[y,x+1,1]:
                    A[countConstraints,indA] = 1
                    # print('right')
                    # print((fg[y,x+1,channel]),(fg[y,x,channel]),(bg[y,x+1,channel]))
                    b[countConstraints,0] = -(fg[y,x+1,channel]) + (fg[y,x,channel]) + (bg[y,x+1,channel])
                    countConstraints += 1

                # Check if pixel to below is outside mask, then add edge gradient constraint
                if (y+1 < imh) and not mask[y+1,x,1]:
                    # print('down')
                    # print((fg[y+1,x,channel]),(fg[y,x,channel]),(bg[y+1,x,channel]))
                    A[countConstraints,indA] = 1
                    b[countConstraints,0] = -(fg[y+1,x,channel]) + (fg[y,x,channel]) + (bg[y+1,x,channel])
                    countConstraints += 1

                # Check if pixel to left is outside mask, then add edge gradient constraint
                if (x-1 >= 0) and not mask[y,x-1,1]:
                    A[countConstraints,indA] = 1
                    b[countConstraints,0] = -(fg[y,x-1,channel]) + (fg[y,x,channel]) + (bg[y,x-1,channel])
                    countConstraints += 1

                # Check if pixel above is outside mask, then add edge gradient constraint
                if (y-1 >= 0) and not mask[y-1,x,1]:
                    A[countConstraints,indA] = 1
                    b[countConstraints,0] = -(fg[y-1,x,channel]) + (fg[y,x,channel]) + (bg[y-1,x,channel])
                    countConstraints += 1

        # Remove unused rows of A and b
        A = A[:countConstraints,:]
        b = b[:countConstraints,:]
        
        # perform sparse linear least squares
        sol = scisparse.linalg.lsqr(A,b)[0]

        # build answer. Could probably do better with vstack
        channelIm = np.zeros([pixelCount])
        channelIm[maskLinear] = sol
        channelIm = channelIm.reshape([imh,imw])
        
        fgConstructed[:,:,channel] = channelIm
        
    fg = fgConstructed

    return fg * mask + bg * (1 - mask)

def align_source(fg, bg, mask, offset):
    """simple way to warp source to target frame"""
    tform = SimilarityTransform(translation=offset)
    fg = warp(fg, tform, output_shape=bg.shape)
    mask = warp(mask, tform, output_shape=bg.shape)
    return fg, mask

# after alignment (masking_code.py)
ratio = 1
fg = cv2.resize(imageio.imread('source_01.jpg'), (0, 0), fx=ratio, fy=ratio)
bg = cv2.resize(imageio.imread('target_01.jpg'), (0, 0), fx=ratio, fy=ratio)
mask = cv2.resize(imageio.imread('mask_01.jpg'), (0, 0), fx=ratio, fy=ratio)

offset = -np.array([5, 200])
fg, mask = align_source(fg, bg, mask, offset)

fg = fg / fg.max()
bg = bg / bg.max()
mask = (mask > 0)

blend_img = poisson_blend(fg, mask, bg)

print('foreground image')
skio.imshow(fg)
skio.show()

print('target image')
skio.imshow(bg)
skio.show()

print('masked foreground image')
skio.imshow(fg*mask)
skio.show()

print('blended image')
skio.imshow(blend_img)
skio.show()

foreground image

target image

masked foreground image

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

blended image

Example from assignment. Seems to turn bear blue

# after alignment (masking_code.py)
ratio = 1

fg = cv2.resize(imageio.imread('crusherOnGrass_newsource.png'), (0, 0), fx=ratio, fy=ratio)
bg = cv2.resize(imageio.imread('pastoral.jpg'), (0, 0), fx=ratio, fy=ratio)
mask = cv2.resize(imageio.imread('pastoral_mask.png'), (0, 0), fx=ratio, fy=ratio)

offset = np.array([0,0])
fg, mask = align_source(fg, bg, mask, offset)

fg = fg / fg.max()
bg = bg / bg.max()
mask = (mask > 0)

blend_img = poisson_blend(fg, mask, bg)

print('foreground image')
skio.imshow(fg)
skio.show()

print('target image')
skio.imshow(bg)
skio.show()

print('masked foreground image')
skio.imshow(fg*mask)
skio.show()

print('blended image')
skio.imshow(blend_img)
skio.show()

foreground image

target image

masked foreground image

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

blended image

Example of merging a real picture into a painting. At this low resolution, you may not notice the extreme texture difference between the target and patch. This example benefits from finding backgrounds that are an approximate match for structure, with a subject on grass with the observer positioned at person height.

# after alignment (masking_code.py)
ratio = 1

fg = cv2.resize(imageio.imread('Toho_King_Kong_newsource.png'), (0, 0), fx=ratio, fy=ratio)
bg = cv2.resize(imageio.imread('kiniChimp.jpg'), (0, 0), fx=ratio, fy=ratio)
mask = cv2.resize(imageio.imread('kiniChimp_mask.png'), (0, 0), fx=ratio, fy=ratio)

#offset = -np.array([5, 200])
offset = np.array([0,0])
fg, mask = align_source(fg, bg, mask, offset)

fg = fg / fg.max()
bg = bg / bg.max()
mask = (mask > 0)

blend_img = poisson_blend(fg, mask, bg)

print('foreground image')
skio.imshow(fg)
skio.show()

print('target image')
skio.imshow(bg)
skio.show()

print('masked foreground image')
skio.imshow(fg*mask)
skio.show()

print('blended image')
skio.imshow(blend_img)
skio.show()

foreground image

target image

masked foreground image

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

blended image

This is my failure case. In this example the structure of the background of the target image and the patch we're transplaning are quite different. The background of the target includes a wall that is just blurred out, instead of ending. There is also a gray haze that seems to be added between King Kong's hand and head. Perhaps most interesting, an artifact appears below his raised arm, which was part of his background.

Bells and Whistles - Color to gray

In the section below, see one approach to the color shift preserving color2gray problem. In my approach, I chose an alternative colorspace, since HSV is actually nonlinear, and in my opinion, overrated I convert the image in to ycbcr, then try to preserve gradients in all of the dimensions.

rgb = imageio.imread('colorBlindTest35.png')

print('Original, color image')
skio.imshow(rgb)
skio.show()

print('Naive grayscale conversion')
skio.imshow(cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY))
skio.show()

Original, color image

Naive grayscale conversion

print('HSV V channel, showing success on its own')
hsv = skolor.rgb2hsv(rgb) * 255
skio.imshow(hsv[:,:,2]/255)
skio.show()

HSV V channel, showing success on its own

I wasn't able to do better than this. So I took the approach of trying to mix the results of the naive grayscale with the value channel. This performs worse and does something strange to the '3' in the corner.

def color2gray(rgb_image):
    """EC: Convert an RGB image to gray image"""
    
    ycbcr = skolor.rgb2hsv(rgb_image) * 255
    ycbcr = ycbcr[:,:,2]
    ycbcr = np.dstack([ycbcr,cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)])
    
    # set up a coordinate to index lut
    imh, imw, imd = rgb_image.shape
    im2var = np.arange(imh * imw).reshape((imh, imw)).astype(int) # sort of like range or 0:(imh*imw) reshaped into image shape
    
    # pre-allocate linalg
    countConstraints = 3 * (2 * imh * imw - imh - imw) + 1
    A = scisparse.lil_matrix((countConstraints, imh * imw))
    b = np.zeros([countConstraints, 1])

    w = [1,1.]
    e = 0
    # add constraints for x finite difference terms
    for x in range(imw-1):
        for y in range(imh):
            for d in range(2):
                A[e, im2var[y, x + 1]] = w[d]
                A[e, im2var[y, x]] = -w[d]
                diff = ((ycbcr[y, x + 1, d]) - (ycbcr[y, x,d]))
                b[e] = w[d] * diff
                e += 1
    # add constraints for y finite difference terms
    for x in range(imw):
        for y in range(imh-1):
            for d in range(2):
                A[e, im2var[y + 1, x]] = w[d]
                A[e, im2var[y, x]] = -w[d]
                diff = ((ycbcr[y + 1, x, d]) - (ycbcr[y, x, d]))
                b[e] = w[d] * diff
                e += 1
                
    # add constraints for corner
    A[e, im2var[0, 0]] = 1
    b[e] = ycbcr[0, 0, 1]
    
    # solve,  reshape to matrix, demote data type
    image = scisparse.linalg.lsqr(A,b)[0].reshape((imh, imw)).astype('uint8')
    
    return image # cv2.cvtColor(rgb_image, cv2.COLOR_RGB2GRAY)

gray = color2gray(rgb)

skio.imshow(gray)
skio.show()