Spaces:
Running
Running
ArchitKohli
commited on
Commit
•
4999d68
1
Parent(s):
a68a08c
Upload 2 files
Browse filesAdded model file
- app.py +182 -0
- 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
|
|