在这个项目中,我的目标是从科学论文图片中提取某些部分(标题、作者和摘要)。预期提取部分是科学论文中常见的部分,例如标题、摘要和作者。输入与最终结果。我的输入是将第一页纸转换成图像。最终结果是一个 txt 文件,其中包含标题、作者和摘要部分,如下图1和图2所示。我将使用 UNet 来了解在哪里可以找到这些部分,然后将训练学到的信息传递到 OCR 中。完整的项目可以在这里找到。
图1 要提取的论文首页(图片格式)
图2 提取得到的论文标题、作者和摘要
步骤说明
我从数据收集和科学论文开始。由于我只对标题、摘要和作者这三个部分感兴趣,所以我选择了多篇论文中的第一页。然后我将其转换pdf
为图像,因为我的 UNet 模型只接受图像。我已将数据集分别按 80/20 分为训练集和测试集。接下来,我编写了一个 python 脚本mask.py
来屏蔽图像。参见下图。
图3:右侧是科学论文的原始首页,左侧是相应论文的mask.py结果。
蒙版图像是一维图像,其中每个像素的值为 0(黑色)和 1(白色)。 0表示该像素不重要,1表示该像素重要。左侧图像的白色区域掩盖了重要的标题、摘要和作者部分。
我们将此掩模和原始图像传递给我们的 UNet 模型进行训练。 UNet 是一种流行的架构。您可以找到该架构的许多实现。我推荐Milesial的架构。如果您已经克隆了dagshub存储库,则无需克隆此存储库。您可以通过将工作目录更改为Unet-OCR/Pytorch-UNet
然后运行来构建模型,train.py
如下所示:
python train.py --epoch 6 --batch-size 1 --learning-rate 0.000001
显然您可以设置自己的参数。当然您也可以使用预训练模型MODEL.pth
,您将在运行时获得:
dvc pull -r origin
确保按照下面的安装说明将 DVC 源设置为我的