スタンフォード/CS131/宿題4-5 オブジェクト除去

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
from skimage import color,io, util
from time import time
from IPython.display import HTML
from __future__ import print_function
from seam_carving import compute_cost,energy_function,enlarge
from seam_carving import backtrack_seam,reduce,remove_seam

%matplotlib inline
plt.rcParams['figure.figsize'] = (15.0, 12.0) # set default size of plots
plt.rcParams["font.size"] = "17"
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'

# Load image
img_yolo = io.imread('imgs/yolo.jpg')
img_yolo = util.img_as_float(img_yolo)

plt.title('Original Image')
plt.imshow(img_yolo)
plt.show()

def compute_forward_cost(image, energy):
    """Computes forward cost map (vertical) and paths of the seams.
    Starting from the first row, compute the cost of each pixel as the sum of energy along the
    lowest energy path from the top.
    Make sure to add the forward cost introduced when we remove the pixel of the seam.
    We also return the paths, which will contain at each pixel either -1, 0 or 1 depending on
    where to go up if we follow a seam at this pixel.
    Args:
        image: numpy array of shape (H, W, 3) or (H, W)
        energy: numpy array of shape (H, W)
    Returns:
        cost: numpy array of shape (H, W)
        paths: numpy array of shape (H, W) containing values -1, 0 or 1
    """
    image = color.rgb2gray(image)
    H, W = image.shape
    cost = np.zeros((H, W))
    paths = np.zeros((H, W), dtype=np.int)
    # Initialization
    cost[0] = energy[0]
    for j in range(W):
        if j > 0 and j < W - 1:
            cost[0,j]+=np.abs(image[0,j+1]-image[0,j-1])
    paths[0] = 0 # we don't care about the first row of paths
    ### YOUR CODE HERE
    for i in range(1, H):
        a = np.insert(image[i,0:W-1],0,0,axis=0)
        b = np.insert(image[i,1:W],W-1,0,axis=0)
        c = image[i-1]
        d = abs(a-b)
        d[0] = 0
        d[-1] = 0
        e = d + abs(c-a)
        f = d + abs(c-b)
        e[0] = 0
        f[-1] = 0
        i1 = np.insert(cost[i-1,0:W-1],0,1e10,axis=0)
        i2 = cost[i-1]
        i3 = np.insert(cost[i-1,1:W],W-1,1e10,axis=0)
        g = np.r_[i1+e,i2+d,i3+f].reshape(3,-1)
        cost[i] = energy[i]+np.min(g,axis=0)
        paths[i] = np.argmin(g,axis=0)-1
    ### END YOUR CODE
    # Check that paths only contains -1, 0 or 1
    assert np.all(np.any([paths == 1, paths == 0, paths == -1], axis=0)), \
           "paths contains other values than -1, 0 or 1"
    return cost, paths

# Let's first test with a small example
img_test = np.array([[1.0, 1.0, 2.0],
                     [0.5, 0.0, 0.0],
                     [1.0, 0.5, 2.0]])
img_test = np.stack([img_test]*3, axis=2)
assert img_test.shape == (3, 3, 3)

energy = energy_function(img_test)

solution_cost = np.array([[0.5, 2.5, 3.0],
                          [1.0, 2.0, 3.0],
                          [2.0, 4.0, 6.0]])

solution_paths = np.array([[ 0,  0,  0],
                           [ 0, -1,  0],
                           [ 0, -1, -1]])

# Vertical Cost Map
vcost, vpaths = compute_forward_cost(img_test, energy)  # don't need the first argument for compute_cost

print("Image:")
print(color.rgb2grey(img_test))

print("Energy:")
print(energy)

print("Cost:")
print(vcost)
print("Solution cost:")
print(solution_cost)

print("Paths:")
print(vpaths)
print("Solution paths:")
print(solution_paths)

assert np.allclose(solution_cost, vcost)
assert np.allclose(solution_paths, vpaths)

energy = energy_function(img_yolo)

out, _ = compute_cost(img_yolo, energy)
plt.subplot(1, 2, 1)
plt.imshow(out, cmap='inferno')
plt.title("Normal cost function")

out, _ = compute_forward_cost(img_yolo, energy)
plt.subplot(1, 2, 2)
plt.imshow(out, cmap='inferno')
plt.title("Forward cost function")
plt.show()

from seam_carving import reduce
out = reduce(img_yolo, 200, axis=0)
plt.imshow(out)
plt.show()

# This step can take a very long time depending on your implementation.
out = reduce(img_yolo, 200, axis=0, cfunc=compute_forward_cost)
plt.imshow(out)
plt.show()

# Load image
image = io.imread('imgs/wyeth.jpg')
image = util.img_as_float(image)

mask = io.imread('imgs/wyeth_mask.jpg', as_gray=True)
mask = util.img_as_bool(mask)

plt.subplot(1, 2, 1)
plt.title('Original Image')
plt.imshow(image)

plt.subplot(1, 2, 2)
plt.title('Mask of the object to remove')
plt.imshow(mask)
plt.show()

def remove_object(image, mask):
    """Remove the object present in the mask.
    Returns an output image with same shape as the input image, but without the object in the mask.
    Args:
        image: numpy array of shape (H, W, 3)
        mask: numpy boolean array of shape (H, W)
    Returns:
        out: numpy array of shape (H, W, 3)
    """
    out = np.copy(image)
    ### YOUR CODE HERE
    H,W,C = out.shape
    while not np.all(mask == 0):
        e = e_function(out)
        e1 = e+mask*(-1000000)
        cost,path = compute_forward_cost(out,e1)
        end = np.argmin(cost[-1])
        seam = backtrack_seam(path,end)
        out = remove_seam(out,seam)
        mask = remove_seam(mask,seam)
        out = enlarge(out,W,axis=1)
    ### END YOUR CODE
    return out

# Use your function to remove the object
from seam_carving import remove_object
out = remove_object(image, mask)

plt.subplot(2, 2, 1)
plt.title('Original Image')
plt.imshow(image)

plt.subplot(2, 2, 2)
plt.title('Mask of the object to remove')
plt.imshow(mask)

plt.subplot(2, 2, 3)
plt.title('Image with object removed')
plt.imshow(out)
plt.show()