aerospaceboy
ebe33392ca
Preserve document revisions and evidence, harden indexing and export behavior, add explicit material selection and direct-folder import, validate model JSON schemas with safe error messages, and retain student selection while loading scoring data.
210 lines
9.5 KiB
Python
210 lines
9.5 KiB
Python
import subprocess
|
|
import re
|
|
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 text_quality(text: str) -> dict:
|
|
unresolved = sum(text.count(char) for char in ("\ufffd", "\ufffe", "\uffff"))
|
|
return {"unresolved_characters": unresolved,
|
|
"needs_review": bool(unresolved),
|
|
"warning": "文字层含未解析字符，请核对原页；相关片段不用于自动内容评分" if unresolved else ""}
|
|
|
|
|
|
def _extract_tables(page):
|
|
"""Extract grids, then bounded three-line tables without treating prose as a grid."""
|
|
result = []
|
|
try:
|
|
for table in page.find_tables().tables:
|
|
result.append({"rows": table.extract(), "bbox": list(table.bbox), "method": "lines"})
|
|
except Exception:
|
|
pass
|
|
# A three-line table has aligned, long horizontal rules and no vertical rules.
|
|
# Group by x extent; do not infer unbounded tables from whole-page prose.
|
|
groups = {}
|
|
for drawing in page.get_drawings():
|
|
for item in drawing["items"]:
|
|
if item[0] != "l":
|
|
continue
|
|
a, b = item[1], item[2]
|
|
left, right = sorted((a.x, b.x))
|
|
if abs(a.y - b.y) > 1 or right - left < page.rect.width * .4:
|
|
continue
|
|
key = (round(left / 3), round(right / 3))
|
|
groups.setdefault(key, []).append((a.y, left, right))
|
|
for rules in groups.values():
|
|
rules = sorted(set(rules))
|
|
for offset in range(0, len(rules) - 2, 3):
|
|
top, header, bottom = rules[offset:offset + 3]
|
|
if not (5 < header[0] - top[0] < 60 and bottom[0] - header[0] > 15):
|
|
continue
|
|
clip = fitz.Rect(top[1], top[0], top[2], bottom[0])
|
|
header_text = page.get_text(clip=fitz.Rect(top[1], top[0], top[2], header[0]), sort=True)
|
|
if (any(marker in header_text for marker in ("//", "#include", "{", "}", ";"))
|
|
or re.search(r"\bListing\s+\d+", page.get_text(clip=clip), re.I)):
|
|
continue # Code-listing boxes are not tabular data.
|
|
if any(clip.intersects(fitz.Rect(t["bbox"])) for t in result):
|
|
continue
|
|
try:
|
|
found = page.find_tables(clip=clip, strategy="text", min_words_vertical=2)
|
|
for table in found.tables:
|
|
rows = [row for row in table.extract() if any(str(cell or "").strip() for cell in row)]
|
|
if len(rows) >= 2 and table.col_count >= 2:
|
|
result.append({"rows": rows, "bbox": list(clip), "method": "three_line_text",
|
|
"text": page.get_text(clip=clip, sort=True),
|
|
"warning": "列边界由文字对齐推断，请核对表头及跨列内容"})
|
|
except Exception:
|
|
continue
|
|
return result
|
|
|
|
|
|
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 = _extract_tables(page)
|
|
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, "text_quality": text_quality(text)})
|
|
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}
|