hr-ocr-idcard/core/direction.py

import dataclasses

import cv2
import numpy as np

from dataclasses import dataclass
from paddleocr import PaddleOCR


@dataclass
# 角度检测器
class AngleDetector(object):
    ocr: PaddleOCR

    def detect_angle(self, img, result) -> int:
        wc = 0
        hc = 0
        count_0 = 0
        count_180 = 0
        angle = 0
        print(result)
        for res in result:
            txt = res[1][0]
            if '号' not in txt: continue
            a = np.array(res[0])
            l, t = np.min(a, axis=0).tolist()
            r, b = np.max(a, axis=0).tolist()
            l, t, r, b = list(map(int, [l, t, r, b]))
            if b - t > r - l:
                hc += 1
            else:
                wc += 1
            imgb = img[t:b, l:r, :]
            r = self.ocr.ocr(imgb, det=False, rec=False, cls=True)
            print(f'ocr angle: {r}')
            if int(r[0][0]) == 180:
                count_180 += 1
            else:
                count_0 += 1
        if hc >= wc:
            if count_0 >= count_180:
                angle = 90
            else:
                angle = 270
        else:
            if count_0 > count_180:
                angle = 0
            else:
                angle = 180

        return angle


def detect_angle(image):
    mask = np.zeros(image.shape, dtype=np.uint8)
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    blur = cv2.GaussianBlur(gray, (3, 3), 0)
    adaptive = cv2.adaptiveThreshold(blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 15, 4)

    cnts = cv2.findContours(adaptive, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cnts = cnts[0] if len(cnts) == 2 else cnts[1]

    for c in cnts:
        area = cv2.contourArea(c)
        if area < 45000 and area > 20:
            cv2.drawContours(mask, [c], -1, (255, 255, 255), -1)

    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    h, w = mask.shape

    # Horizontal
    if w > h:
        left = mask[0:h, 0:0 + w // 2]
        right = mask[0:h, w // 2:]
        left_pixels = cv2.countNonZero(left)
        right_pixels = cv2.countNonZero(right)
        return 0 if left_pixels >= right_pixels else 180
    # Vertical
    else:
        top = mask[0:h // 2, 0:w]
        bottom = mask[h // 2:, 0:w]
        top_pixels = cv2.countNonZero(top)
        bottom_pixels = cv2.countNonZero(bottom)
        return 90 if bottom_pixels >= top_pixels else 270


if __name__ == '__main__':
    image = cv2.imread('d40.jpg')
    angle = detect_angle(image)
    print(angle)