Python利用PaddleOCR进行扫描版PDF的跨页表格提取与合并

pip install pandas paddleocr pdf2image beautifulsoup4 numpy

import pandas as pd
from paddleocr import PPStructure, save_structure_res
from pdf2image import convert_from_path
from bs4 import BeautifulSoup
import warnings
import numpy as np
import logging
import os

warnings.filterwarnings("ignore")
logging.disable(logging.DEBUG)
logging.disable(logging.WARNING)

path = 'E:/Jobcontent/data/测试/'
output_folder = "E:/table/ex/"
topthreshold = 0.2
dthreshold = 0.8

table_engine = PPStructure(show_log=True)

def top_bottom_table_info(table_result):
    top_table = None
    min_y = 0
    for table in table_result:
        if table['type'] == 'table':
            bbox = table['bbox']
            current_bottom_y = bbox[1]
            if top_table is None or current_bottom_y < min_y:
                top_table = table
                min_y = current_bottom_y

    if top_table is not None:
        soup = BeautifulSoup(top_table['res']['html'], 'html.parser')
        last_row = soup.find_all('tr')[-1]
        columns = last_row.find_all('td')
        top_row_columns = len(columns)
        top_xy = top_table['bbox']
        return top_row_columns, top_xy
    else:
        return None

def find_bottom_table_info(table_result):
    bottom_table = None
    bottom_y = 0
    for table in table_result:
        if table['type'] == 'table':
            bbox = table['bbox']
            current_bottom_y = bbox[3]
            if bottom_table is None or current_bottom_y > bottom_y:
                bottom_table = table
                bottom_y = current_bottom_y

    if bottom_table is not None:
        soup = BeautifulSoup(bottom_table['res']['html'], 'html.parser')
        last_row = soup.find_all('tr')[0]
        columns = last_row.find_all('td')
        last_row_columns = len(columns)
        bottom_xy = bottom_table['bbox']
        return last_row_columns, bottom_xy
    else:
        return None

def is_continuation(top_row_columns, last_row_columns, bottom_xy, top_xy, page_height, dthreshold=0.8, topthreshold=0.2):
    if top_row_columns != last_row_columns:
        return False
    is_last_table_at_bottom = bottom_xy[3] > dthreshold * page_height
    is_first_table_at_top = top_xy[1] < topthreshold * page_height
    return is_last_table_at_bottom and is_first_table_at_top

pdf_files = [f for f in os.listdir(path) if f.endswith('.pdf')]
print(pdf_files)

for pdf_file in pdf_files:
    cross_page_tables = []
    pdf_path = os.path.join(path, pdf_file)
    images = convert_from_path(pdf_path, dpi=200)
    print('正在提取表格,请耐心等待...')
    for page_number, image in enumerate(images):
        table_result = table_engine(np.array(image))
        save_structure_res(table_result, output_folder, f'{page_number+1}')

        _, page_height = image.size
        last_row_columns, bottom_xy = find_bottom_table_info(table_result)
        if page_number+1 < len(images):
            table_result_end = table_engine(np.array(images[page_number+1]))
            top_row_columns, top_xy = top_bottom_table_info(table_result_end)

            if is_continuation(top_row_columns, last_row_columns, bottom_xy, top_xy, page_height):
                cross_page_tables.append((bottom_xy, top_xy, page_number+1, page_number+2))
                print(f"{pdf_file} 的表格在第 {page_number+1} 页和第 {page_number+2} 页之间跨页，并且最后一行和下一页的第一行列数相同")

    if cross_page_tables:
        for (bottom_xy, top_xy, start_page, end_page) in cross_page_tables:
            output_folder_s = output_folder + f'{start_page}/'
            output_folder_t = output_folder + f'{end_page}/'
            file_s = f'{bottom_xy}'+'_0'+'.xlsx'
            file_t = f'{top_xy}'+'_0'+'.xlsx'
            s_path = os.path.join(output_folder_s, file_s)
            e_path = os.path.join(output_folder_t, file_t)
            table_result_start =pd.read_excel(s_path, header=None)
            table_result_end = pd.read_excel(e_path, header=None)
            merged_table = pd.concat([table_result_start, table_result_end], ignore_index=True)
            output_path = os.path.join(output_folder, f'{pdf_file}_merged_{start_page}_{end_page}.csv')
            merged_table.to_csv(output_path, index=False)

print('表格提取完成')