From 90d268d8a607ae2066e9f63d009e9b46c01527cb Mon Sep 17 00:00:00 2001 From: 283375 Date: Wed, 30 Aug 2023 20:52:02 +0800 Subject: [PATCH] impr: apply `FixRects` to default ocr method --- src/arcaea_offline_ocr/b30/chieri/v4/ocr.py | 12 +++++++----- src/arcaea_offline_ocr/ocr.py | 5 ++++- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py index 912889d..6f68adb 100644 --- a/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py +++ b/src/arcaea_offline_ocr/b30/chieri/v4/ocr.py @@ -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) diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index b7b6598..3026e04 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/src/arcaea_offline_ocr/ocr.py @@ -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)