Press "Enter" to skip to content

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练)

让机器像人一样学习如何查看PDF

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第1张

在过去的一年里,我一直在努力解决一个商业案例,即使PDF文档也能被机器读取,至少从某种程度上提取出文档中的标题和内容,从而形成一些类似关系的数据结构。起初,我使用卷积神经网络(CNN)以及卷积神经网络(CNN)和循环神经网络(RNN)的组合来处理这个问题,使用文本和文本特征(字体、字号、字重等)来对文档的结构进行分类。这个框架由Rahman & Finin实现,并不仅推断文档结构,还使用长短期记忆(LSTM)对各个分割部分进行语义分类。我面临的问题是,我没有足够的时间和精力来准备和注释足够的数据,以使用相似的框架制作准确的模型。

为了让这个过程更简单,作为一个人组成的团队,我转向了计算机视觉!我不再尝试从PDF中提取文本和文本特征,然后通过对象检测以视觉方式推断文档结构,然后使用OCR对相应的推断进行转换(例如页面标题)和它们所关联的内容(例如标题的内容)。在本文中,我将向您展示如何使用Detectron2完成同样的工作!

Detectron2是Facebook人工智能研究院的下一代库,简化了构建计算机视觉应用程序的过程。在计算机视觉的各种可能用例(图像识别、语义分割、目标检测和实例分割)中,本文将为目标检测训练一个自定义的Detectron2模型,通过其使用边界框来与其他用例区分开来。训练完成后,我们将把应用程序docker化并部署到Heroku/AWS,同时探讨内存管理和批处理推断等其他主题,以帮助您根据自己的用例定制脚本和模型。

在阅读本文之前,提前掌握以下内容将会很有帮助:

  • 扎实的Python知识
  • 了解Django和Docker的部署过程(在第2部分中介绍)
  • 熟悉AWS

Detectron2安装

如果您是Mac或Linux用户,那很幸运!通过运行以下命令,这个过程将相对简单:

pip install torchvision && pip install "detectron2@git+https://github.com/facebookresearch/detectron2.git@v0.5#egg=detectron2"

请注意,这个命令将编译库,所以您需要等一会儿。如果您想安装带有GPU支持的Detectron2,请参考官方Detectron2的安装说明获取详细信息。

如果您是Windows用户,这个过程可能会有点麻烦,但我自己在Windows上成功完成了这个过程。

请按照Layout Parser包为Python提供的指示仔细操作(如果您不在意训练自己的Detectron2模型来进行PDF结构/内容推断,并希望依赖预注释的数据使用,这也是一个有用的包!这样做在时间上确实更友好,但是在特定的用例中,您会发现,您可以自己训练一个更准确且更小的模型,这有助于部署时的内存管理,我后面会讨论)。确保安装pycocotools和Detectron2,因为这个包将在装载、解析和可视化我们所需的数据格式(如COCO数据)的过程中提供帮助,用于训练Detectron2模型。

本地Detectron2安装将在本系列文章的第2部分中使用,因为我们将在本文的后面部分使用AWS EC2实例进行Detectron2的训练。

Detectron2 自定义训练 — 使用 LabelMe 进行标注

对于图像标注,我们需要两样东西:(1)需要标注的图像和(2)标注工具。首先将所有要标注的图像放入一个文件夹中,但如果您正在按照我的示例进行操作,并希望使用 PDF 图像,请组装一个 PDF 文件夹,并安装 pdftoimage 包:

pip install pdf2image

然后使用以下脚本将每个 PDF 页面转换为图像:

import osfrom pdf2image import convert_from_path# 将 input_dir 指定为 PDF 文件夹路径,例如:"C://Users//user//Desktop//pdfs"input_dir = "##"# 将 output_dir 指定为希望保存图像的文件夹路径"output_dir = "##"dir_list = os.listdir(input_dir)index = 0while index < len(dir_list):    images = convert_from_path(f"{input_dir}//" + dir_list[index])    for i in range(len(images)):        images[i].save(f'{output_dir}//doc' + str(index) +'_page'+ str(i) +'.jpg', 'JPEG')    index += 1

准备好包含图像的文件夹后,我们将使用 LabelMe 工具进行标注,请参阅安装说明。安装完成后,只需在命令行或终端中运行 labelme 命令,即可打开以下布局的窗口:

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第2张

在左侧选择“Open Dir”,打开包含图像的文件夹(我们将该文件夹命名为“train”)。LabelMe 将打开文件夹中的第一张图像,并允许您对每个图像进行标注。右键单击图像可以找到各种标注选项,例如使用创建多边形功能在图像中的指定对象周围点击每个点来创建多边形,或使用创建矩形功能在确保 90 度角度的同时捕捉对象。

确定边界框/多边形后,LabelMe 将要求输入标签。在下面的示例中,我为页面上的每个标题实例提供了标签。您可以使用多个标签,标识图像中的各种对象(对于 PDF 示例,可以是标题、表格、段落、列表等),但对于我的目的,我只是要识别标题,然后在模型推断后以算法方式将每个标题与其相应的内容关联起来(请参阅第二部分)。

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第3张

标注完成后,点击保存按钮,然后点击“Next Image”以标注给定文件夹中的下一张图像。Detectron2 在很少的数据下可以进行出色的检测推断,因此可以放心地标注多达约 100 张图像进行初步训练和测试,然后可以通过继续标注和训练来提高模型的准确性(请记住,在超过一个标签类别上训练模型会稍微降低准确性,需要更大的数据集来提高准确性)。

每张图像在 train 文件夹中标注完成后,将其中约 20% 的图像/标注对移动到一个名为 test 的单独文件夹中。

如果您熟悉机器学习,一个简单的经验法则是需要进行测试/训练/验证数据的分割(60-80% 训练数据,10-20% 验证数据,10-20% 测试数据)。对于这个目的,我们只需要进行 20% 的测试和 80% 的训练数据的分割。

Detectron2 自定义训练 — COCO 格式

现在我们有了标注文件夹,我们需要将 labelme 标注转换为 COCO 格式。您可以使用我这里的 labelme2coco.py 文件来完成。我已经从Tony607的仓库中重新调整了这个脚本,用于将多边形标注和任何矩形标注(因为初始脚本未能正确转换矩形标注为 COCO 格式)转换为 COCO 格式。

一旦您下载了labelme2coco.py文件,可以在终端中使用以下命令运行:

python labelme2coco.py path/to/train/folder

它会输出一个train.json文件。再次运行该命令来处理test文件夹,并在labelme2coco.py的第172行编辑,将默认的输出名称更改为test.json(否则会覆盖train.json文件)。

Detectron2定制训练- EC2

现在,繁琐的注释过程结束了,我们可以开始有趣的部分-训练!

如果您的计算机没有Nvidia GPU功能,我们需要使用AWS启动一个EC2实例。Detectron2模型可以在CPU上进行训练,但是如果您尝试此操作,您会发现它将花费非常长的时间,而使用基于GPU的实例上的Nvidia CUDA将在几分钟内训练模型。

首先,登录AWS控制台。登录后,在搜索栏中搜索EC2以进入EC2仪表板。从这里,点击屏幕左侧的”Instances”,然后点击”Launch Instances”按钮

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第4张

实例需要提供的最基本的详细信息是:

  • 一个名称
  • Amazon Machine Image (AMI) 指定软件配置。确保使用支持GPU和PyTorch功能的AMI,因为它将有CUDA所需的软件包和Detectron2所需的其他依赖项,如Torch。为了跟随本教程,请使用Ubuntu AMI。我使用的是Deep Learning AMI GPU PyTorch 2.1.0(Ubuntu 20.04)。
  • 实例类型 指定硬件配置。请参考此处的指南了解各种实例类型。我们希望使用性能优化的实例,如P系列或G系列的实例。我使用的是p3.2xlarge,它提供了我们需要的所有计算能力,特别是GPU功能。

请注意:从P系列启动的实例需要联系AWS客户服务增加配额(因为它们不允许普通用户立即访问高性能实例以避免相关成本)。如果您使用p3.2xlarge实例类型,则需要请求将配额增加到8个vCPU。

  • 指定密钥对(登录)。如果您还没有密钥对,请创建一个,并随意将其命名为p3key,就像我一样。

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第5张

  • 最后,配置存储。如果您使用与我相同的AMI和实例类型,则会看到默认存储为45GB。根据您的训练数据集大小,可以将其增加到60GB或更多以确保实例具有足够的空间来存储您的图像。

启动实例,并单击实例ID的超链接以在EC2控制台中查看。当实例在运行时,打开命令提示符窗口,使用以下命令通过SSH连接到EC2实例(确保替换粗体文本为(1)您的.pem密钥对的路径和(2)您的EC2实例的地址):

ssh -L 8000:localhost:8888 -i C:\path\to\p3key.pem ubuntu@ec2id.ec2region.compute.amazonaws.com

由于这是一个新的主机,请选择”yes”以接受以下消息:

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第6张

然后Ubuntu将与一个名为PyTorch的预装虚拟环境一起启动(来自AWS AMI)。通过以下两个命令激活venv并启动预安装的jupyter notebook:

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第7张

这将返回一些URL供您复制并粘贴到浏览器中。将其中带有localhost的复制到浏览器中,并将8888改为8000。这将带您进入一个类似于下图的Jupyter Notebook:

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第8张

我的github仓库中将Detectron2_Tutorial.ipynb文件上传到笔记本中。从这里,运行Installation标题下的代码行以完全安装Detectron2。然后,重新启动运行时以确保安装生效。

重新启动笔记本后,我们需要在开始训练过程之前上传一些额外的文件:

  • github仓库中的utils.py文件。这为.ipynb文件提供了Detectron2的配置细节(如果您对配置细节感兴趣,可以参考此处的文档)。此文件中还包括了一个在.ipynb文件中引用的plot_samples函数,但在两个文件中都被注释掉了。您可以取消注释并使用它来绘制训练数据,以便在进程中查看样本的可视化效果。请注意,您需要进一步安装cv2才能使用plot_samples功能。
  • 使用labelme2coco.py脚本制作的train.json和test.json文件。
  • 包含训练图像目录和测试图像目录的压缩文件(压缩目录只需上传一个项目到笔记本中;您可以将labelme注释文件保留在目录中,这不会影响训练)。上传这两个压缩文件后,在笔记本中点击右上角的(1) New然后(2) Terminal,打开一个终端,并使用以下命令解压缩每个文件,创建一个单独的Train和Test图像目录:
! unzip ~/train.zip -d ~/! unzip ~/test.zip -d ~/

最后,在.ipynb文件中的Training部分运行笔记本单元格。最后一个单元格将输出类似以下的响应:

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第9张

这将显示用于训练的图像数量,以及在训练数据集中注释的实例数(这里,在训练之前找到了470个“title”类别的实例)。然后,Detectron2会序列化数据并按照配置文件(utils.py)指定的批次加载数据。

一旦开始训练,您将看到Detectron2打印事件:

训练和部署使用PDF文档进行目标检测的自定义Detectron2模型(第一部分:训练) 四海 第10张

这会提供一些信息,例如:估计的训练剩余时间、Detectron2执行的迭代次数,以及最重要的用于监视准确性的total_loss,它是其他损失计算的指标,表示模型在单个示例上的预测有多差。如果模型的预测完美无缺,则损失为零;否则,损失较大。如果模型并非完美无缺,不必担心!我们始终可以添加更多的注释数据来提高模型的准确性,或者在我们的应用程序中使用具有高得分(表示模型对推断是否准确的自信程度)的最终训练模型的推断结果。

完成后,笔记本中将创建一个名为output的目录,其中包含与训练事件和指标相关的文件,记录模型的检查点的文件,最后是一个名为model_final.pth的.pth文件。这是保存和训练的Detectron2模型,现在可以在部署的应用程序中用于进行推断!在关闭或终止AWS EC2实例之前,请确保下载此文件。

现在我们已经有了model_final.pth,我们将在第二部分的部署文章中详细介绍一个使用机器学习的应用程序的部署过程,同时提供一些关键提示,以帮助您更高效地完成这个过程。

除非另有说明,本文中使用的所有图片均为作者所拍。

Leave a Reply

Your email address will not be published. Required fields are marked *