放假了, 终于可以继续可以静下心写一写 OCR 方面的东西上次谈到文字的切割, 今天打算总结一下我们怎么得到用于训练的文字数据集如果是想训练一个手写体识别的模型, 用一些前人收集好的手写文字集就好了, 比如中科院的这些数据集但是如果我们只是想要训练一个专门用于识别印刷汉字的模型, 那么我们就需要各种印刷字体的训练集, 那怎么获取呢? 借助强大的图像库, 自己生成就行了!
先捋一捋思路, 生成文字集需要什么步骤:
确定你要生成多少字体, 生成一个记录着汉字与 label 的对应表
确定和收集需要用到的字体文件
生成字体图像, 存储在规定的目录下
适当的数据增强
第三步的生成字体图像最为重要, 如果仅仅是生成很正规的文字, 那么用这个正规文字集去训练模型, 第一图像数目有点少, 第二模型泛化能力比较差, 所以我们需要对字体图像做大量的图像处理工作, 以增大我们的印刷体文字数据集
我总结了一下, 我们可以做的一些图像增强工作有这些:
文字扭曲
背景噪声(椒盐)
文字位置(设置文字的中心点)
笔画粘连(膨胀来模拟)
笔画断裂(腐蚀来模拟)
文字倾斜(文字旋转)
多种字体
做完以上增强后, 我们得到的数据集已经非常庞大了
现在开始一步一步生成我们的 3755 个汉字的印刷体文字数据集
一生成汉字与 label 的对应表
这里的汉字 label 映射表的生成我使用了 pickel 模块, 借助它生成一个 id: 汉字的映射文件存储下来
这里举个小例子说明怎么生成这个汉字: id 映射表
首先在一个 txt 文件里写入你想要的汉字, 如果对汉字对应的 ID 没有要求的话, 我们不妨使用该汉字的排位作为其 ID, 比如一二三四五中, 五的 ID 就是 00005 如此类推, 把汉字读入内存, 建立一个字典, 把这个关系记录下来, 再使用 pickle.dump 存入文件保存
二收集字体文件
字体文件上网收集就好了, 但是值得注意的是, 不是每一种字体都支持汉字, 所以我们需要筛选出真正适合汉字生成的字体文件才可以我一共使用了十三种汉字字体作为我们接下来汉字数据集用到的字体, 具体如下图:
当然, 如果需要进一步扩大数据集来增强训练得到的模型的泛化能力, 可以花更多的时间去收集各类汉字字体, 那么模型在面对各种字体时也能从容应对, 给出准确的预测
三文字图像生成
首先是定义好输入参数, 其中包括输出目录字体目录测试集大小图像尺寸图像旋转幅度等等
- def args_parse():
- #解析输入参数
- parser = argparse.ArgumentParser(
- description=description, formatter_class=RawTextHelpFormatter)
- parser.add_argument('--out_dir', dest='out_dir',
- default=None, required=True,
- help='write a caffe dir')
- parser.add_argument('--font_dir', dest='font_dir',
- default=None, required=True,
- help='font dir to to produce images')
- parser.add_argument('--test_ratio', dest='test_ratio',
- default=0.2, required=False,
- help='test dataset size')
- parser.add_argument('--width', dest='width',
- default=None, required=True,
- help='width')
- parser.add_argument('--height', dest='height',
- default=None, required=True,
- help='height')
- parser.add_argument('--no_crop', dest='no_crop',
- default=True, required=False,
- help='', action='store_true')
- parser.add_argument('--margin', dest='margin',
- default=0, required=False,
- help='', )
- parser.add_argument('--rotate', dest='rotate',
- default=0, required=False,
- help='max rotate degree 0-45')
- parser.add_argument('--rotate_step', dest='rotate_step',
- default=0, required=False,
- help='rotate step for the rotate angle')
- parser.add_argument('--need_aug', dest='need_aug',
- default=False, required=False,
- help='need data augmentation', action='store_true')
- args = vars(parser.parse_args())
- return args
接下来需要将我们第一步得到的对应表读入内存, 因为这个表示 ID 到汉字的映射, 我们在做一下转换, 改成汉字到 ID 的映射, 用于后面的字体生成
- # 将汉字的 label 读入, 得到 (ID: 汉字) 的映射表 label_dict
- label_dict = get_label_dict()
- char_list=[] # 汉字列表
- value_list=[] # label 列表
- for (value,chars) in label_dict.items():
- print (value,chars)
- char_list.append(chars)
- value_list.append(value)
- # 合并成新的映射关系表:(汉字: ID)
- lang_chars = dict(zip(char_list,value_list))
- font_check = FontCheck(lang_chars)
我们对旋转的角度存储到列表中, 旋转角度的范围是[-rotate,rotate].
- if rotate < 0:
- roate = - rotate
- if rotate > 0 and rotate <= 45:
- all_rotate_angles = []
- for i in range(0, rotate+1, rotate_step):
- all_rotate_angles.append(i)
- for i in range(-rotate, 0, rotate_step):
- all_rotate_angles.append(i)
- #print(all_rotate_angles)
现在说一下字体图像是怎么生成的, 首先我们使用的工具是 PILPIL 里面有很好用的汉字生成函数, 我们用这个函数再结合我们提供的字体文件, 就可以生成我们想要的数字化的汉字了我们先设定好我们生成的字体颜色为黑底白色, 字体尺寸由输入参数来动态设定
- # 生成字体图像
- class Font2Image(object):
- def __init__(self,
- width, height,
- need_crop, margin):
- self.width = width
- self.height = height
- self.need_crop = need_crop
- self.margin = margin
- def do(self, font_path, char, rotate=0):
- find_image_bbox = FindImageBBox()
- # 黑色背景
- img = Image.new("RGB", (self.width, self.height), "black")
- draw = ImageDraw.Draw(img)
- font = ImageFont.truetype(font_path, int(self.width * 0.7),)
- # 白色字体
- draw.text((0, 0), char, (255, 255, 255),
- font=font)
- if rotate != 0:
- img = img.rotate(rotate)
- data = list(img.getdata())
- sum_val = 0
- for i_data in data:
- sum_val += sum(i_data)
- if sum_val > 2:
- np_img = np.asarray(data, dtype='uint8')
- np_img = np_img[:, 0]
- np_img = np_img.reshape((self.height, self.width))
- cropped_box = find_image_bbox.do(np_img)
- left, upper, right, lower = cropped_box
- np_img = np_img[upper: lower + 1, left: right + 1]
- if not self.need_crop:
- preprocess_resize_keep_ratio_fill_bg = \
- PreprocessResizeKeepRatioFillBG(self.width, self.height,
- fill_bg=False,
- margin=self.margin)
- np_img = preprocess_resize_keep_ratio_fill_bg.do(
- np_img)
- # cv2.imwrite(path_img, np_img)
- return np_img
- else:
- print("img doesn't exist.")
我们写两个循环, 外层循环是汉字列表, 内层循环是字体列表, 对于每个汉字会得到一个 image_list 列表, 里面存储着这个汉字的所有图像
- for (char, value) in lang_chars.items(): # 外层循环是字
- image_list = []
- print (char,value)
- #char_dir = os.path.join(images_dir, "%0.5d" % value)
- for j, verified_font_path in enumerate(verified_font_paths): # 内层循环是字体
- if rotate == 0:
- image = font2image.do(verified_font_path, char)
- image_list.append(image)
- else:
- for k in all_rotate_angles:
- image = font2image.do(verified_font_path, char, rotate=k)
- image_list.append(image)
我们将 image_list 中图像按照比例分为训练集和测试集存储
- test_num = len(image_list) * test_ratio
- random.shuffle(image_list) # 图像列表打乱
- count = 0
- for i in range(len(image_list)):
- img = image_list[i]
- #print(img.shape)
- if count < test_num :
- char_dir = os.path.join(test_images_dir, "%0.5d" % value)
- else:
- char_dir = os.path.join(train_images_dir, "%0.5d" % value)
- if not os.path.isdir(char_dir):
- os.makedirs(char_dir)
- path_image = os.path.join(char_dir,"%d.png" % count)
- cv2.imwrite(path_image,img)
- count += 1
写好代码后, 我们执行如下指令, 开始生成印刷体文字汉字集
python gen_printed_char.py--out_dir. / dataset--font_dir. / chinese_fonts--width 30--height 30--margin 4--rotate 30--rotate_step 1
解析一下上述指令的附属参数:
--out_dir 表示生成的汉字图像的存储目录
--font_dir 表示放置汉字字体文件的路径
--width --height 表示生成图像的高度和宽度
--margin 表示字体与边缘的间隔
--rotate 表示字体旋转的范围,[-rotate,rotate]
--rotate_step 表示每次旋转的间隔
生成这么一个 3755 个汉字的数据集的所需的时间还是很久的, 估计接近一个小时其实这个生成过程可以用多线程多进程并行加速, 但是考虑到这种文字数据集只需生成一次就好, 所以就没做这方面的优化了数据集生成完我们可以发现, 在 dataset 文件夹下得到 train 和 test 两个文件夹, train 和 test 文件夹下都有 3755 个子文件夹, 分别存储着生成的 3755 个汉字对应的图像, 每个子文件的名字就是该汉字对应的 id 随便选择一个 train 文件夹下的一个子文件夹打开, 可以看到所获得的汉字图像, 一共 634 个
dataset 下自动生成测试集和训练集
测试集和训练集下都有 3755 个子文件夹, 用于存储每个汉字的图像
生成出来的汉字图像
额外的图像增强
第三步生成的汉字图像是最基本的数据集, 它所做的图像处理仅有旋转这么一项, 如果我们想在数据增强上再做多点东西, 想必我们最终训练出来的 OCR 模型的性能会更加优秀我们使用 opencv 来完成我们定制的汉字图像增强任务
因为生成的图像比较小, 仅仅是 30*30, 如果对这么小的图像加噪声或者形态学处理, 得到的字体图像会很糟糕, 所以我们在做数据增强时, 把图片尺寸适当增加, 比如设置为 100×100, 再进行相应的数据增强, 效果会更好
噪点增加
- def add_noise(cls,img):
- for i in range(20): #添加点噪声
- temp_x = np.random.randint(0,img.shape[0])
- temp_y = np.random.randint(0,img.shape[1])
- img[temp_x][temp_y] = 255
- return img
适当腐蚀
- def add_erode(cls,img):
- kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(3, 3))
- img = cv2.erode(img,kernel)
- return img
适当膨胀
- def add_dilate(cls,img):
- kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(3, 3))
- img = cv2.dilate(img,kernel)
- return img
然后做随机扰动
- def do(self,img_list=[]):
- aug_list= copy.deepcopy(img_list)
- for i in range(len(img_list)):
- im = img_list[i]
- if self.noise and random.random()<0.5:
- im = self.add_noise(im)
- if self.dilate and random.random()<0.25:
- im = self.add_dilate(im)
- if self.erode and random.random()<0.25:
- im = self.add_erode(im)
- aug_list.append(im)
- return aug_list
输入指令
python gen_printed_char.py--out_dir. / dataset2--font_dir. / chinese_fonts--width 100--height 100--margin 10--rotate 30--rotate_step 1--need_aug
使用这种生成的图像如下图所示, 第一数据集扩大了两倍, 第二图像的丰富性进一步提高, 效果还是明显的当然, 如果要获得最好的效果, 还需要调一下里面的参数, 这里就不再详细说明了
至此, 我们所需的印刷体汉字数据集已经成功生成完毕, 下一步要做的就是利用这些数据集设计一个卷积神经网络做文字识别了! 完整的代码可以在我的 github 获取
来源: https://www.cnblogs.com/skyfsm/p/8436820.html