OCR训练数据生成方法

有的时候我们训练网络的时候，数据集在收集的过程中由于种种原因导致图像收集的不完整，比如某些种类很少，或者没有，这个时候我们就可以考虑自己生成数据集。

这个和data augmentation还不太一样，data augmentation是在数据集上做一些变化，得到更多的数据，数据生成是直接生成一些数据集。

在做OCR识别的时候，因为有些类型的数据量很少，所以就考虑了一下数据生成的方法。

安装 imagemagick 工具

imagemagick 工具，可以用来生成不同字体的图像。

imagemagick安装

生成图像格式

为了让生成的图像更加真实，我做了一下操作

1. 修改字体
2. 添加背景
3. 旋转角度
4. 增加噪点

具体实现

#coding=utf8
import os
import cv2
import pickle
import random
from glob import glob
import numpy as np

# 为图像添加背景
def add_background(img):
	# 从自然场景中选择一些图像，因为我是做OCR，所以我从ICDAR 2015的数据集选择了一些图像
    bg_imgs_path = glob('background/*.jpg')
    bg_img = cv2.imread(bg_imgs_path[random.randint(0, len(bg_imgs_path)-1)])
    bg_img = cv2.cvtColor(bg_img, cv2.COLOR_RGB2GRAY)
    # 这个大概是一个字的像素大小，根据具体场景而定
    start_y = random.randint(0, bg_img.shape[0]-9)
    start_x = random.randint(0, bg_img.shape[1]-9)
    bg_img = cv2.resize(bg_img[start_y:start_y+8, start_x:start_x+8], img.shape)
    beta = 0.5
    # 像素叠加
    return img * (1 - beta) + bg_img * beta

# 随机噪点
def salt_and_pepper(img,noise_num):
    noise_img = img
    for i in range(noise_num):
        randX = random.randint(0, img.shape[0]-1)
        randY = random.randint(0, img.shape[1]-1)
        if random.randint(0, 1) == 0:
            noise_img[randX,randY]=0
        else:
            noise_img[randX,randY]=255
    return noise_img

# 增加噪声，高斯平滑
def add_noise(img):
	# 这里没有增加噪点
    img = salt_and_pepper(img, random.randint(0, 0))
    filter_size = random.randint(0, 3) * 2 + 1
    # 高斯平滑
    img = cv2.GaussianBlur(img, (filter_size, filter_size), sigmaX=0)
    return img

# 旋转图像
def rotate_image(mat, angle):
    # angle in degrees
    # opencv python和c++的接口，这里顺序是相反的
    height, width = mat.shape[:2]
    image_center = (width/2, height/2)
    rotation_mat = cv2.getRotationMatrix2D(image_center, angle, 1.)
    abs_cos = abs(rotation_mat[0,0])
    abs_sin = abs(rotation_mat[0,1])
    bound_w = int(height * abs_sin + width * abs_cos)
    bound_h = int(height * abs_cos + width * abs_sin)
    rotation_mat[0, 2] += bound_w/2 - image_center[0]
    rotation_mat[1, 2] += bound_h/2 - image_center[1]
    rotated_mat = cv2.warpAffine(mat, rotation_mat, (bound_w, bound_h))
    return rotated_mat

# 设置需要生成的数据集列表，格式： index num
# index 表示在种类标签中的索引，也就是第几类，num表示此类图像已有的个数
label_path = 'sup_remain_data.txt'
# 输出目录
local_dir = 'supplement2/'

# 输出生成数据集的标签值
gen_suppl_list = 'supplement2/gen_supplement2.txt'

# 每一类图像都生成总共110张，如果已经有一部分了，只需要生成剩下部分的
pic_num = 110 # 100:10
train_num = 100
test_num = 10

# 设置种类和要生成的图像个数
label_genNum = {}
with open(label_path) as fin:
    lines = fin.readlines()
    for line in lines:
        label, num = line.strip().split(' ')
        label_genNum[label] = pic_num - int(num)

# 读入所有标签种类
std_label = {}
with open('data.txt') as fin:
    lines = fin.readlines()
    for index, line in enumerate(lines):
        std_label[str(index)] = line.strip()

#grav=['NorthWest','North','NorthEast','West','Center','East','SouthWest','South','SouthEast']
bgcolor=['black','white', 'gray', 'gray', 'black', 'white']
fillcolor=['white','black','black', 'white', 'gray', 'gray']
# 这个目录下放当前文字常用字体格式文件
ttfs=glob('ttfs/*')

count = 0
with open(gen_suppl_list, 'w') as fout:
    for label, gen_num in label_genNum.items():
        count += 1
        print count
        for i in range(gen_num):
            img_name = label + '_' + str(i) + '.jpg'
            print str(label) + ' ' + std_label[label] + ' ' + label + '_' + str(i)
            fout.write(str(label) + ' ' + std_label[label] + ' ' + label + '_' + str(i) +'\n')
            if os.path.exists(local_dir+img_name):
                continue
            #index_g = random.randint(0, len(grav)-1)
            index_b = random.randint(0, len(bgcolor)-1)
            index_t = random.randint(0, len(ttfs)-1)
            #font_size = random.randint(20, 40)
            # 设置字体、填充方式、背景颜色、像素点大小、图像输出大小
            system_cmd = 'convert -font {0} -fill {1} -background {2} -pointsize 30 -colorspace Gray -gravity "Center" ' \
            '-size 40x40 label:"{3}" "{4}"'.format(ttfs[index_t], fillcolor[index_b], bgcolor[index_b], std_label[str(label)], local_dir+img_name)
            os.system(system_cmd)
            img = cv2.imread(local_dir+img_name, 0)
            # 几种变化效果叠加，这样生成的图像就是具有多种效果的，不需要单独每种效果生成过多的图像
            img = add_background(img)
            img = add_noise(img)
            img = rotate_image(img, random.randint(-5, 5))
            #cv2.imshow('gen image', img)
            #cv2.waitKey(0)
            cv2.imwrite(local_dir+img_name, img)

print count
print 'language supplement-set 生成完毕...'