Spaces:
Running
Running
| import gradio as gr | |
| from PIL import Image | |
| import numpy as np | |
| from io import BytesIO | |
| import glob | |
| import os | |
| import time | |
| from data.dataset import load_itw_samples, crop_ | |
| import torch | |
| import cv2 | |
| import os | |
| import numpy as np | |
| from models.model import TRGAN | |
| from params import * | |
| from torch import nn | |
| from data.dataset import get_transform | |
| import pickle | |
| from PIL import Image | |
| import tqdm | |
| import shutil | |
| model_path = 'files/iam_model.pth' | |
| batch_size = 1 | |
| print ('(1) Loading model...') | |
| model = TRGAN(batch_size = batch_size) | |
| model.netG.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')) ) | |
| print (model_path+' : Model loaded Successfully') | |
| model.eval() | |
| # Define a function to generate an image based on text and images | |
| def generate_image(text,folder, _ch3, images): | |
| # Your image generation logic goes here (replace with your actual implementation) | |
| # For demonstration purposes, we'll just concatenate the uploaded images horizontally. | |
| if images: | |
| style_inputs, width_length = load_itw_samples(images) | |
| elif folder: | |
| style_inputs, width_length = load_itw_samples(folder) | |
| else: | |
| return None | |
| # Load images | |
| text = text.replace("\n", "").replace("\t", "") | |
| text_encode = [j.encode() for j in text.split(' ')] | |
| eval_text_encode, eval_len_text = model.netconverter.encode(text_encode) | |
| eval_text_encode = eval_text_encode.to('cuda').repeat(batch_size, 1, 1) | |
| input_styles, page_val = model._generate_page(style_inputs.to(DEVICE).clone(), width_length, eval_text_encode, eval_len_text, no_concat = True) | |
| page_val = crop_(page_val[0]*255) | |
| input_styles = crop_(input_styles[0]*255) | |
| max_width = max(page_val.shape[1],input_styles.shape[1]) | |
| if page_val.shape[1]!=max_width: | |
| page_val = np.concatenate([page_val, np.ones((page_val.shape[0],max_width-page_val.shape[1]))*255], 1) | |
| else: | |
| input_styles = np.concatenate([input_styles, np.ones((input_styles.shape[0],max_width-input_styles.shape[1]))*255], 1) | |
| upper_pad = np.ones((45,input_styles.shape[1]))*255 | |
| input_styles = np.concatenate([upper_pad, input_styles], 0) | |
| page_val = np.concatenate([upper_pad, page_val], 0) | |
| page_val = Image.fromarray(page_val).convert('RGB') | |
| input_styles = Image.fromarray(input_styles).convert('RGB') | |
| return input_styles, page_val | |
| # Define Gradio Interface | |
| iface = gr.Interface( | |
| fn=generate_image, | |
| inputs=[ | |
| gr.Textbox(value = "In the quiet hum of everyday life, the dance of existence unfolds. Time, an ever-flowing river, carries the stories of triumph and heartache. Each fleeting moment is a brushstroke on the canvas of our memories. Within the tapestry of human connection, threads of empathy weave a fabric that binds us all. Nature's symphony plays, a harmonious blend of rustling leaves and birdsong. In the labyrinth of possibility, dreams take flight. Beneath the veneer of routine, lies the extraordinary. Embrace the kaleidoscope of experience, for in the ordinary, the extraordinary often reveals itself.",label = "Input text"), | |
| gr.Dropdown(value = "files/example_data/style-30", choices=glob.glob('files/example_data/*'), label="Choose from provided writer styles"), | |
| gr.Markdown("### OR"), | |
| gr.File(label="Upload multiple word images", file_count="multiple") | |
| ], | |
| outputs=[#gr.Markdown("## Output"), | |
| gr.Image(type="pil", label="Style Image"), | |
| gr.Image(type="pil", label="Generated Image")] | |
| ) | |
| # Launch the Gradio Interface | |
| iface.launch(debug=True, share=True) | |