Add label to generator output

Author: Belval

README.md

@@ -62,7 +62,7 @@ generator = GeneratorFromStrings(
     random_blur=True
 )
 
-for img in generator:
+for img, lbl in generator:
     # Do something with the pillow images here.
 ```
 

tests.py

@@ -46,31 +46,31 @@ def test_generator_from_dict(self):
         generator = GeneratorFromDict()
         i = 0
         while i < 100:
-            img = next(generator)
+            img, lbl = next(generator)
             self.assertTrue(img.size[1] == 32, "Shape is not right")
             i += 1
 
     def test_generator_from_random(self):
         generator = GeneratorFromRandom()
         i = 0
         while i < 100:
-            img = next(generator)
+            img, lbl = next(generator)
             self.assertTrue(img.size[1] == 32, "Shape is not right")
             i += 1
 
     def test_generator_from_strings(self):
         generator = GeneratorFromStrings(["TEST TEST TEST"])
         i = 0
         while i < 100:
-            img = next(generator)
+            img, lbl = next(generator)
             self.assertTrue(img.size[1] == 32, "Shape is not right")
             i += 1
 
     def test_generator_from_wikipedia(self):
         generator = GeneratorFromWikipedia()
         i = 0
         while i < 100:
-            img = next(generator)
+            img, lbl = next(generator)
             self.assertTrue(img.size[1] == 32, "Shape is not right")
             i += 1
 

trdg/generators/from_strings.py

@@ -85,4 +85,4 @@ def next(self):
             self.space_width,
             self.margins,
             self.fit,
-        )
+        ), self.strings[(self.generated_count - 1) % len(self.strings)]