impr: apply FixRects to default ocr method

This commit is contained in:
283375 2023-08-30 20:52:02 +08:00
parent c1140803b1
commit 90d268d8a6
Signed by: 283375
SSH Key Fingerprint: SHA256:UcX0qg6ZOSDOeieKPGokA5h7soykG61nz2uxuQgVLSk
2 changed files with 11 additions and 6 deletions

View File

@ -1,11 +1,11 @@
from math import floor
from typing import TYPE_CHECKING, List, Optional, Tuple
from typing import List, Optional, Tuple
import cv2
import numpy as np
from ....crop import crop_xywh
from ....ocr import ocr_digits_by_contour_knn, preprocess_hog
from ....ocr import FixRects, ocr_digits_by_contour_knn, preprocess_hog
from ....sift_db import SIFTDatabase
from ....types import Mat, cv2_ml_KNearest
from ....utils import construct_int_xywh_rect
@ -212,10 +212,12 @@ class ChieriBotV4Ocr:
digit_contours, _ = cv2.findContours(
roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
)
digit_rects = sorted(
[cv2.boundingRect(c) for c in digit_contours],
key=lambda r: r[0],
digit_rects = [cv2.boundingRect(c) for c in digit_contours]
digit_rects = FixRects.connect_broken(
digit_rects, roi.shape[1], roi.shape[0]
)
digit_rects = FixRects.split_connected(roi, digit_rects)
digit_rects = sorted(digit_rects, key=lambda r: r[0])
digits = []
for digit_rect in digit_rects:
digit = crop_xywh(roi, digit_rect)

View File

@ -181,7 +181,10 @@ def ocr_digit_samples_knn(__samples, knn_model: cv2_ml_KNearest, k: int = 4):
def ocr_digits_by_contour_get_samples(__roi_gray: Mat, size: int):
roi = __roi_gray.copy()
contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
rects = sorted([cv2.boundingRect(c) for c in contours], key=lambda r: r[0])
rects = [cv2.boundingRect(c) for c in contours]
rects = FixRects.connect_broken(rects, roi.shape[1], roi.shape[0])
rects = FixRects.split_connected(roi, rects)
rects = sorted(rects, key=lambda r: r[0])
# digit_rois = [cv2.resize(crop_xywh(roi, rect), size) for rect in rects]
digit_rois = [resize_fill_square(crop_xywh(roi, rect), size) for rect in rects]
return preprocess_hog(digit_rois)