OCR识别发票内容

2021-06-19  本文已影响0人  小帅明3号

import pytesseract as pt
from PIL import Image
import os
import fitz
import xlwt
import re

进程调度路径

sb_path = r"C:\Program Files\Tesseract-OCR\tesseract.exe"

指定调度进程

pt.pytesseract.tesseract_cmd = sb_path

pdf路径

file_ph = input('请输入pdf存放路径:') # C:\Users\LKM\Desktop\ocr\新建文件夹\pdf

需解析文件路径

file_ph_deal = file_ph + '\信息提取\'

公共变量-pdf清单提取

pdf_dir = []

公共变量-图片解析内容

r_png_txt = []

公共变量-图片解析清单

r_png_txts = []

提取pdf的图片

try:

# 1.转换个数计数器
is_2pic_ok = 0
# 2.处理路径初始化
if not os.path.isdir(file_ph_deal):
    os.mkdir(file_ph_deal)
# 3.优先删除路径下所有png文件
for file_p, k_file_p, file_dir in os.walk(file_ph_deal):
    for file in os.scandir(file_p):
        if file.name.endswith(".png"):  # 指定文件类型
            os.remove(file_p + "\\" + file.name)
# 4.提取路径下所有pdf文件
docunames = os.listdir(file_ph)
for docuname in docunames:
    if os.path.splitext(docuname)[1] == '.pdf':  # 目录下包含.pdf的文件
        pdf_dir.append(docuname)
pdf_dir.sort()
# 5.遍历pdf,提取生成图片
for pdf in pdf_dir:
    print("处理文件:" + pdf)
    file_pdf = file_ph + '\\' + pdf
    # 解析后图片名前置内容获取,截取用
    file_deal = file_ph_deal + '\\' + pdf
    doc = fitz.open(file_pdf)
    # 循环读取pdf页签的内容
    for pg in range(doc.pageCount):
        page = doc[pg]
        # 每个尺寸的缩放系数为2,这将为我们生成分辨率提高四倍的图像。
        trans = fitz.Matrix(2.0, 2.0).preRotate(int(0))
        png_name = os.path.splitext(file_deal)[0] + str(pg) + '.png'
        pm = page.getPixmap(matrix=trans, alpha=False)
        pm.writePNG(png_name)
        # 获取主图片
        img_main = Image.open(png_name)
        # 进行旋转90
        out = img_main.transpose(Image.ROTATE_270)
        # 处理文件保存
        out.save(png_name)
        # 重新读取
        img_main = Image.open(png_name)

        # 图片拆分截取识别
        size = img_main.size
        # 获取长和宽
        weight = int(size[0])
        height = int(size[1])
        # 1.解析单号
        weight_x1 = int(weight * 1287 / 1684)
        height_y1 = int(height * 90 / 1191)
        weight_x2 = int(weight * 1441 / 1684)
        height_y2 = int(height * 139 / 1191)
        box = (weight_x1, height_y1, weight_x2, height_y2)
        region = img_main.crop(box)
        region.save(os.path.splitext(file_deal)[0] + str(pg) + '-1.png')
        img = Image.open(os.path.splitext(file_deal)[0] + str(pg) + '-1.png')
        img = img.convert('L')  # 灰度
        img.load()
        text = pt.image_to_string(img, lang="chi_sim")
        new_text = text.replace(' ', '').replace("\n", "").replace("\x0c","") # 替换空行及空格
        r_png_txt.append(new_text)
        # 2.解析货物名称
        weight_x1 = int(weight * 290 / 1684)
        height_y1 = int(height * 319 / 1191)
        weight_x2 = int(weight * 541 / 1684)
        height_y2 = int(height * 363 / 1191)
        box = (weight_x1, height_y1, weight_x2, height_y2)
        region = img_main.crop(box)
        region.save(os.path.splitext(file_deal)[0] + str(pg) + '-2.png')
        img = Image.open(os.path.splitext(file_deal)[0] + str(pg) + '-2.png')
        img = img.convert('L')  # 灰度
        img.load()
        text = pt.image_to_string(img, lang="chi_sim")
        new_text = text.replace(' ', '').replace("\n", "").replace("\x0c","")  # 替换空行及空格
        new_text = re.sub('[\W_+]', "", new_text)
        r_png_txt.append(new_text)
        # 3.解析销售方名称
        weight_x1 = int(weight * 460 / 1684)
        height_y1 = int(height * 600 / 1191)
        weight_x2 = int(weight * 900 / 1684)
        height_y2 = int(height * 630 / 1191)
        box = (weight_x1, height_y1, weight_x2, height_y2)
        region = img_main.crop(box)
        region.save(os.path.splitext(file_deal)[0] + str(pg) + '-3.png')
        img = Image.open(os.path.splitext(file_deal)[0] + str(pg) + '-3.png')
        img = img.convert('L')  # 灰度
        img.load()
        text = pt.image_to_string(img, lang="chi_sim")
        new_text = text.replace(' ', '').replace("\n", "").replace("\x0c","")  # 替换空行及空格
        new_text = re.sub('[\W_+]', "", new_text)
        r_png_txt.append(new_text)
        # 4.解析金额
        weight_x1 = int(weight * 1063 / 1684)
        height_y1 = int(height * 321 / 1191)
        weight_x2 = int(weight * 1203 / 1684)
        height_y2 = int(height * 380 / 1191)
        box = (weight_x1, height_y1, weight_x2, height_y2)
        region = img_main.crop(box)
        region.save(os.path.splitext(file_deal)[0] + str(pg) + '-4.png')
        img = Image.open(os.path.splitext(file_deal)[0] + str(pg) + '-4.png')
        img = img.convert('L')  # 灰度
        img.load()
        text = pt.image_to_string(img, lang="chi_sim")
        new_text = text.replace(' ', '').replace("\n", "").replace("\x0c","").split("|")[0] # 替换空行及空格
        r_png_txt.append(new_text)
        r_png_txts.append(r_png_txt)
        # 重置
        r_png_txt=[]
    # 处理计数器+1
    is_2pic_ok += 1
if is_2pic_ok==0:
    print('未检测到pdf文件,未转换!')
# 5.再次删除路径下临时png文件
for file_p, k_file_p, file_dir in os.walk(file_ph_deal):
    for file in os.scandir(file_p):
        if file.name.endswith(".png"):  # 指定文件类型
            os.remove(file_p + "\\" + file.name)
# 6.生成Excel
wfile = file_ph_deal + '\\发票识别结果.xls'
we = xlwt.Workbook()  # 创建一个Excel对象
sh = we.add_sheet('sheet1', cell_overwrite_ok=True)  # 某个单元格可以复写,多次写入不报错
# 初始化表头
sh.write(0,0,"发票单号")
sh.write(0,1,"货物名称")
sh.write(0,2,"销售商名称")
sh.write(0,3,"金额")
for i in range(len(r_png_txts)):  # 读行
    for j in range(4):  # 读列
        sh.write(i+1, j, r_png_txts[i][j])  # 在第i行第1列写内容
we.save(wfile)

except Exception as e:
print(e)

os.system('pause')

上一篇 下一篇

猜你喜欢

热点阅读