174 lines
5 KiB
Python
174 lines
5 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""OCR wrapper using pytesseract."""
|
|
import pytesseract
|
|
import numpy as np
|
|
|
|
from . import image_io
|
|
|
|
_AUTO_LANGS = frozenset(("eng", "rus"))
|
|
_SCRIPT_TO_LANG = {
|
|
"latin": "eng",
|
|
"cyrillic": "rus",
|
|
}
|
|
|
|
|
|
def choose_ocr_lang(img_rgb, boxes, requested_lang='eng+rus', prep_options=None):
|
|
"""
|
|
Pick OCR language for RU/EN inputs.
|
|
|
|
Uses a fast Tesseract OSD pass to detect script (Latin/Cyrillic) on a few
|
|
largest text regions. Falls back to requested_lang when uncertain.
|
|
"""
|
|
requested = (requested_lang or 'eng').strip()
|
|
langs = _split_langs(requested)
|
|
if set(langs) != _AUTO_LANGS:
|
|
return requested, "fixed"
|
|
|
|
detected = _detect_lang_via_osd(img_rgb, boxes, prep_options=prep_options)
|
|
if detected:
|
|
return detected, "auto_script"
|
|
return requested, "auto_fallback"
|
|
|
|
|
|
def run_ocr(img_rgb, boxes, lang='eng', config='', return_boxes=True, prep_options=None):
|
|
if not boxes:
|
|
boxes = [(0, 0, img_rgb.shape[1], img_rgb.shape[0])]
|
|
|
|
blocks = []
|
|
collected_text = []
|
|
prep_options = prep_options or {}
|
|
|
|
for bbox in boxes:
|
|
crop = image_io.crop(img_rgb, bbox)
|
|
if crop.size == 0:
|
|
continue
|
|
prepared = image_io.prepare_for_ocr(
|
|
crop,
|
|
use_threshold=prep_options.get("use_threshold", True),
|
|
upscale=prep_options.get("upscale", 1.0),
|
|
clahe=prep_options.get("clahe", False),
|
|
bilateral=prep_options.get("bilateral", False),
|
|
blur=prep_options.get("blur", 0),
|
|
invert=prep_options.get("invert", False),
|
|
)
|
|
try:
|
|
data = pytesseract.image_to_data(
|
|
prepared,
|
|
lang=lang,
|
|
config=config,
|
|
output_type=pytesseract.Output.DICT,
|
|
)
|
|
texts = []
|
|
confs = []
|
|
for txt, conf in zip(data.get('text', []), data.get('conf', [])):
|
|
if txt and txt.strip():
|
|
texts.append(txt.strip())
|
|
try:
|
|
confs.append(float(conf))
|
|
except Exception:
|
|
confs.append(-1.0)
|
|
block_text = ' '.join(texts).strip()
|
|
avg_conf = float(np.mean(confs)) if confs else -1.0
|
|
except Exception:
|
|
# Fallback to simple string extraction
|
|
block_text = pytesseract.image_to_string(prepared, lang=lang, config=config)
|
|
avg_conf = -1.0
|
|
|
|
if block_text:
|
|
collected_text.append(block_text)
|
|
if return_boxes:
|
|
blocks.append({
|
|
"bbox": [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])],
|
|
"text": block_text,
|
|
"conf": avg_conf,
|
|
})
|
|
|
|
combined_text = '\n'.join(collected_text)
|
|
if not return_boxes:
|
|
blocks = []
|
|
return combined_text, blocks
|
|
|
|
|
|
def _split_langs(lang):
|
|
parts = []
|
|
for item in (lang or '').split('+'):
|
|
item = item.strip()
|
|
if item and item not in parts:
|
|
parts.append(item)
|
|
return parts
|
|
|
|
|
|
def _bbox_area(bbox):
|
|
try:
|
|
x1, y1, x2, y2 = bbox
|
|
return max(0, int(x2) - int(x1)) * max(0, int(y2) - int(y1))
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def _probe_boxes(img_rgb, boxes, limit=3):
|
|
if not boxes:
|
|
h, w = img_rgb.shape[:2]
|
|
return [(0, 0, w, h)]
|
|
ranked = sorted(list(boxes), key=_bbox_area, reverse=True)
|
|
return ranked[:limit]
|
|
|
|
|
|
def _prepare_probe_crop(crop, prep_options):
|
|
prep_options = prep_options or {}
|
|
upscale = prep_options.get("upscale", 1.0)
|
|
try:
|
|
upscale = float(upscale)
|
|
except Exception:
|
|
upscale = 1.0
|
|
if upscale < 1.0:
|
|
upscale = 1.0
|
|
if upscale > 1.5:
|
|
upscale = 1.5
|
|
|
|
return image_io.prepare_for_ocr(
|
|
crop,
|
|
use_threshold=False,
|
|
upscale=upscale,
|
|
clahe=prep_options.get("clahe", False),
|
|
bilateral=False,
|
|
blur=0,
|
|
invert=False,
|
|
)
|
|
|
|
|
|
def _detect_lang_via_osd(img_rgb, boxes, prep_options=None):
|
|
votes = {"eng": 0, "rus": 0}
|
|
for bbox in _probe_boxes(img_rgb, boxes):
|
|
crop = image_io.crop(img_rgb, bbox)
|
|
if crop is None or crop.size == 0:
|
|
continue
|
|
prepared = _prepare_probe_crop(crop, prep_options)
|
|
if prepared is None or prepared.size == 0:
|
|
continue
|
|
lang = _lang_from_osd(prepared)
|
|
if lang in votes:
|
|
votes[lang] += 1
|
|
|
|
if votes["eng"] == votes["rus"]:
|
|
return None
|
|
if votes["eng"] > votes["rus"]:
|
|
return "eng"
|
|
if votes["rus"] > votes["eng"]:
|
|
return "rus"
|
|
return None
|
|
|
|
|
|
def _lang_from_osd(prepared):
|
|
try:
|
|
osd_text = pytesseract.image_to_osd(prepared, config='--psm 0')
|
|
except Exception:
|
|
return None
|
|
|
|
for line in osd_text.splitlines():
|
|
line_low = line.lower()
|
|
if not line_low.startswith("script:"):
|
|
continue
|
|
script = line.split(':', 1)[1].strip().lower()
|
|
return _SCRIPT_TO_LANG.get(script)
|
|
return None
|