OCR训练数据生成方法

安装 imagemagick 工具

生成图像格式

具体实现

#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 生成完毕...'