| import logging |
| import uuid |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| logger = logging.getLogger(__name__) |
|
|
| import fitz |
| from django.conf import settings |
| from django.core.files.uploadedfile import UploadedFile |
| from docx import Document |
|
|
|
|
| def get_doc_config() -> dict[str, int]: |
| """获取文档处理配置""" |
| from django.conf import settings |
|
|
| |
| try: |
| if getattr(settings, "CONFIG_MANAGER_AVAILABLE", False): |
| get_unified_config = getattr(settings, "get_unified_config", None) |
| if get_unified_config: |
| return { |
| "DEFAULT_TEXT_LIMIT": get_unified_config("features.document_processing.default_text_limit", 1500), |
| "DEFAULT_PREVIEW_PAGE": get_unified_config("features.document_processing.default_preview_page", 1), |
| "MAX_TEXT_LIMIT": get_unified_config("features.document_processing.max_text_limit", 10000), |
| "MAX_PREVIEW_PAGES": get_unified_config("features.document_processing.max_preview_pages", 5), |
| } |
| except Exception: |
| pass |
|
|
| |
| return getattr( |
| settings, |
| "DOCUMENT_PROCESSING", |
| { |
| "DEFAULT_TEXT_LIMIT": 1500, |
| "DEFAULT_PREVIEW_PAGE": 1, |
| "MAX_TEXT_LIMIT": 10000, |
| "MAX_PREVIEW_PAGES": 5, |
| }, |
| ) |
|
|
|
|
| def extract_text_from_image_with_rapidocr(file_path: str) -> str: |
| """ |
| 使用 OCR 从图片中提取文字 |
| |
| 自动根据 SystemConfig 中的 OCR_PROVIDER 配置选择本地 RapidOCR 或 PaddleOCR API。 |
| """ |
| from apps.automation.services.ocr.ocr_service import OCRService |
|
|
| ocr_service = OCRService() |
| return ocr_service.recognize(file_path) |
|
|
|
|
| def render_pdf_page_to_image(file_path: str, page_num: int = 0) -> str: |
| """ |
| 将PDF指定页面渲染为图片 |
| |
| Args: |
| file_path: PDF文件路径 |
| page_num: 页码（从0开始），默认为0（第一页） |
| |
| Returns: |
| 图片URL |
| """ |
| |
| if page_num < 0: |
| page_num = 0 |
|
|
| p = Path(file_path) |
| with fitz.open(p) as doc: |
| |
| if page_num >= doc.page_count: |
| page_num = min(page_num, doc.page_count - 1) |
|
|
| page = doc.load_page(page_num) |
| pix = page.get_pixmap() |
| out_dir = Path(settings.MEDIA_ROOT) / "automation" / "processed" |
| out_dir.mkdir(parents=True, exist_ok=True) |
| out_name = f"{uuid.uuid4().hex}_page{page_num + 1}.png" |
| out_path = out_dir / out_name |
| pix.save(out_path.as_posix()) |
| return f"{settings.MEDIA_URL}automation/processed/{out_name}" |
|
|
|
|
| def render_pdf_first_page_to_image(file_path: str) -> str: |
| """保持向后兼容性的函数""" |
| return render_pdf_page_to_image(file_path, page_num=0) |
|
|
|
|
| def extract_docx_text(file_path: str, limit: int | None = None) -> str: |
| """提取 .docx 文件的文本""" |
| |
| if limit is None: |
| config = get_doc_config() |
| limit = config["DEFAULT_TEXT_LIMIT"] |
|
|
| |
| config = get_doc_config() |
| if limit > config["MAX_TEXT_LIMIT"]: |
| limit = config["MAX_TEXT_LIMIT"] |
|
|
| p = Path(file_path) |
| d = Document(p.as_posix()) |
| parts = [] |
| for para in d.paragraphs: |
| if para.text: |
| parts.append(para.text) |
| if limit is not None and sum(len(x) for x in parts) >= int(limit): |
| break |
| text = "\n".join(parts) |
| if limit is not None: |
| return text[: int(limit)] |
| return text |
|
|
|
|
| def extract_pdf_text(file_path: str, limit: int | None = None, max_pages: int | None = None) -> str: |
| |
| if limit is None: |
| config = get_doc_config() |
| limit = config["DEFAULT_TEXT_LIMIT"] |
|
|
| |
| config = get_doc_config() |
| if limit > config["MAX_TEXT_LIMIT"]: |
| limit = config["MAX_TEXT_LIMIT"] |
|
|
| p = Path(file_path) |
| parts: list[str] = [] |
| with fitz.open(p) as doc: |
| for i in range(doc.page_count): |
| if max_pages is not None and max_pages > 0 and i >= max_pages: |
| break |
| page = doc.load_page(i) |
| t = page.get_text() |
| if t: |
| parts.append(t) |
| if limit is not None and sum(len(x) for x in parts) >= int(limit): |
| break |
| text = "".join(parts) |
| if limit is not None: |
| return text[: int(limit)] |
| return text |
|
|
|
|
| def _apply_pdf_limits(limit: int | None, preview_page: int | None, config: dict[str, Any]) -> tuple[int, int]: |
| """应用配置默认值和范围限制""" |
| lim = limit if limit is not None else config["DEFAULT_TEXT_LIMIT"] |
| page = preview_page if preview_page is not None else config["DEFAULT_PREVIEW_PAGE"] |
| lim = min(lim, config["MAX_TEXT_LIMIT"]) |
| page = min(page, config["MAX_PREVIEW_PAGES"]) |
| return lim, page |
|
|
|
|
| def _ocr_pdf_page(file_path: str, page_num_1based: int, limit: int) -> str | None: |
| """将 PDF 指定页 OCR，返回文字或 None""" |
| try: |
| p = Path(file_path) |
| with fitz.open(p) as doc: |
| page_num = min(max(page_num_1based - 1, 0), doc.page_count - 1) |
| page = doc.load_page(page_num) |
| pix = page.get_pixmap() |
|
|
| temp_dir = Path(settings.MEDIA_ROOT) / "automation" / "processed" |
| temp_dir.mkdir(parents=True, exist_ok=True) |
| temp_path = temp_dir / f"temp_{uuid.uuid4().hex}_page{page_num + 1}.png" |
| pix.save(temp_path.as_posix()) |
|
|
| ocr_text = extract_text_from_image_with_rapidocr(temp_path.as_posix()) |
| temp_path.relative_to(Path(settings.MEDIA_ROOT)) |
| temp_path.unlink(missing_ok=True) |
|
|
| if ocr_text.strip(): |
| return ocr_text[:limit] |
| except Exception as e: |
| logger.info(f"OCR处理PDF失败: {e}") |
| return None |
|
|
|
|
| def process_pdf( |
| file_path: str, limit: int | None = None, preview_page: int | None = None |
| ) -> tuple[str | None, str | None]: |
| """处理PDF文件：先提取文字，失败则 OCR，再失败则返回预览图""" |
| config = get_doc_config() |
| lim, page = _apply_pdf_limits(limit, preview_page, config) |
|
|
| text = extract_pdf_text(file_path, lim) |
| if text.strip(): |
| return None, text |
|
|
| ocr_text = _ocr_pdf_page(file_path, page, lim) |
| if ocr_text: |
| return None, ocr_text |
|
|
| image_url = render_pdf_page_to_image(file_path, page - 1) |
| return image_url, None |
|
|
|
|
| @dataclass |
| class DocumentExtraction: |
| file_path: str |
| text: str | None |
| image_url: str | None |
| kind: str |
|
|
|
|
| def save_uploaded_document(upload: UploadedFile) -> Path: |
| out_dir = Path(settings.MEDIA_ROOT) / "automation" / "uploads" |
| out_dir.mkdir(parents=True, exist_ok=True) |
| fname = f"{uuid.uuid4().hex}_{upload.name}" |
| dest = out_dir / fname |
| with dest.open("wb") as w: |
| for chunk in upload.chunks(): |
| w.write(chunk) |
| return dest |
|
|
|
|
| def extract_document_content( |
| file_path: str, limit: int | None = None, preview_page: int | None = None |
| ) -> DocumentExtraction: |
| """ |
| 提取文档内容 |
| |
| Args: |
| file_path: 文件路径 |
| limit: 文字提取限制，None时使用配置默认值 |
| preview_page: PDF预览页码，None时使用配置默认值 |
| |
| Returns: |
| DocumentExtraction对象 |
| """ |
| ext = Path(file_path).suffix.lower() |
| supported_image_exts = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"] |
|
|
| if ext == ".pdf": |
| image_url, text = process_pdf(file_path, limit, preview_page) |
| return DocumentExtraction(file_path=file_path, text=text, image_url=image_url, kind="pdf") |
|
|
| if ext == ".docx": |
| text = extract_docx_text(file_path, limit=limit) |
| return DocumentExtraction(file_path=file_path, text=text, image_url=None, kind="docx") |
|
|
| |
| if ext in supported_image_exts: |
| text = extract_text_from_image_with_rapidocr(file_path) |
|
|
| |
| if limit is None: |
| config = get_doc_config() |
| limit = config["DEFAULT_TEXT_LIMIT"] |
| if limit > get_doc_config()["MAX_TEXT_LIMIT"]: |
| limit = get_doc_config()["MAX_TEXT_LIMIT"] |
|
|
| if limit is not None: |
| text = text[: int(limit)] |
| return DocumentExtraction(file_path=file_path, text=text, image_url=None, kind="image") |
|
|
| raise ValueError(f"不支持的文件类型 {ext}，支持的格式：PDF、DOCX、图片({', '.join(supported_image_exts)})") |
|
|
|
|
| def process_uploaded_document( |
| upload: UploadedFile, limit: int | None = None, preview_page: int | None = None |
| ) -> DocumentExtraction: |
| """ |
| 统一的上传文件处理接口： |
| 1. 保存上传文件 |
| 2. 抽取可用文本 / 首图 |
| 3. 返回 DocumentExtraction 结果 |
| |
| Args: |
| upload: 上传的文件对象 |
| limit: 文字提取限制，None时使用配置默认值 |
| preview_page: PDF预览页码，None时使用配置默认值 |
| |
| Returns: |
| DocumentExtraction对象 |
| """ |
| dest = save_uploaded_document(upload) |
| extraction = extract_document_content(dest.as_posix(), limit=limit, preview_page=preview_page) |
| return extraction |
|
|
