UNIT-Tensorflow/utils.py at master · EmptyOps/UNIT-Tensorflow · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import tensorflow as tf
from tensorflow.contrib import slim
from scipy import misc
import os, random
import numpy as np
from PIL import Image
import cv2


class ImagePool:
    """ History of generated images
        Same logic as https://github.com/junyanz/CycleGAN/blob/master/util/image_pool.lua
    """

    def __init__(self, pool_size):
        self.pool_size = pool_size
        self.images = []

    def query(self, image):
        if self.pool_size == 0:
            return image

        if len(self.images) < self.pool_size:
            self.images.append(image)
            return image
        else:
            p = random.random()
            if p > 0.5:
                # use old image
                random_id = random.randrange(0, self.pool_size)
                tmp = self.images[random_id].copy()
                self.images[random_id] = image.copy()
                return tmp
            else:
                return image


def prepare_data(dataset_name, size):
    data_path = os.path.join("./dataset", dataset_name)

    trainA = []
    trainB = []
    for path, dir, files in os.walk(data_path):
        for file in files:
            if file == '.DS_Store':
                continue

            image = os.path.join(path, file)
            if path.__contains__('trainA') :
                #hiren fixed this to support the new scipy version
                # trainA.append(misc.imresize(misc.imread(image, mode='RGB'), [size, size]))
                print( "image ", image, " ", size )
                trainA.append( np.array( Image.fromarray( cv2.imread(image) ).resize( [size, size] ) ) )
            if path.__contains__('trainB') :
                #hiren fixed this to support the new scipy version
                # trainB.append(misc.imresize(misc.imread(image, mode='RGB'), [size, size]))
                trainB.append( np.array( Image.fromarray( cv2.imread(image) ).resize( [size, size] ) ) )


    trainA = preprocessing(np.asarray(trainA))
    trainB = preprocessing(np.asarray(trainB))

    np.random.shuffle(trainA)
    np.random.shuffle(trainB)

    return trainA, trainB

def test_data(dataset_name, size) :
    data_path = os.path.join("./dataset", dataset_name)
    testA = []
    testB = []
    for path, dir, files in os.walk(data_path) :
        for file in files :
            image = os.path.join(path, file)
            if path.__contains__('testA') :
                testA.append(misc.imresize(misc.imread(image, mode='RGB'), [size, size]))
            if path.__contains__('testB') :
                testB.append(misc.imresize(misc.imread(image, mode='RGB'), [size, size]))

    testA = preprocessing(np.asarray(testA))
    testB = preprocessing(np.asarray(testB))

    return testA, testB

def load_test_data(image_path, size=256):
    img = misc.imread(image_path, mode='RGB')
    img = misc.imresize(img, [size, size])
    img = np.expand_dims(img, axis=0)
    img = preprocessing(img)

    return img

def preprocessing(x):
    """
    # Create Normal distribution
    x = x.astype('float32')
    x[:, :, :, 0] = (x[:, :, :, 0] - np.mean(x[:, :, :, 0])) / np.std(x[:, :, :, 0])
    x[:, :, :, 1] = (x[:, :, :, 1] - np.mean(x[:, :, :, 1])) / np.std(x[:, :, :, 1])
    x[:, :, :, 2] = (x[:, :, :, 2] - np.mean(x[:, :, :, 2])) / np.std(x[:, :, :, 2])
    """
    x = x/127.5 - 1 # -1 ~ 1
    return x

def augmentation(image, augment_size):
    seed = random.randint(0, 2 ** 31 - 1)
    ori_image_shape = tf.shape(image)
    image = tf.image.resize_images(image, [augment_size, augment_size])
    image = tf.random_crop(image, ori_image_shape, seed=seed)
    image = tf.map_fn(lambda x: tf.image.random_flip_left_right(x, seed), image)
    return image

def save_images(images, size, image_path):
    return imsave(inverse_transform(images), size, image_path)

def inverse_transform(images):
    return (images+1.) / 2

def imsave(images, size, path):
    return misc.imsave(path, merge(images, size))

def merge(images, size):
    h, w = images.shape[1], images.shape[2]
    img = np.zeros((h * size[0], w * size[1], 3))
    for idx, image in enumerate(images):
        i = idx % size[1]
        j = idx // size[1]
        img[h*j:h*(j+1), w*i:w*(i+1), :] = image

    return img

def show_all_variables():
    model_vars = tf.trainable_variables()
    slim.model_analyzer.analyze_vars(model_vars, print_info=True)

def check_folder(log_dir):
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    return log_dir