OpenVINO+DBNet实时OCR文本检测

OpenVINO+DBNet实时OCR文本检测

1. 简介

DBNet算法是今年来非常流行的文本检测算法，该算法在众多的实际场景中具有很好的检测性能，如百度的ppocr也是使用了DBNet作为文本检测阶段的算法。
更详细的介绍：
（1）知乎-DBNet阅读笔记
（2）论文：Real-time Scene Text Detection with Differentiable Binarization
（3）作者实现的代码
（4）WenmuZhou实现的代码-推荐

2.模型结构

模型结构：经典的FPN网络+后处理

3.模型转换

为了演示，我白piao一个pth模型(未完全训练好的模型，拒接白piao)，我使用了BaofengZan，先感谢BaofengZan大佬的模型。本人在BaofengZan/DBNet.pytorch基础上略微修改实现的DBNet-OpenVINO, one more,欢迎star，star，star哈…

（1） 第一步是讲pytorch模型转为onnx模型

# 来源：https://github.com/BaofengZan/DBNet.pytorch
def export_onnx(self):
    img = torch.zeros((1, 3, 480, 720))#.cuda()  # image size(1,3,320,192) iDetection

    y = self.model(img)
    try:
        import onnx

        print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
        f = ("./onnx-xml/dbnet.onnx") # filename
        #self.model.fuse()  # only for ONNX
        torch.onnx.export(self.model, img, f, verbose=False, opset_version=12, input_names=['images'],
                            output_names=['classes', 'boxes'] if y is None else ['output'])
        # Checks
        onnx_model = onnx.load(f)  # load onnx model
        onnx.checker.check_model(onnx_model)  # check onnx model
        print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model
        print('ONNX export success, saved as %s' % f)
    except Exception as e:
        print('ONNX export failure: %s' % e)

（2） 验证onnx模型是否转换正确
这一点很重要，如果转换不正确，都是白搭！

     # 此处省略...
    def init_onnx_session(self):
        self.ort_session = ort.InferenceSession("./onnx-xml/dbnet.onnx")
     # 此处省略...
    preds = self.ort_session.run(None, {'images': tensor.cpu().numpy()})

（3） 直接使用onnx模型进行openvino推理

    # 此处省略...
    self.ie = IECore()
    self.dbnet = self.ie.read_network(model="./onnx-xml/dbnet.onnx")

    self.input_blob = next(iter(self.dbnet.input_info))
    self.out_blob = next(iter(self.dbnet.outputs))
    self.exec_net = self.ie.load_network(network=self.dbnet, device_name="CPU")

    # Read and pre-process input images
    self.b, self.c, self.h, self.w = self.dbnet.input_info[self.input_blob].input_data.shape
    # 此处省略...
    preds = self.exec_net.infer(inputs={self.input_blob: tensor})

也可将onnx转为ir模型

python3.6 /opt/intel/openvino/deployment_tools/model_optimizer/mo_onnx.py \
                        -m ./onnx-xml/dbnet.onnx \
                        -o ./onnx-xml/ \
                        --log_level ERROR \
                        --data_type FP32

• 不包含预处理和后处理时间（不同的CPU推理时间会存在差异）

4.结果

输入：

输出二值化图：

后处理：

参考：
（1）知乎-DBNet阅读笔记
（2）论文：Real-time Scene Text Detection with Differentiable Binarization
（3）作者实现的代码
（4）WenmuZhou实现的代码-推荐
（5）BaofengZan
（6）DBNet-OpenVINO, one more