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 math import floor
from typing import TYPE_CHECKING, List, Optional, Tuple from typing import List, Optional, Tuple
import cv2 import cv2
import numpy as np import numpy as np
from ....crop import crop_xywh 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 ....sift_db import SIFTDatabase
from ....types import Mat, cv2_ml_KNearest from ....types import Mat, cv2_ml_KNearest
from ....utils import construct_int_xywh_rect from ....utils import construct_int_xywh_rect
@ -212,10 +212,12 @@ class ChieriBotV4Ocr:
digit_contours, _ = cv2.findContours( digit_contours, _ = cv2.findContours(
roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
) )
digit_rects = sorted( digit_rects = [cv2.boundingRect(c) for c in digit_contours]
[cv2.boundingRect(c) for c in digit_contours], digit_rects = FixRects.connect_broken(
key=lambda r: r[0], 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 = [] digits = []
for digit_rect in digit_rects: for digit_rect in digit_rects:
digit = crop_xywh(roi, digit_rect) 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): def ocr_digits_by_contour_get_samples(__roi_gray: Mat, size: int):
roi = __roi_gray.copy() roi = __roi_gray.copy()
contours, _ = cv2.findContours(roi, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) 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 = [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] digit_rois = [resize_fill_square(crop_xywh(roi, rect), size) for rect in rects]
return preprocess_hog(digit_rois) return preprocess_hog(digit_rois)