159 lines
6.8 KiB
Python
159 lines
6.8 KiB
Python
import subprocess
|
|
import tempfile
|
|
from functools import lru_cache
|
|
from importlib.metadata import version
|
|
from pathlib import Path
|
|
from threading import Lock
|
|
|
|
import pymupdf as fitz
|
|
from docx import Document
|
|
from pptx import Presentation
|
|
|
|
from .config import get_settings
|
|
|
|
|
|
_ocr_lock = Lock()
|
|
|
|
|
|
def _render(page, output: Path):
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
page.get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False).save(output)
|
|
return str(output)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _ocr_engine(model_dir: str):
|
|
root = Path(model_dir)
|
|
detection, recognition = root / "detection", root / "recognition"
|
|
if not detection.is_dir() or not recognition.is_dir():
|
|
raise RuntimeError(f"扫描PDF需要离线OCR模型目录: {detection} 和 {recognition}")
|
|
try:
|
|
from paddleocr import PaddleOCR
|
|
except ImportError as exc:
|
|
raise RuntimeError("扫描PDF需要PaddleOCR；请安装OCR依赖或交由教师处理") from exc
|
|
if int(version("paddleocr").split(".")[0]) >= 3:
|
|
def model_name(directory):
|
|
config = directory / "inference.yml"
|
|
if not config.is_file():
|
|
raise RuntimeError(f"离线OCR模型缺少 inference.yml: {directory}")
|
|
for line in config.read_text(encoding="utf-8").splitlines():
|
|
if line.strip().startswith("model_name:"):
|
|
return line.split(":", 1)[1].strip().strip("'\"")
|
|
raise RuntimeError(f"离线OCR模型未声明 model_name: {config}")
|
|
|
|
return PaddleOCR(text_detection_model_name=model_name(detection), text_detection_model_dir=str(detection),
|
|
text_recognition_model_name=model_name(recognition), text_recognition_model_dir=str(recognition),
|
|
use_textline_orientation=False, use_doc_orientation_classify=False, use_doc_unwarping=False,
|
|
device="cpu", enable_mkldnn=False)
|
|
return PaddleOCR(lang="ch", det_model_dir=str(detection), rec_model_dir=str(recognition), use_angle_cls=False, show_log=False)
|
|
|
|
|
|
def _ocr(image_path: str):
|
|
engine = _ocr_engine(str(get_settings().ocr_model_dir))
|
|
with _ocr_lock:
|
|
return _ocr_unlocked(engine, image_path)
|
|
|
|
|
|
def _ocr_unlocked(engine, image_path: str):
|
|
if hasattr(engine, "predict"):
|
|
result = engine.predict(image_path)
|
|
lines = []
|
|
for page in result:
|
|
payload = page.json if hasattr(page, "json") else page
|
|
payload = payload() if callable(payload) else payload
|
|
payload = payload.get("res", payload) if isinstance(payload, dict) else {}
|
|
lines.extend(payload.get("rec_texts", []))
|
|
return "\n".join(str(line) for line in lines)
|
|
result = engine.ocr(image_path, cls=True)
|
|
return "\n".join(line[1][0] for page in result for line in (page or []))
|
|
|
|
|
|
def parse_pdf(path: Path, asset_dir: Path, *, ocr=True):
|
|
pages = []
|
|
with fitz.open(path) as pdf:
|
|
if not len(pdf):
|
|
raise ValueError("PDF没有页面")
|
|
for index, page in enumerate(pdf):
|
|
image = _render(page, asset_dir / f"page-{index + 1}.png")
|
|
text = page.get_text(sort=True).strip()
|
|
if ocr and len(text) < 30:
|
|
text = _ocr(image)
|
|
tables = []
|
|
try:
|
|
tables = [{"rows": table.extract(), "bbox": list(table.bbox)} for table in page.find_tables().tables]
|
|
except Exception:
|
|
pass
|
|
images = [{"xref": item[0], "width": item[2], "height": item[3]} for item in page.get_images(full=True)]
|
|
pages.append({"page_number": index + 1, "text": text, "tables": tables, "images": images,
|
|
"section": next((line.strip() for line in text.splitlines() if line.strip()), ""),
|
|
"rendered_image_path": image})
|
|
return pages
|
|
|
|
|
|
def _convert_to_pdf(path: Path, output_dir: Path):
|
|
with tempfile.TemporaryDirectory() as profile:
|
|
command = ["libreoffice", f"-env:UserInstallation=file://{profile}", "--headless", "--convert-to", "pdf", "--outdir", str(output_dir), str(path)]
|
|
result = subprocess.run(command, capture_output=True, text=True, timeout=180, check=False)
|
|
pdf = output_dir / f"{path.stem}.pdf"
|
|
if result.returncode or not pdf.exists():
|
|
raise RuntimeError(f"LibreOffice转换失败: {(result.stderr or result.stdout).strip()[:500]}")
|
|
return pdf
|
|
|
|
|
|
def _docx_structure(path: Path):
|
|
doc = Document(path)
|
|
headings = [{"text": p.text, "style": p.style.name} for p in doc.paragraphs if p.style and p.style.name.startswith("Heading")]
|
|
tables = [[list(cell.text for cell in row.cells) for row in table.rows] for table in doc.tables]
|
|
images = len(doc.inline_shapes)
|
|
return {"headings": headings, "tables": tables, "image_count": images}
|
|
|
|
|
|
def _pptx_structure(path: Path):
|
|
slides = []
|
|
deck = Presentation(path)
|
|
for number, slide in enumerate(deck.slides, start=1):
|
|
texts, tables, images = [], [], 0
|
|
for shape in slide.shapes:
|
|
if shape.has_text_frame:
|
|
texts.append(shape.text)
|
|
if shape.has_table:
|
|
tables.append([[cell.text for cell in row.cells] for row in shape.table.rows])
|
|
if shape.shape_type == 13:
|
|
images += 1
|
|
try:
|
|
notes = slide.notes_slide.notes_text_frame.text if slide.has_notes_slide else ""
|
|
except Exception:
|
|
notes = ""
|
|
slides.append({"page_number": number, "text": "\n".join(texts), "tables": tables, "images": [{"count": images}],
|
|
"notes": notes, "section": texts[0][:200] if texts else ""})
|
|
if not slides:
|
|
raise ValueError("PPTX没有幻灯片")
|
|
return slides
|
|
|
|
|
|
def parse_document(path: Path, category: str, asset_dir: Path, *, ocr=True):
|
|
suffix = path.suffix.lower()
|
|
asset_dir.mkdir(parents=True, exist_ok=True)
|
|
if suffix == ".pdf":
|
|
pages = parse_pdf(path, asset_dir, ocr=ocr)
|
|
extra = {}
|
|
elif suffix in {".docx", ".doc"}:
|
|
extra = _docx_structure(path) if suffix == ".docx" else {}
|
|
pdf = _convert_to_pdf(path, asset_dir)
|
|
pages = parse_pdf(pdf, asset_dir, ocr=ocr)
|
|
pdf.unlink(missing_ok=True)
|
|
elif suffix == ".pptx":
|
|
slides = _pptx_structure(path)
|
|
pdf = _convert_to_pdf(path, asset_dir)
|
|
rendered = parse_pdf(pdf, asset_dir, ocr=False)
|
|
pdf.unlink(missing_ok=True)
|
|
if len(rendered) != len(slides):
|
|
raise ValueError("PPTX渲染页数与幻灯片数不一致")
|
|
pages = [{**slide, "rendered_image_path": rendered[index]["rendered_image_path"]} for index, slide in enumerate(slides)]
|
|
extra = {}
|
|
else:
|
|
raise ValueError("仅支持PDF、DOCX和PPTX")
|
|
if not any(page["text"].strip() for page in pages):
|
|
raise ValueError("未能提取文档内容")
|
|
return {"document_type": category, "source_path": str(path), "pages": pages, **extra}
|