import cv2
import random
import numpy as np
import glob
import os
from multiprocessing import Queue, Process, Event, Pool, freeze_support
import imutils
import time
import sys
import math
from imgaug import augmenters as iaa
from PIL import Image, ImageDraw, ImageFont, ImageFilter
keys = " 0123456789"
nclass = len(keys)
imgW = 1000
imgH = 48
from concurrent.futures import ThreadPoolExecutor
class Data_Generator:
def __init__(self, batch_size, num_workers=1, train = True):
self.batch_size = batch_size
self.train = train
self.queue = Queue(100)
self.bg_root = './bg'
self.bg_list = []
for path in os.listdir(self.bg_root):
self.bg_list.append(np.clip(cv2.imread(os.path.join(self.bg_root, path)) , 0, 255))
self.font_root = './font'
self.font_list = [i for i in os.listdir(self.font_root) if i.endswith('ttf') or i.endswith('TTF')]
self.theard_pool = ThreadPoolExecutor(max_workers=num_workers)
for i in range(num_workers):
self.theard_pool.submit(self.proc_worker)
def gen(self, font):
if random.randint(0,100) >= 20:
label = [random.choice("0123456789") for i in range(random.randint(4, 18))]
img = self.StereoscopicText(label, font)
elif random.randint(0,100) >= 50:
label = [random.choice("0123456789") for i in range(random.randint(4, 18))]
img = self.PrintText(label, font)
label = "".join(label)
seq = iaa.SomeOf(3, [
iaa.GaussianBlur((0, 1.0)),
iaa.AdditiveGaussianNoise(),
iaa.JpegCompression(),
iaa.ChangeColorTemperature(),
iaa.AddToBrightness(),
iaa.Invert(),
iaa.MultiplyBrightness(),
iaa.MotionBlur(),
], True)
img = seq.augment_images(np.expand_dims(img, 0).astype('uint8'))[0]
return img, label
def StereoscopicText(self, label, font):
for i in range(random.randint(1,3)):
label.insert(random.randint(0, len(label)), ' ')
label = ''.join(label)
size = font.getsize(''.join(label))
pad_x = random.randint(0,25)
pad_y = random.randint(0,25)
img = Image.fromarray(np.zeros((size[1] + pad_y + 20, size[0] + pad_x + 20)).astype('uint8') + 255)
draw = ImageDraw.Draw(img)
draw.text((random.randint(10, 10 + pad_x), random.randint(10, 10 + pad_y)), ''.join(label), 0, font=font)
img = img.filter(ImageFilter.EMBOSS)
img = np.array(img).astype('float32')
img = img[10:-10,10:-10]
bg = random.choice(self.bg_list)
if bg.shape[0] < img.shape[0] or bg.shape[1] < img.shape[1]:
bg = cv2.resize(bg, (max(bg.shape[1], img.shape[1]),max(bg.shape[0], img.shape[0])))
x = random.randint(0, bg.shape[1] - img.shape[1])
y = random.randint(0, bg.shape[0] - img.shape[0])
bg = bg[ y:img.shape[0] +y, x:x+img.shape[1] ,: ]
img = img * (random.uniform(1,1.5) / 255 )
img = np.expand_dims(img, -1)
res = img * bg
# res = np.clip(res, 0, 255)
return res
def PrintText(self, label, font):
for i in range(random.randint(1, 3)):
label.insert(random.randint(0, len(label)), ' ')
label = ''.join(label)
size = font.getsize(''.join(label))
pad_x = random.randint(0, 25)
pad_y = random.randint(0, 25)
img = Image.fromarray(np.zeros((size[1] + pad_y + 20, size[0] + pad_x + 20)).astype('uint8') + 255)
draw = ImageDraw.Draw(img)
draw.text((random.randint(10, 10 + pad_x), random.randint(10, 10 + pad_y)), ''.join(label), 0, font=font)
img = np.array(img).astype('float32')
img = img[10:-10, 10:-10]
bg = random.choice(self.bg_list)
if bg.shape[0] < img.shape[0] or bg.shape[1] < img.shape[1]:
bg = cv2.resize(bg, (max(bg.shape[1], img.shape[1]), max(bg.shape[0], img.shape[0])))
x = random.randint(0, bg.shape[1] - img.shape[1])
y = random.randint(0, bg.shape[0] - img.shape[0])
bg = bg[y:img.shape[0] + y, x:x + img.shape[1], :]
img = img / 255
img = np.expand_dims(img, -1)
res = img * bg
return res
def next_batch(self):
imgs = []
labels = []
font = ImageFont.truetype(os.path.join(self.font_root, random.choice(self.font_list)), random.randint(32, 64))
for i in range(self.batch_size):
img, label = self.gen(font)
imgs.append(img)
labels.append(label)
return imgs, labels
def proc_worker(self):
while True:
images_, labels_ = self.next_batch()
self.queue.put([images_, labels_])
def get_batch(self):
return self.queue.get()
if __name__=='__main__':
test = Data_Generator(64, 4)
t = time.time()
for i in range(10):
test.get_batch()
t = time.time() - t
print(64 * 10/ t)
