app.py

import torch
import argparse
import numpy as np
from helper import *
from config.GlobalVariables import *
from SynthesisNetwork import SynthesisNetwork
from DataLoader import DataLoader
import convenience
import gradio as gr

device = 'cpu'
num_samples = 10

net = SynthesisNetwork(weight_dim=256, num_layers=3).to(device)

if not torch.cuda.is_available():
    net.load_state_dict(torch.load('./model/250000.pt', map_location=torch.device(device))["model_state_dict"])


dl = DataLoader(num_writer=1, num_samples=10, divider=5.0, datadir='./data/writers')


writer_options = [5, 14, 15, 16, 17, 22, 25, 80, 120, 137, 147, 151]
all_loaded_data = []
chosen_writers = [120, 80]
avail_char = "0 1 2 3 4 5 6 7 8 9 a b c d e f g h i j k l m n o p q r s t u v w x y z A B C D E F G H I J K L M N O P Q R S T U V W X Y Z ! ? \" ' * + - = : ; , . < > \ / [ ] ( ) # $ % &"
avail_char_list = avail_char.split(" ")
for writer_id in chosen_writers:
    loaded_data = dl.next_batch(TYPE='TRAIN', uid=writer_id, tids=list(range(num_samples)))
    all_loaded_data.append(loaded_data)

default_loaded_data = all_loaded_data[-1]

# data for writer interpolation
writer_words = ["hello", "world"]
writer_mean_Ws = []
all_word_writer_Ws = []
all_word_writer_Cs = []
writer_weight = 0.7


# data for char interpolation
blend_chars = ["y", "s"]
char_mean_global_W = None
char_weight = 0.7
default_mean_global_W = convenience.get_mean_global_W(net, default_loaded_data, device)
char_Ws = default_mean_global_W.reshape(1, 1, convenience.L)
char_Cs = all_Cs = torch.zeros(1, 2, convenience.L, convenience.L)


# data for MDN
mdn_words = ["hello", "world"]
mdn_mean_global_W = None
all_word_mdn_Ws = []
all_word_mdn_Cs = []


def update_writer_word(target_word):
    writer_words.clear()
    for word in target_word.split(" "):
        writer_words.append(word)

    all_word_writer_Ws.clear()
    all_word_writer_Cs.clear()
    for word in writer_words:
        all_writer_Ws, all_writer_Cs = convenience.get_DSD(net, word, writer_mean_Ws, all_loaded_data, device)
        all_word_writer_Ws.append(all_writer_Ws)
        all_word_writer_Cs.append(all_writer_Cs)

    return update_writer_slider(writer_weight)


# for writer interpolation
def update_writer_slider(val):
    global writer_weight
    writer_weight = val
    weights = [1 - writer_weight, writer_weight]

    net.clamp_mdn = 0
    im = convenience.draw_words(writer_words, all_word_writer_Ws, all_word_writer_Cs, weights, net)
    return im.convert("RGB")


def update_chosen_writers(writer1, writer2):
    net.clamp_mdn = 0
    chosen_writers[0], chosen_writers[1] = int(writer1.split(" ")[1]), int(writer2.split(" ")[1])

    all_loaded_data.clear()
    for writer_id in chosen_writers:
        loaded_data = dl.next_batch(TYPE='TRAIN', uid=writer_id, tids=list(range(num_samples)))
        all_loaded_data.append(loaded_data)

    writer_mean_Ws.clear()
    for loaded_data in all_loaded_data:
        mean_global_W = convenience.get_mean_global_W(net, loaded_data, device)
        writer_mean_Ws.append(mean_global_W)

    return gr.Slider.update(label=f"{writer1} vs. {writer2}"), update_writer_slider(writer_weight)

# for character blend


def update_char_slider(weight):
    """Generates an image of handwritten text based on target_sentence"""
    net.clamp_mdn = 0

    global char_weight
    char_weight = weight
    character_weights = [1 - weight, weight]

    all_W_c = convenience.get_character_blend_W_c(character_weights, char_Ws, char_Cs)
    all_commands = convenience.get_commands(net, blend_chars[0], all_W_c)
    im = convenience.commands_to_image(all_commands, 160, 750, 375, 30)

    return im.convert("RGB")


def update_blend_chars(c1, c2):
    global blend_chars
    blend_chars[0], blend_chars[1] = c1, c2

    for i in range(2):  # get corners of grid
        _, char_matrix = convenience.get_DSD(net, blend_chars[i], default_mean_global_W, [default_loaded_data], device)
        char_Cs[:, i, :, :] = char_matrix

    return gr.Slider.update(label=f"'{c1}' vs. '{c2}'")

# for MDN


def update_mdn_word(target_word):
    mdn_words.clear()
    for word in target_word.split(" "):
        mdn_words.append(word)

    all_word_mdn_Ws.clear()
    all_word_mdn_Cs.clear()
    for word in mdn_words:
        all_writer_Ws, all_writer_Cs = convenience.get_DSD(net, word, default_mean_global_W, [default_loaded_data], device)
        all_word_mdn_Ws.append(all_writer_Ws)
        all_word_mdn_Cs.append(all_writer_Cs)

    return sample_mdn(net.scale_sd, net.clamp_mdn)


def sample_mdn(maxs, maxr):
    net.clamp_mdn = maxr
    net.scale_sd = maxs
    im = convenience.draw_words(mdn_words, all_word_mdn_Ws, all_word_mdn_Cs, [1], net)
    return im.convert("RGB")


update_writer_word(" ".join(writer_words))
update_chosen_writers(f"Writer {chosen_writers[0]}", f"Writer {chosen_writers[1]}")

update_mdn_word(" ".join(writer_words))

update_blend_chars(*blend_chars)

with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.TabItem("Blend Writers"):
            target_word = gr.Textbox(label="Target Word", value=" ".join(writer_words), max_lines=1)
            with gr.Row():
                left_ratio_options = ["Style " + str(id) for i, id in enumerate(writer_options) if i % 2 == 0]
                right_ratio_options = ["Style " + str(id) for i, id in enumerate(writer_options) if i % 2 == 1]
                with gr.Column():
                    writer1 = gr.Radio(left_ratio_options, value="Style 120", label="Style for first writer")
                with gr.Column():
                    writer2 = gr.Radio(right_ratio_options, value="Style 80", label="Style for second writer")
            with gr.Row():
                writer_slider = gr.Slider(0, 1, value=writer_weight, label="Style 120 vs. Style 80")
            with gr.Row():
                writer_submit = gr.Button("Submit")
            with gr.Row():
                writer_default_image = update_writer_slider(writer_weight)
                writer_output = gr.Image(writer_default_image)

            writer_submit.click(fn=update_writer_slider, inputs=[writer_slider], outputs=[writer_output], show_progress=False)
            writer_slider.change(fn=update_writer_slider, inputs=[writer_slider], outputs=[writer_output], show_progress=False)
            target_word.submit(fn=update_writer_word, inputs=[target_word], outputs=[writer_output], show_progress=False)

            writer1.change(fn=update_chosen_writers, inputs=[writer1, writer2], outputs=[writer_slider, writer_output], show_progress=False)
            writer2.change(fn=update_chosen_writers, inputs=[writer1, writer2], outputs=[writer_slider, writer_output], show_progress=False)

        with gr.TabItem("Blend Characters"):
            with gr.Row():
                with gr.Column():
                    char1 = gr.Dropdown(choices=avail_char_list, value=blend_chars[0], label="Character 1")
                with gr.Column():
                    char2 = gr.Dropdown(choices=avail_char_list, value=blend_chars[1], label="Character 2")
            with gr.Row():
                char_slider = gr.Slider(0, 1, value=char_weight, label=f"'{blend_chars[0]}' vs. '{blend_chars[1]}'")
            with gr.Row():
                char_default_image = update_char_slider(char_weight)
                char_output = gr.Image(char_default_image)

            char_slider.change(fn=update_char_slider, inputs=[char_slider], outputs=[char_output], show_progress=False)

            char1.change(fn=update_blend_chars, inputs=[char1, char2], outputs=[char_slider], show_progress=False)
            char2.change(fn=update_blend_chars, inputs=[char1, char2], outputs=[char_slider], show_progress=False)

        with gr.TabItem("Add Randomness"):
            mdn_word = gr.Textbox(label="Target Word", value=" ".join(mdn_words), max_lines=1)
            with gr.Row():
                with gr.Column():
                    max_rand = gr.Slider(0, 1, value=net.clamp_mdn, label="Maximum Randomness")
                with gr.Column():
                    scale_rand = gr.Slider(0, 3, value=net.scale_sd, label="Scale of Randomness")
            with gr.Row():
                mdn_sample_button = gr.Button(value="Resample!")
            with gr.Row():
                default_im = sample_mdn(net.scale_sd, net.clamp_mdn)
                mdn_output = gr.Image(default_im)

            max_rand.change(fn=sample_mdn, inputs=[scale_rand, max_rand], outputs=[mdn_output], show_progress=False)
            scale_rand.change(fn=sample_mdn, inputs=[scale_rand, max_rand], outputs=[mdn_output], show_progress=False)
            mdn_sample_button.click(fn=sample_mdn, inputs=[scale_rand, max_rand], outputs=[mdn_output], show_progress=False)
            mdn_word.submit(fn=update_mdn_word, inputs=[mdn_word], outputs=[mdn_output], show_progress=False)

demo.launch()