"""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))