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