ArchitKohli commited on
Commit
4999d68
1 Parent(s): a68a08c

Upload 2 files

Browse files

Added model file

Files changed (2) hide show
  1. app.py +182 -0
  2. requirements.txt +0 -0
app.py ADDED
@@ -0,0 +1,182 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import tensorflow as tf
4
+ import numpy as np
5
+
6
+ BUFFER_SIZE = 1000
7
+ BATCH_SIZE = 1
8
+ IMG_WIDTH = 256
9
+ IMG_HEIGHT = 256
10
+
11
+ def random_crop(image):
12
+ cropped_image = tf.image.random_crop(
13
+ image[0], size=[IMG_HEIGHT, IMG_WIDTH, 3])
14
+
15
+ return cropped_image
16
+
17
+ # normalizing the images to [-1, 1]
18
+ def normalize(image):
19
+ image = tf.cast(image, tf.float32)
20
+ image = (image / 127.5) - 1
21
+ return image
22
+
23
+ def random_jitter(image):
24
+ # resizing to 286 x 286 x 3
25
+ image = tf.image.resize(image, [286, 286],
26
+ method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
27
+
28
+ # randomly cropping to 256 x 256 x 3
29
+ image = random_crop(image)
30
+
31
+ # random mirroring
32
+ image = tf.image.random_flip_left_right(image)
33
+
34
+ return image
35
+
36
+ def preprocess_image_train(image):
37
+ image = random_jitter(image)
38
+ image = normalize(image)
39
+ return image
40
+
41
+ def preprocess_image_test(image):
42
+ image = tf.image.resize(image, [286, 286],method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
43
+ image = random_crop(image)
44
+ image = normalize(image)
45
+ return image
46
+
47
+ OUTPUT_CHANNELS=3
48
+
49
+ def downsample(filters, size, strides=2, apply_batchnorm=True, leaky_relu=None, padding=True):
50
+ initializer = tf.random_normal_initializer(0., 0.02)
51
+
52
+ result = tf.keras.Sequential()
53
+ if padding:
54
+ result.add(
55
+ tf.keras.layers.Conv2D(filters, size, strides=strides, padding='same',
56
+ kernel_initializer=initializer, use_bias=False))
57
+ else:
58
+ result.add(
59
+ tf.keras.layers.Conv2D(filters, size, strides=strides, padding='valid',
60
+ kernel_initializer=initializer, use_bias=False))
61
+
62
+ if apply_batchnorm:
63
+ result.add(tf.keras.layers.BatchNormalization())
64
+ if leaky_relu is True:
65
+ result.add(tf.keras.layers.LeakyReLU())
66
+ elif leaky_relu is False:
67
+ result.add(tf.keras.layers.ReLU())
68
+ return result
69
+
70
+ def upsample(filters, size, apply_dropout=False):
71
+ initializer = tf.random_normal_initializer(0., 0.02)
72
+
73
+ result = tf.keras.Sequential()
74
+ result.add(
75
+ tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
76
+ padding='same',
77
+ kernel_initializer=initializer,
78
+ use_bias=False))
79
+
80
+ result.add(tf.keras.layers.BatchNormalization())
81
+
82
+ if apply_dropout:
83
+ result.add(tf.keras.layers.Dropout(0.5))
84
+
85
+ result.add(tf.keras.layers.ReLU())
86
+
87
+ return result
88
+
89
+ def resnet_block():
90
+ result = tf.keras.models.Sequential()
91
+ result.add(downsample(256,3,strides=1,apply_batchnorm=True,leaky_relu=False))
92
+ result.add(downsample(256,3,strides=1,apply_batchnorm=True,leaky_relu=None))
93
+ return result
94
+
95
+ def Generator():
96
+ inputs = tf.keras.layers.Input(shape=[256, 256, 3])
97
+ down_stack = [
98
+ downsample(64,7,strides=1,apply_batchnorm=False,leaky_relu=True),
99
+ downsample(128,3,strides=2,apply_batchnorm=True,leaky_relu=True),
100
+ downsample(256,3,strides=2,apply_batchnorm=True,leaky_relu=True),
101
+ ]
102
+ resnet_stack = [
103
+ resnet_block()
104
+ ]*9
105
+ up_stack = [
106
+ upsample(128,3),
107
+ upsample(64,3),
108
+ ]
109
+ final_layer = downsample(3,7,strides=1,apply_batchnorm=False)
110
+
111
+ x = inputs
112
+
113
+ for down in down_stack:
114
+ x = down(x)
115
+
116
+ for res in resnet_stack:
117
+ x = x + res(x)
118
+
119
+ for up in up_stack:
120
+ x = up(x)
121
+
122
+ x = final_layer(x)
123
+ return tf.keras.Model(inputs=inputs, outputs=x)
124
+
125
+
126
+ def Discriminator():
127
+ initializer = tf.random_normal_initializer(0., 0.02)
128
+
129
+ inputs = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image')
130
+ x = inputs
131
+ x = downsample(64,4,strides=2,apply_batchnorm=False,leaky_relu=True)(x)
132
+ x = downsample(128,4,strides=2,apply_batchnorm=False,leaky_relu=True)(x)
133
+ x = downsample(256,4,strides=2,apply_batchnorm=False,leaky_relu=True)(x)
134
+
135
+ x = tf.keras.layers.ZeroPadding2D()(x)
136
+ x = downsample(512,4,strides=1,apply_batchnorm=False,leaky_relu=True,padding=False)(x)
137
+ x = tf.keras.layers.ZeroPadding2D()(x)
138
+ x = downsample(1,4,strides=1,apply_batchnorm=False,leaky_relu=True,padding=False)(x)
139
+
140
+ return tf.keras.Model(inputs=inputs,outputs=x)
141
+
142
+
143
+ generator_m2f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
144
+ generator_f2m_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
145
+
146
+ discriminator_m_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
147
+ discriminator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
148
+
149
+ generator_m2f = Generator()
150
+ generator_f2m = Generator()
151
+
152
+ discriminator_f = Discriminator()
153
+ discriminator_m = Discriminator()
154
+
155
+ checkpoint_dir = "/AnimizerGAN/checkpoint_dir"
156
+ checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
157
+ checkpoint = tf.train.Checkpoint(generator_m2f=generator_m2f,
158
+ generator_f2m=generator_f2m,
159
+ discriminator_m=discriminator_m,
160
+ discriminator_f=discriminator_f,
161
+ generator_m2f_optimizer=generator_m2f_optimizer,
162
+ generator_f2m_optimizer=generator_f2m_optimizer,
163
+ discriminator_m_optimizer=discriminator_m_optimizer,
164
+ discriminator_f_optimizer=discriminator_f_optimizer)
165
+
166
+ checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
167
+
168
+ def animize(img):
169
+ img = preprocess_image_test(np.reshape(img,(1,np.shape(img)[0],np.shape(img)[1],np.shape(img)[2])))
170
+ img = tf.expand_dims(img,0)
171
+ prediction = generator_f2m(img)
172
+
173
+ # converting to -1 to 1
174
+ prediction = (prediction - np.min(prediction)) / (np.max(prediction) - np.min(prediction))
175
+ prediction = prediction * 2 - 1
176
+ prediction = (prediction[0]*0.5+0.5).numpy()
177
+ # out = tf.image.resize(image, [128, 128],method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
178
+
179
+ return prediction
180
+
181
+ interface = gr.Interface(fn=animize, inputs="image", outputs="image")
182
+ interface.launch(share=True)
requirements.txt ADDED
Binary file (3.56 kB). View file