pengxiao
6395ee3b49
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>
141 lines
4.2 KiB
Python
141 lines
4.2 KiB
Python
"""测试 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()
|