Update ocr_core.py

ocr_core.py

@@ -19,32 +19,80 @@ def load_bgr(source):
     img = ImageOps.exif_transpose(Image.open(source)).convert("RGB")
     return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
 
-
 def preprocess(bgr, size=2200):
     """Redresse la photo, neutralise les ombres, efface les traits de la grille."""
     h, w = bgr.shape[:2]
+
     s = size / max(h, w)
-    bgr = cv2.resize(bgr, None, fx=s, fy=s, interpolation=cv2.INTER_CUBIC)
+    bgr = cv2.resize(
+        bgr, None, fx=s, fy=s,
+        interpolation=cv2.INTER_CUBIC
+    )
+
     gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
-    bw = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 41, 15)
+
+    bw = cv2.adaptiveThreshold(
+        gray, 255,
+        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+        cv2.THRESH_BINARY_INV,
+        41, 15
+    )
 
     angle = 0.0
-    segs = cv2.HoughLinesP(bw, 1, np.pi / 720, threshold=200,
-                           minLineLength=bw.shape[1] // 4, maxLineGap=20)
+
+    segs = cv2.HoughLinesP(
+        bw, 1, np.pi / 720,
+        threshold=200,
+        minLineLength=bw.shape[1] // 4,
+        maxLineGap=20
+    )
+
     if segs is not None:
-        angs = [np.degrees(np.arctan2(y2 - y1, x2 - x1)) for x1, y1, x2, y2 in segs[:, 0]]
-        angs = [a for a in angs if abs(a) < 20]
+        angs = []
+
+        for segment in segs.reshape(-1, 4):
+            x1, y1, x2, y2 = map(int, segment)
+
+            a = np.degrees(
+                np.arctan2(y2 - y1, x2 - x1)
+            )
+
+            if abs(a) < 20:
+                angs.append(float(a))
+
         if angs:
             angle = float(np.median(angs))
-    M = cv2.getRotationMatrix2D((bw.shape[1] / 2, bw.shape[0] / 2), angle, 1)
-    bw = cv2.warpAffine(bw, M, (bw.shape[1], bw.shape[0]), flags=cv2.INTER_LINEAR)
 
-    hk = cv2.getStructuringElement(cv2.MORPH_RECT, (80, 1))
-    vk = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 80))
-    grid = cv2.add(cv2.morphologyEx(bw, cv2.MORPH_OPEN, hk), cv2.morphologyEx(bw, cv2.MORPH_OPEN, vk))
-    bw = cv2.subtract(bw, cv2.dilate(grid, np.ones((3, 3), np.uint8)))
-    return 255 - bw, angle
+    M = cv2.getRotationMatrix2D(
+        (bw.shape[1] / 2, bw.shape[0] / 2),
+        angle, 1
+    )
+
+    bw = cv2.warpAffine(
+        bw, M,
+        (bw.shape[1], bw.shape[0]),
+        flags=cv2.INTER_LINEAR
+    )
+
+    hk = cv2.getStructuringElement(
+        cv2.MORPH_RECT, (80, 1)
+    )
 
+    vk = cv2.getStructuringElement(
+        cv2.MORPH_RECT, (1, 80)
+    )
+
+    grid = cv2.add(
+        cv2.morphologyEx(bw, cv2.MORPH_OPEN, hk),
+        cv2.morphologyEx(bw, cv2.MORPH_OPEN, vk)
+    )
+
+    bw = cv2.subtract(
+        bw,
+        cv2.dilate(grid, np.ones((3, 3), np.uint8))
+    )
+
+    return 255 - bw, angle
 
 def read_words(bw, lang="eng", psm=11):
     d = pytesseract.image_to_data(bw, lang=lang, config=f"--psm {int(psm)}",