|
@@ -1,163 +1,41 @@
|
|
|
-import math
|
|
|
+from core.anchor import *
|
|
|
+from dataclasses import dataclass
|
|
|
+from paddleocr import PaddleOCR
|
|
|
|
|
|
-import cv2
|
|
|
import numpy as np
|
|
|
-from pathlib import Path
|
|
|
|
|
|
-BASE_DIR = Path(__file__).parent.parent
|
|
|
-logo_path = str(BASE_DIR/'logo.png')
|
|
|
-print(f'logo_path: {logo_path}')
|
|
|
-template_img = cv2.imread(logo_path)
|
|
|
-template_img = cv2.cvtColor(template_img, cv2.COLOR_BGR2RGB)
|
|
|
-
|
|
|
-
|
|
|
-def format_img(img, max_height=480):
|
|
|
- max_width = img.shape[1] * max_height // img.shape[0]
|
|
|
- return cv2.resize(img, (max_width, max_height))
|
|
|
-
|
|
|
-
|
|
|
-def search_best(des1, copy_target, sift):
|
|
|
- (kp2, des2) = sift.detectAndCompute(copy_target, None)
|
|
|
- bf = cv2.BFMatcher()
|
|
|
- matches1 = bf.knnMatch(des1, des2, k=2)
|
|
|
- ratio1 = 0.5
|
|
|
- good1 = []
|
|
|
- for m1, n1 in matches1:
|
|
|
- # 如果最接近和次接近的比值大于一个既定的值,那么我们保留这个最接近的值,认为它和其匹配的点为good_match
|
|
|
- if m1.distance < ratio1 * n1.distance:
|
|
|
- good1.append(m1)
|
|
|
- goodx = [(m1, m1.distance / n1.distance, m1.distance, n1.distance) for m1, n1 in matches1]
|
|
|
- goodx = sorted(goodx, key=lambda x: x[1])
|
|
|
- for gx, _, _, _ in goodx:
|
|
|
- if len(good1) >= 20: break
|
|
|
- if gx not in good1:
|
|
|
- good1.append(gx)
|
|
|
- return [[p] for p in good1], kp2
|
|
|
-
|
|
|
-
|
|
|
-def transform2stardard(ori_logo, ori_target):
|
|
|
- sift = cv2.xfeatures2d.SIFT_create()
|
|
|
- copy_logo = ori_logo.copy()
|
|
|
- (kp1, des1) = sift.detectAndCompute(copy_logo, None)
|
|
|
- match_pairs = []
|
|
|
- angle_rotate = {-1: 0, cv2.ROTATE_90_CLOCKWISE: 90} # ,cv2.ROTATE_180:180}
|
|
|
- for rotate in angle_rotate.keys():
|
|
|
- copy_target = cv2.rotate(ori_target, rotate) if rotate != -1 else ori_target.copy()
|
|
|
- goodx, kpx = search_best(des1, copy_target, sift)
|
|
|
- match_pairs.append((goodx, rotate, kpx))
|
|
|
- if len(goodx) >= 20:
|
|
|
- break
|
|
|
-
|
|
|
- max_match = max(match_pairs, key=lambda y: len(y[0]))
|
|
|
- good1 = max_match[0]
|
|
|
- kp2 = max_match[2]
|
|
|
-
|
|
|
- copy_target = cv2.rotate(ori_target, max_match[1]) if max_match[1] != -1 else ori_target.copy()
|
|
|
- if len(good1) > 4:
|
|
|
- ptsA = np.float32([kp1[m[0].queryIdx].pt for m in good1]).reshape(-1, 1, 2)
|
|
|
- ptsB = np.float32([kp2[m[0].trainIdx].pt for m in good1]).reshape(-1, 1, 2)
|
|
|
- ransacReprojThreshold = min(max(4, int(len(good1))), 20)
|
|
|
- # RANSAC算法选择其中最优的四个点
|
|
|
- H, status = cv2.findHomography(ptsA, ptsB, cv2.RANSAC, ransacReprojThreshold)
|
|
|
- imgout = cv2.warpPerspective(copy_target, H, (copy_logo.shape[1], copy_logo.shape[0]),
|
|
|
- flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)
|
|
|
- angle = math.atan2(H[1, 0], H[0, 0]) * 180 / math.pi - angle_rotate[max_match[1]]
|
|
|
- print(f'------------angle: {angle}')
|
|
|
- return imgout, angle
|
|
|
- # return np.hstack((format_img(ori_target.copy()),
|
|
|
- # cv2.putText(img=format_img(imgout.copy()), text='{}=>{}'.format(int(angle), tag),
|
|
|
- # org=(20, 70), fontFace=cv2.FONT_HERSHEY_TRIPLEX, fontScale=3, color=(0, 0, 255),
|
|
|
- # thickness=3)))
|
|
|
- else:
|
|
|
- return None, None
|
|
|
-
|
|
|
-
|
|
|
-def detect_angle(img):
|
|
|
- img_copy = np.copy(img)
|
|
|
- imgout, angle = transform2stardard(template_img, img)
|
|
|
- if imgout is not None:
|
|
|
- angle = np.argmin(np.abs(np.array([0, 90, 180, -180, -90]) - angle))
|
|
|
- if angle == 2 or angle == 3:
|
|
|
- angle = 2
|
|
|
- if angle == 4:
|
|
|
- angle = 3
|
|
|
- # return format_img(imgout), int(angle)
|
|
|
- return img_copy, int(angle)
|
|
|
- else:
|
|
|
- raise Exception('无法识别方向, 请把银行卡摆放正确')
|
|
|
-
|
|
|
-
|
|
|
-# @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_angle2(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
|
|
|
|
|
|
+def detect_angle(result, ocr_anchor: OcrAnchor):
|
|
|
+ lp = LineParser(result)
|
|
|
+ res = lp.parse()
|
|
|
+ print('------ angle ocr -------')
|
|
|
+ print(res)
|
|
|
+ print('------ angle ocr -------')
|
|
|
+ is_horizontal = lp.is_horizontal
|
|
|
+ return ocr_anchor.locate_anchor(res, is_horizontal)
|
|
|
+
|
|
|
+
|
|
|
+@dataclass
|
|
|
+class AngleDetector(object):
|
|
|
+ """
|
|
|
+ 角度检测器
|
|
|
+ """
|
|
|
+ ocr: PaddleOCR
|
|
|
+
|
|
|
+ def detect_angle(self, img):
|
|
|
+ ocr_anchor = BankCardAnchor('银行卡号', [Direction.BOTTOM])
|
|
|
+
|
|
|
+ result = self.ocr.ocr(img, cls=True)
|
|
|
+
|
|
|
+ try:
|
|
|
+ angle = detect_angle(result, ocr_anchor)
|
|
|
+ return angle, result
|
|
|
+
|
|
|
+ except Exception as e:
|
|
|
+ print(e)
|
|
|
+ # 如果第一次识别不到,旋转90度再识别
|
|
|
+ img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
|
|
|
+ result = self.ocr.ocr(img, cls=True)
|
|
|
+ angle = detect_angle(result, ocr_anchor)
|
|
|
+ # 旋转90度之后要重新计算角度
|
|
|
+ return (angle - 1 + 4) % 4, result
|