前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >tensorflow LSTM + CTC实现端到端OCR

tensorflow LSTM + CTC实现端到端OCR

作者头像
机器学习AI算法工程
发布2018-03-15 14:25:22
4.2K1
发布2018-03-15 14:25:22
举报

本文github源码地址:

在公众号 datadw 里 回复 OCR 即可获取。

最近在做OCR相关的东西,关于OCR真的是有悠久了历史了,最开始用tesseract,然而效果总是不理想,其中字符分割真的是个博大精深的问题,那么多年那么多算法,然而应用到实际总是有诸多问题。比如说非等间距字体的分割,汉字的分割,有光照阴影的图片的字体分割等等,针对特定的问题,特定的算法能有不错的效果,但也仅限于特定问题,很难有一些通用的结果。于是看了Xlvector的博客之后,发现可以端到端来实现OCR,他是基于mxnet的,于是我想把它转到tensorflow这个框架来,顺便还能熟悉一下这个框架。更加细节的实现方法见另一篇 http://ilovin.me/2017-04-23/tensorflow-lstm-ctc-input-output/

生成数据

利用captcha来生成验证码,具体生成验证码的代码,

在公众号 datadw 里 回复 OCR 即可获取。

共生成4-6位包含数字和英文大小写的训练图片128000张和测试图片400张。命名规则就是num_label.png,生成的图片如下图

关于生成数据,再多说一点,可以像Xlvector那样一边生成一边训练,这样样本是无穷的,效果更好。但是实际应用中有限样本的情况还是更多的。

载入数据

两种载入数据方式

pipeline

最开始想通过一个tf.train.string_input_producer来读入所有的文件名,然后以pipline的方式读入,但是由于标签的是不定长的,想通过正则来生成label,一开始是想用py_func来实现,后来发现传入string会有问题,所以最后还是选择生成tf.record文件,关于不定长问题,把比较短的标签在后面补零(0是blank的便签,就是说自己的类别中不能出现0这个类),然后读出每个batch后,再把0去掉。

一次性载入

我这里给一个目录,然后遍历里面所有的文件,等到训练的时候,每一个epoch循环把文件的index给手动shuffle一下,然后就可以每次截取出一个batch来用作输入了

代码语言:javascript
复制
class DataIterator:
    def __init__(self, data_dir):
        self.image_names = []
        self.image = []
        self.labels=[]
        for root, sub_folder, file_list in os.walk(data_dir):
            for file_path in file_list:
                image_name = os.path.join(root,file_path)
                self.image_names.append(image_name)
                im = cv2.imread(image_name,0).astype(np.float32)/255.
                im = cv2.resize(im,(image_width,image_height))
                # transpose to (160*60) and the step shall be 160
                # in this way, each row is a feature vector
                im = im.swapaxes(0,1)
                self.image.append(np.array(im))
                #image is named as ./<folder>/00000_abcd.png
                code = image_name.split('/')[2].split('_')[1].split('.')[0]
                code = [SPACE_INDEX if code == SPACE_TOKEN else maps[c] for c in list(code)]
                self.labels.append(code)
                print(image_name,' ',code)
    @property
    def size(self):
        return len(self.labels)
    def input_index_generate_batch(self,index=None):
        if index:
            image_batch=[self.image[i] for i in index]
            label_batch=[self.labels[i] for i in index]
        else:
            # get the whole data as input
            image_batch=self.image
            label_batch=self.labels
        def get_input_lens(sequences):
            lengths = np.asarray([len(s) for s in sequences], dtype=np.int64)
            return sequences,lengths
        batch_inputs,batch_seq_len = get_input_lens(np.array(image_batch))
        batch_labels = sparse_tuple_from_label(label_batch)
        return batch_inputs,batch_seq_len,batch_labels

需要注意的是tensorflow lstm输入格式的问题,其label tensor应该是稀疏矩阵,所以读取图片和label之后,还要进行一些处理,具体可以看代码

在公众号 datadw 里 回复 OCR 即可获取。

关于载入图片,发现12.8w张图一次读进内存,内存也就涨了5G,如果训练数据加大,还是加一个pipeline来读比较好。

网络结构

然后是网络结构

代码语言:javascript
复制
graph = tf.Graph()
with graph.as_default():
    inputs = tf.placeholder(tf.float32, [None, None, num_features])
    labels = tf.sparse_placeholder(tf.int32)
    seq_len = tf.placeholder(tf.int32, [None])
    # Stacking rnn cells
    stack = tf.contrib.rnn.MultiRNNCell([tf.contrib.rnn.LSTMCell(FLAGS.num_hidden,state_is_tuple=True) for i in range(FLAGS.num_layers)] , state_is_tuple=True)
    # The second output is the last state and we will no use that
    outputs, _ = tf.nn.dynamic_rnn(stack, inputs, seq_len, dtype=tf.float32)
    shape = tf.shape(inputs)
    batch_s, max_timesteps = shape[0], shape[1]
    # Reshaping to apply the same weights over the timesteps
    outputs = tf.reshape(outputs, [-1, FLAGS.num_hidden])
    # Truncated normal with mean 0 and stdev=0.1
    W = tf.Variable(tf.truncated_normal([FLAGS.num_hidden,
                                         num_classes],
                                        stddev=0.1),name='W')
    b = tf.Variable(tf.constant(0., shape=[num_classes],name='b'))
    # Doing the affine projection
    logits = tf.matmul(outputs, W) + b
    # Reshaping back to the original shape
    logits = tf.reshape(logits, [batch_s, -1, num_classes])
    # Time major
    logits = tf.transpose(logits, (1, 0, 2))
    global_step = tf.Variable(0,trainable=False)
    loss = tf.nn.ctc_loss(labels=labels,inputs=logits, sequence_length=seq_len)
    cost = tf.reduce_mean(loss)
    #optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate,
    #        momentum=FLAGS.momentum).minimize(cost,global_step=global_step)
    optimizer = tf.train.AdamOptimizer(learning_rate=FLAGS.initial_learning_rate,
            beta1=FLAGS.beta1,beta2=FLAGS.beta2).minimize(loss,global_step=global_step)
    # Option 2: tf.contrib.ctc.ctc_beam_search_decoder
    # (it's slower but you'll get better results)
    #decoded, log_prob = tf.nn.ctc_greedy_decoder(logits, seq_len,merge_repeated=False)
    decoded, log_prob = tf.nn.ctc_beam_search_decoder(logits, seq_len,merge_repeated=False)
    # Inaccuracy: label error rate
    lerr = tf.reduce_mean(tf.edit_distance(tf.cast(decoded[0], tf.int32), labels))

这里我参考了stackoverflow的一篇帖子写的,

https://stackoverflow.com/questions/38059247/using-tensorflows-connectionist-temporal-classification-ctc-implementation

根据tensorflow 1.0.1的版本做了微调,使用了Adam作为optimizer。

需要注意的是ctc_beam_search_decoder是非常耗时的.

和greedy_decoder的区别是,greedy_decoder根据当前序列预测下一个字符,并且取概率最高的作为结果,再此基础上再进行下一次预测。而beam_search_decoder每次会保存取k个概率最高的结果,以此为基础再进行预测,并将下一个字符出现的概率与当前k个出现的概率相乘,这样就可以减缓贪心造成的丢失好解的情况,当k=1的时候,二者就一样了。

结果

稍微调一调,网络可以跑到85%以上。 把网络用在识别身份证号,试了73张网上爬的(不同分辨率下的)真实图片,错了一张,准确率在98%左右(不过毕竟身份证号比较简单)

大概14个epoch后,准确率过了50%,现在跑到了73%的正确率。

最后,代码托管在Github上。

后记

百度出了一个warpCTC可以加速CTC的计算,试用了一下CPU的版本发现好像没什么速度的提升,不知道是不是姿势不对,回头再试试GPU的版本。 对于更加细节的实现方法(输入输出的构造,以及warpCTC和内置ctc_loss的异同)放在了另一篇博客。

http://ilovin.me/2017-04-23/tensorflow-lstm-ctc-input-output/

  • warpCTC的GPU版本试过之后发现速度差不多,但是能极大的减少CPU的占用
  • 对于不同的优化器,数据,同样的参数是不能普适的。往往之前的参数可以收敛,换个optimizer,数据,网络就不能收敛了。这个时候要微调参数。对于不同的优化器之间区别,文末有一篇神文可以参考

http://ilovin.me/2017-04-06/tensorflow-lstm-ctc-ocr/

本文参与 腾讯云自媒体分享计划,分享自微信公众号。
原始发表:2017-11-16,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 大数据挖掘DT数据分析 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体分享计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 生成数据
  • 载入数据
    • pipeline
      • 一次性载入
      • 网络结构
      • 结果
      • 后记
      相关产品与服务
      验证码
      腾讯云新一代行为验证码(Captcha),基于十道安全栅栏, 为网页、App、小程序开发者打造立体、全面的人机验证。最大程度保护注册登录、活动秒杀、点赞发帖、数据保护等各大场景下业务安全的同时,提供更精细化的用户体验。
      领券
      问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档