使用神经网络提取PDF表格工具来了，支持图片，关键是能白嫖谷歌GPU资源

贾浩楠 发自 凹非寺
量子位 报道 | 公众号 QbitAI

“表哥表姐”们还在为大量PDF文件中的表格发愁吗&＃xff1f;

百度一下&＃xff0c;网上有大量提取PDF表格的工具&＃xff0c;但是&＃xff0c;它们都只支持文本格式PDF。

但扫描生成的图片表格怎么办&＃xff1f;

别着急&＃xff0c;一种使用深度神经网络识别提取表格的开源工具可以帮助你。

兼容图片、高准确率、还不占用本地运算资源&＃xff0c;如此实用的工具值得你拥有。

测试实例

如果在输入的PDF文件中检测的表格&＃xff0c;模型会在边界框&＃xff08;bounding box&＃xff09;标出表格边框&＃xff1a;

然后&＃xff0c;表格数据会被转化为Panda数据框架&＃xff0c;方便后续处理&＃xff1a;

怎么样&＃xff0c;是不是很实用&＃xff1f;那这个工具如何使用呢&＃xff1f;

使用姿势详解

神经网络算法&＃xff0c;还不占用本地运算资源&＃xff1f;

对&＃xff0c;你没听错&＃xff0c;这个工具的所有代码都可以在谷歌Colab上运行。也就是说你可以利用Colab云端资源完成训练和推理&＃xff0c;无需本地安装。

这个工具使用到的深度神经网络是Keras-RetinaNet&＃xff0c;首先要在Colab上安装Keras-RetinaNet&＃xff0c;通过一下一行代码就可以完成&＃xff1a;

git clone https://github.com/fizyr/keras-retinanet


同时需要安装必要的库&＃xff1a;

pip install .
python setup.py build_ext — inplace


训练Keras-RetinaNet识别表格

首先要构建或一个训练使用的数据库。

这里要用到Colab工具PDF2Img&＃xff0c;将PDF文件转为JPG格式供算法学习。

然后将转好的图片保存在Images文件夹中。接下来需要手动给这些训练数据打标签&＃xff0c;这里推荐使用在线标记工具makesense.ai。

将XML注释文件保存在注释文件夹中&＃xff0c;并创建用于培训和测试的PDF文件列表&＃xff0c;将该列表导入train.txt和test.txt中。

接下来&＃xff0c;克隆Github项目https://github.com/ferrygun/PDFTableExtract&＃xff0c;并安装beautifulsoup。

运行以下命令以将PASCALVOC格式转换为Keras-RetinaNet所需的格式&＃xff1a;

python build_logos.py


运行上述命令后&＃xff0c;会得到retinanet_classes.csv&＃xff0c;retinanet_test.csv和retinanet_train.csv。

在retinanet_classses.csv中&＃xff0c;由于只识别PDF文档中的表&＃xff0c;所以只有到一个class&＃xff0c;即class 0。

如果你在打标签的过程中加入页眉页脚等标签&＃xff0c;相应能得到多个class。

完整的文件和文件夹结构的列表&＃xff1a;

然后&＃xff0c;将retinanet_classes.csv&＃xff0c;retinanet_test.csv&＃xff0c;retinanet_train.csv&＃xff0c;train.txt和test.txt 导入keras-retinanet的根文件夹中&＃xff1a;

接下来&＃xff0c;运行Colab TrainOCR&＃xff0c;可以根据要训练的JPG文件数量来调整训练的epoch数量。

训练完成后&＃xff0c;就会得到权重文件output.h5&＃xff0c;下载此文件并将其保存到本地主机。后面将使用该文件来运行测试。

这里需要注意的一点是&＃xff0c;在Colab中&＃xff0c;已将文件上传到Git并进行了Git克隆。

运行测试

运行测试前&＃xff0c;还需要安装处理PDF文件必要的库。

PyPDF2是一个python工具库&＃xff0c;能够提取文档信息&＃xff0c;裁剪页面等。

使用以下命令安装此库&＃xff1a;

pip install PyPDF2


Camelot是专门用于解析PDF页面表格的python库。

使用以下命令安装此库&＃xff1a;

pip install camelot-py[cv]


PDF2IMG是将Python转换为PIL Image对象的python库。

使用以下命令安装此库&＃xff1a;

pip install pdf2imag


在运行预测之前&＃xff0c;需要使用从训练中获得的权重文件output.h5加载模型&＃xff0c;并从retinanet_classes.csv定义标签class 0。

model_path &＃61; ‘output.h5’
model &＃61; models.load_model(model_path, backbone_name&＃61;’resnet50’)
labels &＃61; ‘retinanet_classes.csv’
LABELS &＃61; open(labels).read().strip().split(‘\n’)
LABELS &＃61; {int(L.split(‘,’)[1]): L.split(‘,’)[0] for L in LABELS}
print(LABELS)
{0: ‘tabel’}


接下来运行测试

import numpy as np
import cv2
import matplotlib.pyplot as pltimage_path &＃61; imgfname
#image &＃61; cv2.imread(image_path)
image &＃61; read_image_bgr(image_path)
image &＃61; cv2.cvtColor(image, cv2.COLOR_BGR2RGB)output &＃61; image.copy()
output &＃61; cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
print(output.shape) # row (height) x column (width) x color (3)image &＃61; preprocess_image(image)
(image, scale) &＃61; resize_image(image)
image &＃61; np.expand_dims(image, axis&＃61;0)# detect objects in the input image and correct for the image scale
(boxes, scores, labels) &＃61; model.predict_on_batch(image)
boxes /&＃61; scaleconfidence &＃61; 0.2
from google.colab.patches import cv2_imshow
import matplotlib.pyplot as plt
label_out &＃61; []
result &＃61; ""# loop over the detections
for (box, score, label) in zip(boxes[0], scores[0], labels[0]):# filter out weak detectionsif score plt.imshow(output)
plt.show()


如果检测的表格&＃xff0c;模型会在边界框&＃xff08;bounding box&＃xff09;标出表格边框&＃xff1a;

边界框坐标&＃xff08;x1&＃xff0c;y1&＃xff0c;x2&＃xff0c;y2&＃xff09;&＃xff0c;将table_area输入到Camelot read_pdf函数中&＃xff0c;table_area是已标准化的边界框。

interesting_areas&＃61;[]output &＃61; [[x1, y1, x2, y2]]
for x in output:[x1, y1, x2, y2] &＃61; bboxes_pdf(img, pdf_page, x)bbox_camelot &＃61; [",".join([str(x1), str(y1), str(x2), str(y2)])][0] # x1,y1,x2,y2 where (x1, y1) -> left-top and (x2, y2) -> right-bottom in PDF coordinate space#print(bbox_camelot)interesting_areas.append(bbox_camelot)print(interesting_areas)
output_camelot &＃61; camelot.read_pdf(filepath&＃61;pdf_file, pages&＃61;str(pg), flavor&＃61;"stream", table_areas&＃61;interesting_areas
)output_camelot[0].df


这样就可以将表格数据转化为Panda数据框架&＃xff0c;方便后续处理&＃xff1a;

图片怎么办

针对扫描图片格式的PDF表格&＃xff0c;可以使用Ocrmypdf包来将图片格式PDF转为文本格式。

Ocrmypdf是一个python软件包&＃xff0c;可将基于图像的pdf转换为基于文本的PDF。

安装ocrmypdf&＃xff0c;可以通过以下命令行将其用于macOS和Linux&＃xff1a;

brew install ocrmypdf


通过以下命令调用&＃xff1a;

ocrmypdf input_file.pdf output_file.pdf


之后就可以按照上面的方法进行表格提取了。

怎么样&＃xff0c;这款兼容图片&＃xff0c;又能白嫖谷歌GPU的PDF表格工具对你有帮助吗&＃xff1f;

传送门

项目地址&＃xff1a;
https://github.com/ferrygun/PDFTableExtract

在线打标签工具&＃xff1a;
https://www.makesense.ai/

— 完 —

本文系网易新闻•网易号特色内容激励计划签约账号【量子位】原创内容&＃xff0c;未经账号授权&＃xff0c;禁止随意转载。

喜欢就点「在看」吧 !