EAST 算法超详细源码解析(二)、模型推断与测试结果生成
Date: 2020/05/10
Author: CW
前言:
继上篇 EAST 算法超详细源码解析(一)、数据预处理与标签生成 后,本文就模型推断与测试结果生成的代码部分进行解读,咦,这是不是还有点“首尾呼应”的 feel,哈哈!
对一张测试图片或整个测试集进行推断
我们可以对一张图片进行推断,也可以对整个测试集(有多张测试图片)进行推断,下图是对一张图片的推断。
detect再来看看在整个测试集上的推断。
detect_dataset(i) detect_dataset(ii)plot_boxes:画出图像中的所有文本框
这部分相对简单,即绘制多边形,只要知道了每个多边形的4个顶点坐标位置即可。
plot_boxesget_boxes:对模型输出进行后处理,生成最终的预测结果(文本框)
这是核心部分,让我们来一起来仔细研读。
注意这个方法是对一张图片进行处理的,于是 score map 和 geo map 都是3维,在输入该方法前去掉了最外面对应 batch size 那一维(详细可以看前面部分 'detect' 那张图中的代码)。
get_boxes(i)有了位于文本框内部的点的位置以及它们的geo map,我们就可以转换出对应的文本框位置。
get_boxes(ii)restore_polys:计算图像中所有文本框4个顶点的坐标
restore_polys(i)注意,由于模型输出结果是对应在下采样1/4的特征图上的,于是我们需要在坐标转换到输入图像的尺寸。
restore_polys(ii)以上将角度取反,是因为在生成gt时,d值是将文本框旋转到与水平轴夹角为0的情况下计算出来的(在上一篇解读 EAST 算法超详细源码解析(一)、数据预处理与标签生成 中有详细说明),于是模型学到的d对应的也是文本框与水平轴夹角为0的情况,从而我们根据d计算出文本框顶点坐标后需要将文本框旋转回去,这样得到的顶点坐标才是真实的值。
光说可能不太好理解,画个图以示吾意(如果还没get到 CW 的意思那就只能yy了):
将文本框旋转为真实位置 restore_polys(iii)最后,记得要判断下旋转后得到的文本框顶点位置是否在输入图像尺寸范围内,如果是,才将其记录下来,作为合法的文本框。
到此,我们已经生成一批预测结果了,但是这其中有许多重复的文本框,需要将其去除,化繁为简。这种后处理通常都会用NMS,但这里的NMS有点不一样哦,做了一些小改动,称为局部感知的 NMS(Locality-Aware NMS)。
locality_aware_nms:后处理,局部感知的非极大值抑制
局部感知的 NMS 的不同之处在于,对于两个框,在它们的 IoU 大于阀值的时候,不是直接去掉一个,而是将它们进行合并。这样,根据 IoU 阀值,相似的框不会去除而是相互合并。
locality_aware_nms(i)合并完毕后使用标准的NMS。
locality_aware_nms(ii)iou:计算多边形之间的交并比
这里计算的 IoU 不是常规的bbox而是多边形之间的,但原理其实一致。
iouweighted_merge:文本框之间的合并
weighted_merge结语
标准 NMS 的代码就不在这里解读,这已经是常规操作了,大家应该也没什么胃口,模型推断与测试结果生成的代码解读就到此为止。
总地来看,重点部分有两个,一个是从 geo map 转换出旋转的文本框位置,另一个就是局部感知的NMS。