Explorar o código

fix direction

zeke-chin %!s(int64=2) %!d(string=hai) anos
pai
achega
a0a373e932
Modificáronse 1 ficheiros con 17 adicións e 16 borrados
  1. 17 16
      core/direction.py

+ 17 - 16
core/direction.py

@@ -1,7 +1,7 @@
 import re
 from dataclasses import dataclass
 from enum import Enum
-from typing import Tuple
+from typing import Tuple, List
 
 import cv2
 import numpy as np
@@ -16,19 +16,26 @@ class Direction(Enum):
     BOTTOM = 2
     LEFT = 3
 
+
 # 父类
 class OcrAnchor(object):
     # 输入识别anchor的名字, 如身份证号
-    def __init__(self, name: str, d: List[Directoin]):
+    def __init__(self, name: str, d: List[Direction]):
         self.name = name
         # anchor位置
         self.direction = d
 
         def t_func(anchor, c, is_horizontal):
             if is_horizontal:
-                return 2 if anchor[1] > c[1] else 0
+                return 0 if anchor[1] < c[1] else 2
             else:
-                return 3 if anchor[0] < c[0] else 1
+                return 1 if anchor[0] > c[0] else 3
+
+        def r_func(anchor, c, is_horizontal):
+            if is_horizontal:
+                return 0 if anchor[0] > c[0] else 2
+            else:
+                return 1 if anchor[1] > c[1] else 3
 
         def b_func(anchor, c, is_horizontal):
             if is_horizontal:
@@ -39,12 +46,6 @@ class OcrAnchor(object):
         def l_func(anchor, c, is_horizontal):
             if is_horizontal:
                 return 0 if anchor[0] < c[0] else 2
-            else:
-                return 1 if anchor[1] > c[1] else 3
-
-        def r_func(anchor, c, is_horizontal):
-            if is_horizontal:
-                return 0 if anchor[0] > c[0] else 2
             else:
                 return 1 if anchor[1] < c[1] else 3
 
@@ -77,7 +78,7 @@ class OcrAnchor(object):
     # 找 锚点 -> 锚点坐标
     def find_anchor(self, res) -> Tuple[bool, float, float]:
         """
-        寻找身份证号的识别区域以及中心点
+        寻找锚点 中心点坐标
         """
         for row in res:
             for r in row:
@@ -101,7 +102,7 @@ class OcrAnchor(object):
         # print(f'cx: {cx}, cy: {cy}')
         pre = None
         for d in self.direction:
-            f = self.direction_funcs.get(self.direction, None)
+            f = self.direction_funcs.get(d, None)
             angle = f((id_cx, id_cy), (cx, cy), is_horizontal)
             if pre is None:
                 pre = angle
@@ -110,7 +111,6 @@ class OcrAnchor(object):
                     raise Exception('angle is not compatiable')
         return pre
 
-
         # if is_horizontal:
         #     # 如果是水平的,身份证号的位置在相对识别区域的下方,方向则为0度,否则是180度
         #     return 0 if id_cy > cy else 2
@@ -121,7 +121,7 @@ class OcrAnchor(object):
 
 # 子类1 人像面
 class FrontSideAnchor(OcrAnchor):
-    def __init__(self, name: str, d: List[Directoin]):
+    def __init__(self, name: str, d: List[Direction]):
         super(FrontSideAnchor, self).__init__(name, d)
 
     def is_anchor(self, txt, box) -> bool:
@@ -136,7 +136,7 @@ class FrontSideAnchor(OcrAnchor):
 
 # 子类2 国徽面
 class BackSideAnchor(OcrAnchor):
-    def __init__(self, name: str, d: List[Directoin]):
+    def __init__(self, name: str, d: List[Direction]):
         super(BackSideAnchor, self).__init__(name, d)
 
     def is_anchor(self, txt, box) -> bool:
@@ -177,7 +177,8 @@ class AngleDetector(object):
     # -> angle       result(ocr生)
     def detect_angle(self, img, image_type):
         image_type = int(image_type)
-        ocr_anchor = BackSideAnchor('有效期', [Directoin.BOTTOM]) if image_type != 0 else FrontSideAnchor('身份证号', [Directoin.BOTTOM])
+        ocr_anchor = BackSideAnchor('有效期', [Direction.BOTTOM]) if image_type != 0 else FrontSideAnchor('身份证号', [
+            Direction.BOTTOM])
 
         result = self.ocr.ocr(img, cls=True)