generator.py

# -*- coding: utf-8 -*-
"""
 @File    : generator.py
 @Time    : 2019/12/22 下午8:22
 @Author  : yizuotian
 @Description    :  中文数据生成器
"""
import random

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from torch.utils.data.dataset import Dataset

from fontutils import FONT_CHARS_DICT


def random_color(lower_val, upper_val):
    return [random.randint(lower_val, upper_val),
            random.randint(lower_val, upper_val),
            random.randint(lower_val, upper_val)]


def put_text(image, x, y, text, font, color=None):
    """
    写中文字
    :param image:
    :param x:
    :param y:
    :param text:
    :param font:
    :param color:
    :return:
    """

    im = Image.fromarray(image)
    draw = ImageDraw.Draw(im)
    color = (255, 0, 0) if color is None else color
    draw.text((x, y), text, color, font=font)
    return np.array(im)


class Generator(Dataset):
    def __init__(self, alpha, direction='horizontal'):
        """

        :param alpha: 所有字符
        :param direction: 文字方向：horizontal|vertical
        """
        super(Generator, self).__init__()
        self.alpha = alpha
        self.direction = direction
        self.alpha_list = list(alpha)
        self.min_len = 5
        self.max_len_list = [16, 19, 24, 26]
        self.max_len = max(self.max_len_list)
        self.font_size_list = [30, 25, 20, 18]
        self.font_path_list = list(FONT_CHARS_DICT.keys())
        self.font_list = []  # 二位列表[size,font]
        for size in self.font_size_list:
            self.font_list.append([ImageFont.truetype(font_path, size=size)
                                   for font_path in self.font_path_list])
        if self.direction == 'horizontal':
            self.im_h = 32
            self.im_w = 512
        else:
            self.im_h = 512
            self.im_w = 32

    def gen_background(self):
        """
        生成背景;随机背景|纯色背景|合成背景
        :return:
        """
        a = random.random()
        pure_bg = np.ones((self.im_h, self.im_w, 3)) * np.array(random_color(0, 100))
        random_bg = np.random.rand(self.im_h, self.im_w, 3) * 100
        if a < 0.1:
            return random_bg
        elif a < 0.8:
            return pure_bg
        else:
            b = random.random()
            mix_bg = b * pure_bg + (1 - b) * random_bg
            return mix_bg

    def horizontal_draw(self, draw, text, font, color, char_w, char_h):
        """
        水平方向文字合成
        :param draw:
        :param text:
        :param font:
        :param color:
        :param char_w:
        :param char_h:
        :return:
        """
        text_w = len(text) * char_w
        h_margin = max(self.im_h - char_h, 1)
        w_margin = max(self.im_w - text_w, 1)
        x_shift = np.random.randint(0, w_margin)
        y_shift = np.random.randint(0, h_margin)
        i = 0
        while i < len(text):
            draw.text((x_shift, y_shift), text[i], color, font=font)
            i += 1
            x_shift += char_w
            y_shift = np.random.randint(0, h_margin)
            # 如果下个字符超出图像，则退出
            if x_shift + char_w > self.im_w:
                break
        return text[:i]

    def vertical_draw(self, draw, text, font, color, char_w, char_h):
        """
        锤子方向文字生成
        :param draw:
        :param text:
        :param font:
        :param color:
        :param char_w:
        :param char_h:
        :return:
        """
        text_h = len(text) * char_h
        h_margin = max(self.im_h - text_h, 1)
        w_margin = max(self.im_w - char_w, 1)
        x_shift = np.random.randint(0, w_margin)
        y_shift = np.random.randint(0, h_margin)
        i = 0
        while i < len(text):
            draw.text((x_shift, y_shift), text[i], color, font=font)
            i += 1
            x_shift = np.random.randint(0, w_margin)
            y_shift += char_h
            # 如果下个字符超出图像，则退出
            if y_shift + char_h > self.im_h:
                break
        return text[:i]

    def draw_text(self, draw, text, font, color, char_w, char_h):
        if self.direction == 'horizontal':
            return self.horizontal_draw(draw, text, font, color, char_w, char_h)
        return self.vertical_draw(draw, text, font, color, char_w, char_h)

    def gen_image(self):
        idx = np.random.randint(len(self.max_len_list))
        image = self.gen_background()
        image = image.astype(np.uint8)
        target_len = int(np.random.uniform(self.min_len, self.max_len_list[idx], size=1))

        # 随机选择size,font
        size_idx = np.random.randint(len(self.font_size_list))
        font_idx = np.random.randint(len(self.font_path_list))
        font = self.font_list[size_idx][font_idx]
        font_path = self.font_path_list[font_idx]
        # 在选中font字体的可见字符中随机选择target_len个字符
        text = np.random.choice(FONT_CHARS_DICT[font_path], target_len)
        text = ''.join(text)
        # 计算字体的w和h
        w, char_h = font.getsize(text)
        char_w = int(w / len(text))

        # 写文字，生成图像
        im = Image.fromarray(image)
        draw = ImageDraw.Draw(im)
        color = tuple(random_color(105, 255))
        text = self.draw_text(draw, text, font, color, char_w, char_h)
        target_len = len(text)  # target_len可能变小了
        # 对应的类别
        indices = np.array([self.alpha.index(c) for c in text])
        # 转为灰度图
        image = np.array(im)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # 亮度反转
        if random.random() > 0.5:
            image = 255 - image
        return image, indices, target_len

    def __getitem__(self, item):
        image, indices, target_len = self.gen_image()
        if self.direction == 'horizontal':
            image = np.transpose(image[:, :, np.newaxis], axes=(2, 1, 0))  # [H,W,C]=>[C,W,H]
        else:
            image = np.transpose(image[:, :, np.newaxis], axes=(2, 0, 1))  # [H,W,C]=>[C,H,W]
        # 标准化
        image = image.astype(np.float32) / 255.
        image -= 0.5
        image /= 0.5

        target = np.zeros(shape=(self.max_len,), dtype=np.long)
        target[:target_len] = indices
        input_len = self.im_w // 4 - 3
        return image, target, input_len, target_len

    def __len__(self):
        return len(self.alpha) * 100


def test_image_gen(direction='vertical'):
    from config import cfg
    gen = Generator(cfg.word.get_all_words()[:10], direction=direction)
    for i in range(10):
        im, indices, target_len = gen.gen_image()
        # cv2.imwrite('output/{}-{:03d}.jpg'.format(direction, i + 1), im)
        print(''.join([gen.alpha[j] for j in indices]))


def test_gen():
    from data.words import Word
    gen = Generator(Word().get_all_words())
    for x in gen:
        print(x[1])


def test_font_size():
    font = ImageFont.truetype('fonts/simsun.ttc')
    print(font.size)
    font.size = 20
    print(font.size)


if __name__ == '__main__':
    test_image_gen('horizontal')
    # test_image_gen('vertical')
    # test_gen()
    # test_font_size()