utils.py

# Utility

import numpy as np

from os import listdir, mkdir, sep
from os.path import join, exists, splitext
from scipy.misc import imread, imsave, imresize
import skimage
import skimage.io
import skimage.transform
import tensorflow as tf
from PIL import Image
from functools import reduce
import cv2

def list_images(directory):
    images = []
    for file in listdir(directory):
        name = file.lower()
        if name.endswith('.png'):
            images.append(join(directory, file))
        elif name.endswith('.jpg'):
            images.append(join(directory, file))
        elif name.endswith('.jpeg'):
            images.append(join(directory, file))
    return images


def get_train_images(paths, gray, resize_len=512, crop_height=256, crop_width=256, flag = True):
    if isinstance(paths, str):
        paths = [paths]

    images = []
    ny = 0
    nx = 0
    for path in paths:
        if gray:
            input_ir = cv2.imread(path, 0)
            input_ir = input_ir.astype(np.float32)
            input_ir = input_ir.reshape([256, 256, 1])
            return  input_ir
        else:
            input_vi = cv2.imread(path)
            input_vi = input_vi.astype(np.float32)
            input_vi_YCrCb = cv2.cvtColor(input_vi, cv2.COLOR_BGR2YCrCb)
            Y, Cr, Cb = cv2.split(input_vi_YCrCb)
            Y = Y.reshape([256, 256, 1])
            return Y, Cr, Cb





def get_train_images_rgb(path, resize_len=512, crop_height=256, crop_width=256, flag = True):

    # image = imread(path, mode='L')
    image = imread(path, mode='RGB')

    return image


def get_images(paths, height=None, width=None):
    if isinstance(paths, str):
        paths = [paths]

    images = []
    for path in paths:
        image = imread(path, mode='RGB')

        if height is not None and width is not None:
            image = imresize(image, [height, width], interp='nearest')

        images.append(image)

    images = np.stack(images, axis=0)
    print('images shape gen:', images.shape)
    return images


def save_images(paths, datas, save_path, prefix=None, suffix=None):
    if isinstance(paths, str):
        paths = [paths]

    assert(len(paths) == len(datas))

    if not exists(save_path):
        mkdir(save_path)

    if prefix is None:
        prefix = ''
    if suffix is None:
        suffix = ''

    for i, path in enumerate(paths):
        data = datas[i]
        # print('data ==>>\n', data)
        if data.shape[2] == 1:
            data = data.reshape([data.shape[0], data.shape[1]])
        # print('data reshape==>>\n', data)

        name, ext = splitext(path)
        name = name.split(sep)[-1]
        
        path = join(save_path, prefix + suffix + ext)
        print('data path==>>', path)


        # new_im = Image.fromarray(data)
        # new_im.show()

        imsave(path, data)

def get_l2_norm_loss(diffs):
    shape = diffs.get_shape().as_list()
    size = reduce(lambda x, y: x * y, shape) ** 2
    sum_of_squared_diffs = tf.reduce_sum(tf.square(diffs))
    return sum_of_squared_diffs / size