跳转到主要内容
公告社区

docTR 加入 PyTorch 生态系统:从像素到数据,使用 PyTorch 和 docTR 构建识别管道

作者: 2024 年 12 月 18 日2025 年 5 月 3 日无评论
docTR logo

我们很高兴地宣布 docTR 项目已集成到 PyTorch 生态系统中!此次集成确保 docTR 符合 PyTorch 的标准和实践,为开发人员提供了一个可靠、社区支持的强大 OCR 工作流解决方案。

有关成为 PyTorch 生态系统项目的更多信息,请参阅 PyTorch 生态系统工具页面

关于 docTR

docTR 是一个 Apache 2.0 项目,由 Mindee 开发和分发,旨在帮助开发人员将 OCR 功能集成到应用程序中,无需任何先验知识。

为了快速有效地提取文本信息,docTR 采用两阶段方法

  • 首先,它执行文本检测以定位单词。
  • 然后,它进行文本识别以识别单词中的所有字符。

检测识别由用 PyTorch 编写的最新模型执行。要了解有关此方法的更多信息,您可以参考 docTR 文档

docTR 通过提供高性能的 OCR 功能,开箱即用,从而增强了 PyTorch 项目中的用户体验。其专门设计的模型对于常见用例几乎不需要进行微调,使开发人员能够快速集成高级文档分析功能。

本地安装

docTR 需要 Python >= 3.10,并支持 Windows、Mac 和 Linux。有关配备 M1 芯片的 MacBook 的必要依赖项,请参阅我们的 README

pip3 install -U pip
pip3 install "python-doctr[torch,viz]"

这将安装 docTR 以及最新版本的 PyTorch。

Note: docTR also provides docker images for an easy deployment, such as a part of Kubernetes cluster.

文本识别

现在,让我们在这个示例上尝试 docTR 的 OCR 识别

OCR sample

OCR 识别模型需要一个只包含一个单词的图像,并将输出预测的单词和置信度分数。您可以使用以下代码片段测试 docTR 的 OCR 功能

python
from doctr.io import DocumentFile
from doctr.models import recognition_predictor

doc = DocumentFile.from_images("/path/to/image")

# Load the OCR model
# This will download pre-trained models hosted by Mindee
model = recognition_predictor(pretrained=True)

result = model(doc)
print(result)

在这里,最重要的代码行是 model = recognition_predictor(pretrained=True)。这将加载默认的文本识别模型 crnn_vgg16_bn,但您可以通过 arch 参数选择其他模型。您可以查看 可用架构

在示例上运行时,识别预测器检索以下数据:[('MAGAZINE', 0.9872216582298279)]

Note: using the DocumentFile object docTR provides an easy way to manipulate PDF or Images.

文本检测

上一个示例是单个单词的裁剪。现在,如果图像上有多个单词,比如这个呢?

photo of magazines

在文本识别之前使用文本检测模型,以输出表示文本位置的分割图。之后,对每个检测到的补丁应用文本识别。

下面是仅运行检测部分的代码片段

from doctr.io import DocumentFile
from doctr.models import detection_predictor
from matplotlib import pyplot as plt
from doctr.utils.geometry import detach_scores
from doctr.utils.visualization import draw_boxes

doc = DocumentFile.from_images("path/to/my/file")
model = detection_predictor(pretrained=True)

result = model(doc)

draw_boxes(detach_scores([result[0]["words"]])[0][0], doc[0])
plt.axis('off')
plt.show()

在完整示例上运行它会产生以下结果

photo of magazines

与文本识别类似,detection_predictor 将加载默认模型(此处为 fast_base)。您也可以通过 arch 参数提供另一个模型来加载它。

完整实现

现在,让我们将这两个组件插入到同一个管道中。

方便的是,docTR 提供了一个完全满足我们需求的包装器

from doctr.io import DocumentFile
from doctr.models import ocr_predictor

doc = DocumentFile.from_images("/path/to/image")

model = ocr_predictor(pretrained=True, assume_straight_pages=False)

result = model(doc)
result.show()
photo of magazines

最后一行应显示一个 matplotlib 窗口,其中显示检测到的补丁。将鼠标悬停在它们上面会显示它们的内容。

您还可以用这个输出做更多事情,例如这样重构一个合成文档

import matplotlib.pyplot as plt

synthetic_pages = result.synthesize()
plt.imshow(synthetic_pages[0])
plt.axis('off')
plt.show()
black text on white

该管道是高度可定制的,您可以通过向 ocr_predictor 传递参数来修改检测或识别模型的行为。请参阅 文档 以了解更多信息。

结论

我们很高兴欢迎 docTR 加入 PyTorch 生态系统,它与 PyTorch 管道无缝集成,开箱即用地提供最先进的 OCR 功能。

通过让开发人员能够使用熟悉的工具快速从图像或 PDF 中提取文本,docTR 简化了复杂的文档分析任务,并增强了整体 PyTorch 体验。

我们邀请您探索 docTR GitHub 存储库,加入 Slack 上的 docTR 社区,并发送电子邮件至 contact@mindee.com 进行咨询或寻求合作机会。

我们可以一起继续突破文档理解的界限,并为 PyTorch 社区中的每个人开发更强大、更易于访问的工具。