计算机视觉/视频分析

自定义 AI 模型:使用 NVIDIA Triton 部署字符检测和识别模型

 

NVIDIA Triton Inference Server 通过使团队能够在任何基于 GPU 或 CPU 的基础设施上部署、运行和扩展经过训练的 ML 或 DL 模型,简化和标准化 AI 推理。它帮助开发人员在云端、本地、边缘和嵌入式设备上提供高性能推理。

nvOCDR 库已集成到 Triton 中进行推理。nvOCDR 库封装了用于光学字符检测和识别(OCD/OCR)的整个推理管道。该库使用在 TAO Toolkit 上训练的 OCDNet 和 OCRNet 模型。想要了解更多详细信息,请参阅 nvOCDR 文档

本文是关于使用 NVIDIA TAO 和预训练模型创建和部署自定义 AI 模型以准确检测和识别手写文本的系列文章的一部分。第一部分 解释了如何使用 TAO 对字符检测和识别模型进行训练和微调。本部分将引导您完成使用 NVIDIA Triton 部署模型的步骤。所提供的步骤可用于任何其他 OCR 任务。

使用 OCD/OCR 模型构建 Triton 样本

以下步骤显示了在 Triton 推理服务器中使用 Docker 图像构建和使用 OCD/OCR 模型的简单推荐方法。

步骤 1:准备 ONNX 模型

一旦你按照ocdnet.ipynbocrnet.ipynb要完成模型训练和导出,可以获得两个 ONNX 模型,例如ocdnet.onnxocrnet.onnx.(在ocdnet.ipynb 里,导出的 ONNX 命名为model_best.onnx在里面ocrnet.ipynb,导出的 ONNX 命名为best_accuracy.onnx

# bash commands
$ mkdir onnx_models
$ cd onnx_models
$ cp <ocd results dir>/export/model_best.onnx ./ocdnet.onnx
$ cp <ocr results dir>/export/best_accuracy.onnx ./ocrnet.onnx

字符列表文件,生成于ocrnet.ipynb,也是必需的:

$ cp <ocr DATA_DIR>/character_list ./

步骤 2:获取 nvOCDR 存储库

要获取 nvOCDR 存储库,请使用以下脚本:

$ git clone https://github.com/NVIDIA-AI-IOT/NVIDIA-Optical-Character-Detection-and-Recognition-Solution.git

步骤 3:构建 Triton 服务器 Docker 映像

Triton 服务器和客户端 Docker 镜像的构建过程可以通过运行相关脚本自动启动:

$ cd NVIDIA-Optical-Character-Detection-and-Recognition-Solution/triton

# bash setup_triton_server.sh [input image height] [input image width] [OCD input max batchsize] [DEVICE] [ocd onnx path] [ocr onnx path] [ocr character list path]
$ bash setup_triton_server.sh 1024 1024 4 0 ~/onnx_models/ocd.onnx ~/onnx_models/ocr.onnx ~/onnx_models/ocr_character_list

步骤 4:构建 Triton 客户端 Docker 镜像

使用以下脚本构建 Triton 客户端 Docker 映像:

$ cd NVIDIA-Optical-Character-Detection-and-Recognition-Solution/triton
$ bash setup_triton_client.sh

步骤 5:运行 nvOCDR Triton 服务器

在构建 Triton 服务器和 Triton client Docker 镜像后,创建一个容器并启动 Triton’服务器:

$ docker run -it --net=host --gpus all --shm-size 8g nvcr.io/nvidian/tao/nvocdr_triton_server:v1.0 bash

接下来,修改 nvOCDR-lib 的配置文件。nvOCDR-lib 可以支持高分辨率输入图像(4000 x 4000 或更大)。如果您的输入图像很大,您可以在 Triton 服务器容器中将配置文件更改为/opt/nvocdr/ocdr/triton/models/nvOCDR/spec.json,以支持高分辨率图像推断。

# to support high resolution images
$ vim /opt/nvocdr/ocdr/triton/models/nvOCDR/spec.json

   "is_high_resolution_input": true,
   "resize_keep_aspect_ratio": true,

这个resize_keep_aspect_ratio将自动设置为 Trueis_high_resolution_input为 True。如果您要推断分辨率较小的图像(例如 640 x 640 或 960 x 1280),您可以设置 is_high_resolution_input为 False。

在容器中,运行以下命令以启动 Triton 服务器:

$ CUDA_VISIBLE_DEVICES=<gpu idx> tritonserver --model-repository /opt/nvocdr/ocdr/triton/models/

步骤 6:发送推理请求

在一个单独的控制台中,从 Triton 客户端容器启动 nvOCDR 示例:

$ docker run -it --rm -v <path to images dir>:<path to images dir>  --net=host nvcr.io/nvidian/tao/nvocdr_triton_client:v1.0 bash

启动推理:

$ python3 client.py -d <path to images dir> -bs 1
The predicted output from OCDNet and OCRNet on a sample handwritten image, with bounding boxes around words such as ‘stairs’ and ‘rushed.’
图 1。OCDNet 和 OCRNet 对样本手写图像的预测输出

结论

NVIDIA TAO 5.0 介绍了光学字符检测(OCD)和光学字符识别(OCR)的几个功能和模型。本文介绍了自定义和微调预训练模型的步骤,以准确识别 IAM 数据集上的手写文本。该模型实现了 90%的字符检测准确率和 80%的字符识别准确率。文章中提到的所有步骤都可以从提供的 Jupyter 笔记本上运行,从而可以轻松地用最少的代码创建自定义的人工智能模型。

有关详细信息,请参阅:

 

Tags