yolov3使用ICDAR2013数据集训练字符区域识别

1. OCR任务简述

最近项目需求要做一个OCR的demo给客户展示一下，之前也没有涉及过ocr，做了一些调研。感觉整体思路和两步法的目标检测比较类似。
首先我们需要从一张自然图像中定位出有文字的区域并把这个区域提取出来；然后对提取出文字的区域进行识别，识别出文字的内容，再后续处理利用（比如交给nlp进行分析处理）。如下图：

上图显示的是ocr的第一步，及把整张图像中有文字的区域分离出来。

然后对提取出的有文字的区域进行识别，如上图所示。如上即完成了ocr的全部过程，下面介绍具体的实现方式。

2.OCR开源数据集

因为ocr的第一步和目标检测比较类似，加之之前对yolov3 比较熟悉，所以优先考虑使用yolov3对有文字的区域进行检测。下面介绍几个ocr领域的常用数据集。
2.1 ICDAR(International Conference on Document Analysis and Recognition) 2011 - About Robust Reading Competition Reading Text in Born-Digital Images
(Web and Email) http://www.cvc.uab.es/icdar2011competition/
Reading Text in Scene Images http://robustreading.opendfki.de/wiki/SceneText

2.2 ICDAR 2013 - About Robust Reading Competition Reading Text in Born-Digital Images (Web and Email) http://dag.cvc.uab.es/icdar2013competition/?ch=1&com=downloads
这个数据集也是我选用的数据集，外网下载可能比较慢，网盘链接：
链接: https://pan.baidu.com/s/1rwlXWhTtEq6-op4HBKBcOQ 密码: mcbi
2.4 MNIST - Data Base http://yann.lecun.com/exdb/mnist/
2.5 USPS - Handwritten digit data http://www.gaussianprocess.org/gpml/data/

3.ICDAR2013数据集整理为yolov3可用格式

看一下标注的文件内容，每一个图片对应一个txt，打开后每一行前4个数字对应字母区域在图片中的位置， 最后的“”中的内容是字母的内容。

3.1 ICDAR2013数据集标注格式转换

但是在yolov3检测文字区域这步我们并不需要知道文字区域中的内容是什么，也就是说我们只需要检测区域，所以不需要最后的“”中的信息；此外，还要给原始txt标注更改成voc数据集的.xml标注格式方便进一步使用。
如何整理可以参考博文yolov2训练ICDAR2011数据集
将txt标注转换成xml格式标注代码如下：

#! /usr/bin/python
#-*-coding:utf8-*-

import os, sys
import glob
from PIL import Image

# ICDAR 图像存储位置
src_img_dir = "img"
# ICDAR 图像的 ground truth 的 txt 文件存放位置，新生成的xml格式标注和txt在一起
src_txt_dir = "rot_label"

img_Lists = glob.glob(src_img_dir + '/*.jpg')

img_basenames = [] # e.g. 100.jpg
for item in img_Lists:
    img_basenames.append(os.path.basename(item))

img_names = [] # e.g. 100
for item in img_basenames:
    temp1, temp2 = os.path.splitext(item)
    img_names.append(temp1)

for img in img_names:
    im = Image.open((src_img_dir + '/' + img + '.jpg'))
    width, height = im.size

    # open the crospronding txt file
    gt = open(src_txt_dir + '/gt_' + img + '.txt').read().splitlines()

    # write in xml file
    #os.mknod(src_txt_dir + '/' + img + '.xml')
    xml_file = open((src_txt_dir + '/' + img + '.xml'), 'w')
    xml_file.write('\n')
    xml_file.write('    VOC2007\n')
    xml_file.write('    ' + str(img) + '.jpg' + '\n')
    xml_file.write('    \n')
    xml_file.write('        ' + str(width) + '\n')
    xml_file.write('        ' + str(height) + '\n')
    xml_file.write('        3\n')
    xml_file.write('    \n')

    # write the region of text on xml file
    for img_each_label in gt:
        spt = img_each_label.split(' ')   
        #''中间有一个空格，该txt中每行数据用空格分离，不同的数据集可能需要改动，
        #主要看原始的标注txt中每行数据的间隔方式
        
        xml_file.write('    \n')

    xml_file.write('')

代码运行完成会在存放txt的文件夹中生成.xml的标注文件，标注格式与voc数据集格式相同。

3.2 按照darkent yolov3需要模式整理数据

可以参照windows10+keras下的yolov3的快速使用及自己数据集的训练
如下是我的文件夹结构

运行gen_train_val.py程序会自动划分训练集和测试集，代码：

import os
import random
 
trainval_percent = 0.05
train_percent = 0.95       #划分训练集和测试集的比例，数据量小建议训练集比重适当调大
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
total_xml = os.listdir(xmlfilepath)
 
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
 
ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')
 
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
 
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

运行结束会生成如下文件

3.3 运行voc_label.py生成yolov3可以使用的标注

程序位置：/home/your user name/darknet-gpu/scripts/voc_label.py
代码如下：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

sets=[ ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]

classes = ["text"]  #注意：此处["  "]中内容必须和xml中 text
                    #内容一致，即“text”，因为程序是按照“text”关键字在xml中
                    #索引对应信息后写入新的txt中，如果不一致会导致生成的txt
                    #内容为空
def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(year, image_id):
    in_file = open('VOCdevkit-text/VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOCdevkit-text/VOC%s/labels/%s.txt'%(year, image_id), 'w')
    tree=ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult)==1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()

for year, image_set in sets:
    if not os.path.exists('VOCdevkit-text/VOC%s/labels/'%(year)):
        os.makedirs('VOCdevkit-text/VOC%s/labels/'%(year))
    image_ids = open('VOCdevkit-text/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOCdevkit-text/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()

os.system("cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt > train.txt")
os.system("cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt > train.all.txt")

运行结束会在label文件夹中生成yolov3可用的txt格式标注数据

txt中内容如下：第一个0代表区域属于第一类（本例也只有一类）

将标注处理成如上格式即可以开始yolov3训练。

4.使用yolov3进行训练，找出有文字的区域

安装darkent，使用darknet 下的yolov3进行训练，安装方式网上很多。

4.1 修改yolov3-voc.cfg

以下是cfg中需要修改的位置

[net]
# Testing
 batch=128                  #每一个批次送进网络图像的数量，主要由显存空间决定。
 subdivisions=32            #只影响训练速度，不影响最终精度
# Training
# batch=32
# subdivisions=16
width=416                   #图像被统一resize到416*416
height=416
channels=3
momentum=0.9                #动量
decay=0.0005
angle=0                     #旋转角度，图像增强使用，慎重使用，旋转可能会导致方框超出
saturation = 1.5
exposure = 1.5
hue=.1

learning_rate=0.001        #学习率，训练效果长期不收敛考虑调低
burn_in=10000
max_batches = 40000        #最多训练次数
policy=steps
steps=40000,45000
scales=.1,.1

然后在cfg文档中搜索yolo，一共可以找到3处，3处均按照如下修改。

[convolutional]
size=1
stride=1
pad=1
filters=18    #计算公式为3 * （5 + classes），本例classes=1，故filters数量为3 *（1+5）=18
activation=linear

[yolo]
mask = 6,7,8
anchors = 10,13,  16,30,  33,23,  30,61,  62,45,  59,119,  116,90,  156,198,  373,326
# anchors通常不需要修改，如果为了进一步提升效果也可以通过看
#kmeans产生所用训练数据匹配的anchor box值
classes=1    #本例只有一类，故为1
num=9
jitter=.3
ignore_thresh = .5
truth_thresh = 1
random=1

4.2 修改voc.data和voc.names

voc.data修改如下

classes= 1
train  = /home/your user name/darknet-gpu/2007_train.txt  
valid  = /home/your user name/darknet-gpu/2007_test.txt
names = /home/your user name/darknet-gpu/data/voc.names
backup = backup

voc.names修改如下

text   #因为只有一类，类名为text

4.3 开始训练

darknet 框架下yolov3的训练方式：

./darknet detector train cfg/voc-text.data
 cfg/yolov3-voc-text.cfg backup/yolov3-voc-text.backup -gpus 0,1

训练开始如图：

Loading weights from backup/yolov3-voc-text.backup...Done!
Learning Rate: 0.002, Momentum: 0.9, Decay: 0.0005
Resizing
512
Loaded: 0.000041 seconds
Region 82 Avg IOU: 0.799022, Class: 0.999835, Obj: 0.990258, No Obj: 0.008770, .5R: 1.000000, .75R: 0.777778,  count: 9
Region 94 Avg IOU: 0.783882, Class: 0.999762, Obj: 0.962048, No Obj: 0.001786, .5R: 1.000000, .75R: 0.500000,  count: 6
Region 82 Avg IOU: 0.724783, Class: 0.999746, Obj: 0.996917, No Obj: 0.003994, .5R: 1.000000, .75R: 0.500000,  count: 2
Region 106 Avg IOU: -nan, Class: -nan, Obj: -nan, No Obj: 0.000008, .5R: -nan, .75R: -nan,  count: 0
Region 94 Avg IOU: 0.785549, Class: 0.999787, Obj: 0.991846, No Obj: 0.002388, .5R: 1.000000, .75R: 0.875000,  count: 8
Region 106 Avg IOU: 0.757414, Class: 0.997299, Obj: 0.622859, No Obj: 0.000092, .5R: 1.000000, .75R: 0.500000,  count: 2
Region 82 Avg IOU: 0.814417, Class: 0.999138, Obj: 0.985323, No Obj: 0.009260, .5R: 1.000000, .75R: 0.857143,  count: 7
Region 94 Avg IOU: 0.794187, Class: 0.998478, Obj: 0.978485, No Obj: 0.001021, .5R: 1.000000, .75R: 0.500000,  count: 4
Region 106 Avg IOU: 0.447214, Class: 0.998129, Obj: 0.712037, No Obj: 0.000188, .5R: 0.375000, .75R: 0.125000,  count: 8
Region 82 Avg IOU: 0.868999, Class: 0.999798, Obj: 0.996820, No Obj: 0.003046, .5R: 1.000000, .75R: 1.000000,  count: 3
Region 94 Avg IOU: 0.804574, Class: 0.999718, Obj: 0.913576, No Obj: 0.000859, .5R: 1.000000, .75R: 0.500000,  count: 4
Region 106 Avg IOU: 0.701177, Class: 0.999827, Obj: 0.423744, No Obj: 0.000053, .5R: 1.000000, .75R: 0.500000,  count: 2

训练好的模型权重文件保存再/backup文件夹中，如图

训练结束后使用得到的weights文件执行推断验证效果。

4.4 推断验证训练效果

./darknet detector test cfg/voc-text.data.data cfg/yolov3-voc-text.cfg 
yolov3-voc-text_20000.weights 240.jpg

得到结果保存成图片predictions.jpg,如图：

截止此时已经完成了ocr的第一步，把文字区域从整张图片中分离出来。

5.使用pytesseract识别文字区域

import cv2
import numpy as np
from PIL import Image
import pytesseract as tess

def recognize_text(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    cv2.imshow("binimg", gray)
    ret, binnary = cv2.threshold(gray, 100, 255, cv2.THRESH_OTSU)
    cv2.imshow("binmg", binnary)
    kerhel1 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    bin1 = cv2.morphologyEx(binnary, cv2.MORPH_OPEN, kerhel1, iterations=1)
    kerhel2 = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 1))
    bin2 = cv2.morphologyEx(binnary, cv2.MORPH_OPEN, kerhel2, iterations=1)
    cv2.imshow("binary_img",bin2)
    text = tess.image_to_string(bin2)
    print("识别结果："%text)

img = cv2.imread('12.jpg',1)
recognize_text(img)
cv2.waitKey(0)
cv2.destroyAllWindows()

注意: pytesseract的安装方式

sudo apt install tesseract
pip install pytesseract