博客生态系统

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

作者 2024年12月18日2025年9月4日暂无评论
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 以获取搭载 M1 芯片的 MacBook 所需的依赖项。

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 社区中的每个人开发更强大、更易用的工具。