由 Olivier Dulcy & Sebastian Olivera,Mindee 编写

docTR logo

我们非常高兴地宣布,docTR 项目已集成到 PyTorch 生态系统中!这次集成确保 docTR 与 PyTorch 的标准和规范保持一致,为开发者提供了一个可靠、社区支持的强大 OCR 工作流程解决方案。

想了解更多关于成为 PyTorch 生态系统项目的意义,请参阅 PyTorch 生态系统工具页面。

关于 docTR

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

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

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

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

docTR 通过提供开箱即用的高性能 OCR 功能,增强了 PyTorch 项目的用户体验。其特别设计的模型对常见用例仅需最小化或无需微调,使开发者能够快速集成高级文档分析功能。

本地安装

docTR 需要 Python >= 3.10,并支持 Windows、Mac 和 Linux。请参阅我们的 README 以获取 MacBook M1 芯片所需的依赖项。

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 仓库,加入 docTR 的 Slack 社区,或通过 contact@mindee.com 联系,以获取咨询或合作机会。

一起,我们可以继续拓展文档理解的边界,并为 PyTorch 社区中的每个人开发更强大、更易用的工具。