Update app.py

app.py

@@ -1,3 +1,6 @@
+import os
+os.system("pip install git+https://github.com/ai-forever/ScrabbleGAN")
+
 import numpy as np
 import cv2
 
@@ -27,10 +30,66 @@ def get_text_from_image(img):
     return bg
 
 
+def split_text_to_rows(text, n=4):
+    # https://stackoverflow.com/a/6187258
+    l = text.split()
+    return [' '.join(l[x:x+n]) for x in range(0, len(l), n)]
+
+
+def remove_right_padding(img, len_text, char_w=32):
+    # char_w for a standard ScrabbleGAN char width
+    return img[:, :len_text*char_w]
+
+
+def split_list2batches(lst, batch_size):
+    """Split list of images to list of bacthes."""
+    return [lst[i:i+batch_size] for i in range(0, len(lst), batch_size)]
+
+
+def get_canvas_size(images, row_width, left_pad):
+    canvas_width = 0
+    canvas_height = 0
+    for image in images:
+        h, w = image.shape[:2]
+        canvas_height += h*row_width
+        if w > canvas_width:
+            canvas_width = w
+    canvas_width += left_pad
+    # expand canvas to the height of the last image
+    # (to correct the effect of rows shrinking)
+    h = images[-1].shape[0]
+    canvas_height += h - h*row_width
+    return int(canvas_height), canvas_width
+
+
 def predict(text):
-    imgs, texts = GENERATOR.generate(word_list=[text])
-    image_on_white = get_text_from_image(imgs[0])
-    return image_on_white
+    if text.find(NEW_LINE_SYMB) == -1:
+        texts = split_text_to_rows(text)
+    else:
+        texts = [row.strip() for row in text.split(NEW_LINE_SYMB)]
+
+    texts_batches = split_list2batches(texts, BATCH_SIZE)
+
+    images_on_white = []
+    for texts_batch in texts_batches:
+        imgs, texts_on_image = GENERATOR.generate(word_list=texts_batch)
+        for img, text_on_image in zip(imgs, texts_on_image):
+            cropped_image = remove_right_padding(
+                img, len(text_on_image))
+            images_on_white.append(
+                get_text_from_image(cropped_image))
+
+    canvas_height, canvas_width = get_canvas_size(
+        images_on_white, ROW_WIDTH, LEFT_PAD)
+    canvas = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8)
+    canvas.fill(255)
+    
+    start_draw = 0
+    for image_on_white in images_on_white:
+        h, w = image_on_white.shape[:2]
+        canvas[start_draw:start_draw+h, LEFT_PAD:LEFT_PAD+w] = image_on_white
+        start_draw += int(h * ROW_WIDTH)
+    return canvas
 
 
 CHAR_MAP_PATH, WEIGHTS_PATH = download_weights("sberbank-ai/scrabblegan-peter")
@@ -41,9 +100,14 @@ GENERATOR = ImgGenerator(
     char_map_path=CHAR_MAP_PATH
 )
 
+BATCH_SIZE = 3
+ROW_WIDTH = 0.7
+LEFT_PAD = 10
+NEW_LINE_SYMB = '{n}'
+
 gr.Interface(
     predict,
-    inputs=gr.Textbox(label="Type your text (RU) to generate it on an image"),
+    inputs=gr.Textbox(label=f"Type your text (RU) to generate it on an image. The text will be automatically splitted on lines, or you can use a new line symbol {NEW_LINE_SYMB}"),
     outputs=gr.Image(label="Generated image"),
     title="Peter the Great handwritten image generation",
 ).launch()