89 lines
4.4 KiB
Python
89 lines
4.4 KiB
Python
"""
|
|
===============================================================================
|
|
FILE: services/office/document_extractor.py
|
|
PROJECT: SCUD Orion AI (Office Domain)
|
|
ROLE: Постраничное извлечение текста и OCR сканов документов через Vision LLM.
|
|
===============================================================================
|
|
"""
|
|
|
|
import os
|
|
import base64
|
|
import logging
|
|
from typing import List, Dict, Any, Optional
|
|
|
|
import fitz # PyMuPDF
|
|
from modules.web_api.llm.core.ollama_client import call_ollama_chat
|
|
|
|
logger = logging.getLogger("OFFICE_EXTRACTOR")
|
|
|
|
|
|
def ocr_image_b64(image_b64: str, page_num: int = 1) -> str:
|
|
"""
|
|
Распознает текст с одного растрового изображения через Qwen 2.5 VL.
|
|
"""
|
|
system_prompt = (
|
|
"Ты — высокоточный профессиональный модуль OCR для канцелярии и документооборота.\n"
|
|
"Твоя задача — точно переписать весь текст с предоставленного изображения документа.\n"
|
|
"ПРАВИЛА:\n"
|
|
"1. Переписывай текст дословно, сохраняя структуру, заголовки, списки, нумерацию и таблицы.\n"
|
|
"2. Запрещено добавлять вводные слова, приветствия, комментарии ('Спасибо за обращение', 'Вот текст' и т.д.).\n"
|
|
"3. Выводи СТРОГО чистый распознанный текст документа."
|
|
)
|
|
user_message = {
|
|
"role": "user",
|
|
"content": f"Распознай весь печатный и рукописный текст страницы №{page_num} без пропусков.",
|
|
"images": [image_b64]
|
|
}
|
|
try:
|
|
res = call_ollama_chat(
|
|
messages=[{"role": "system", "content": system_prompt}, user_message],
|
|
is_vision=True,
|
|
timeout=300 # 5 минут на тяжелые страницы
|
|
)
|
|
return (res.get("content") or "").strip()
|
|
except Exception as e:
|
|
logger.error(f"[Office OCR] Ошибка распознавания страницы {page_num}: {e}")
|
|
return f"[Ошибка распознавания страницы {page_num}: {e}]"
|
|
|
|
|
|
def process_pdf_full(file_path: str, max_pages: Optional[int] = None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Обрабатывает PDF целиком:
|
|
- Если есть качественный цифровой текст — мгновенно извлекает его со всех страниц.
|
|
- Если страница является сканом/картинкой — рендерит ее в 200 DPI и прогоняет через Vision OCR.
|
|
"""
|
|
results = []
|
|
doc = fitz.open(file_path)
|
|
total_pages = len(doc)
|
|
limit = min(total_pages, max_pages) if max_pages else total_pages
|
|
|
|
logger.info(f"[Office] Начало обработки PDF: {os.path.basename(file_path)} (всего страниц: {total_pages})")
|
|
|
|
for idx in range(limit):
|
|
page_num = idx + 1
|
|
page = doc[idx]
|
|
extracted_text = (page.get_text("text") or "").strip()
|
|
|
|
# Если на странице есть хороший машинный текст (не скан)
|
|
if len(extracted_text) > 80:
|
|
logger.info(f"[Office] Страница {page_num}/{limit}: извлечен цифровой текст ({len(extracted_text)} симв.)")
|
|
results.append({
|
|
"page": page_num,
|
|
"method": "DIGITAL_TEXT",
|
|
"text": extracted_text
|
|
})
|
|
else:
|
|
# Чистый скан — рендерим страницу в PNG и передаем в Vision LLM
|
|
logger.info(f"[Office] Страница {page_num}/{limit}: распознавание скана через Ollama Vision...")
|
|
pix = page.get_pixmap(dpi=200)
|
|
img_b64 = base64.b64encode(pix.tobytes("png")).decode("utf-8")
|
|
|
|
ocr_text = ocr_image_b64(img_b64, page_num=page_num)
|
|
results.append({
|
|
"page": page_num,
|
|
"method": "VISION_OCR",
|
|
"text": ocr_text
|
|
})
|
|
|
|
doc.close()
|
|
return results |