Add text_color support for handwritten

Author: Belval

TextRecognitionDataGenerator/computer_text_generator.py

@@ -30,9 +30,9 @@ def __generate_horizontal_text(cls, text, font, text_color, font_size, space_wid
         c1, c2 = colors[0], colors[-1]
 
         fill = (
-            random.randint(c1[0], c2[0]),
-            random.randint(c1[1], c2[1]),
-            random.randint(c1[2], c2[2])
+            random.randint(min(c1[0], c2[0]), max(c1[0], c2[0])),
+            random.randint(min(c1[1], c2[1]), max(c1[1], c2[1])),
+            random.randint(min(c1[2], c2[2]), max(c1[2], c2[2]))
         )
 
         for i, w in enumerate(words):

TextRecognitionDataGenerator/data_generator.py

@@ -30,7 +30,7 @@ def generate(cls, index, text, font, out_dir, size, extension, skewing_angle, ra
         if is_handwritten:
             if orientation == 1:
                 raise ValueError("Vertical handwritten text is unavailable")
-            image = HandwrittenTextGenerator.generate(text)
+            image = HandwrittenTextGenerator.generate(text, text_color)
         else:
             image = ComputerTextGenerator.generate(text, font, text_color, size, orientation, space_width)
 

TextRecognitionDataGenerator/handwritten_text_generator.py

@@ -1,12 +1,13 @@
 import os
 import pickle
 import numpy as np
+import random
 import tensorflow as tf
 import matplotlib.pyplot as plt
 import matplotlib.cm as cm
 import matplotlib.mlab as mlab
 import seaborn
-from PIL import Image
+from PIL import Image, ImageColor
 from collections import namedtuple
 
 class HandwrittenTextGenerator(object):
@@ -115,7 +116,7 @@ def __join_images(cls, images):
         return compound_image
 
     @classmethod
-    def generate(cls, text):
+    def generate(cls, text, text_color):
         with open(os.path.join('handwritten_model', 'translation.pkl'), 'rb') as file:
             translation = pickle.load(file)
 
@@ -127,6 +128,15 @@ def generate(cls, text):
             saver = tf.train.import_meta_graph('handwritten_model/model-29.meta')
             saver.restore(sess, 'handwritten_model/model-29')
             images = []
+            colors = [ImageColor.getrgb(c) for c in text_color.split(',')]
+            c1, c2 = colors[0], colors[-1]
+
+            color = '#{:02x}{:02x}{:02x}'.format(
+                random.randint(min(c1[0], c2[0]), max(c1[0], c2[0])),
+                random.randint(min(c1[1], c2[1]), max(c1[1], c2[1])),
+                random.randint(min(c1[2], c2[2]), max(c1[2], c2[2]))
+            )
+
             for word in text.split(' '):
                 _, window_data, kappa_data, stroke_data, coords = cls.__sample_text(sess, word, translation)
 
@@ -140,7 +150,7 @@ def generate(cls, text):
                 ax.axis('off')
 
                 for stroke in cls.__split_strokes(cls.__cumsum(np.array(coords))):
-                    plt.plot(stroke[:, 0], -stroke[:, 1], color='#080808')
+                    plt.plot(stroke[:, 0], -stroke[:, 1], color=color)
 
                 fig.patch.set_alpha(0)
                 fig.patch.set_facecolor('none')