RobotJelly's picture
app.py
15a6641
raw history blame
No virus
3.41 kB
import os
import random
import numpy as np
import tensorflow as tf
from PIL import Image
import gradio as gr
from huggingface_hub import from_pretrained_keras
model = from_pretrained_keras("keras-io/GauGAN-Image-generation")
def predict(image_file, segmentation_png, bitmap_img):
image_list = [segmentation_png, image_file, bitmap_img]
image = tf.image.decode_png(tf.io.read_file(image_list[1]), channels=3)
image = tf.cast(image, tf.float32) / 127.5 - 1
segmentation_file = tf.image.decode_png(tf.io.read_file(image_list[0]), channels=3)
segmentation_file = tf.cast(segmentation_file, tf.float32)/127.5 - 1
label_file = tf.image.decode_bmp(tf.io.read_file(image_list[2]), channels=0)
# label_file = tf.image.rgb_to_grayscale(tf.image.decode_bmp(tf.io.read_file(image_list[2]), channels=3))
# print("after decode_bmp --> ", label_file.shape, type(label_file))
label_file = tf.squeeze(label_file)
image_list = [segmentation_file, image, label_file]
crop_size = tf.convert_to_tensor((256, 256))
image_shape = tf.shape(image_list[1])[:2]
margins = image_shape - crop_size
y1 = tf.random.uniform(shape=(), maxval=margins[0], dtype=tf.int32)
x1 = tf.random.uniform(shape=(), maxval=margins[1], dtype=tf.int32)
y2 = y1 + crop_size[0]
x2 = x1 + crop_size[1]
cropped_images = []
for img in image_list:
cropped_images.append(img[y1:y2, x1:x2])
final_img_list = [tf.expand_dims(cropped_images[0], axis=0), tf.expand_dims(cropped_images[1], axis=0), tf.expand_dims(tf.one_hot(cropped_images[2], 12), axis=0)]
# print(final_img_list[0].shape)
# print(final_img_list[1].shape)
# print(final_img_list[2].shape)
latent_vector = tf.random.normal(shape=(1, 256), mean=0.0, stddev=2.0)
# Generate fake images
fake_image = model.predict([latent_vector, final_img_list[2]])
fake_img = tf.squeeze(fake_image, axis=0)
return np.array((fake_img+1)/2)
# input
input = [gr.inputs.Image(type="filepath", label="Ground Truth - Real Image (jpg)"),
gr.inputs.Image(type="filepath", label="Corresponding Segmentation (png)"),
gr.inputs.Image(type="filepath", label="Corresponding bitmap image (bmp)", image_mode="L")]
examples = [["facades_data/cmp_b0010.jpg", "facades_data/cmp_b0010.png", "facades_data/cmp_b0010.bmp"],
["facades_data/cmp_b0020.jpg", "facades_data/cmp_b0020.png", "facades_data/cmp_b0020.bmp"],
["facades_data/cmp_b0030.jpg", "facades_data/cmp_b0030.png", "facades_data/cmp_b0030.bmp"],
["facades_data/cmp_b0040.jpg", "facades_data/cmp_b0040.png", "facades_data/cmp_b0040.bmp"],
["facades_data/cmp_b0050.jpg", "facades_data/cmp_b0050.png", "facades_data/cmp_b0050.bmp"]]
# output
output = [gr.outputs.Image(type="numpy", label="Generated - Conditioned Images")]
title = "GauGAN For Conditional Image Generation"
description = "Upload an Image or take one from examples to generate realistic images that are conditioned on cue images and segmentation maps"
gr.Interface(fn=predict, inputs = input, outputs = output, examples=examples, allow_flagging=False, analytics_enabled=False,
title=title, description=description, article="<center>Space By: <u><a href='https://github.com/robotjellyzone'><b>Kavya Bisht</b></a></u> \n Based on <a href='https://keras.io/examples/generative/gaugan/'><b>this notebook</b></a></center>").launch(enable_queue=True, debug=True)