feat: replace pymupdf image extraction with pdf2image + Heron layout detection

PDF image extraction now renders pages with pdf2image, detects figure regions using docling-layout-heron (RT-DETRv2), and crops only the detected pictures. Removes full-page fallback for text-only pages. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-02 22:38:31 +08:00
parent bbf6b921af
commit 6395ee3b49
7 changed files with 381 additions and 55 deletions
@@ -0,0 +1,140 @@
+"""测试 Heron 版面检测 + 裁剪 + VLM 图片描述流程"""
+import argparse
+import asyncio
+import logging
+import sys
+import time
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent.parent))
+
+from PIL import Image
+from src.rag.document_loaders import PDFImageExtractor
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+PROJECT_ROOT = Path(__file__).parent.parent.parent  # course-agent-od/
+TEST_PDF = (
+    PROJECT_ROOT
+    / "data" / "uploads" / "testuser" / "knowledge_bases" / "植物知识图谱"
+    / "ab810d4a-b030-45c0-94d2-bd6e2a23053b.pdf"
+)
+OUTPUT_DIR = PROJECT_ROOT / "data" / "images" / "_test_pdf_extract"
+
+
+def test_extract(max_pages: int | None = None):
+    """阶段1：Heron 版面检测 + 裁剪"""
+    if not TEST_PDF.exists():
+        logger.error(f"测试PDF不存在: {TEST_PDF}")
+        sys.exit(1)
+
+    # 清理旧输出
+    if OUTPUT_DIR.exists():
+        for f in OUTPUT_DIR.iterdir():
+            f.unlink()
+
+    print(f"\n{'='*60}")
+    print(f"[阶段1] Heron 版面检测 + 裁剪测试")
+    print(f"  PDF: {TEST_PDF.name}")
+    print(f"  输出: {OUTPUT_DIR}")
+    print(f"{'='*60}\n")
+
+    t0 = time.time()
+    images = PDFImageExtractor.extract_images(
+        str(TEST_PDF), str(OUTPUT_DIR),
+    )
+    elapsed = time.time() - t0
+
+    if not images:
+        print("  ❌ 未提取到任何图片")
+        return []
+
+    if max_pages:
+        images = [img for img in images if img["page"] <= max_pages]
+
+    # 按页分组统计
+    from collections import defaultdict
+    by_page = defaultdict(list)
+    for img in images:
+        by_page[img["page"]].append(img)
+
+    total_size = 0
+    fig_count = 0
+    full_count = 0
+    for page_num in sorted(by_page.keys()):
+        page_imgs = by_page[page_num]
+        for img in page_imgs:
+            path = Path(img["path"])
+            with Image.open(path) as pil:
+                w, h = pil.size
+            size_kb = img["size"] / 1024
+            total_size += img["size"]
+            is_full = "_full." in img["filename"]
+            tag = "FULL" if is_full else "FIG"
+            if is_full:
+                full_count += 1
+            else:
+                fig_count += 1
+            print(
+                f"  page {img['page']:>3d}  [{tag}] {img['filename']:<25s}  "
+                f"{w}x{h}  {size_kb:>7.1f} KB"
+            )
+
+    print(
+        f"\n  汇总: {len(images)} 张 ({fig_count} 裁剪 + {full_count} 整页兜底), "
+        f"{total_size/1024/1024:.1f} MB, 耗时 {elapsed:.1f}s\n"
+    )
+    return images
+
+
+async def test_vlm(images: list[dict], max_images: int = 3):
+    """阶段2：VLM图片描述"""
+    from src.llm.siliconflow import get_llm_client
+
+    # 优先选裁剪图
+    fig_images = [img for img in images if "_full." not in img["filename"]]
+    targets = (fig_images or images)[:max_images]
+
+    print(f"\n{'='*60}")
+    print(f"[阶段2] VLM图片描述测试 ({len(targets)} 张)")
+    print(f"{'='*60}\n")
+
+    client = get_llm_client()
+    sem = asyncio.Semaphore(3)
+
+    async def describe_one(idx: int, img: dict):
+        async with sem:
+            t0 = time.time()
+            desc = await client.describe_image(img["path"], img.get("context_text", ""))
+            elapsed = time.time() - t0
+            return idx, desc, elapsed
+
+    tasks = [describe_one(i, img) for i, img in enumerate(targets)]
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    for r in results:
+        if isinstance(r, Exception):
+            print(f"  ❌ 失败: {r}\n")
+            continue
+        idx, desc, elapsed = r
+        img = targets[idx]
+        print(f"  --- page {img['page']} ({img['filename']}) {elapsed:.1f}s ---")
+        print(f"  {desc}\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="测试PDF图片提取")
+    parser.add_argument("--vlm", action="store_true", help="启用VLM图片描述")
+    parser.add_argument("--max-pages", type=int, default=None, help="限制提取页数")
+    parser.add_argument("--max-vlm", type=int, default=3, help="VLM描述最大图片数")
+    args = parser.parse_args()
+
+    images = test_extract(max_pages=args.max_pages)
+
+    if args.vlm and images:
+        asyncio.run(test_vlm(images, max_images=args.max_vlm))
+
+
+if __name__ == "__main__":
+    main()