# CNN-LSTM-CTC-OCR
# Copyright (C) 2017 Jerod Weinman
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
# model.py -- Constructs the graph representing the network
# model. Inputs start from convnet_layers(), whose outputs hook into
# rnn_layers(), which produces the logits for CTC loss (for
# training) and decoding (for prediction/evaluation).
import tensorflow as tf
from tensorflow.contrib import learn
# Layer params: Filts K Padding Name BatchNorm?
layer_params = [ [ 64, 3, 'valid', 'conv1', False],
[ 64, 3, 'same', 'conv2', True], # pool
[ 128, 3, 'same', 'conv3', False],
[ 128, 3, 'same', 'conv4', True], # hpool
[ 256, 3, 'same', 'conv5', False],
[ 256, 3, 'same', 'conv6', True], # hpool
[ 512, 3, 'same', 'conv7', False],
[ 512, 3, 'same', 'conv8', True] ] # hpool 3
rnn_size = 2**9 # Dimensionality of all RNN elements' hidden layers
dropout_rate = 0.5 # For RNN layers (currently not used--uncomment below)
def conv_layer( bottom, params, training ):
"""Build a convolutional layer using entry from layer_params)"""
batch_norm = params[4] # Boolean
if batch_norm:
activation = None
else:
activation = tf.nn.relu
kernel_initializer = tf.contrib.layers.variance_scaling_initializer()
bias_initializer = tf.constant_initializer( value=0.0 )
top = tf.layers.conv2d( bottom,
filters=params[0],
kernel_size=params[1],
padding=params[2],
activation=activation,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
name=params[3] )
if batch_norm:
top = norm_layer( top, training, params[3]+'/batch_norm' )
top = tf.nn.relu( top, name=params[3]+'/relu' )
return top
def pool_layer( bottom, wpool, padding, name ):
"""Short function to build a pooling layer with less syntax"""
top = tf.layers.max_pooling2d( bottom,
2,
[2, wpool],
padding=padding,
name=name )
return top
def norm_layer( bottom, training, name):
"""Short function to build a batch normalization layer with less syntax"""
top = tf.layers.batch_normalization( bottom,
axis=3, # channels last
training=training,
name=name )
return top
def convnet_layers( inputs, widths, mode ):
"""
Build convolutional network layers attached to the given input tensor
"""
training = (mode == learn.ModeKeys.TRAIN)
# inputs should have shape [ ?, 32, ?, 1 ]
with tf.variable_scope( "convnet" ): # h,w
conv1 = conv_layer( inputs, layer_params[0], training ) # 30,30
conv2 = conv_layer( conv1, layer_params[1], training ) # 30,30
pool2 = pool_layer( conv2, 2, 'valid', 'pool2' ) # 15,15
conv3 = conv_layer( pool2, layer_params[2], training ) # 15,15
conv4 = conv_layer( conv3, layer_params[3], training ) # 15,15
pool4 = pool_layer( conv4, 1, 'valid', 'pool4' ) # 7,14
conv5 = conv_layer( pool4, layer_params[4], training ) # 7,14
conv6 = conv_layer( conv5, layer_params[5], training ) # 7,14
pool6 = pool_layer( conv6, 1, 'valid', 'pool6') # 3,13
conv7 = conv_layer( pool6, layer_params[6], training ) # 3,13
conv8 = conv_layer( conv7, layer_params[7], training ) # 3,13
pool8 = tf.layers.max_pooling2d( conv8, [3, 1], [3, 1],
padding='valid',
name='pool8' ) # 1,13
# squeeze row dim
features = tf.squeeze( pool8, axis=1, name='features' )
sequence_length = get_sequence_lengths( widths )
# Vectorize
sequence_length = tf.reshape( sequence_length, [-1], name='seq_len' )
return features, sequence_length
def get_sequence_lengths( widths ):
"""Tensor calculating output sequence length from original image widths"""
kernel_sizes = [params[1] for params in layer_params]
with tf.variable_scope("sequence_length"):
conv1_trim = tf.constant( 2 * (kernel_sizes[0] // 2),
dtype=tf.int32,
name='conv1_trim' )
one = tf.constant( 1, dtype=tf.int32, name='one' )
two = tf.constant( 2, dtype=tf.int32, name='two' )
after_conv1 = tf.subtract( widths, conv1_trim, name='after_conv1' )
after_pool2 = tf.floor_div( after_conv1, two, name='after_pool2' )
after_pool4 = tf.subtract( after_pool2, one, name='after_pool4' )
after_pool6 = tf.subtract( after_pool4, one, name='after_pool6' )
after_pool8 = tf.identity( after_pool6, name='after_pool8' )
return after_pool8
def rnn_layer( bottom_sequence, sequence_length, rnn_size, scope ):
"""Build bidirectional (concatenated output) RNN layer"""
weight_initializer = tf.truncated_normal_initializer( stddev=0.01 )
# Default activation is tanh
cell_fw = tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell( rnn_size )
cell_bw = tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell( rnn_size )
# Pre-CUDNN (slower) alternatve. Default activation is tanh .
#cell_fw = tf.contrib.rnn.LSTMCell( rnn_size,
# initializer=weight_initializer)
#cell_bw = tf.contrib.rnn.LSTMCell( rnn_size,
# initializer=weight_initializer)
# Include?
#cell_fw = tf.contrib.rnn.DropoutWrapper( cell_fw,
# input_keep_prob=dropout_rate )
#cell_bw = tf.contrib.rnn.DropoutWrapper( cell_bw,
# input_keep_prob=dropout_rate )
rnn_output,_ = tf.nn.bidirectional_dynamic_rnn(
cell_fw,
cell_bw,
bottom_sequence,
sequence_length=sequence_length,
time_major=True,
dtype=tf.float32,
scope=scope )
# Concatenation allows a single output op because [A B]*[x;y] = Ax+By
# [ paddedSeqLen batchSize 2*rnn_size]
rnn_output_stack = tf.concat( rnn_output, 2, name='output_stack' )
return rnn_output_stack
def rnn_layers( features, sequence_length, num_classes ):
"""Build a stack of RNN layers from input features"""
# Input features is [batchSize paddedSeqLen numFeatures]
logit_activation = tf.nn.relu
weight_initializer = tf.contrib.layers.variance_scaling_initializer()
bias_initializer = tf.constant_initializer( value=0.0 )
with tf.variable_scope( "rnn" ):
# Transpose to time-major order for efficiency
rnn_sequence = tf.transpose( features,
perm=[1, 0, 2],
name='time_major' )
rnn1 = rnn_layer( rnn_sequence, sequence_length, rnn_size, 'bdrnn1' )
rnn2 = rnn_layer( rnn1, sequence_length, rnn_size, 'bdrnn2' )
rnn_logits = tf.layers.dense( rnn2,
num_classes+1,
activation=logit_activation,
kernel_initializer=weight_initializer,
bias_initializer=bias_initializer,
name='logits' )
return rnn_logits
def ctc_loss_layer( rnn_logits, sequence_labels, sequence_length,
reduce_mean=True ):
"""Build CTC Loss layer for training"""
losses = tf.nn.ctc_loss( sequence_labels,
rnn_logits,
sequence_length,
time_major=True,
ignore_longer_outputs_than_inputs=True )
if (reduce_mean):
loss = tf.reduce_mean( losses )
else:
loss = tf.reduce_sum( losses )
return loss
