NIM简单实践-图像分割

项目背景

我正在学习一个图像分割的 Demo，使用 NVIDIA 提供的预训练大模型进行光学字符检测 (OCDNet) 和光学字符识别 (OCRNet)。这些模型专门为光学字符检测和识别设计，能够自动将图像中的字符进行分割和识别。

预训练模型介绍

• OCDNet (Optical Character Detection Network)：这是一个专门用于光学字符检测的预训练模型。该模型能够自动检测图像中的字符区域，为后续的字符识别提供基础。
• OCRNet (Optical Character Recognition Network)：OCRNet 是一个用于光学字符识别的预训练模型，能够识别由 OCDNet 检测到的字符区域，并将其转换为可读的文本。

实现步骤

1. 环境准备

   • 需要安装 NVIDIA 提供的工具包和相关依赖。
   • 创建并配置 Python 环境，确保能够加载和使用预训练模型。

2. 模型加载与初始化

   • 使用 NVIDIA 的预训练模型，通过调用 OCDNet 和 OCRNet 进行字符区域检测与识别。
   • 通过加载图像数据，模型能够自动进行字符区域的分割，识别图像中的文字信息。

3. 图像分割与识别

   • 首先通过 OCDNet 检测图像中的字符位置，分割出包含字符的区域。
   • 接着通过 OCRNet 对这些分割出的字符区域进行识别，输出最终的文本结果。

使用场景

• 图像中的文字检测：用于自动识别图像中的文字，如扫描文档、街道标识、产品包装等。
• 文字识别应用：在各类文档自动化处理中，可以帮助高效提取文本信息。

总结

NVIDIA 提供的 OCDNet 和 OCRNet 预训练模型极大简化了图像字符检测和识别的开发过程，适合快速搭建光学字符检测与识别的应用。在该 Demo 中，通过简单的步骤即可实现字符区域的检测与分割，帮助我们在图像处理中实现高效的文字识别功能。

代码如下：

import os
import sys
import uuid
import zipfileimport requests# NVAI endpoint for the ocdrnet NIM
nvai_url="https://ai.api.nvidia.com/v1/cv/nvidia/ocdrnet"
API_KEY_REQUIRED_IF_EXECUTING_OUTSIDE_NGC = "***"header_auth = f"Bearer {API_KEY_REQUIRED_IF_EXECUTING_OUTSIDE_NGC}"def _upload_asset(input, description):"""Uploads an asset to the NVCF API.:param input: The binary asset to upload:param description: A description of the asset"""assets_url = "https://api.nvcf.nvidia.com/v2/nvcf/assets"headers = {"Authorization": header_auth,"Content-Type": "application/json","accept": "application/json",}s3_headers = {"x-amz-meta-nvcf-asset-description": description,"content-type": "image/jpeg",}payload = {"contentType": "image/jpeg", "description": description}response = requests.post(assets_url, headers=headers, json=payload, timeout=30)response.raise_for_status()asset_url = response.json()["uploadUrl"]asset_id = response.json()["assetId"]response = requests.put(asset_url,data=input,headers=s3_headers,timeout=300,)response.raise_for_status()return uuid.UUID(asset_id)if __name__ == "__main__":"""Uploads an image of your choosing to the NVCF API and sends arequest to the Optical character detection and recognition model.The response is saved to a local directory.Note: You must set up an environment variable, NGC_PERSONAL_API_KEY."""if len(sys.argv) != 3:print("Usage: python test.py <image> <output_dir>")sys.exit(1)asset_id = _upload_asset(open(sys.argv[1], "rb"), "Input Image")inputs = {"image": f"{asset_id}", "render_label": True}asset_list = f"{asset_id}"headers = {"Content-Type": "application/json","NVCF-INPUT-ASSET-REFERENCES": asset_list,"NVCF-FUNCTION-ASSET-IDS": asset_list,"Authorization": header_auth,}response = requests.post(nvai_url, headers=headers, json=inputs)with open(f"{sys.argv[2]}.zip", "wb") as out:out.write(response.content)with zipfile.ZipFile(f"{sys.argv[2]}.zip", "r") as z:z.extractall(sys.argv[2])print(f"Output saved to {sys.argv[2]}")print(os.listdir(sys.argv[2]))

然后运行：

python .\ocdrnet.py .\ocdrnet_output_v1\food_02_.jpg .\ocdrnet_output_v1\

图片切割示例（分析中文不适合这个模型）