Spaces:
Running
Running
import os | |
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
BUFFER_SIZE = 1000 | |
BATCH_SIZE = 1 | |
IMG_WIDTH = 256 | |
IMG_HEIGHT = 256 | |
def random_crop(image): | |
cropped_image = tf.image.random_crop( | |
image[0], size=[IMG_HEIGHT, IMG_WIDTH, 3]) | |
return cropped_image | |
# normalizing the images to [-1, 1] | |
def normalize(image): | |
image = tf.cast(image, tf.float32) | |
image = (image / 127.5) - 1 | |
return image | |
def random_jitter(image): | |
# resizing to 286 x 286 x 3 | |
image = tf.image.resize(image, [286, 286], | |
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) | |
# randomly cropping to 256 x 256 x 3 | |
image = random_crop(image) | |
# random mirroring | |
image = tf.image.random_flip_left_right(image) | |
return image | |
def preprocess_image_train(image): | |
image = random_jitter(image) | |
image = normalize(image) | |
return image | |
def preprocess_image_test(image): | |
image = tf.image.resize(image, [286, 286],method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) | |
image = random_crop(image) | |
image = normalize(image) | |
return image | |
OUTPUT_CHANNELS=3 | |
def downsample(filters, size, strides=2, apply_batchnorm=True, leaky_relu=None, padding=True): | |
initializer = tf.random_normal_initializer(0., 0.02) | |
result = tf.keras.Sequential() | |
if padding: | |
result.add( | |
tf.keras.layers.Conv2D(filters, size, strides=strides, padding='same', | |
kernel_initializer=initializer, use_bias=False)) | |
else: | |
result.add( | |
tf.keras.layers.Conv2D(filters, size, strides=strides, padding='valid', | |
kernel_initializer=initializer, use_bias=False)) | |
if apply_batchnorm: | |
result.add(tf.keras.layers.BatchNormalization()) | |
if leaky_relu is True: | |
result.add(tf.keras.layers.LeakyReLU()) | |
elif leaky_relu is False: | |
result.add(tf.keras.layers.ReLU()) | |
return result | |
def upsample(filters, size, apply_dropout=False): | |
initializer = tf.random_normal_initializer(0., 0.02) | |
result = tf.keras.Sequential() | |
result.add( | |
tf.keras.layers.Conv2DTranspose(filters, size, strides=2, | |
padding='same', | |
kernel_initializer=initializer, | |
use_bias=False)) | |
result.add(tf.keras.layers.BatchNormalization()) | |
if apply_dropout: | |
result.add(tf.keras.layers.Dropout(0.5)) | |
result.add(tf.keras.layers.ReLU()) | |
return result | |
def resnet_block(): | |
result = tf.keras.models.Sequential() | |
result.add(downsample(256,3,strides=1,apply_batchnorm=True,leaky_relu=False)) | |
result.add(downsample(256,3,strides=1,apply_batchnorm=True,leaky_relu=None)) | |
return result | |
def Generator(): | |
inputs = tf.keras.layers.Input(shape=[256, 256, 3]) | |
down_stack = [ | |
downsample(64,7,strides=1,apply_batchnorm=False,leaky_relu=True), | |
downsample(128,3,strides=2,apply_batchnorm=True,leaky_relu=True), | |
downsample(256,3,strides=2,apply_batchnorm=True,leaky_relu=True), | |
] | |
resnet_stack = [ | |
resnet_block() | |
]*9 | |
up_stack = [ | |
upsample(128,3), | |
upsample(64,3), | |
] | |
final_layer = downsample(3,7,strides=1,apply_batchnorm=False) | |
x = inputs | |
for down in down_stack: | |
x = down(x) | |
for res in resnet_stack: | |
x = x + res(x) | |
for up in up_stack: | |
x = up(x) | |
x = final_layer(x) | |
return tf.keras.Model(inputs=inputs, outputs=x) | |
def Discriminator(): | |
initializer = tf.random_normal_initializer(0., 0.02) | |
inputs = tf.keras.layers.Input(shape=[256, 256, 3], name='input_image') | |
x = inputs | |
x = downsample(64,4,strides=2,apply_batchnorm=False,leaky_relu=True)(x) | |
x = downsample(128,4,strides=2,apply_batchnorm=False,leaky_relu=True)(x) | |
x = downsample(256,4,strides=2,apply_batchnorm=False,leaky_relu=True)(x) | |
x = tf.keras.layers.ZeroPadding2D()(x) | |
x = downsample(512,4,strides=1,apply_batchnorm=False,leaky_relu=True,padding=False)(x) | |
x = tf.keras.layers.ZeroPadding2D()(x) | |
x = downsample(1,4,strides=1,apply_batchnorm=False,leaky_relu=True,padding=False)(x) | |
return tf.keras.Model(inputs=inputs,outputs=x) | |
generator_m2f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) | |
generator_f2m_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) | |
discriminator_m_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) | |
discriminator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) | |
generator_m2f = Generator() | |
generator_f2m = Generator() | |
discriminator_f = Discriminator() | |
discriminator_m = Discriminator() | |
checkpoint_dir = "/AnimizerGAN/checkpoint_dir" | |
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt") | |
checkpoint = tf.train.Checkpoint(generator_m2f=generator_m2f, | |
generator_f2m=generator_f2m, | |
discriminator_m=discriminator_m, | |
discriminator_f=discriminator_f, | |
generator_m2f_optimizer=generator_m2f_optimizer, | |
generator_f2m_optimizer=generator_f2m_optimizer, | |
discriminator_m_optimizer=discriminator_m_optimizer, | |
discriminator_f_optimizer=discriminator_f_optimizer) | |
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir)) | |
def animize(img): | |
img = preprocess_image_test(np.reshape(img,(1,np.shape(img)[0],np.shape(img)[1],np.shape(img)[2]))) | |
img = tf.expand_dims(img,0) | |
prediction = generator_f2m(img) | |
# converting to -1 to 1 | |
prediction = (prediction - np.min(prediction)) / (np.max(prediction) - np.min(prediction)) | |
prediction = prediction * 2 - 1 | |
prediction = (prediction[0]*0.5+0.5).numpy() | |
# out = tf.image.resize(image, [128, 128],method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) | |
return prediction | |
interface = gr.Interface(fn=animize, inputs="image", outputs="image") | |
interface.launch(share=True) |