99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from importlib import import_module
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OcrResult:
|
|
text: str
|
|
confidence: float
|
|
bbox: tuple[int, int, int, int] | None
|
|
debug: dict[str, Any]
|
|
|
|
|
|
def preprocess_for_ocr(image: np.ndarray, *, scale: float = 2.0) -> np.ndarray:
|
|
if image.size == 0:
|
|
raise ValueError("image is empty")
|
|
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
|
|
if scale != 1.0:
|
|
gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
|
|
denoised = cv2.GaussianBlur(gray, (3, 3), 0)
|
|
return cv2.adaptiveThreshold(
|
|
denoised,
|
|
255,
|
|
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
|
cv2.THRESH_BINARY,
|
|
31,
|
|
7,
|
|
)
|
|
|
|
|
|
def run_tesseract_ocr(
|
|
image: np.ndarray,
|
|
*,
|
|
lang: str = "eng",
|
|
psm: int = 6,
|
|
tesseract_cmd: str | None = None,
|
|
) -> OcrResult:
|
|
try:
|
|
pytesseract = import_module("pytesseract")
|
|
except ModuleNotFoundError as exc:
|
|
raise RuntimeError(
|
|
"pytesseract is not installed; run install.bat or install requirements.txt"
|
|
) from exc
|
|
|
|
if tesseract_cmd:
|
|
pytesseract.pytesseract.tesseract_cmd = tesseract_cmd
|
|
|
|
processed = preprocess_for_ocr(image)
|
|
config = f"--psm {psm}"
|
|
text = pytesseract.image_to_string(processed, lang=lang, config=config).strip()
|
|
confidences = _read_confidences(
|
|
pytesseract.image_to_data(
|
|
processed,
|
|
lang=lang,
|
|
config=config,
|
|
output_type=pytesseract.Output.DICT,
|
|
)
|
|
)
|
|
confidence = sum(confidences) / len(confidences) if confidences else 0.0
|
|
return OcrResult(
|
|
text=text,
|
|
confidence=confidence,
|
|
bbox=None,
|
|
debug={
|
|
"backend": "pytesseract",
|
|
"lang": lang,
|
|
"psm": psm,
|
|
"preprocessed_shape": tuple(map(int, processed.shape)),
|
|
"word_confidences": confidences,
|
|
},
|
|
)
|
|
|
|
|
|
def save_ocr_preprocess_debug(image: np.ndarray, output_path: str | Path) -> Path:
|
|
output = Path(output_path)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
if not cv2.imwrite(str(output), preprocess_for_ocr(image)):
|
|
raise RuntimeError(f"failed to write OCR debug image: {output}")
|
|
return output
|
|
|
|
|
|
def _read_confidences(data: dict[str, list[Any]]) -> list[float]:
|
|
values: list[float] = []
|
|
for raw in data.get("conf", []):
|
|
try:
|
|
confidence = float(raw)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if confidence >= 0:
|
|
values.append(confidence / 100.0)
|
|
return values
|