项目背景
我正在学习一个图像分割的 Demo,使用 NVIDIA 提供的预训练大模型进行光学字符检测 (OCDNet) 和光学字符识别 (OCRNet)。这些模型专门为光学字符检测和识别设计,能够自动将图像中的字符进行分割和识别。
预训练模型介绍
- OCDNet (Optical Character Detection Network):这是一个专门用于光学字符检测的预训练模型。该模型能够自动检测图像中的字符区域,为后续的字符识别提供基础。
- OCRNet (Optical Character Recognition Network):OCRNet 是一个用于光学字符识别的预训练模型,能够识别由 OCDNet 检测到的字符区域,并将其转换为可读的文本。
实现步骤
-
环境准备
- 需要安装 NVIDIA 提供的工具包和相关依赖。
- 创建并配置 Python 环境,确保能够加载和使用预训练模型。
-
模型加载与初始化
- 使用 NVIDIA 的预训练模型,通过调用
OCDNet
和OCRNet
进行字符区域检测与识别。 - 通过加载图像数据,模型能够自动进行字符区域的分割,识别图像中的文字信息。
- 使用 NVIDIA 的预训练模型,通过调用
-
图像分割与识别
- 首先通过
OCDNet
检测图像中的字符位置,分割出包含字符的区域。 - 接着通过
OCRNet
对这些分割出的字符区域进行识别,输出最终的文本结果。
- 首先通过
使用场景
- 图像中的文字检测:用于自动识别图像中的文字,如扫描文档、街道标识、产品包装等。
- 文字识别应用:在各类文档自动化处理中,可以帮助高效提取文本信息。
总结
NVIDIA 提供的 OCDNet
和 OCRNet
预训练模型极大简化了图像字符检测和识别的开发过程,适合快速搭建光学字符检测与识别的应用。在该 Demo 中,通过简单的步骤即可实现字符区域的检测与分割,帮助我们在图像处理中实现高效的文字识别功能。
代码如下:
import os
import sys
import uuid
import zipfile
import 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 a
request 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\
图片切割示例(分析中文不适合这个模型)