python实现图片方向纠正

目前有两种方案：

1、图片种类(发票、表格、单据...)较少，可直接使用基于模版的方向纠正

2、图片种类较多，基于图片分类模型预测 ，然后进行图片纠正

实现：

1、基于模版的方向纠正代码如下：

根据模版的方向纠正代码

# -*- coding: utf-8 -*-

import numpy as np

import cv2

"""

模式匹配，匹配出身份证信息，并且矫正身份证信息；

单应性指的是图像在投影发生了畸变后仍然能够有较高的检测和匹配准确率；

"""

class findidcard:

    def __init__(self):

        pass

    # 身份证正面截取，img1为身份证模板, img2为需要识别的图像

    def front_find(self, img2_name):

        #身份证背面识别

        img1_name = './picture_templates/invoice_mask.png'

        # 最小匹配数量设为10个， 大于这个数量从中筛选出10个最好的

        MIN_MATCH_COUNT = 10

        # 读入两幅图片 图片中有相同部分

        img1 = cv2.UMat(cv2.imread(img1_name, 0))  # queryImage in Gray

        img1 = self.img_resize(img1, 640)

        img2 = cv2.UMat(cv2.imread(img2_name, 0))  # trainImage in Gray

        img2 = self.img_resize(img2, 1920)

        img_org = cv2.UMat(cv2.imread(img2_name))

        img_org = self.img_resize(img_org, 1920)

        # 获取sift特征检测器

        sift = cv2.xfeatures2d.SIFT_create()

        # 使用SIFT查找关键点和描述符

        kp1, des1 = sift.detectAndCompute(img1, None)

        kp2, des2 = sift.detectAndCompute(img2, None)

        # kdtree建立索引方式的常量参数

        FLANN_INDEX_KDTREE = 0

        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)

        search_params = dict(checks=10)  # checks指定索引树要被遍历的次数

        flann = cv2.FlannBasedMatcher(index_params, search_params)

        # 进行匹配搜索

        matches = flann.knnMatch(des1, des2, k=2)

        # 两个最佳匹配之间距离需要大于ratio 0.7,距离过于相似可能是噪声点

        # 寻找距离近的放入good列表

        good = []

        for m, n in matches:

            if m.distance < 0.7 * n.distance:

                good.append(m)

        # reshape为(x,y)数组

        # 如果足够多，就筛选

        if len(good) > MIN_MATCH_COUNT:

            # 通过距离近的描述符 找到两幅图片的关键点

            src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)

            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

            # 单应性匹配图关键点匹配线。。不懂啥意思

            # 用HomoGraphy计算图像与图像之间映射关系, M为转换矩阵

            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)

            # 使用转换矩阵M计算出img1在img2的对应形状

            h, w = cv2.UMat.get(img1).shape

            M_r = np.linalg.inv(M)

            im_r = cv2.warpPerspective(img_org, M_r, (w, h))

            # cv2.imwrite("middle.jpg", im_r)

        else:

            print("Not enough matches are found - %d/%d" % (len(good), MIN_MATCH_COUNT))

        return im_r

    def showimg(img):

        cv2.namedWindow("contours", 0)

        cv2.resizeWindow("contours", 1280, 720)

        cv2.imshow("contours", img)

        cv2.waitKey()

    # 身份证背面截取，img1为身份证模板, img2为需要识别的图像

    def back_find(self, img2_name):

        # 身份证背面识别

        img1_name = './picture_templates/idcard_back_mask.jpg'

        # 最小匹配数量设为10个， 大于这个数量从中筛选出10个最好的

        MIN_MATCH_COUNT = 10

        # 读入两幅图片 图片中有相同部分

        img1 = cv2.UMat(cv2.imread(img1_name, 0))  # queryImage in Gray

        img1 = self.img_resize(img1, 640)

        img2 = cv2.UMat(cv2.imread(img2_name, 0))  # trainImage in Gray

        img2 = self.img_resize(img2, 1920)

        img_org = cv2.UMat(cv2.imread(img2_name))

        img_org = self.img_resize(img_org, 1920)

        # 获取sift特征检测器

        sift = cv2.xfeatures2d.SIFT_create()

        # 使用SIFT查找关键点和描述符

        kp1, des1 = sift.detectAndCompute(img1, None)

        kp2, des2 = sift.detectAndCompute(img2, None)

        # kdtree建立索引方式的常量参数

        FLANN_INDEX_KDTREE = 0

        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)

        search_params = dict(checks=10)  # checks指定索引树要被遍历的次数

        flann = cv2.FlannBasedMatcher(index_params, search_params)

        # 进行匹配搜索

        matches = flann.knnMatch(des1, des2, k=2)

        # 两个最佳匹配之间距离需要大于ratio 0.7,距离过于相似可能是噪声点

        # 寻找距离近的放入good列表

        good = []

        for m, n in matches:

            if m.distance < 0.7 * n.distance:

                good.append(m)

        # reshape为(x,y)数组

        # 如果足够多，就筛选

        if len(good) > MIN_MATCH_COUNT:

            # 通过距离近的描述符 找到两幅图片的关键点

            src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)

            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

            # 单应性匹配图关键点匹配线。。不懂啥意思

            # 用HomoGraphy计算图像与图像之间映射关系, M为转换矩阵

            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)

            # 使用转换矩阵M计算出img1在img2的对应形状

            h, w = cv2.UMat.get(img1).shape

            M_r = np.linalg.inv(M)

            im_r = cv2.warpPerspective(img_org, M_r, (w, h))

        else:

            print("Not enough matches are found - %d/%d" % (len(good), MIN_MATCH_COUNT))

        return im_r

    #重置图片大小

    def img_resize(self, imggray, dwidth):

        crop = imggray

        size = crop.get().shape

        height = size[0]

        width = size[1]

        height = height * dwidth / width

        crop = cv2.resize(src=crop, dsize=(dwidth, int(height)), interpolation=cv2.INTER_CUBIC)

        return crop

if __name__ == "__main__":

    idfind = findidcard()

    result = idfind.front_find('./testimages/invoice5.png')

    # idfind.showimg(result)

    cv2.namedWindow("contours", 0)

    cv2.resizeWindow("contours", 1280, 720)

    cv2.imshow("contours", result)

    cv2.waitKey()

2、基于方向模型预测实现方向纠正

方向纠正模型训练使用百度飞浆paddle开源的paddleClas进行模型训练。训练样本可使用业务中的具体数据，准备0度图片样本1000张，然后使用脚本对其进行90、180、270度方向旋转，分别放到0、90、180、270 文件夹,使用脚本生成coco数据集。修改paddleClas的ppcls/configs/ImageNet/PPLCNetV2/PPLCNetV2_base_direction.yaml 配置文件，更改样本地址路径和类别，进行训练。最终使用模型进行方向预测，通过预测能得出图片的旋转方向，程序再次对图片进行旋转纠正。

预测结果：

进行图片旋转纠正：

               try:
                    im = Image.open(img_path)
                except Exception as e:
                    print(str(e))
                else:
                    if im.mode == "RGBA":
                        im = im.convert('RGB')
                        os.remove(img_path)
                        im.save(img_path)
                        im = Image.open(img_path)
                    if rorate_radio == 90:
                        rimg = im.transpose(Image.ROTATE_90)
                    if rorate_radio == 180:
                        rimg = im.transpose(Image.ROTATE_180)
                    if rorate_radio == 270:
                        rimg = im.transpose(Image.ROTATE_270)
                    if rorate_radio == 0:
                        rimg = im
                    img = dst+"_"+img
                    print("target img:",img)
                    rimg.save(os.path.join(dst, img))

效果：