作者: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。有关配备 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 社区中的每个人开发更强大、更易于访问的工具。