OCR训练数据生成方法

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

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

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

安装 imagemagick 工具

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

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

巴特西

OCR训练数据生成方法

安装 imagemagick 工具

生成图像格式

具体实现

最新文章

热门文章