switched to SVG rendering

app.py

@@ -76,8 +76,8 @@ def update_writer_slider(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")
+    svg = convenience.draw_words_svg(writer_words, all_word_writer_Ws, all_word_writer_Cs, weights, net)
+    return svg
 
 
 def update_chosen_writers(writer1, writer2):
@@ -109,9 +109,9 @@ def update_char_slider(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)
+    svg = convenience.commands_to_svg(all_commands, 750, 160, 375)
 
-    return im.convert("RGB")
+    return svg
 
 
 def update_blend_chars(c1, c2):
@@ -145,8 +145,8 @@ def update_mdn_word(target_word):
 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")
+    svg = convenience.draw_words_svg(mdn_words, all_word_mdn_Ws, all_word_mdn_Cs, [1], net)
+    return svg
 
 
 update_writer_word(" ".join(writer_words))
@@ -173,7 +173,7 @@ with gr.Blocks() as demo:
                 writer_submit = gr.Button("Submit")
             with gr.Row():
                 writer_default_image = update_writer_slider(writer_weight)
-                writer_output = gr.Image(writer_default_image)
+                writer_output = gr.HTML(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)
@@ -192,7 +192,7 @@ with gr.Blocks() as demo:
                 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_output = gr.HTML(char_default_image)
 
             char_slider.change(fn=update_char_slider, inputs=[char_slider], outputs=[char_output], show_progress=False)
 
@@ -210,7 +210,7 @@ with gr.Blocks() as demo:
                 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)
+                mdn_output = gr.HTML(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)

convenience.py

@@ -14,10 +14,12 @@ from config.GlobalVariables import *
 from tensorboardX import SummaryWriter
 from SynthesisNetwork import SynthesisNetwork
 from DataLoader import DataLoader
+import svgwrite
 # import ffmpeg # for problems with ffmpeg uninstall ffmpeg and then install ffmpeg-python
 
 L = 256
 
+
 def get_mean_global_W(net, loaded_data, device):
     """gets the mean global style vector for a given writer"""
     [_, _, _, _, _, _, all_word_level_stroke_in, all_word_level_stroke_out, all_word_level_stroke_length, all_word_level_term, all_word_level_char, all_word_level_char_length, all_segment_level_stroke_in, all_segment_level_stroke_out,
@@ -231,14 +233,14 @@ def get_character_blend_W_c(character_weights, all_Ws, all_Cs):
     W_vector = all_Ws[0, 0, :].unsqueeze(-1)
 
     weights_tensor = torch.tensor(character_weights).repeat_interleave(L * L).reshape(1, M, L, L)  # repeat accross remaining dimensions
-    char_matrix = (weights_tensor * all_Cs).sum(axis=1).squeeze() # take weighted sum accross characters axis
+    char_matrix = (weights_tensor * all_Cs).sum(axis=1).squeeze()  # take weighted sum accross characters axis
 
     W_c = char_matrix @ W_vector
 
     return W_c.reshape(1, 1, L)
 
 
-def get_commands(net, target_word, all_W_c): # seems like target_word is only used for length
+def get_commands(net, target_word, all_W_c):  # seems like target_word is only used for length
     """converts character-dependent style-dependent DSDs to a list of commands for drawing"""
     all_commands = []
     current_id = 0
@@ -285,6 +287,7 @@ def get_commands(net, target_word, all_W_c): # seems like target_word is only us
 
     return commands
 
+
 def mdn_video(target_word, num_samples, scale_sd, clamp_mdn, net, all_loaded_data, device):
     '''
     Method creating gif of mdn samples
@@ -306,7 +309,7 @@ def mdn_video(target_word, num_samples, scale_sd, clamp_mdn, net, all_loaded_dat
             writer_Ws, writer_Cs = get_DSD(net, word, [mean_global_W], [all_loaded_data[0]], device)
             word_Ws.append(writer_Ws)
             word_Cs.append(writer_Cs)
-            
+
         im = draw_words(words, word_Ws, word_Cs, [1], net)
         im.convert("RGB").save(f'results/{us_target_word}_mdn_samples/sample_{i}.png')
     # Convert fromes to video using ffmpeg
@@ -314,6 +317,7 @@ def mdn_video(target_word, num_samples, scale_sd, clamp_mdn, net, all_loaded_dat
     videos = photos.output(f'results/{us_target_word}_video.mov', vcodec="libx264", pix_fmt="yuv420p")
     videos.run(overwrite_output=True)
 
+
 def sample_blended_writers(writer_weights, target_sentence, net, all_loaded_data, device="cpu"):
     """Generates an image of handwritten text based on target_sentence"""
     words = target_sentence.split(' ')
@@ -329,7 +333,7 @@ def sample_blended_writers(writer_weights, target_sentence, net, all_loaded_data
         writer_Ws, writer_Cs = get_DSD(net, word, writer_mean_Ws, all_loaded_data, device)
         word_Ws.append(writer_Ws)
         word_Cs.append(writer_Cs)
-        
+
     return draw_words(words, word_Ws, word_Cs, writer_weights, net)
 
 
@@ -356,10 +360,10 @@ def sample_character_grid(letters, grid_size, net, all_loaded_data, device="cpu"
             wx = i / (grid_size - 1)
             wy = j / (grid_size - 1)
 
-            character_weights = [(1 - wx) * (1 - wy), # top left is 1 at (0, 0)
-                                 wx       * (1 - wy), # top right is 1  at (1, 0)
+            character_weights = [(1 - wx) * (1 - wy),  # top left is 1 at (0, 0)
+                                 wx * (1 - wy),  # top right is 1  at (1, 0)
                                  (1 - wx) * wy,       # bottom left is 1 at (0, 1)
-                                 wx       * wy]       # bottom right is 1 at (1, 1)
+                                 wx * wy]       # bottom right is 1 at (1, 1)
             all_W_c = get_character_blend_W_c(character_weights, all_Ws, all_Cs)
             all_commands = get_commands(net, letters[0], all_W_c)
 
@@ -377,6 +381,7 @@ def sample_character_grid(letters, grid_size, net, all_loaded_data, device="cpu"
 
     return im
 
+
 def writer_interpolation_video(target_sentence, transition_time, net, all_loaded_data, device="cpu"):
     """
     Generates a video of interpolating between each provided writer
@@ -416,6 +421,7 @@ def writer_interpolation_video(target_sentence, transition_time, net, all_loaded
     videos = photos.output(f"results/{target_sentence}_blend_video.mov", vcodec="libx264", pix_fmt="yuv420p")
     videos.run(overwrite_output=True)
 
+
 def mdn_single_sample(target_word, scale_sd, clamp_mdn, net, all_loaded_data, device):
     '''
     Method creating gif of mdn samples
@@ -462,7 +468,7 @@ def sample_blended_chars(character_weights, letters, net, all_loaded_data, devic
 def char_interpolation_video(letters, transition_time, net, all_loaded_data, device="cpu"):
     """Generates an image of handwritten text based on target_sentence"""
 
-    os.makedirs(f"./results/{''.join(letters)}_frames", exist_ok=True) # make a folder for the frames
+    os.makedirs(f"./results/{''.join(letters)}_frames", exist_ok=True)  # make a folder for the frames
 
     M = len(letters)
     mean_global_W = get_mean_global_W(net, all_loaded_data[0], device)
@@ -507,6 +513,25 @@ def draw_words(words, word_Ws, word_Cs, writer_weights, net):
 
     return im
 
+
+def draw_words_svg(words, word_Ws, word_Cs, writer_weights, net):
+    dwg = svgwrite.Drawing("output.svg", size=(750, 160), style="background-color: black;")
+    width = 50
+    for word, all_writer_Ws, all_writer_Cs in zip(words, word_Ws, word_Cs):
+        all_W_c = get_writer_blend_W_c(writer_weights, all_writer_Ws, all_writer_Cs)
+        all_commands = get_commands(net, word, all_W_c)
+
+        for [x, y, t] in all_commands:
+            if t == 0:
+                path.push("L", x + width, y)
+            else:
+                path = svgwrite.path.Path(stroke="white", stroke_width="1")
+                dwg.add(path)
+            path.push("M", x + width, y)
+        width += np.max(all_commands[:, 0]) + 25
+    return dwg.tostring()
+
+
 def commands_to_image(commands, imW, imH, xoff, yoff):
     im = Image.fromarray(np.zeros([imW, imH]))
     dr = ImageDraw.Draw(im)
@@ -519,4 +544,15 @@ def commands_to_image(commands, imW, imH, xoff, yoff):
                 y - yoff), 255, 1)
         px, py = x, y
     return im
-    
+
+
+def commands_to_svg(commands, imW, imH, xoff):
+    dwg = svgwrite.Drawing("output.svg", size=(imW, imH), style="background-color:black")
+    for [x, y, t] in commands:
+        if t == 0:
+            path.push("L", x + xoff, y)
+        else:
+            path = svgwrite.path.Path(stroke="white", stroke_width="1")
+            dwg.add(path)
+        path.push("M", x + xoff, y)
+    return dwg.tostring()

interpolation.py

@@ -20,11 +20,10 @@ def main(params):
     net = SynthesisNetwork(weight_dim=256, num_layers=3).to(device)
 
     if not torch.cuda.is_available():
-        try: # retrained model also contains loss in dict 
+        try:  # retrained model also contains loss in dict
             net.load_state_dict(torch.load('./model/250000.pt', map_location=torch.device('cpu'))["model_state_dict"])
         except:
             net.load_state_dict(torch.load('./model/250000.pt', map_location=torch.device('cpu')))
-        
 
     dl = DataLoader(num_writer=1, num_samples=10, divider=5.0, datadir='./data/writers')
 
@@ -34,7 +33,6 @@ def main(params):
         loaded_data = dl.next_batch(TYPE='TRAIN', uid=writer_id, tids=list(range(params.num_samples)))
         all_loaded_data.append(loaded_data)
 
-    
     if params.output == "image":
 
         if params.interpolate == "writer":
@@ -78,6 +76,7 @@ def main(params):
     else:
         raise ValueError("Invalid output")
 
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Arguments for generating samples with the handwriting synthesis model.')
 
@@ -89,25 +88,25 @@ if __name__ == '__main__':
     parser.add_argument('--interpolate', type=str, default="randomness", choices=["writer", "character", "randomness"])
 
     # PARAMS FOR BOTH WRITER AND CHARACTER INTERPOLATION:
-        # IF IMAGE - weights to use for a single sample of interpolation
+    # IF IMAGE - weights to use for a single sample of interpolation
     parser.add_argument('--blend_weights', type=float, nargs="+", default=[0.5, 0.5])
-        # IF VIDEO - the number of frames for each character/writer
+    # IF VIDEO - the number of frames for each character/writer
     parser.add_argument('--frames_per_step', type=int, default=10)
 
     # PARAMS IF WRITER INTERPOLATION:
     parser.add_argument('--target_word', type=str, default="hello world")
     parser.add_argument('--writer_ids', type=int, nargs="+", default=[80, 120])
-    
+
     # PARAMS IF CHARACTER INTERPOLATION:
-        # IF VIDEO OR BLEND
-    parser.add_argument('--blend_chars', type=str, nargs="+", default = ["a", "b", "c", "d", "e"])
-        # IF GRID
-    parser.add_argument('--grid_chars', type=str, nargs="+", default= ["y", "s", "u", "n"])
+    # IF VIDEO OR BLEND
+    parser.add_argument('--blend_chars', type=str, nargs="+", default=["a", "b", "c", "d", "e"])
+    # IF GRID
+    parser.add_argument('--grid_chars', type=str, nargs="+", default=["y", "s", "u", "n"])
     parser.add_argument('--grid_size', type=int, default=10)
 
     # PARAMS IF RANDOMNESS ITERPOLATION (--output will be ignored):
-    parser.add_argument('--max_randomness', type=float, default=1) 
-    parser.add_argument('--scale_randomness', type=float, default=0.5) 
+    parser.add_argument('--max_randomness', type=float, default=1)
+    parser.add_argument('--scale_randomness', type=float, default=0.5)
     parser.add_argument('--num_random_samples', type=int, default=10)
 
     main(parser.parse_args())

requirements.txt

@@ -8,4 +8,5 @@ torch==1.11.0
 typing_extensions==4.1.1
 ffmpeg-python
 gradio
+svgwrite