总论中对PPStructure核心类做了概要介绍,本文旨在详解核心类之一:LayoutPredictor类。
回顾下总论中的内容,LayoutPredictor类负责版面分析,用于获取图片中各版面要素的边界框、类别标签、置信度等信息。代码在paddleocr/ppstructure/layout/predict_layout.py文件中。
本文首先详细解释LayoutPredictor类的__init__与__call__方法,接着通过应用场景代码实践,进一步了解该类的实现原理与用法。
构造函数__init__方法
构造函数用于构造LayoutPredictor类的实例,是使用该类的前提。构造函数位于predict_layout.py文件的第39行,定义如下:
def __init__(self, args)
可见,只要传入一个参数对象args即可,系统并没有对参数对象的类进行任何约束,看起来非常简单。但是,如果仔细阅读源码,可以看出构造函数中对args对象的属性有最小集要求。亦即传入的args参数对象,至少应该包含源码中使用到的属性,否则会报错。以下是args必须包含的属性列表:
主要参数包括:
- args.use_onnx 代表是否使用ONNX模型格式,可以为True或者False。
- args.layout_dict_path 代表版面分析字典文件所在路径,字符串str类型,默认值为layout_cdla_dict.txt。字典对形成正确的分析结果非常重要,必须保持与推理模型配套。推理模型运行的一般过程是将图像张量输入到模型,执行模型并获取输出张量。然后,将输出张量通过字典中单词序号的简单映射关系,转换为单词的列表。举例来说,假设通过版面分析模型,得到输出张量为[0,0,3,4],假设模型的配套字典为默认值,那么返回的解析结果是[‘text’,‘text’,‘figure_caption’,‘table’]。同样的输出张量,如果模型的配套字典为PubLayNet,那么返回的解析结果是[‘text’,‘text’,‘table’,‘figure’]。差之毫厘谬之千里也。
- args.layout_score_threshold代表版面分析置信度阈值,浮点float类型。置信度低于此阈值者筛选掉,不输出到解析结果。如果使用paddleocr的parse_args来构造参数对象,那么layout_score_threshold的默认值被设置为0.5。
- args.layout_nms_threshold代表版面分析nms算法的iou阈值,浮点float类型。开源团队使用了非极大值抑制算法,来处理计算机视觉识别到的候选框序列,高于此阈值者会被视为重叠框从而从结果中排除。 nms的概念理解,可以看看这篇文章。如果使用paddleocr的parse_args来构造参数对象,那么layout_nms_threshold的默认值被设置为0.5。
上述必选参数,数量不多,可以自定义一个类来设置这些属性。为了保持与源码的一致性,推荐使用parse_args函数,该函数的解释参照PPStructure核心源码研究(二)。
魔法函数__call__
上节的构造函数完成了实例基础设置,真正体现功能的地方在魔法函数__call__。函数位于predict_layout.py文件的第70行,定义如下:
def __call__(self, img)
传入参数解释如下:
- img参数,图像数据三维张量,形状为[h,w,3],分别代表图像高度、图像宽度、通道数(RGB)
魔法函数__call__函数执行过程中,首先对原始图像进行归一化等预处理操作,然后传入推理模型,执行模型,获取输出张量,并通过字典翻译,形成结果列表res_list和耗时elapse两项返回值。
- res_list
第一项返回值是一个结果列表,列表中的每一项代表一个版面要素。每个识别到的版面要素是一个字典,包含以下属性:- bbox 要素边界框,列表list类型,包含四个数值,分别是左上角x、左上角y、右下角x、右下角y。
- label 要素类型标签,字符串str类型,值来自于版面分析模型配套字典。
- score 要素置信度,浮点float类型,置信度越高结果越可信。
- elapse
第二项返回结果是耗时,代表整个模型的执行时间,单位为秒。
实验一:获取版面结果
第一个实验的目标,是直接使用LayoutPredictor类,识别官网中示例图片(假设已经存储到本地,名为layout1.jpg):
通过程序运行,得到每个版面分析的类别、置信度等信息。主要代码如下:
args = construct_engine_params()
lp = LayoutPredictor(args)
source_img = './img/layout1.jpg'
img = cv2