diff --git a/src/arcaea_offline_ocr/ocr.py b/src/arcaea_offline_ocr/ocr.py index b15f0d1..83a30b4 100644 --- a/src/arcaea_offline_ocr/ocr.py +++ b/src/arcaea_offline_ocr/ocr.py @@ -3,7 +3,6 @@ from typing import Optional, Sequence, Tuple import cv2 import numpy as np -from numpy.linalg import norm from .crop import crop_xywh from .types import Mat, cv2_ml_KNearest @@ -142,28 +141,11 @@ def resize_fill_square(img: Mat, target: int = 20): def preprocess_hog(digit_rois): - # https://github.com/opencv/opencv/blob/f834736307c8328340aea48908484052170c9224/samples/python/digits.py + # https://learnopencv.com/handwritten-digits-classification-an-opencv-c-python-tutorial/ samples = [] for digit in digit_rois: - gx = cv2.Sobel(digit, cv2.CV_32F, 1, 0) - gy = cv2.Sobel(digit, cv2.CV_32F, 0, 1) - mag, ang = cv2.cartToPolar(gx, gy) - bin_n = 16 - _bin = np.int32(bin_n * ang / (2 * np.pi)) - bin_cells = _bin[:10, :10], _bin[10:, :10], _bin[:10, 10:], _bin[10:, 10:] - mag_cells = mag[:10, :10], mag[10:, :10], mag[:10, 10:], mag[10:, 10:] - hists = [ - np.bincount(b.ravel(), m.ravel(), bin_n) - for b, m in zip(bin_cells, mag_cells) - ] - hist = np.hstack(hists) - - # transform to Hellinger kernel - eps = 1e-7 - hist /= hist.sum() + eps - hist = np.sqrt(hist) - hist /= norm(hist) + eps - + hog = cv2.HOGDescriptor((20, 20), (10, 10), (5, 5), (10, 10), 9) + hist = hog.compute(digit) samples.append(hist) return np.float32(samples)