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 | |
| wellcomingMessage = """ | |
| <h1>💥 Handwriting Synthesis - Generate text in anyone's handwriting 💥 </h1> | |
| <p>🚀 This app is a demo for the ICCV'21 paper "Handwriting Transformer". Visit our github paper for more information - <a href="https://github.com/ankanbhunia/Handwriting-Transformers" target="_blank">https://github.com/ankanbhunia/Handwriting-Transformers</a></p> | |
| <p>🚀 You can either choose from an existing style gallery or upload your own handwriting. If you choose to upload, please ensure that you provide a sufficient number of (~15) cropped handwritten word images for the model to work effectively. The demo is made available for research purposes, and any other use is not intended.</p> | |
| """ | |
| 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(DEVICE).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.",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")], | |
| description = wellcomingMessage, | |
| thumbnail = "Handwriting Synthesis - Mimic anyone's handwriting!" | |
| ) | |
| # Launch the Gradio Interface | |
| iface.launch(debug=True, share=True) | |