main.py

#from src.model.models import generator_unet_upsampling
from PIL import Image
import h5py
import cv2
import sys
import os
import numpy as np
import cv2
from keras.applications.resnet50 import preprocess_input
from keras.models import load_model
from config import CONFIG
from inf_cycle import test, test_dep

img_dim = 224
output_path = 'static/results'

# Make dictionary
if not CONFIG['development']:
    # print 'Loading p2d_model'
    # p2d_model = load_model('weights/model_resglass.h5')

    # print 'Loading d2p_model'
    # d2p_model = load_model('weights/model_resglass.h5')
    model_list = {  
            'pix2depth':{ 
                'pix2pix' : load_model('weights/p2d_pix2pix.h5'),
                'CycleGAN':load_model('weights/p2d_cycle.h5'),
                'CNN': load_model('weights/p2d_cnn.h5'),
                },
            'depth2pix':{ 
                'pix2pix' : load_model('weights/d2p_pix2pix.h5'),
                'CycleGAN':load_model('weights/d2p_cycle.h5'),
                }
             }

def pix2depth(path, model):
    model_name = 'p2d'
    originalImage = cv2.imread(path)
    loaded_model =  model_list['pix2depth'][model]
    file_name = model+'_'+path.split('/')[-1]
    output_file = os.path.join(output_path,file_name)
    if model =='CNN':
        originalImage = cv2.resize(originalImage,(img_dim,img_dim))
        x = preprocess_input(originalImage/1.)
    elif model == 'CycleGAN':
        test(path)
        os.system('cp gautam/inf_results/imgs/fakeA_0_0.jpg %s' % output_file)
    else:
        originalImage = cv2.resize(originalImage,(256,256))
        x = originalImage/255.
    if not model == 'CycleGAN':
        p1 = get_depth_map(x, loaded_model)
        cv2.imwrite(output_file,p1)
    return output_file

def depth2pix(path,model):
    model_name = 'd2p'
    originalImage = cv2.imread(path)
    loaded_model =  model_list['depth2pix'][model]
    file_name = model+'_'+path.split('/')[-1]
    output_file = os.path.join(output_path,file_name)
    if model =='CNN':
        img_dim = 256
        originalImage = cv2.resize(originalImage,(img_dim,img_dim))
        x = preprocess_input(originalImage/1.)
    elif model == 'CycleGAN':
        test_dep(path)
        os.system('cp gautam/inf_results/imgs/fakeB_0_0.jpg %s' % output_file)
    else:
        originalImage = cv2.resize(originalImage,(256,256))
        x = originalImage/255.
    if not model == 'CycleGAN':
        p1 = get_depth_map(x, loaded_model)
        cv2.imwrite(output_file,p1)
    return output_file

def blur_effect(image, depthImage, outputPath):
    try:
        if len(depthImage.shape) == 3:
            depthImage = np.mean(depthImage, axis=-1).astype(int)
        print depthImage.shape
        (h, w) = depthImage.shape
        image = cv2.resize(image, (h,w))
        blurredImage = cv2.GaussianBlur(image,(5,5),0)
        print 'b'
        # Need path to depth Image
        thresh = 200
        maskImage = cv2.threshold(depthImage, thresh, 255, cv2.THRESH_BINARY)[1]
        print maskImage.shape
        print h,w
        new_image = np.zeros((h, w, 3),dtype=np.int)
        for i in range(len(maskImage)):
            for j in range(len(maskImage[i])):
                if maskImage[i][j] ==255.0:
                    new_image[i,j,0] = blurredImage[i,j,0]
                    new_image[i,j,1] = blurredImage[i,j,1]
                    new_image[i,j,2] = blurredImage[i,j,2]
                else:
                    new_image[i,j,0] = image[i,j,0]
                    new_image[i,j,1] = image[i,j,1]
                    new_image[i,j,2] = image[i,j,2]
        cv2.imwrite(outputPath, new_image)
        return True
    except Exception as e:
        print e
        return False


def get_depth_map(input_image, model):
    pred_dep = model.predict(np.array([input_image]), batch_size=1)[0]*255.
    return pred_dep

def portrait_mode(path, model):
    originalImage = cv2.imread(path)
    file_name = model+'_'+path.split('/')[-1]
    output_file = pix2depth(path, model)

    portrait_out_path = os.path.join(output_path, 'portrait_'+file_name)
    p1 = cv2.imread(output_file, 0)
    if blur_effect(originalImage, p1, portrait_out_path):
        return portrait_out_path