单字ocr-cnn效果实现与测试

目标

1. 在无数据集基础上,训练基于CNN的ocr模型
2. 使用vae增强数据

环境

1. python3.6
2. numpy=1.18.0
3. opencv-python=4.1.1.26
4. pandas=1.1.5
5. Pillow=5.4.1
6. python-igraph=0.8.0
7. PythonMagick=0.9.19
8. tensorflow-gpu=2.0.0
9. torch=1.8.0

大概就这些

思路

1. 数据集生成
2. 训练VAE并生成数据
3. 训练CNN模型
4. 查看效果

数据集

在一张空白图片上绘制单个字并保存
字体是从操作系统字体库中找的
基本参数如下

# 图片大小统一为28*28
self.imgsize = kwargs.get('imgsize', (28, 28))
# 太小的字看不清
self.fontsize = kwargs.get('fontsize', range(25, 27))
# 字体有这些,最后有效的只有两个
self.fonttype = kwargs.get('fonttype', ['simfang.ttf',
                                        'STSONG.TTF',
                                        'SIMYOU.TTF',
                                        'STZHONGS.TTF',
                                        'STKAITI.TTF'])

汉字编码范围我选择的是(0x4E00, 0x9FA5),再加上0-9,a-z差不多足够了

# 绘图的代码,绘制完成后将图片二值化为0,1
with Image.new('RGB', self.imgsize) as img:
     draw = ImageDraw.Draw(img)
     draw.text((1, 1),
     	kwargs.get('label'),
     	'white', font=kwargs.get('font')[-1])
     img = cv2.cvtColor(np.asarray(img, np.uint8),
     			       cv2.COLOR_RGB2GRAY)
     img = np.where(img > 0, 255, 0)
     img = self.std(img)
# 随机点的位置
def ptr(start=0, stop=28):
    return random.randint(start, stop)
# 在绘制了汉字的图片上添加噪音
with Image.new('RGB', self.imgsize) as img:
    draw = ImageDraw.Draw(img)
    draw.text((1, 1), kwargs.get('label'), 'white', font=kwargs.get('font')[-1])
    if kwargs.get('drawtype') == 'drawline':
        draw.line([(ptr(), ptr()), (ptr(), ptr())], fill='white', width=1)
    if kwargs.get('drawtype') == 'drawpoint':
        draw.point((ptr(), ptr()), fill='white')
    img = cv2.cvtColor(np.asarray(img, np.uint8), cv2.COLOR_RGB2GRAY)

    img = np.where(img > 0, 255, 0)
    img = self.std(img)
    # self.show(img)
    return img

示例:

VAE

使用变分自编码器在上述生成图片的基础上伪造更丰富的图片
代码参考生成式深度学习中提到的,摘录其中易懂的部分

"""
模型训练
"""
import codecs
import cv2
import os
import numpy as np
from sklearn.model_selection import train_test_split
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model
from config import get_variable


class Generate(object):
    input_img = keras.Input(shape=get_variable('img_shape'))
    x = layers.Conv2D(128, 3, padding='same', activation='relu')(input_img)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(64, 3, padding=