Working Laplacian Creator

laplacian.py

# Standalone version of the Matting Laplacian code from here: https://github.com/martinbenson/deep-photo-styletransfer
# Usage: python3 laplacian.py -in_dir <directory> -lap_dir <directory> -width <value>
# Install the depdendencies with: pip3 install numpy scipy Pillow
# This script is intended for use with artistic style transfer neural networks, and Deep Photo Style Transfer.

import argparse
import glob
import os
import shutil
import multiprocessing
import math
import subprocess
import scipy.misc as spm
import scipy.ndimage as spi
import scipy.sparse as sps
import numpy as np

def getlaplacian1(i_arr: np.ndarray, consts: np.ndarray, epsilon: float = 0.0000001, win_size: int = 1):
    neb_size = (win_size * 2 + 1) ** 2
    h, w, c = i_arr.shape
    img_size = w * h
    consts = spi.morphology.grey_erosion(consts, footprint=np.ones(shape=(win_size * 2 + 1, win_size * 2 + 1)))

    indsM = np.reshape(np.array(range(img_size)), newshape=(h, w), order='F')
    tlen = int((-consts[win_size:-win_size, win_size:-win_size] + 1).sum() * (neb_size ** 2))

    row_inds = np.zeros(tlen)
    col_inds = np.zeros(tlen)
    vals = np.zeros(tlen)
    l = 0
    for j in range(win_size, w - win_size):
        for i in range(win_size, h - win_size):
            if consts[i, j]:
                continue
            win_inds = indsM[i - win_size:i + win_size + 1, j - win_size: j + win_size + 1]
            win_inds = win_inds.ravel(order='F')
            win_i = i_arr[i - win_size:i + win_size + 1, j - win_size: j + win_size + 1, :]
            win_i = win_i.reshape((neb_size, c), order='F')
            win_mu = np.mean(win_i, axis=0).reshape(1, win_size * 2 + 1)
            win_var = np.linalg.inv(
                np.matmul(win_i.T, win_i) / neb_size - np.matmul(win_mu.T, win_mu) + epsilon / neb_size * np.identity(
                    c))

            win_i2 = win_i - win_mu
            tvals = (1 + np.matmul(np.matmul(win_i2, win_var), win_i2.T)) / neb_size

            ind_mat = np.broadcast_to(win_inds, (neb_size, neb_size))
            row_inds[l: (neb_size ** 2 + l)] = ind_mat.ravel(order='C')
            col_inds[l: neb_size ** 2 + l] = ind_mat.ravel(order='F')
            vals[l: neb_size ** 2 + l] = tvals.ravel(order='F')
            l += neb_size ** 2

    vals = vals.ravel(order='F')
    row_inds = row_inds.ravel(order='F')
    col_inds = col_inds.ravel(order='F')
    a_sparse = sps.csr_matrix((vals, (row_inds, col_inds)), shape=(img_size, img_size))

    sum_a = a_sparse.sum(axis=1).T.tolist()[0]
    a_sparse = sps.diags([sum_a], [0], shape=(img_size, img_size)) - a_sparse

    return a_sparse


def im2double(im):
    min_val = np.min(im.ravel())
    max_val = np.max(im.ravel())
    return (im.astype('float') - min_val) / (max_val - min_val)


def reshape_img(in_img, l=512):
    in_h, in_w, _ = in_img.shape
    if in_h > in_w:
        h2 = l
        w2 = int(in_w * h2 / in_h)
    else:
        w2 = l
        h2 = int(in_h * w2 / in_w)
	
    return spm.imresize(in_img, (h2, w2))
    
    
    
if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("-in_dir", "--in_directory", help="Path to inputs", required=True)
    parser.add_argument("-lap_dir", "--laplacian_directory", help="Path to laplacians", required=True)
    parser.add_argument("-width", "--width", help="Image width", default=512)
    args = parser.parse_args()
    width = int(args.width)
	
    if not os.path.exists("/tmp/deep_photo/"):
        os.makedirs("/tmp/deep_photo/")
	
    if not os.path.exists("/tmp/deep_photo/in"):
        os.makedirs("/tmp/deep_photo/in")
	
    if not os.path.exists(args.laplacian_directory):
        os.makedirs(args.laplacian_directory)
		
    files = []
    for f in glob.iglob(os.path.join(args.in_directory, '*.png')):
        files.append(f)
		
    good_images = []
    for f in files:
        image_name = os.path.basename(f)
        good_images.append(image_name)
		
    def process_image(image_name):
        filename = os.path.join(args.in_directory, image_name)
        lap_name = os.path.join(args.laplacian_directory,
        image_name.replace(".png", "") + "_" + str(args.width) + ".csv")

        img = spi.imread(filename, mode="RGB")
        resized_img = reshape_img(img, width)
        spm.imsave("/tmp/deep_photo/in/" + image_name, resized_img)
		
	#if not os.path.exists(lap_name):
        print("Calculating matting laplacian for " + str(image_name) + "...")
        img = im2double(resized_img)
        h, w, c = img.shape
        csr = getlaplacian1(img, np.zeros(shape=(h, w)), 1e-7, 1)
        coo = csr.tocoo()
        zipped = zip(coo.row + 1, coo.col + 1, coo.data)
        with open(lap_name, 'w') as out_file:
            out_file.write(str(len(coo.data))+"\n")
            for row, col, val in zipped:
                 out_file.write("%d,%d,%.15f\n" % (row, col, val))

pool = multiprocessing.Pool(multiprocessing.cpu_count())
pool.map(process_image, good_images)

shutil.rmtree("/tmp/deep_photo/", ignore_errors=True)