This commit is contained in:
+42
-18
@@ -2,13 +2,17 @@ from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import logging
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from numbers import Real
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from core.models import Bounds, SceneElement
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OCRBox:
|
||||
@@ -61,9 +65,10 @@ def run_ocr(
|
||||
) -> list[SceneElement]:
|
||||
try:
|
||||
boxes = (engine or PaddleOCREngine()).extract(image)
|
||||
except ImportError:
|
||||
except Exception:
|
||||
if strict:
|
||||
raise
|
||||
logger.warning("OCR extraction failed; continuing without OCR", exc_info=True)
|
||||
return []
|
||||
return [box.to_scene_element(f"ocr-{index:03d}") for index, box in enumerate(boxes)]
|
||||
|
||||
@@ -105,20 +110,18 @@ def _flatten_pages(raw: Any) -> Iterable[Any]:
|
||||
return flattened
|
||||
json_attr = getattr(raw, "json", None)
|
||||
if callable(json_attr):
|
||||
return _flatten_pages(json_attr())
|
||||
if isinstance(json_attr, (dict, list)):
|
||||
return _flatten_pages(json_attr)
|
||||
return []
|
||||
|
||||
|
||||
def _dict_lines(data: dict[str, Any]) -> list[Any]:
|
||||
payload = data.get("res") if isinstance(data.get("res"), dict) else data
|
||||
texts = payload.get("rec_texts") or payload.get("texts") or []
|
||||
scores = payload.get("rec_scores") or payload.get("scores") or []
|
||||
boxes = (
|
||||
payload.get("rec_boxes")
|
||||
or payload.get("rec_polys")
|
||||
or payload.get("dt_polys")
|
||||
or []
|
||||
)
|
||||
result = data.get("res")
|
||||
payload = result if isinstance(result, dict) else data
|
||||
texts = _first_nonempty(payload, "rec_texts", "texts")
|
||||
scores = _first_nonempty(payload, "rec_scores", "scores")
|
||||
boxes = _first_nonempty(payload, "rec_boxes", "rec_polys", "dt_polys")
|
||||
return [
|
||||
{
|
||||
"text": text,
|
||||
@@ -136,11 +139,13 @@ def _looks_like_ocr_line(value: Any) -> bool:
|
||||
def _parse_line(line: Any) -> OCRBox | None:
|
||||
if isinstance(line, dict):
|
||||
text = line.get("text")
|
||||
points = line.get("points") or line.get("box") or line.get("bounds")
|
||||
points = _first_nonempty(line, "points", "box", "bounds")
|
||||
confidence = line.get("confidence")
|
||||
if not text or not points:
|
||||
if not _has_items(text) or not _has_items(points):
|
||||
return None
|
||||
return OCRBox(str(text), _bounds_from_points(points), _float_or_none(confidence))
|
||||
return OCRBox(
|
||||
str(text), _bounds_from_points(points), _float_or_none(confidence)
|
||||
)
|
||||
|
||||
if not _looks_like_ocr_line(line):
|
||||
return None
|
||||
@@ -161,11 +166,7 @@ def _parse_line(line: Any) -> OCRBox | None:
|
||||
def _bounds_from_points(points: Any) -> Bounds:
|
||||
if isinstance(points, dict):
|
||||
return Bounds.from_dict(points)
|
||||
if (
|
||||
isinstance(points, (list, tuple))
|
||||
and len(points) == 4
|
||||
and all(isinstance(value, (int, float)) for value in points)
|
||||
):
|
||||
if _has_length(points, 4) and all(isinstance(value, Real) for value in points):
|
||||
x1, y1, x2, y2 = [float(value) for value in points]
|
||||
return Bounds(x1, y1, x2 - x1, y2 - y1)
|
||||
|
||||
@@ -186,3 +187,26 @@ def _float_or_none(value: Any) -> float | None:
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
def _first_nonempty(data: dict[str, Any], *keys: str) -> Any:
|
||||
for key in keys:
|
||||
value = data.get(key)
|
||||
if _has_items(value):
|
||||
return value
|
||||
return []
|
||||
|
||||
|
||||
def _has_items(value: Any) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
try:
|
||||
return len(value) > 0
|
||||
except TypeError:
|
||||
return bool(value)
|
||||
|
||||
|
||||
def _has_length(value: Any, length: int) -> bool:
|
||||
try:
|
||||
return len(value) == length
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
Reference in New Issue
Block a user