跳转到主要内容
博客生态系统

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

作者： 2024 年 12 月 18 日2025 年 9 月 4 日暂无评论
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 社区中的每个人开发更强大、更易于访问的工具。