app.py

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
from datetime import datetime


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.
    # try:
        text_copy = text
        if images:
          style_log = images
          style_inputs, width_length = load_itw_samples(images)
        elif folder:
          style_log = 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')
        
        current_datetime = datetime.now()
        formatted_datetime = current_datetime.strftime("%Y-%m-%d %H:%M:%S")

        print (f'{formatted_datetime}: input_string - {text_copy}, style_input - {style_log}')
    # except:
        print ('ERROR! Try again.')

        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)