这篇文章解析以下CTPN的模型结构。
不多说,先上代码,用的是tensorflow1中的slim库。
def model(image):
image = mean_image_subtraction(image)
with slim.arg_scope(vgg.vgg_arg_scope()):
conv5_3 = vgg.vgg_16(image)
rpn_conv = slim.conv2d(conv5_3, 512, 3)
lstm_output = Bilstm(rpn_conv, 512, 128, 512, scope_name='BiLSTM')
bbox_pred = lstm_fc(lstm_output, 512, 10 * 4, scope_name="bbox_pred")
cls_pred = lstm_fc(lstm_output, 512, 10 * 2, scope_name="cls_pred")
cls_pred_shape = tf.shape(cls_pred)
cls_pred_reshape = tf.reshape(cls_pred, [cls_pred_shape[0], cls_pred_shape[1], -1, 2])
cls_pred_reshape_shape = tf.shape(cls_pred_reshape)
cls_prob = tf.reshape(tf.nn.softmax(tf.reshape(cls_pred_reshape, [-1, cls_pred_reshape_shape[3]])),
[-1, cls_pred_reshape_shape[1], cls_pred_reshape_shape[2], cls_pred_reshape_shape[3]],
name="cls_prob")
return bbox_pred, cls_pred, cls_prob
然后我们一句句来看。
image = mean_image_subtraction(image)
这一句,是因为一般使用VGG,用的是在Imgnet上预训练好的,需要减去Imgnet上图片的均值,做数据特征标准化。
with slim.arg_scope(vgg.vgg_arg_scope()):
conv5_3 = vgg.vgg_16(image)
这一步得到VGG16提取的特征图,取的是conv5_3那一层。
rpn_conv = slim.conv2d(conv5_3, 512, 3)
这一句进行卷积核大小是3,步长为1的卷积进一步提取特征。
lstm_output = Bilstm(rpn_conv, 512, 128, 512, scope_name='BiLSTM')
这一步将上面得到的特征输入到一个双向的LSTM中。BiLSTM函数后面会进一步解析。
bbox_pred = lstm_fc(lstm_output, 512, 10 * 4, scope_name="bbox_pred")
cls_pred = lstm_fc(lstm_output, 512, 10 * 2, scope_name="cls_pred")
这两句是将LSTM层的得到的结果输入到全连接层得到最终输出。
其中bbox_pred得到的是检测框中心点坐标x,y的偏移以及宽和高的偏移。这里其实和原论文中描述的不一样,这份代码是根据Faster RCNN改过来的。原论文中,不需要预测x的偏移以及w的偏移,因为在CTPN设计的时候w固定死了为16个像素(因为vgg提取的特征图大小是原图1/16),x就是特征图上的点映射到原图上的区域的中心点的x坐标。所以原CTPN只需要预测y的偏移以及h的偏移就好了。当然像这里这么做也有好处,之后会提到,有side-refinement的作用。
cls_pred得到的是这个框中 文本/非文本的得分(置信度)。也就是是文本的可能性和不是文本的可能性。
cls_pred_shape = tf.shape(cls_pred)
cls_pred_reshape = tf.reshape(cls_pred, [cls_pred_shape[0], cls_pred_shape[1], -1, 2])
cls_pred_reshape_shape = tf.shape(cls_pred_reshape)
cls_prob = tf.reshape(tf.nn.softmax(tf.reshape(cls_pred_reshape, [-1, cls_pred_reshape_shape[3]])),
[-1, cls_pred_reshape_shape[1], cls_pred_reshape_shape[2], cls_pred_reshape_shape[3]],
name="cls_prob")
这些操作是为了计算 文本/非文本的概率。主要是为了使用softmax,所以需要先对cls_pred进行reshape,然后再reshape回去。最后cls_prob就是 文本/非文本的概率。
下面介绍解析Bilstm函数,具体可以看注解。
def Bilstm(net, input_channel, hidden_unit_num, output_channel, scope_name):
# width--->time step
with tf.variable_scope(scope_name) as scope:
# reshape 输入
shape = tf.shape(net)
N, H, W, C = shape[0], shape[1], shape[2], shape[3]
net = tf.reshape(net, [N * H, W, C])
net.set_shape([None, None, input_channel])
# 声明前向LSTM lstm_fw_cell 和后向LSTM lstm_bw_cell
lstm_fw_cell = tf.contrib.rnn.LSTMCell(hidden_unit_num, state_is_tuple=True)
lstm_bw_cell = tf.contrib.rnn.LSTMCell(hidden_unit_num, state_is_tuple=True)
# 建立双向LSTM
lstm_out, last_state = tf.nn.bidirectional_dynamic_rnn(lstm_fw_cell, lstm_bw_cell, net, dtype=tf.float32)
# concat LSTM的输出
lstm_out = tf.concat(lstm_out, axis=-1)
# reshape LSTM的输出,为了之后输入到全连接层
lstm_out = tf.reshape(lstm_out, [N * H * W, 2 * hidden_unit_num])
# 这里建立全连接层,初始化weight,biases
init_weights = tf.contrib.layers.variance_scaling_initializer(factor=0.01, mode='FAN_AVG', uniform=False)
init_biases = tf.constant_initializer(0.0)
weights = make_var('weights', [2 * hidden_unit_num, output_channel], init_weights)
biases = make_var('biases', [output_channel], init_biases)
outputs = tf.matmul(lstm_out, weights) + biases
# 全连接结束, 再reshape回去
outputs = tf.reshape(outputs, [N, H, W, output_channel])
return outputs
论文中其实还提到了side-refinement。这个东西是用来修正文本行的边界,这里没有实现,作者提供的caffe代码也没有实现(甚至没有训练代码)。但是这里我们预测框的时候对框的水平位置和宽度也进行了修正,所以这也起来了side-refinement的作用。
到这里model就解析完了,之后是reshape的细节,不感兴趣的可以不看。
先介绍一下tf.reshape()函数。 tf.reshape(输入,[reshape之后的维度])。
首先一般图像输入都是四维的 NHW*C,N代表几张图片也就是batch_size,W的图片的宽度,H的图片的高度,C是图片的通道数。各个模型这几维的排列顺序可能不同。
还有一个地方,四个维度中如果有一个填的是-1,那代表的意思就是这个维度是多少是自动计算的。还有就是因为tensorflow是静态图,所以你在运行的时候如果想得到shape只能调用tf.shape()。
之前我们得到的cls_pred是 文本/非文本的得分,我们需要的是 文本/非文本的概率,那么自然而然就想到用softmax。比如说 文本/非文本分数是 [3,1],那么softmax之后应该是[0.5 , 0.5],softmax在这里就不介绍了,具体的可以看这个。https://www.jianshu.com/p/4666213d20c7](https://www.jianshu.com/p/4666213d20c7
这里想用tf.nn.softmax,默认计算输入张量的最后一个维度。那我们就需要把我们之前得到的cls_pred reshape成最后一个维度是 文本/非文本 得分,也就是说最后一个维度大小是2。
cls_pred_reshape = tf.reshape(cls_pred, [cls_pred_shape[0], cls_pred_shape[1], -1, 2])
tf.reshape(cls_pred_reshape, [-1, cls_pred_reshape_shape[3]])
上面两句话完成了这一部分工作。
tf.nn.softmax(tf.reshape(cls_pred_reshape, [-1, cls_pred_reshape_shape[3]]))
上面的代码,对reshape之后的张量进行softmax操作
tf.reshape(tf.nn.softmax(tf.reshape(cls_pred_reshape, [-1, cls_pred_reshape_shape[3]])),
[-1, cls_pred_reshape_shape[1], cls_pred_reshape_shape[2], cls_pred_reshape_shape[3]],
name="cls_prob")
最后再reshape回去。