OCR-文本检测后的文字纠正

偏执的太偏执、 2022-12-27 15:23 99阅读 0赞

# -*- encoding: utf-8 -*-
    """
    @File    : 文本角度检测以及角度矫正.py
    @Time    : 2020/12/12 10:56
    @Author  : Johnson
    
    
    参考链接：https://blog.csdn.net/qq_39056987/article/details/108671180
    """
    '''
    1：霍夫变换预测倾斜角度
    2：图像二值化用于分割文本区域
    3：ROI区域获取；
    4：得到纠正后的文本区域截图
    
    具体实现如下：
    首先通过霍夫变换找到图片中的直线，并通过这些直线求得近似真实值的旋转角度
    再使用旋转函数对图片进行角度旋转。这里说明的是，其实文本图像中的直线比较多。
    文本越清晰，求出的角度越准确。
    
    旋转后为了不缺边，使用了退昂的RGB均值对边缘进行填充
    
    第二步，是将图片灰度后进行二值化处理，二值化用到了闭合运算这样能够保证
    相连的文本不会被划分成几个部分，最后得到的二值图像如下，对比上图，可以看到有文字的区域都被区分的很开，这里尝试了后面接着用开运算，但是效果不理想，所以果断放弃这个操作。
    
    
    得到文本区域后，接下来对这些区域进行ROI提取，提取出面积最大的最小外接矩形区域，其他的剔除，最后得到的纠正效果如下：
    
    
    '''
    # -*- coding: utf-8 -*-
    import cv2
    import numpy as np
    import math
    
    def compute(img):
        "求图像每个通道的均值"
        per_image_Rmean = []
        per_image_Gmean = []
        per_image_Bmean = []
        per_image_Bmean.append(np.mean(img[:, :, 0]))
        per_image_Gmean.append(np.mean(img[:, :, 1]))
        per_image_Rmean.append(np.mean(img[:, :, 2]))
        R_mean = np.mean(per_image_Rmean)
        G_mean = np.mean(per_image_Gmean)
        B_mean = np.mean(per_image_Bmean)
        return B_mean, G_mean, R_mean
    
    
    def rotate_bound(image, angle):
        """
        旋转任意角度不缺边缘
        """
        (h, w) = image.shape[:2]  # 获取图像的尺寸
        (cX, cY) = (w // 2, h // 2)  # 旋转中心
    
        # 设置旋转矩阵
        M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
        cos = np.abs(M[0, 0])
        sin = np.abs(M[0, 1])
    
        # 计算图像旋转后的新边界
        nW = int((h * sin) + (w * cos))
        nH = int((h * cos) + (w * sin))
    
        # 调整旋转矩阵的移动距离（t_{x}, t_{y}）
        M[0, 2] += (nW / 2) - cX
        M[1, 2] += (nH / 2) - cY
    
        # perform the actual rotation and return the image
        b, g, r = compute(image)
        return cv2.warpAffine(image, M, (nW, nH), borderValue=(b, g, r))
    
    
    def Hough(img):
        """
        霍夫变换角度检测并纠正
        :param img:
        :return: 矫正后的图片
        """
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 25, 150, apertureSize=3)
    
        lines = cv2.HoughLines(edges, 1, np.pi / 180, 0)
        for rho, theta in lines[0]:
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a * rho
            y0 = b * rho
            x1 = int(x0 + 1000 * (-b))
            y1 = int(y0 + 1000 * (a))
            x2 = int(x0 - 1000 * (-b))
            y2 = int(y0 - 1000 * (a))
        if x1 == x2 or y1 == y2:
            return img
        t = float(y2 - y1) / (x2 - x1)
        rotate_angle = math.degrees(math.atan(t))
        print(rotate_angle)
        if 12 > rotate_angle > 2 or -12 < rotate_angle < -2:    # 角度纠正的范围
            rotate_img = rotate_bound(img, -rotate_angle)
            afterimg=find_license(rotate_img)
            cv2.imshow("afterimg",afterimg)
            return afterimg
        else:
            return img
    
    
    def stretch(img):
        '''
        图像拉伸函数
        '''
        maxi=float(img.max())
        mini=float(img.min())
    
        for i in range(img.shape[0]):
            for j in range(img.shape[1]):
                img[i,j]=(255/(maxi-mini)*img[i,j]-(255*mini)/(maxi-mini))
    
        return img
    
    def dobinaryzation(img):
        '''
        二值化处理函数
        '''
        maxi=float(img.max())
        mini=float(img.min())
    
        x=maxi-((maxi-mini)/2)
        ret,thresh=cv2.threshold(img,x,255,cv2.THRESH_BINARY)
        #返回二值化后的黑白图像
        return thresh
    
    def find_license(img):
        '''
        预处理函数
        '''
        gray_img=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
        #灰度拉伸
        stretchedimg=stretch(gray_img)
    
        #图像二值化
        binaryimg=dobinaryzation(stretchedimg)
    
        #canny边缘检测
        canny=cv2.Canny(binaryimg,binaryimg.shape[0],binaryimg.shape[1])
    
        #进行闭运算
        kernel=np.ones((3,19),np.uint8)
        closingimg=cv2.morphologyEx(canny,cv2.MORPH_CLOSE,kernel)
        # cv2.imshow("closingimg",closingimg)
        # 提取图像区域
        img_crop = ROI_extract(img,closingimg)
    
        return img_crop
    
    def findVertices(points):
        # 获取最小外接矩阵，中心点坐标，宽高，旋转角度
        rect = cv2.minAreaRect(points)
        # 获取矩形四个顶点，浮点型
        box = cv2.boxPoints(rect)
        # 取整
        box = np.int0(box)
        # 获取四个顶点坐标
        left_point_x = np.min(box[:, 0])
        right_point_x = np.max(box[:, 0])
        top_point_y = np.min(box[:, 1])
        bottom_point_y = np.max(box[:, 1])
    
        left_point_y = box[:, 1][np.where(box[:, 0] == left_point_x)][0]
        right_point_y = box[:, 1][np.where(box[:, 0] == right_point_x)][0]
        top_point_x = box[:, 0][np.where(box[:, 1] == top_point_y)][0]
        bottom_point_x = box[:, 0][np.where(box[:, 1] == bottom_point_y)][0]
        # 上下左右四个点坐标
        vertices = np.array([[top_point_x, top_point_y], [bottom_point_x, bottom_point_y], [left_point_x, left_point_y],
                          [right_point_x, right_point_y]])
        return vertices, rect
    
    def ROI_extract(image,thresh):
        """
        提取目标区域
        :param image:
        :return:
        """
        contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    
        max_area = 0
        img = None
        size = image.shape
        if len(contours) != 0:
            for cnt in range(len(contours)):
                x, y, w, h = cv2.boundingRect(contours[cnt])
                box, _ = findVertices(contours[cnt])
                area = w * h
                if area > max_area:  # 过滤小区域，只保留最大的区域
                    y1 = max(0,y-2)
                    x1 = max(0,x-2)
                    y2 = min(size[0],y + h + 2)
                    x2 = min(size[1],x + w + 2)
                    newimage = image[y1 :y2, x1 :x2]
                    max_area = area
                    img = newimage
    
                    # 框出文字区域
                    image_copy = image.copy()
                    cv2.rectangle(image_copy, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.imshow('image_copy', image_copy)
                else:
                    continue
        else:
            return img
        return img
    
    
    if __name__=='__main__':
        import os
    
        path = r"G:\CRNN_data\diff_img\test"
        for name in os.listdir(path):
            img=cv2.imread(os.path.join(path,name),cv2.IMREAD_COLOR)
            #预处理图像
            cv2.imshow("img0",img)
            img = Hough(img)
            # cv2.imwrite(r"G:\CRNN_data\diff_img\test_save\{}".format(name),img)
            cv2.imshow("hough_img",img)
            cv2.waitKey()