197 lines
6.8 KiB
Python
197 lines
6.8 KiB
Python
"""Normalize PP-OCRv6 results into a stable project-owned schema."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
SCHEMA_VERSION = 1
|
|
OCR_VERSION = "PP-OCRv6"
|
|
DEFAULT_MODEL_SIZE = "medium"
|
|
MODEL_VARIANTS = {
|
|
"tiny": ("PP-OCRv6_tiny_det", "PP-OCRv6_tiny_rec"),
|
|
"small": ("PP-OCRv6_small_det", "PP-OCRv6_small_rec"),
|
|
"medium": ("PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),
|
|
}
|
|
DEFAULT_DETECTION_MODEL, DEFAULT_RECOGNITION_MODEL = MODEL_VARIANTS[DEFAULT_MODEL_SIZE]
|
|
|
|
|
|
def model_names_for_size(model_size: str) -> tuple[str, str]:
|
|
try:
|
|
return MODEL_VARIANTS[model_size]
|
|
except KeyError as exc:
|
|
supported = ", ".join(MODEL_VARIANTS)
|
|
raise ValueError(f"不支持的 PP-OCRv6 模型规格: {model_size}；可选: {supported}") from exc
|
|
|
|
|
|
def _payload(result: Any) -> dict[str, Any]:
|
|
if isinstance(result, dict):
|
|
return result
|
|
try:
|
|
return dict(result)
|
|
except (TypeError, ValueError) as exc:
|
|
raise TypeError(f"不支持的 PP-OCRv6 结果类型: {type(result).__name__}") from exc
|
|
|
|
|
|
def _python_value(value: Any) -> Any:
|
|
if hasattr(value, "tolist"):
|
|
value = value.tolist()
|
|
if isinstance(value, dict):
|
|
return {str(key): _python_value(item) for key, item in value.items()}
|
|
if isinstance(value, (list, tuple)):
|
|
return [_python_value(item) for item in value]
|
|
if hasattr(value, "item"):
|
|
try:
|
|
return value.item()
|
|
except (TypeError, ValueError):
|
|
pass
|
|
return value
|
|
|
|
|
|
def _polygon(value: Any) -> list[list[float | int]]:
|
|
points = _python_value(value) or []
|
|
normalized: list[list[float | int]] = []
|
|
for point in points:
|
|
if isinstance(point, (list, tuple)) and len(point) >= 2:
|
|
normalized.append([point[0], point[1]])
|
|
return normalized
|
|
|
|
|
|
def _box(value: Any, polygon: list[list[float | int]]) -> list[float | int]:
|
|
box = _python_value(value)
|
|
if isinstance(box, (list, tuple)) and len(box) >= 4:
|
|
return [box[0], box[1], box[2], box[3]]
|
|
if not polygon:
|
|
return []
|
|
xs = [point[0] for point in polygon]
|
|
ys = [point[1] for point in polygon]
|
|
return [min(xs), min(ys), max(xs), max(ys)]
|
|
|
|
|
|
def _image_size(result: dict[str, Any], input_path: Path | None) -> tuple[int | None, int | None]:
|
|
image = result.get("doc_preprocessor_res", {}).get("output_img")
|
|
shape = getattr(image, "shape", None)
|
|
if shape is not None and len(shape) >= 2:
|
|
return int(shape[1]), int(shape[0])
|
|
if input_path is not None and input_path.is_file():
|
|
try:
|
|
with Image.open(input_path) as opened:
|
|
return opened.size
|
|
except OSError:
|
|
pass
|
|
return None, None
|
|
|
|
|
|
def adapt_ocr_result(
|
|
result: Any,
|
|
*,
|
|
input_path: Path | str | None,
|
|
source_type: str,
|
|
language: str,
|
|
detection_model: str = DEFAULT_DETECTION_MODEL,
|
|
recognition_model: str = DEFAULT_RECOGNITION_MODEL,
|
|
model_size: str | None = None,
|
|
page_index: int | None = None,
|
|
page_number: int | None = None,
|
|
) -> dict[str, Any]:
|
|
raw = _payload(result)
|
|
path = Path(input_path).expanduser().resolve() if input_path is not None else None
|
|
def as_list(value: Any) -> list[Any]:
|
|
if value is None:
|
|
return []
|
|
if hasattr(value, "tolist"):
|
|
value = value.tolist()
|
|
return list(value)
|
|
|
|
texts = as_list(raw.get("rec_texts"))
|
|
scores = as_list(raw.get("rec_scores"))
|
|
polygon_values = raw.get("rec_polys")
|
|
if polygon_values is None:
|
|
polygon_values = raw.get("dt_polys")
|
|
polygons = as_list(polygon_values)
|
|
boxes = as_list(raw.get("rec_boxes"))
|
|
angles = as_list(raw.get("textline_orientation_angles"))
|
|
line_count = max(len(texts), len(scores), len(polygons), len(boxes))
|
|
|
|
lines: list[dict[str, Any]] = []
|
|
for index in range(line_count):
|
|
text = str(texts[index]) if index < len(texts) else ""
|
|
try:
|
|
score = float(scores[index]) if index < len(scores) else None
|
|
except (TypeError, ValueError):
|
|
score = None
|
|
polygon = _polygon(polygons[index]) if index < len(polygons) else []
|
|
box = _box(boxes[index] if index < len(boxes) else None, polygon)
|
|
orientation = angles[index] if index < len(angles) else None
|
|
lines.append(
|
|
{
|
|
"index": index + 1,
|
|
"text": text,
|
|
"score": round(score, 6) if score is not None and math.isfinite(score) else None,
|
|
"polygon": polygon,
|
|
"box": box,
|
|
"orientation": _python_value(orientation),
|
|
}
|
|
)
|
|
|
|
valid_scores = [line["score"] for line in lines if line["score"] is not None]
|
|
width, height = _image_size(raw, path)
|
|
resolved_page_index = page_index if page_index is not None else raw.get("page_index")
|
|
resolved_model_size = model_size
|
|
if resolved_model_size is None:
|
|
for size, names in MODEL_VARIANTS.items():
|
|
if names == (detection_model, recognition_model):
|
|
resolved_model_size = size
|
|
break
|
|
|
|
payload: dict[str, Any] = {
|
|
"schema_version": SCHEMA_VERSION,
|
|
"source_type": source_type,
|
|
"input_path": str(path) if path is not None else str(raw.get("input_path") or ""),
|
|
"page_index": resolved_page_index,
|
|
"model": {
|
|
"ocr_version": OCR_VERSION,
|
|
"model_size": resolved_model_size,
|
|
"detection_model": detection_model,
|
|
"recognition_model": recognition_model,
|
|
"language": language,
|
|
},
|
|
"image": {"width": width, "height": height},
|
|
"lines": lines,
|
|
"summary": {
|
|
"detected_lines": len(lines),
|
|
"non_empty_lines": sum(bool(line["text"].strip()) for line in lines),
|
|
"mean_score": round(sum(valid_scores) / len(valid_scores), 6) if valid_scores else None,
|
|
"min_score": min(valid_scores) if valid_scores else None,
|
|
"max_score": max(valid_scores) if valid_scores else None,
|
|
},
|
|
}
|
|
if page_number is not None:
|
|
payload["page_number"] = page_number
|
|
return payload
|
|
|
|
|
|
def result_plain_text(payload: dict[str, Any]) -> str:
|
|
return "\n".join(
|
|
str(line.get("text", "")).strip()
|
|
for line in payload.get("lines", [])
|
|
if str(line.get("text", "")).strip()
|
|
)
|
|
|
|
|
|
def result_markdown(payload: dict[str, Any], *, title: str | None = None) -> str:
|
|
text = result_plain_text(payload)
|
|
if title:
|
|
return f"# {title}\n\n{text}".rstrip() + "\n"
|
|
return text.rstrip() + ("\n" if text else "")
|
|
|
|
|
|
def raw_result_json(result: Any) -> dict[str, Any]:
|
|
raw_json = getattr(result, "json", None)
|
|
if raw_json is not None:
|
|
return _python_value(raw_json)
|
|
return _python_value(_payload(result))
|