imseldrith's picture
Update app.py
78724b5
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
from io import BytesIO
from tensorflow.keras.applications.inception_v3 import InceptionV3, preprocess_input
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from transformers import TFAutoModelWithLMHead, AutoTokenizer
from stylegan2.tf_api import G_synthesis as StyleGAN2
# Load pre-trained image captioning model
tokenizer = AutoTokenizer.from_pretrained("t5-large")
model = TFAutoModelWithLMHead.from_pretrained("t5-large")
# Load pre-trained StyleGAN2 model
g = StyleGAN2()
g.load_weights('models/stylegan2-ffhq-config-f.pkl')
# Load pre-trained InceptionV3 model for image preprocessing
inception_v3 = InceptionV3(weights='imagenet')
# Define function to preprocess image for GAN
def preprocess_image(image):
image = image.resize((256, 256))
image_array = img_to_array(image)
image_array = preprocess_input(image_array)
image_array = np.expand_dims(image_array, axis=0)
return image_array
# Define function to generate image from text using StyleGAN2
def generate_image(description):
z = tf.random.normal([1, g.input_shape[1]])
text = "generate image of a " + description
input_ids = tokenizer.encode(text, return_tensors='tf')
output = model.generate(input_ids=input_ids)
caption = tokenizer.decode(output[0], skip_special_tokens=True)
image = g(z, caption)
image = (image.numpy()[0] * 255).astype(np.uint8)
image = Image.fromarray(image, mode='RGB')
return image
# Define function to generate text description of uploaded image using InceptionV3 and T5
def generate_description(image_file):
image = Image.open(BytesIO(image_file.read()))
image = preprocess_image(image)
features = inception_v3.predict(image)
features = tf.keras.backend.flatten(features)
input_text = tokenizer.encode("generate a description of an image", return_tensors="tf")
output = model.generate(input_ids=input_text, attention_mask=tf.ones(input_text.shape), max_length=50)
caption = tokenizer.decode(output[0], skip_special_tokens=True)
return caption
# Define functionto create the web application using Gradio
def image_generation(text_input, image_file):
if image_file is not None:
# Generate text description of uploaded image
description = generate_description(image_file)
# Generate image from text description
generated_image = generate_image(description)
else:
# Generate image from user input text
generated_image = generate_image(text_input)
return generated_image
#Define Gradio interface
inputs = [gr.inputs.Textbox(label="Input text"),
gr.inputs.Image(label="Upload an image (optional)")
]
outputs = gr.outputs.Image(label="Generated Image")
gr.Interface(
fn=image_generation,
inputs=inputs,
outputs=outputs,
title="Image Generation from Text",
description="Generate high-quality images from text descriptions.",
theme="default",
layout="vertical",
examples=[
["a red sports car on a mountain road"],
["a cute puppy"],
["an elegant woman with a hat and a scarf"],
["a scenic beach with palm trees and blue water"],
["a golden retriever sitting on a couch"],
["a delicious pizza with pepperoni and cheese"],
["a futuristic city with tall buildings and flying cars"],
["an adorable kitten playing with a ball of yarn"],
],
).launch(debug=True)