from __future__ import annotations import base64 import json import shutil import subprocess from dataclasses import dataclass, field from pathlib import Path from uuid import uuid4 from sqlalchemy.orm import Session from app.core.config import SERVER_DIR, get_settings from app.schemas.ocr import OcrRecognizeBatchRead, OcrRecognizeDocumentRead, OcrRecognizeFieldRead, OcrRecognizeLineRead from app.services.document_intelligence import DocumentIntelligenceService WORKER_JSON_PREFIX = "__OCR_JSON__=" SUPPORTED_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".pdf"} @dataclass(slots=True) class PreparedOcrInput: input_path: Path source_key: str filename: str media_type: str page_index: int | None = None preview_kind: str = "" preview_data_url: str = "" @dataclass(slots=True) class AggregatedOcrDocument: filename: str media_type: str source_key: str engine: str = "paddleocr_mobile" model: str = "PP-OCRv5_mobile" summary_fragments: list[str] = field(default_factory=list) text_fragments: list[str] = field(default_factory=list) score_values: list[float] = field(default_factory=list) warnings: list[str] = field(default_factory=list) lines: list[OcrRecognizeLineRead] = field(default_factory=list) page_count: int = 0 preview_kind: str = "" preview_data_url: str = "" class OcrService: def __init__(self, db: Session | None = None) -> None: self.settings = get_settings() self.document_intelligence_service = DocumentIntelligenceService(db) def recognize_files( self, files: list[tuple[str, bytes, str | None]], ) -> OcrRecognizeBatchRead: if not files: raise ValueError("至少需要上传一个文件。") temp_root = self.settings.resolved_ocr_temp_dir temp_root.mkdir(parents=True, exist_ok=True) documents: list[OcrRecognizeDocumentRead] = [] prepared_inputs: list[PreparedOcrInput] = [] cleanup_paths: list[Path] = [] python_bin = self._resolve_python_bin() worker_path = self._resolve_worker_path() worker_payload: dict = {} try: for filename, content, media_type in files: normalized_name = Path(str(filename or "").strip()).name or "upload.bin" suffix = Path(normalized_name).suffix.lower() resolved_media_type = str(media_type or "application/octet-stream") if not content: documents.append( OcrRecognizeDocumentRead( filename=normalized_name, media_type=resolved_media_type, warnings=["文件内容为空，未执行 OCR。"], ) ) continue if suffix not in SUPPORTED_SUFFIXES: documents.append( OcrRecognizeDocumentRead( filename=normalized_name, media_type=resolved_media_type, warnings=["当前仅支持图片和 PDF 文件进行 OCR。"], ) ) continue if len(content) > self.settings.ocr_max_file_size_mb * 1024 * 1024: documents.append( OcrRecognizeDocumentRead( filename=normalized_name, media_type=resolved_media_type, warnings=[ f"文件超过 {self.settings.ocr_max_file_size_mb} MB，未执行 OCR。" ], ) ) continue temp_path = temp_root / f"{uuid4().hex}{suffix}" temp_path.write_bytes(content) cleanup_paths.append(temp_path) if suffix == ".pdf": try: prepared_inputs.extend( self._prepare_pdf_inputs( pdf_path=temp_path, filename=normalized_name, media_type=resolved_media_type, cleanup_paths=cleanup_paths, ) ) except RuntimeError as exc: documents.append( OcrRecognizeDocumentRead( filename=normalized_name, media_type=resolved_media_type, warnings=[str(exc)], ) ) continue prepared_inputs.append( PreparedOcrInput( input_path=temp_path, source_key=uuid4().hex, filename=normalized_name, media_type=resolved_media_type, preview_kind="image" if resolved_media_type.startswith("image/") else "", preview_data_url=( self._build_preview_data_url(temp_path, media_type=resolved_media_type) if resolved_media_type.startswith("image/") else "" ), ) ) if prepared_inputs: worker_payload = self._invoke_worker( python_bin=python_bin, worker_path=worker_path, input_paths=[item.input_path for item in prepared_inputs], ) documents.extend( self._build_documents( worker_documents=worker_payload.get("documents", []), prepared_inputs=prepared_inputs, ) ) success_count = sum( 1 for item in documents if item.line_count > 0 or not item.warnings ) engine = ( str(worker_payload.get("engine", "paddleocr_mobile")) if prepared_inputs else "paddleocr_mobile" ) model = ( str(worker_payload.get("model", "PP-OCRv5_mobile")) if prepared_inputs else "PP-OCRv5_mobile" ) return OcrRecognizeBatchRead( engine=engine, model=model, total_file_count=len(files), success_count=success_count, documents=documents, ) finally: self._cleanup_temp_paths(cleanup_paths) def _resolve_python_bin(self) -> str: candidates = [] configured = str(self.settings.ocr_python_bin or "").strip() if configured: candidates.append(configured) candidates.append(str(SERVER_DIR / ".venv-ocr312" / "bin" / "python")) candidates.append("/usr/local/bin/python3.12") resolved = shutil.which("python3.12") if resolved: candidates.append(resolved) for candidate in candidates: if candidate and Path(candidate).exists(): return candidate raise RuntimeError( "未找到可用的 OCR Python 运行时。请先执行 scripts/bootstrap_paddleocr_mobile.sh " "或通过 OCR_PYTHON_BIN 指向已安装 PaddleOCR 的 Python 3.12。" ) @staticmethod def _resolve_worker_path() -> str: worker_path = SERVER_DIR / "scripts" / "paddle_ocr_worker.py" if not worker_path.exists(): raise RuntimeError(f"OCR worker 不存在：{worker_path}") return str(worker_path) def _invoke_worker( self, *, python_bin: str, worker_path: str, input_paths: list[Path], ) -> dict: command = [ python_bin, worker_path, "--lang", self.settings.ocr_language, "--text-detection-model", self.settings.ocr_text_detection_model, "--text-recognition-model", self.settings.ocr_text_recognition_model, ] for path in input_paths: command.extend(["--input", str(path)]) completed = subprocess.run( command, capture_output=True, text=True, timeout=self.settings.ocr_timeout_seconds, check=False, ) if completed.returncode != 0: detail = (completed.stderr or completed.stdout or "").strip() raise RuntimeError(f"OCR 执行失败：{detail or 'worker 返回非 0 状态码。'}") payload = self._parse_worker_stdout(completed.stdout) if payload is None: raise RuntimeError("OCR worker 未返回可解析的 JSON 结果。") return payload @staticmethod def _parse_worker_stdout(stdout: str) -> dict | None: for line in reversed(stdout.splitlines()): normalized = line.strip() if normalized.startswith(WORKER_JSON_PREFIX): return json.loads(normalized[len(WORKER_JSON_PREFIX) :]) return None def _prepare_pdf_inputs( self, *, pdf_path: Path, filename: str, media_type: str, cleanup_paths: list[Path], ) -> list[PreparedOcrInput]: output_dir = pdf_path.with_suffix("") output_dir.mkdir(parents=True, exist_ok=True) cleanup_paths.append(output_dir) image_paths = self._convert_pdf_to_images(pdf_path=pdf_path, output_dir=output_dir) if not image_paths: raise RuntimeError("PDF 转图片后未生成可识别页面。") preview_data_url = self._build_preview_data_url(image_paths[0], media_type="image/png") source_key = uuid4().hex descriptors: list[PreparedOcrInput] = [] for page_index, image_path in enumerate(image_paths): descriptors.append( PreparedOcrInput( input_path=image_path, source_key=source_key, filename=filename, media_type=media_type, page_index=page_index, preview_kind="image" if page_index == 0 else "", preview_data_url=preview_data_url if page_index == 0 else "", ) ) return descriptors def _convert_pdf_to_images(self, *, pdf_path: Path, output_dir: Path) -> list[Path]: prefix = output_dir / "page" completed = subprocess.run( [ "pdftoppm", "-png", "-r", "160", str(pdf_path), str(prefix), ], capture_output=True, text=True, timeout=self.settings.ocr_timeout_seconds, check=False, ) if completed.returncode != 0: detail = (completed.stderr or completed.stdout or "").strip() raise RuntimeError(f"PDF 转图片失败：{detail or 'pdftoppm 返回非 0 状态码。'}") return sorted(output_dir.glob("page-*.png"), key=self._extract_pdf_page_sort_key) @staticmethod def _extract_pdf_page_sort_key(path: Path) -> tuple[int, str]: suffix = path.stem.rsplit("-", 1)[-1] try: return int(suffix), path.name except ValueError: return 0, path.name @staticmethod def _build_preview_data_url(path: Path, *, media_type: str) -> str: encoded = base64.b64encode(path.read_bytes()).decode("ascii") return f"data:{media_type};base64,{encoded}" def _build_documents( self, *, worker_documents: list[dict], prepared_inputs: list[PreparedOcrInput], ) -> list[OcrRecognizeDocumentRead]: descriptor_by_path = {str(item.input_path): item for item in prepared_inputs} source_order: list[str] = [] seen_sources: set[str] = set() for item in prepared_inputs: if item.source_key in seen_sources: continue seen_sources.add(item.source_key) source_order.append(item.source_key) aggregated_by_source: dict[str, AggregatedOcrDocument] = {} for payload in worker_documents: if not isinstance(payload, dict): continue input_path = str(payload.get("input_path") or "") descriptor = descriptor_by_path.get(input_path) if descriptor is None: continue aggregated = aggregated_by_source.get(descriptor.source_key) if aggregated is None: aggregated = AggregatedOcrDocument( filename=descriptor.filename, media_type=descriptor.media_type, source_key=descriptor.source_key, engine=str(payload.get("engine", "paddleocr_mobile")), model=str(payload.get("model", "PP-OCRv5_mobile")), ) aggregated_by_source[descriptor.source_key] = aggregated aggregated.page_count = max( aggregated.page_count, (descriptor.page_index + 1) if descriptor.page_index is not None else int(payload.get("page_count", 1) or 1), ) if descriptor.preview_kind and not aggregated.preview_kind: aggregated.preview_kind = descriptor.preview_kind if descriptor.preview_data_url and not aggregated.preview_data_url: aggregated.preview_data_url = descriptor.preview_data_url page_summary = str(payload.get("summary", "") or "").strip() if page_summary: aggregated.summary_fragments.append(page_summary) page_text = str(payload.get("text", "") or "").strip() if page_text: aggregated.text_fragments.append(page_text) lines = self._build_lines( payload.get("lines", []), page_index_override=descriptor.page_index, ) aggregated.lines.extend(lines) aggregated.score_values.extend(line.score for line in lines if line.score > 0) if not lines: avg_score = float(payload.get("avg_score", 0.0) or 0.0) if avg_score > 0: aggregated.score_values.append(avg_score) for warning in payload.get("warnings", []): normalized_warning = str(warning or "").strip() if normalized_warning and normalized_warning not in aggregated.warnings: aggregated.warnings.append(normalized_warning) documents: list[OcrRecognizeDocumentRead] = [] for source_key in source_order: descriptors = [item for item in prepared_inputs if item.source_key == source_key] if not descriptors: continue aggregated = aggregated_by_source.get(source_key) if aggregated is None: first_descriptor = descriptors[0] documents.append( OcrRecognizeDocumentRead( filename=first_descriptor.filename, media_type=first_descriptor.media_type, page_count=max(1, len(descriptors)), preview_kind=first_descriptor.preview_kind, preview_data_url=first_descriptor.preview_data_url, warnings=["OCR worker 未返回该文件的识别结果。"], ) ) continue documents.append(self._finalize_document(aggregated)) return documents @staticmethod def _build_lines( items: list[dict], *, page_index_override: int | None = None, ) -> list[OcrRecognizeLineRead]: lines: list[OcrRecognizeLineRead] = [] for item in items: if not isinstance(item, dict): continue page_index = page_index_override if page_index is None and item.get("page_index") is not None: page_index = int(item["page_index"]) lines.append( OcrRecognizeLineRead( text=str(item.get("text", "")), score=float(item.get("score", 0.0) or 0.0), box=[ [int(point[0]), int(point[1])] for point in item.get("box", []) if isinstance(point, list) and len(point) == 2 ], page_index=page_index, ) ) return lines @staticmethod def _truncate_summary(parts: list[str]) -> str: summary = "；".join([part for part in parts if part][:3]) if len(summary) > 180: return f"{summary[:177]}..." return summary def _finalize_document(self, aggregated: AggregatedOcrDocument) -> OcrRecognizeDocumentRead: full_text = "\n".join(fragment for fragment in aggregated.text_fragments if fragment).strip() summary = self._truncate_summary(aggregated.summary_fragments or aggregated.text_fragments) insight = self.document_intelligence_service.build_document_insight( filename=aggregated.filename, summary=summary, text=full_text, preview_data_url=aggregated.preview_data_url, ) warnings = list(aggregated.warnings) for warning in insight.warnings: normalized_warning = str(warning or "").strip() if normalized_warning and normalized_warning not in warnings: warnings.append(normalized_warning) return OcrRecognizeDocumentRead( filename=aggregated.filename, media_type=aggregated.media_type, engine=aggregated.engine, model=aggregated.model, text=full_text, summary=summary, avg_score=( sum(aggregated.score_values) / len(aggregated.score_values) if aggregated.score_values else 0.0 ), line_count=len(aggregated.lines), page_count=max(1, aggregated.page_count), document_type=insight.document_type, document_type_label=insight.document_type_label, scene_code=insight.scene_code, scene_label=insight.scene_label, classification_source=insight.classification_source, classification_confidence=insight.classification_confidence, classification_evidence=list(insight.evidence), document_fields=[ OcrRecognizeFieldRead( key=field.key, label=field.label, value=field.value, ) for field in insight.fields ], preview_kind=aggregated.preview_kind, preview_data_url=aggregated.preview_data_url, warnings=warnings, lines=sorted( aggregated.lines, key=lambda item: item.page_index if item.page_index is not None else -1, ), ) @staticmethod def _cleanup_temp_paths(paths: list[Path]) -> None: for path in reversed(paths): if path.is_dir(): shutil.rmtree(path, ignore_errors=True) continue path.unlink(missing_ok=True)