PaddleOCR文字检测模型的预处理
目录
- 简介
- 1 预处理流程初探
- 2 算子处理过程
-
- 2.1 DecodeImage
- 2.2 DetResizeForTest
-
- 2.2.1 构造方法
- 2.2.2 调用方法
- 2.3 NormalizeImage
- 2.4 ToCHWImage
- 2.5 KeepKeys
- 3 预处理结果
- 4 实际部署时的预处理
- 5 总结
简介
导出ONNX格式的模型后，在部署模型时，需要对模型的输入进行预处理，转换成符合模型输入维度的张量；模型输出张量结果后，也需要通过后处理，将张量转换成需要的预测结果。
在Java中尝试通过OnnxRuntime部署PaddleOCR文字检测模型时，发现上述预处理和后处理过程还是有些东西值得学习的，于是花了些时间学习了项目源码，经过一些实践调试后，总结出了下面的预处理流程步骤。
整个内容比较新手向，适合对文字检测模型的图像预处理不了解的新手。
本文使用的文字检测模型相关信息如下：
| 模型名称 | 对应配置文件 |
|---|---|
| ch_ppocr_mobile_v2.0_det_train | det_mv3_db.yml |
注意：在项目源码中，模型推理过程的预处理实际上有两个：
-
一个是在训练过程中，为了快速检验模型训练效果，在 “训练模型” 上进行推理时所做的预处理。
-
另一个是实际部署模型时，将模型转换成 “推理模型” 后，在 “推理模型” 上进行推理时所做的预处理。
上述的两个预处理的步骤方法基本上是一致的，只是细节上有些差别。
本文的内容，是从第一种预处理的流程开始的。其实还是一开始没太注意这个细节，所以在后面补充说明了 实际部署时的预处理方法 。
1 预处理流程初探
根据官方文档的介绍，可以通过项目中的tools/infer_det.py直接在文字检测模型的 训练模型 上进行推理。
在tools/infer_det.py中可以看到如下代码：
# create data ops
transforms = []
for op in config['Eval']['dataset']['transforms']:
op_name = list(op)[0]
if 'Label' in op_name:
continue
elif op_name == 'KeepKeys':
op[op_name]['keep_keys'] = ['image', 'shape']
transforms.append(op)
ops = create_operators(transforms, global_config)
这一步目的是读取预处理算子的配置参数到transforms列表中，然后通过方法create_operators()，构造算子类的对象并添加到列表ops。
从上面的代码中可以看到，算子的相关参数在Eval.dataset.transforms下，查看模型的对应配置文件，相关参数如下：
Eval:
dataset:
transforms:
- DecodeImage: # load image
img_mode: BGR
channel_first: False
- DetLabelEncode: # Class handling label
- DetResizeForTest:
image_shape: [736, 1280]
- NormalizeImage:
scale: 1./255.
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]
order: 'hwc'
- ToCHWImage:
- KeepKeys:
keep_keys: ['image', 'shape', 'polys', 'ignore_tags']
因为模型训练和推理使用的是同一个配置文件，而推理和训练在具体细节上又有所区别。因此，上面的代码中读取算子配置参数时，做了两个调整：
- 跳过了
DetLabelEncode这个算子； - 将
KeepKeys算子的参数keep_keys的值修改成了['image', 'shape']
因此从上述代码和配置文件中，可以得出结论，文字检测模型在推理时，预处理共包含5个算子：
- DecodeImage
- DetResizeForTest
- NormalizeImage
- ToCHWImage
- KeepKeys
继续往下看代码：
for file in get_image_file_list(config['Global']['infer_img']): # 从配置文件读取图片路径列表
logger.info("infer_img: {}".format(file))
with open(file, 'rb') as f:
img = f.read()
data = {'image': img}
batch = transform(data, ops) # 预处理方法
images = np.expand_dims(batch[0], axis=0)
shape_list = np.expand_dims(batch[1], axis=0)
images = paddle.to_tensor(images) # 飞桨框架的张量转换
preds = model(images) # 模型调用
post_result = post_process_class(preds, shape_list) # 对模型输出张量的后处理
boxes = post_result[0]['points']
# 写入json和可视化结果
dt_boxes_json = []
for box in boxes:
tmp_json = {"transcription": ""}
tmp_json['points'] = box.tolist()
dt_boxes_json.append(tmp_json)
otstr = file + "\t" + json.dumps(dt_boxes_json) + "\n"
fout.write(otstr.encode())
src_img = cv2.imread(file)
draw_det_res(boxes, config, src_img, file)
可以看到一开始的循环，读取了配置文件中的需要进行推理的图片路径列表。
随后，在循环内对每张图片进行了预处理，然后输入到模型进行了推理预测，最后将模型输出写入了json，并绘制到图片上给出了可视化结果。
预处理的部分在上述代码的第6行，此处调用了transform()方法，该方法参数如下：
| 参数 | 类型 | 说明 |
|---|---|---|
data |
dict |
key是'image'；value是从图片读取的bytes类型数据 |
ops |
list |
包含算子对象的列表 |
跳转到该方法，代码如下：
def transform(data, ops=None):
""" transform """
if ops is None:
ops = []
for op in ops:
data = op(data)
if data is None:
return None
return data
方法很简单，就是循环调用算子对象列表中的每一个算子，依次对数据进行处理。从这里可以看出，在前面的配置文件中，算子的先后顺序也是有意义的，因为每个算子处理后的输出都是下个算子处理的输入。
此处不能直接跳转到算子类的__call__()方法，经过调试，发现预处理算子类的定义在ppocr/data/imaug/operators.py中，接下来按照顺序，具体了解每个算子是如何处理的。
2 算子处理过程
2.1 DecodeImage
因为是第一个算子，此时的输入data是如前所述的bytes类型的字节值，因此第一个算子的任务，就是在此基础上进行解码，得到图像的像素矩阵。该算子类的代码如下：
class DecodeImage(object):
""" decode image """
def __init__(self, img_mode='RGB', channel_first=False, **kwargs):
self.img_mode = img_mode
self.channel_first = channel_first
def __call__(self, data):
img = data['image']
if six.PY2:
assert type(img) is str and len(
img) > 0, "invalid input 'img' in DecodeImage"
else:
assert type(img) is bytes and len(
img) > 0, "invalid input 'img' in DecodeImage"
img = np.frombuffer(img, dtype='uint8')
img = cv2.imdecode(img, 1)
if img is None:
return None
if self.img_mode == 'GRAY':
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
elif self.img_mode == 'RGB':
assert img.shape[2] == 3, 'invalid shape of image[%s]' % (img.shape)
img = img[:, :, ::-1]
if self.channel_first:
img = img.transpose((2, 0, 1))
data['image'] = img
return data
解码的方法很简单，通过numpy的frombuffer()方法，将数据转换成uint8类型的矩阵，然后直接使用openCV的imdecode()方法进行解码，解码后的图像颜色格式为BGR，矩阵维度按 HWC 顺序排列。
随后根据img_mode的参数值，将图像转换成灰度图、RGB格式或保持BGR格式不变；根据channel_first参数决定是否将channel变换到第一个维度。
该算子总结如下：
作用：将图片从bytes类型解码成所需格式的图像矩阵
| 参数 | 参数值 | 参数说明 |
|---|---|---|
| img_mode | ‘BGR’ / ‘GRAY’ / ‘RGB’ | 图像颜色格式 |
| channel_first | True / False | 是否将channel变换到第一个维度 |
2.2 DetResizeForTest
这个算子的任务是将图片进行缩放，想象中代码应该很简短，但实际查看源码后发现，其中内容还不少。
2.2.1 构造方法
先来看看这个算子类的构造方法：
def __init__(self, **kwargs):
super(DetResizeForTest, self).__init__()
self.resize_type = 0
if 'image_shape' in kwargs:
self.image_shape = kwargs['image_shape']
self.resize_type = 1
elif 'limit_side_len' in kwargs:
self.limit_side_len = kwargs['limit_side_len']
self.limit_type = kwargs.get('limit_type', 'min')
elif 'resize_long' in kwargs:
self.resize_type = 2
self.resize_long = kwargs.get('resize_long', 960)
else:
self.limit_side_len = 736
self.limit_type = 'min'
在PaddleOCR中推荐的文字检测算法是DB算法，也就是本文选取模型采用的算法，但实际上在项目内还支持了EAST和SAST两种检测算法。使用不同的算法或者在不同的训练数据下，图片的输入大小也有一定差别，因此在DetResizeForTest这个算子内包含了不同类型的缩放方法。
而这里源码处理的方式也比较简单：在构造算子对象时，直接从参数列表中判断是否包含某个参数，来决定此时的缩放类别，并赋值到数据成员self.resize_type中。
从前面的配置文件可以看到，在使用 训练模型 进行快速推理时，配置参数为image_shape: [736, 1280]，此时对象中的数据成员self.resize_type的值应该为1。
由不同参数决定的缩放类型resize_type具体代表了什么含义呢？从调用方法上可以一探究竟。
2.2.2 调用方法
以下是调用方法的代码：
def __call__(self, data):
img = data['image']
src_h, src_w, _ = img.shape
if self.resize_type == 0:
img, [ratio_h, ratio_w] = self.resize_image_type0(img)
elif self.resize_type == 2:
img, [ratio_h, ratio_w] = self.resize_image_type2(img)
else:
img, [ratio_h, ratio_w] = self.resize_image_type1(img)
data['image'] = img
data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
return data
和想象中一样，根据不同的缩放类别，调用对应方法。这里需要注意的是，缩放方法给出了两个返回值。
-
第一个值
img，是缩放后的图像； -
第二个值
[ratio_h, ratio_w]，是高、宽缩放比例的列表
同时，算子返回的结果中则增加了key为'shape'的1X4的矩阵，存储了图像的原始高、宽和对应缩放比例。此处存储缩放的shape数据是为了在后处理的过程中对图像进行还原。
下面依次来看看不同缩放方法具体是怎么操作的。
1. 当self.resize_type为0时
def resize_image_type0(self, img):
limit_side_len = self.limit_side_len
h, w, _ = img.shape
if self.limit_type == 'max':
if max(h, w) > limit_side_len:
if h > w:
ratio = float(limit_side_len) / h
else:
ratio = float(limit_side_len) / w
else:
ratio = 1.
else:
if min(h, w) < limit_side_len:
if h < w:
ratio = float(limit_side_len) / h
else:
ratio = float(limit_side_len) / w
else:
ratio = 1.
resize_h = int(h * ratio)
resize_w = int(w * ratio)
resize_h = max(int(round(resize_h / 32) * 32), 32)
resize_w = max(int(round(resize_w / 32) * 32), 32)
try:
if int(resize_w) <= 0 or int(resize_h) <= 0:
return None, (None, None)
img = cv2.resize(img, (int(resize_w), int(resize_h)))
except:
print(img.shape, resize_w, resize_h)
sys.exit(0)
ratio_h = resize_h / float(h)
ratio_w = resize_w / float(w)
return img, [ratio_h, ratio_w]
此时的配置参数为limit_side_len和limit_type，从代码上来看，limit_side_len实际上是定义一个图片边长的最值，根据limit_type来确定定义的是最大值还是最小值。
这种方法最后将把图片宽高缩放成均在参数限制内的、32的整数倍（这里取32是模型结构决定的）。
2. 当self.resize_type为1时
def resize_image_type1(self, img):
resize_h, resize_w = self.image_shape
ori_h, ori_w = img.shape[:2]
ratio_h = float(resize_h) / ori_h
ratio_w = float(resize_w) / ori_w
img = cv2.resize(img, (int(resize_w), int(resize_h)))
return img, [ratio_h, ratio_w]
此时的配置参数为image_shape，为缩放后图片的宽高。
这种方法直接调用了openCV的resize()方法进行缩放，然后计算了缩放比例。
3. 当self.resize_type为2时
def resize_image_type2(self, img):
h, w, _ = img.shape
resize_w = w
resize_h = h
if resize_h > resize_w:
ratio = float(self.resize_long) / resize_h
else:
ratio = float(self.resize_long) / resize_w
resize_h = int(resize_h * ratio)
resize_w = int(resize_w * ratio)
max_stride = 128
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
img = cv2.resize(img, (int(resize_w), int(resize_h)))
ratio_h = resize_h / float(h)
ratio_w = resize_w / float(w)
return img, [ratio_h, ratio_w]
此时的配置参数为resize_long，为缩放后图片的最长边的值。
这种方法首先是从图片宽高中找出最长边，计算缩放比例，然后保持图片的宽高比不变进行缩放。
但到这里并没有结束，接下来代码中定义了一个名为max_stride的变量，值为128，然后在此基础上做了一个计算：
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
其实就是将边长做了一个向上取max_stride整数倍的处理。
随后根据最终取整的宽高值计算了缩放比例，返回结果。
从项目配置中来看，这种缩放方法主要是用于SAST这个检测算法的预处理中，因此猜测max_stride变量的值以及取整的操作和具体算法有关联，有兴趣的读者可以阅读 SAST paper 进一步研究。
最后，该算子总结如下：
作用：对图像进行缩放，在返回结果中添加缩放比例
| resize_type（缩放类型） | 缩放类型说明 | 参数 | 参数值 | 参数说明 |
|---|---|---|---|---|
| 0 | 在限定的边长范围内，缩放边长到32的整数倍 | limit_side_len | 736 | 边长最值 |
| 0 | limit_type | min / max | 边长限值为最大值还是最小值 | |
| 1 | 根据参数缩放图像到指定宽高 | image_shape | [736, 1280] | 缩放后的图像高、宽 |
| 2 | 根据最长边确定比例，缩放边长到128的整数倍 | resize_long | 960 | 最长边缩放长度 |
注：该算子参数可以为None，从代码上可以看到，此时resize_type为0，limit_side_len为736，limit_type为min。
2.3 NormalizeImage
从配置参数上看：
NormalizeImage:
scale: 1./255.
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]
order: 'hwc'
这里使用的图像归一化是常见的方法，先乘上scale进行线性变换，再减去对应通道的平均值，最后除以对应通道的标准差。算子的调用方法如下：
def __call__(self, data):
img = data['image']
from PIL import Image
if isinstance(img, Image.Image):
img = np.array(img)
assert isinstance(img,
np.ndarray), "invalid input 'img' in NormalizeImage"
data['image'] = (
img.astype('float32') * self.scale - self.mean) / self.std
return data
该算子总结如下：
作用：对图像进行归一化
| 参数 | 参数值 | 参数说明 |
|---|---|---|
| scale | 1./255. | 线性变换参数 |
| mean | [0.485, 0.456, 0.406] | BGR三通道对应平均值 |
| std | [0.229, 0.224, 0.225] | BGR三通道对应标准差 |
| order | ‘hwc’ | 图像矩阵维度顺序 |
2.4 ToCHWImage
这个算子非常简单，代码如下：
class ToCHWImage(object):
""" convert hwc image to chw image
"""
def __init__(self, **kwargs):
pass
def __call__(self, data):
img = data['image']
from PIL import Image
if isinstance(img, Image.Image):
img = np.array(img)
data['image'] = img.transpose((2, 0, 1))
return data
即将图像矩阵维度变换为 CHW ，没有配置参数。
2.5 KeepKeys
最后一个算子也很简单，直接看源码：
class KeepKeys(object):
def __init__(self, keep_keys, **kwargs):
self.keep_keys = keep_keys
def __call__(self, data):
data_list = []
for key in self.keep_keys:
data_list.append(data[key])
return data_list
其实就是将data从dict类型，变成了list。个人认为叫做RemoveKeys更合适，这个算子也没有参数。
3 预处理结果
为了便于描述，将代码再贴一遍（更为完整的部分在最前）：
batch = transform(data, ops)
images = np.expand_dims(batch[0], axis=0)
shape_list = np.expand_dims(batch[1], axis=0)
images = paddle.to_tensor(images)
经过上面的一系列处理，transform()方法返回了一个list结果batch。
batch中的第一个元素为图像矩阵，颜色格式为BGR，维度顺序依次为CHW ；
batch中的第二个元素为图像缩放数据的1X4矩阵，分别为：原始高、原始宽、高度缩放比例和宽度缩放比例。
可以看到在转换成张量前，两个元素均扩充了一个维度，按照上文给出的训练模型推理时的配置参数，最终输入到模型的张量维度为[1, 3, 736, 1280]
4 实际部署时的预处理
根据项目导出的可用于部署的模型来看，实际部署推理时，文字检测模型运行的是项目下tools/infer/predict_det.py这个脚本，脚本内定义了一个TextDetector的类，其构造方法中包含了预处理的配置列表，如下：
pre_process_list = [{
'DetResizeForTest': None
}, {
'NormalizeImage': {
'std': [0.229, 0.224, 0.225],
'mean': [0.485, 0.456, 0.406],
'scale': '1./255.',
'order': 'hwc'
}
}, {
'ToCHWImage': None
}, {
'KeepKeys': {
'keep_keys': ['image', 'shape']
}
}]
很明显和上文的预处理基本一致，首先是少了DecodeImage这个图像解码的过程，因为在实际部署时，这个步骤将在把图像输入到文字检测模型前完成。
另一个变化是DetResizeForTest的参数为None，采用了和上文预处理中不同的缩放方法（具体参考上文2.2节）。
最后，KeepKeys的参数相比上文预处理的参数虽然少了两个，但算子内实际只用到了这两个参数，可以理解为没有变化。
5 总结
其实整体学习研究一遍后，发现整个预处理过程并不复杂，处理方法也比较常见，同时也对整个项目有了更深入的了解，接下来将尝试继续学习文字检测模型的后处理流程。