python使用fitz和opencv库提取pdf中的表格

首先说明下，如果能在网上找到开源的库并适用于自己的项目，那最好直接使用开源库，避免重复造轮子。

目前在网上可以找到一些python提取处理pdf表格的库，但是总有一些不令人满意的地方，于是打算自己利用图像处理方式来实现，提取pdf文件中的表格数据

tabula-py库，这个库源码是用java实现的可以参考https://github.com/tabulapdf/tabula-java。提取内容能力感觉是最强的，但是在运行过程中总是报错，例如java.io.EOFException，UnicodeDecodeError: 'utf-8' codec can't decode byte 0xc6 in position 0: invalid continuation byte

pdfplumber库，对我手上很多pdf中的表格无法提取，所以放弃了

camelot库在pip安装过程中有问题

既然上面的开源项目不能满足自己的要求，于是打算使用机器视觉的方式来提取表格相关的信息。大致处理流程如下：

1.获取pdf的page文档的长宽

2.创建一块相同大小cv2的mat画布

3.获取page的所有线条，并在画布上绘制线条

4.使用轮廓包围，查找到所有各自表格cell的框

5.使用page的get_text_selection方法获取到没cell的文字信息

其中有几个需要注意的细节部分：

1.可以创建一个单通道的画布，这样可以避免灰度化和二值化操作

2.使用白底黑线，并使用漫水填充边缘，这样可以避免轮廓分析

3.如果表格线条是双实线，可以用开闭运算去掉双实线

4.使用get_text_selection方法时需要注意文本是否已经超出cell框的边界，如果超过边界，则只能获取到边框内的文字。如果存在这种表格，可以根据判断文字区域的中心坐标是否在cell中来提取文字。

所有的代码如下

import fitz
import numpy as np
import cv2


def to_int(*kwargs):
    v = []
    for k in kwargs:
        v.append(int(k))
    return v


def page_to_words_list(page: fitz.fitz.Page) -> list:
    '''
    将每一页中的textWords信息使用list封装，这样方便后续使用
    :param page:
    :return:
    '''
    # 获取文字及坐标信息
    words = page.getTextWords()

    # 将元素转为list
    # 因为list[0],list[1]....对于不熟悉代码的人很容易忘记含义，所以用字典封装
    # words = [[w[0], w[1], w[2], w[3], w[4]] for w in words]
    # 此处更适合用实体类，但是调试print的时候不方便，虽然可以重写__str__
    word_list = [{'rect': [w[0], w[1], w[2], w[3]], 'text': w[4]} for w in words]
    return word_list


def draw_pdf_tables(page: fitz.fitz.Page):
    assert isinstance(page, fitz.fitz.Page), '必须传入fitz.Page对象'
    # 创建一个白色的画布
    pixmap = page.getPixmap(matrix=fitz.Matrix(1, 1))
    # 二进制数据，宽，高
    img = np.zeros([pixmap.h, pixmap.w], dtype=np.uint8) + 255
    draws = page.getDrawings()
    # 在白色的画布上，画上黑色的线条
    for draw in draws:
        items = draw['items']
        for item in items:
            # 线条
            if 'l' in item:
                p1, p2 = to_int(*item[1]), to_int(*item[2])
                img = cv2.rectangle(img, (p1[0], p1[1]), (p2[0], p2[1]), (0))
    # 使用漫水填充算法，将周围变为黑色
    # 这样也可以去掉单独的线条
    cv2.floodFill(img, None, (0, 0), (0), cv2.FLOODFILL_FIXED_RANGE)
    return img


def get_pdf_table(page: fitz.fitz.Page, words=None):
    '''
    :param page:
    :param words:
    :return:
    '''
    assert isinstance(page, fitz.fitz.Page), '必须传入fitz.Page对象'
    if words is None:
        words = page_to_words_list(page)
    img = draw_pdf_tables(page)
    # 查找相应的轮廓，得到每个表格cell的矩形框
    contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 如果少于10个cell，认为这不是一个表格，这个判可以根据实际情况去掉
    if len(contours) < 10:
        return [], img
    tabel_cell = []
    # 判断文字是否在表格cell中
    for c in contours:
        r = cv2.boundingRect(c)
        r = [r[0], r[1], r[0] + r[2], r[1] + r[3]]
        ws = []
        for word in words[:]:
            w = word['rect']
            center = [(w[0] + w[2]) / 2, (w[1] + w[3]) / 2]
            if inside_rectangle(center, r):
                ws.append(word)
        tabel_cell.append({'rect': r, 'words': ws})
    # 闭运算
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    morp = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel, iterations=2)
    # 查找相应的轮廓，得到每个表格cell的矩形框
    contours, hierarchy = cv2.findContours(morp, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    tabels = []
    for c in contours:
        r = cv2.boundingRect(c)
        r = [r[0], r[1], r[0] + r[2], r[1] + r[3]]
        tabels.append(r)
    # 排序
    tabel_cell = sorted(tabel_cell, key=lambda data: (data['rect'][1], data['rect'][0]))
    tabels = sorted(tabels, key=lambda data: (data[1], data[0]))
    # 将cell合并到表中
    pdf_tabels = []
    for t in tabels:
        tabel = {'rect': t, 'cell': []}
        for cell in tabel_cell:
            c = cell['rect']
            center = [(c[0] + c[2]) / 2, (c[1] + c[3]) / 2]
            if inside_rectangle(center, t):
                tabel['cell'].append(cell)
        pdf_tabels.append(tabel)
    return pdf_tabels, img


def inside_rectangle(point, rect):
    '''
    判断点是否在框内
    :param point:
    :param rect:
    :return:
    '''
    x, y = point[0], point[1]
    x1, y1, x2, y2 = rect
    if x1 <= x <= x2 and y1 <= y <= y2:
        return True
    return False

最后看下效果图吧

返回的pdf_tabels是一个list列表，表示一页pdf中包含多个表格。每个表格是一个字典，有rect和cell两个字段，rect表示表格的区域位置。cell也是一个列表，表示一个表格格子，里面含有格子的坐标信息，以及文字信息。

最开始的时候我觉得返回这些信息比较繁琐，啰嗦，但是后续将pdf文件中的文档格式化输出的时候可以利用里面的信息做参考