专栏名称: 机器学习研究会
机器学习研究会是北京大学大数据与机器学习创新中心旗下的学生组织,旨在构建一个机器学习从事者交流的平台。除了及时分享领域资讯外,协会还会举办各种业界巨头/学术神牛讲座、学术大牛沙龙分享会、real data 创新竞赛等活动。
目录
51好读  ›  专栏  ›  机器学习研究会

tensorflow LSTM + CTC实现端到端OCR

机器学习研究会  · 公众号  · AI  · 2017-11-16 22:34

正文

最近在做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来用作输入了

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 .//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来读比较好。

网络结构

然后是网络结构

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'






请到「今天看啥」查看全文