基於KNN的發票識別

項目概況:

有一個PDF文件，里面的每頁都是一張發票，把每頁的發票單獨存為一個PDF並用該發票的的發票號碼進行文件的命名，發票號碼需要OCR識別，即識別下圖中紅色方塊的內容。

 

 一:拆分PDF

現有一個PDF文件，里面有很多張發票圖片，每張發票占一頁

 

我們先把這整個PDF拆分為單獨的PDF

使用PyPDF2這個包

代碼如下，基本上每句都寫了注釋

from PyPDF2 import PdfFileWriter,PdfFileReader

def test1(file_path,folder_path,num,end_page,start_page=0):
    """
    :param file_path: pdf文件路徑
    :param folder_path: 存放路徑
    :param num: 拆分后的pdf存在幾個原pdf頁數
    :param end_page: 拆分到的最后一頁
    :param start_page: 起始的頁數，默認為0
    :return:
    """
    # 打開PDF文件
    pdf_file = PdfFileReader(open(file_path, 'rb'))
    # 獲取pdf的頁數
    pdf_file_num = pdf_file.getNumPages()
    # 如果輸入的end_page頁數比pdf文件的頁數大或者小於等於0，讓停止的頁數為pdf最大的頁數
    if end_page>pdf_file_num or end_page<=0:
        end_page=pdf_file_num
    # 從起始頁到最后一頁進行遍歷
    for i in range(start_page,end_page,num):
        #創建一個PdfFileWriter的對象
        out_put = PdfFileWriter()
        # 給out_put這個對象傳num數的頁，項目中每個發票都只占了1頁，所以num為1，如果發票占據2頁，那么num為2
        for k in range(num):
            out_put.addPage(pdf_file.getPage(i))
        # 設置保存的路徑
        out_file = folder_path + "\\" + f"{i}.pdf"
        # 把out_put里面的數據寫入到文件中
        out_put.write(open(out_file, 'wb'))

運行結果如下:

 

 

 二:把PDF變成圖片，並進行切分

現在發票是PDF格式，我們需要轉為圖片格式，而且我需要的發票號碼在發票的右上角，所以對圖片進行大致的切分有助於提高后面的識別速率。

這里解釋一下rect = page.rect，rect可以獲取頁面的大小，rect.tl,tl為topleft的縮寫，也就是左上角的意思，所以有tl(左上),tf(右上),bl(左下)，bf(右下)等坐標

import fitz

def my_fitz(pdfPath, imagePath):
    """
    :param pdfPath: pdf的路徑
    :param imagePath: 圖片文件夾的路徑，不是圖片路徑
    :return:
    """
    # 打開pdf文件
    pdfDoc = fitz.open(pdfPath)
    for pg in range(pdfDoc.pageCount):
        page = pdfDoc[pg]
        rotate = int(0)
        # 每個尺寸的縮放系數為2，生成的圖像的分辨率會提高，參數也可以自由設置，沒有硬性要求
        zoom_x = 2
        zoom_y = 2
        # 這個函數可以理解為，把zoom_x，zoom_y這兩個參數保存起來
        mat = fitz.Matrix(zoom_x, zoom_y).preRotate(rotate)
        rect = page.rect  # 頁面大小
        # mp為截取矩形的左上角坐標
        mp=rect.tr-(500/zoom_x,0)
        # tem為截取矩形的右下角坐標
        tem=rect.tr+(0,200/zoom_y)
        # clip為截取的矩形
        clip = fitz.Rect(mp, tem)
        # 進行圖片的截取
        pix = page.getPixmap(matrix=mat, alpha=False,clip=clip)
        if not os.path.exists(imagePath):  # 判斷存放圖片的文件夾是否存在
            os.makedirs(imagePath)  # 若圖片文件夾不存在就創建
        new_img_path = imagePath + '/' + '0.png'
        pix.writePNG(new_img_path)  # 將圖片寫入指定的文件夾內

        return new_img_path

運行結果如圖所示:

 

 

 

 

 

三:檢測邊緣，把中間的數字截取出來

邊緣檢測我使用的CV2模塊，注意使用cv2.threshold函數時，里面的圖片必須為灰度圖，不然會報錯

import cv2

def my_croping(imgpath):
    # 讀取圖片的路徑
    img = cv2.imread(imgpath)
    # 把該圖片轉換為灰度圖
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    #設置固定級別的閾值應用於矩陣
    ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
    # 尋找邊緣，返回的contours為邊緣數據的集合
    _, contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_TC89_L1)
    # 畫出邊緣,-1為畫出所有的邊緣，如果為任意自然數那么為contours的索引,(0,0,255)為顏色，最后的2是線條的粗細，數值越大，線條越粗
    cv2.drawContours(img, contours, -1, (0, 0, 255), 2)
    # 展示圖片
    cv2.imshow("pic", img)
    # 等待，當參數為0時，為無限等待，直到有鍵盤指令
    cv2.waitKey(0)

運行結果:

 

可見上一步驟的圖片中的發票號碼已經被圈起來了，但是有很多不必要的東西也被圈進來了，所以我們需要對初始的的contours進行篩選。

contours是一個包含多個列表的列表，我們需要的中間的數字，觀察可知，中間數字的邊緣比較大，所以我們只需要通過len()方法就可以進行初步的過濾

contours.sort(key=lambda x: len(x), reverse=True)
for i in range(len(contours)):
        if len(contours[i]) > 10:
            continue
        else:
            contours = contours[:i]
            break

加入過濾后運行結果:

 

 我們初步的縮小了范圍，下面需要制定具體的規則來確定想要獲得的對象

首先，我們先獲取各個邊緣所組成的矩形的坐標

rect_list=[]
for i in range(len(contours)):
        cont_ = contours[i]
        # 找到boundingRect
        rect = cv2.boundingRect(cont_)
        print(rect)
        rect_list.append(rect)
        

運行結果如下:

從左到右分別是x,y,寬度，高度

很明顯，我們要找的坐標是8個，寬度，高度差不多的坐標,n為閾值，初始為10，當兩個矩陣的寬和高直接的差的絕對值在閾值范圍內,填入集合，如果這樣的元素超過8個，那么則找到號碼對應的矩陣，在傳入之前，用X坐標的大小進行排序，能減少很多時間

def xyhw(li):
    n=10
    while n<30:
        for i in range(len(li)):
            tem_li=[li[i]]
            for k in range(i+1,len(li)):
                if abs(li[i][1]-li[k][1])+abs(li[i][2]-li[k][2])+abs(li[i][3]-li[k][3])<n:
                    tem_li.append(li[k])
            if len(tem_li)>=8:
                return tem_li
        n+=1

但是這個篩選完，還有一個問題，有時候會出現分割后NO沒有分割掉的情況，所以需要過濾掉NO

 

def filter_li(li):
    if len(li)>8:
        li = li[:9]
    interval=li[0][0]-li[1][0]
    test_interval=li[-2][0]-li[-1][0]
    if test_interval/interval>1.5:
        li=li[:-1]
    return li

這樣我們就可以獲得號碼的八個矩陣坐標，我們只需要把這八個矩陣融合即可

#進行排序
rect_list.sort(key= lambda x:x[0],reverse=True)
#進行篩選
rect_list=filter_li(rect_list)
#x0,y0為矩陣的左上角，x1,y1為矩陣的右下角
y0=rect_list[0][1]
y1=rect_list[0][1]+rect_list[0][3]
x0=rect_list[-1][0]
x1=rect_list[0][0]+rect_list[0][2]
print(y0,y1,x0,x1)
#進行圖片切割
cropImg = img2[y0:y1,x0:x1]
#寫入圖片
cv2.imwrite(img_path,cropImg)

可以獲得這樣的圖片:

 

 

四:把圖片中的數字分別截取出來

第四步和第三步的原理一樣，先邊緣檢測，然后獲取矩形坐標后進行截圖，比第三步簡單不少，這里就不多贅述了

 

import cv2
import numpy as np


def xyhw(li):
    n=10
    tem_li=[]
    while n<30:
        for i in range(len(li)):
            tem_li=[li[i]]
            for k in range(i+1,len(li)):
                if abs(li[i][1]-li[k][1])+abs(li[i][2]-li[k][2])+abs(li[i][3]-li[k][3])<n:
                    tem_li.append(li[k])
            if len(tem_li)>=8:
                return tem_li
        n+=1
    else:
        return tem_li


# 將img的高度調整為28，先后對圖像進行如下操作：直方圖均衡化，形態學，閾值分割
def pre_treat(img):
    height_ = 28
    ratio_ = float(img.shape[1]) / float(img.shape[0])
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.resize(gray, (int(ratio_ * height_), height_))
    gray = cv2.equalizeHist(gray)
    _, binary = cv2.threshold(gray, 190, 255, cv2.THRESH_BINARY)
    img_ = 255 - binary  # 反轉：文字置為白色，背景置為黑色
    return img_


def get_roi(contours):
    rect_list = []
    for i in range(len(contours)):
        rect = cv2.boundingRect(contours[i])
        if rect[3] > 10:
            rect_list.append(rect)
    return rect_list


def get_rect(img):
    _, contours, hierarchy = cv2.findContours(img,cv2.RETR_TREE, cv2.CHAIN_APPROX_TC89_L1)
    rect_list = get_roi(contours)
    rect_list.sort(key= lambda x:x[0],reverse=True)
    rect_list=xyhw(rect_list)
    
    return rect_list

def change_(img):
    length = 28
    h,w = img.shape
    H = np.float32([[1,0,(length-w)/2],[0,1,(length-h)/2]])
    img = cv2.warpAffine(img,H,(length,length))
    M = cv2.getRotationMatrix2D((length/2,length/2),0,26/float(img.shape[0]))
    return cv2.warpAffine(img,M,(length,length))

def fenge(img_path):
    cont = 0
    img = cv2.imread(img_path)
    img = pre_treat(img)
    contours = get_rect(img)
    folder_path=r"C:\Users\86173\Desktop\jetbrains2019.2\new\tem"
    file_list=[]
    # img=cv2.drawContours(img,contours,2,(0, 0, 255),3)
    print("*********************%s*************" %contours)
    for i in range(len(contours)):
        y0 = contours[i][1]
        y1 = contours[i][1] + contours[i][3]
        x0 = contours[i][0]
        x1 = contours[i][0] + contours[i][2]
        print(y0, y1, x0, x1)
        cropImg = img[y0:y1, x0:x1]
        cropImg = change_(cropImg)
        fenge_img=rf"{folder_path}\{cont}.png"
        cv2.imwrite(fenge_img, cropImg)
        cont += 1
        file_list.append(fenge_img)
    return file_list

五:苦力活

通過第四步的分割，我們可以得到分割后的數字，那么第一步就是給這些分割后的數字命名,類似這樣:

 

建議在分割的時候，用input輸入來命名嗷

 第二步就是把這些圖片轉為矩陣存入txt中:

from PIL import Image
import numpy

def noise_remove_pil(image_name, k):
    """
    8鄰域降噪
    Args:
        image_name: 圖片文件命名
        k: 判斷閾值

    Returns:

    """

    def calculate_noise_count(img_obj, w, h):
        """
        計算鄰域非白色的個數
        Args:
            img_obj: img obj
            w: width
            h: height
        Returns:
            count (int)
        """
        count = 0
        width, height = img_obj.size
        for _w_ in [w - 1, w, w + 1]:
            for _h_ in [h - 1, h, h + 1]:
                if _w_ > width - 1:
                    continue
                if _h_ > height - 1:
                    continue
                if _w_ == w and _h_ == h:
                    continue
                if img_obj.getpixel((_w_, _h_)) < 190:  # 這里因為是灰度圖像，設置小於230為非白色
                    count += 1
        return count

    img = Image.open(image_name)
    # 灰度
    gray_img = img.convert('L')

    w, h = gray_img.size
    for _w in range(w):
        for _h in range(h):
            if _w == 0 or _h == 0:
                gray_img.putpixel((_w, _h), 255)
                continue
            # 計算鄰域非白色的個數
            pixel = gray_img.getpixel((_w, _h))
            if pixel == 255:
                continue

            if calculate_noise_count(gray_img, _w, _h) < k:
                gray_img.putpixel((_w, _h), 255)
    # gray_img = gray_img.resize((32, 32), Image.LANCZOS)
    gray_img.save(image_name)
    # gray_img.show()
    im = numpy.array(gray_img)
    for i in range(im.shape[0]):  # 轉化為二值矩陣
        for j in range(im.shape[1]):
            if im[i, j] <190:
                im[i, j] = 1
            else:
                im[i, j] = 0
    return im




if __name__ == '__main__':
    for i in range(0,10):
        for k in range(0,100):
            png_file_path=rf"C:\Users\86173\Desktop\jetbrains2019.2\model_test\{i}_{k}.png"
            txt_file_path=rf"C:\Users\86173\Desktop\jetbrains2019.2\model_test\txt_folder\{i}_{k}.txt"
            try:
                im = noise_remove_pil(png_file_path, 4)
                with open(txt_file_path,'at',encoding='utf-8')as f:
                    for n in im:
                        f.writelines(str(n).replace("[","").replace("]","").replace(" ","")+"\n")
            except Exception as e:
                continue

運行結果:

 

 

 

獲得這樣的文件，那么准備工作就結束了

六:KNN模型的使用

 導入sklearn使用knn模型非常簡單，代碼量很少

import numpy as np
from os import listdir
from sklearn.neighbors import KNeighborsClassifier as kNN

def np2vector(im):
    returnVect = np.zeros((1, 784))
    for i in range(28):
        # 讀一行數據
        lineStr = im[i]
        # 每一行的前28個元素依次添加到returnVect中
        for j in range(28):
            returnVect[0, 28 * i + j] = int(lineStr[j])
    # 返回轉換后的1x784向量
    return returnVect
def img2vector(filename):
    #創建1x784零向量
    returnVect = np.zeros((1, 784))
    #打開文件
    fr = open(filename)
    #按行讀取
    for i in range(28):
        #讀一行數據
        lineStr = fr.readline()
        #每一行的前28個元素依次添加到returnVect中
        for j in range(28):

            returnVect[0,28*i+j] = int(lineStr[j])
    #返回轉換后的1x784向量
    return returnVect

def handwritingClassTest(im):
    #測試集的Labels
    hwLabels = []
    #返回trainingDigits目錄下的文件名
    trainingFileList = listdir(r"C:\Users\86173\Desktop\jetbrains2019.2\model_test\txt_folder")
    #返回文件夾下文件的個數
    m = len(trainingFileList)
    #初始化訓練的Mat矩陣,測試集
    trainingMat = np.zeros((m, 784))
    #從文件名中解析出訓練集的類別
    for i in range(m):
        #獲得文件的名字
        fileNameStr = trainingFileList[i]
        #獲得分類的數字
        classNumber = int(fileNameStr.split('_')[0])
        #將獲得的類別添加到hwLabels中
        hwLabels.append(classNumber)
        trainingMat[i,:] = img2vector(r'C:\Users\86173\Desktop\jetbrains2019.2\model_test\txt_folder\%s' % (fileNameStr))
    #構建kNN分類器
    neigh = kNN(n_neighbors = 4, algorithm = 'auto')
    #擬合模型, trainingMat為測試矩陣,hwLabels為對應的標簽
    neigh.fit(trainingMat, hwLabels)
    
    vectorUnderTest = np2vector(im)

    classifierResult = neigh.predict(vectorUnderTest)
    return classifierResult

有這個模型，我們調用一下，就可以獲取到對應的發票號碼了

最終運行結果:

 

 

 

 

 最后:

knn的原理比較簡單，但是因為是在工作之余寫的，寫的比較匆忙，有些步驟說的不夠詳細，如果有什么問題歡迎在評論區留言，如果有改進方案那就更好了，博主只是一個初入機器學習的小學生，歡迎各位大佬的指點，謝謝