目标
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=