100d1a06eb
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
496 lines
14 KiB
Python
496 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
PDF 图片提取模块
|
||
|
||
从 PDF 文档中提取图片，支持：
|
||
- 使用 PyMuPDF (fitz) 提取嵌入式图片
|
||
- 保存到指定目录
|
||
- 生成图片元数据供 RAG 系统使用
|
||
"""
|
||
|
||
import os
|
||
import hashlib
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional, Tuple, Any
|
||
from dataclasses import dataclass, asdict
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class ImageInfo:
|
||
"""图片信息"""
|
||
image_id: str # 图片唯一 ID
|
||
original_name: str # 原始文件名
|
||
storage_path: str # 存储路径（相对路径）
|
||
page: int # 所在页码
|
||
width: int # 宽度
|
||
height: int # 高度
|
||
format: str # 格式 (png, jpg, etc.)
|
||
size_bytes: int # 文件大小
|
||
caption: str = "" # 图片说明（可选）
|
||
bbox: Optional[List[float]] = None # 边界框坐标
|
||
|
||
|
||
def extract_images_from_pdf(
|
||
pdf_path: str,
|
||
output_dir: str,
|
||
min_width: int = 100,
|
||
min_height: int = 100,
|
||
max_width: int = 2000, # 新增：最大宽度阈值（过滤跨页底纹）
|
||
max_height: int = 2000, # 新增：最大高度阈值
|
||
max_images: int = 50
|
||
) -> List[ImageInfo]:
|
||
"""
|
||
从 PDF 中提取图片
|
||
|
||
Args:
|
||
pdf_path: PDF 文件路径
|
||
output_dir: 图片输出目录
|
||
min_width: 最小宽度阈值（过滤小图标）
|
||
min_height: 最小高度阈值
|
||
max_width: 最大宽度阈值（过滤跨页底纹、背景横幅）
|
||
max_height: 最大高度阈值
|
||
max_images: 最大提取图片数量
|
||
|
||
Returns:
|
||
图片信息列表
|
||
"""
|
||
try:
|
||
import fitz # PyMuPDF
|
||
except ImportError:
|
||
logger.warning("PyMuPDF 未安装，无法提取图片。请运行: pip install PyMuPDF")
|
||
return []
|
||
|
||
pdf_path = Path(pdf_path)
|
||
if not pdf_path.exists():
|
||
raise FileNotFoundError(f"PDF 文件不存在: {pdf_path}")
|
||
|
||
# 创建输出目录
|
||
output_dir = Path(output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
images = []
|
||
image_count = 0
|
||
|
||
try:
|
||
doc = fitz.open(str(pdf_path))
|
||
filename = pdf_path.stem
|
||
|
||
for page_num in range(len(doc)):
|
||
page = doc[page_num]
|
||
image_list = page.get_images(full=True)
|
||
|
||
for img_index, img_info in enumerate(image_list):
|
||
if image_count >= max_images:
|
||
break
|
||
|
||
try:
|
||
# 获取图片引用
|
||
xref = img_info[0]
|
||
|
||
# 提取图片
|
||
base_image = doc.extract_image(xref)
|
||
if not base_image:
|
||
continue
|
||
|
||
image_bytes = base_image.get("image")
|
||
if not image_bytes:
|
||
continue
|
||
|
||
# 获取图片属性
|
||
width = base_image.get("width", 0)
|
||
height = base_image.get("height", 0)
|
||
image_ext = base_image.get("ext", "png")
|
||
|
||
# 过滤太小的图片（通常是图标）
|
||
if width < min_width or height < min_height:
|
||
continue
|
||
|
||
# 过滤太大的图片（跨页底纹、背景横幅）
|
||
if width > max_width or height > max_height:
|
||
continue
|
||
|
||
# 生成唯一 ID
|
||
image_id = f"{filename}_p{page_num + 1}_img{img_index + 1}"
|
||
|
||
# 确定文件扩展名
|
||
if image_ext in ["jpeg", "jpg"]:
|
||
ext = ".jpg"
|
||
elif image_ext == "png":
|
||
ext = ".png"
|
||
else:
|
||
ext = f".{image_ext}"
|
||
|
||
# 保存图片
|
||
image_filename = f"{image_id}{ext}"
|
||
image_path = output_dir / image_filename
|
||
|
||
with open(image_path, "wb") as f:
|
||
f.write(image_bytes)
|
||
|
||
# 记录图片信息
|
||
images.append(ImageInfo(
|
||
image_id=image_id,
|
||
original_name=filename,
|
||
storage_path=f"images/{image_filename}",
|
||
page=page_num + 1,
|
||
width=width,
|
||
height=height,
|
||
format=image_ext,
|
||
size_bytes=len(image_bytes),
|
||
caption=f"图片 {img_index + 1}"
|
||
))
|
||
|
||
image_count += 1
|
||
|
||
except Exception as e:
|
||
# 单个图片提取失败不影响其他图片
|
||
logger.warning(f"提取图片失败 (页 {page_num + 1}, 图片 {img_index + 1}): {e}")
|
||
continue
|
||
|
||
if image_count >= max_images:
|
||
break
|
||
|
||
doc.close()
|
||
|
||
except Exception as e:
|
||
logger.error(f"PDF 图片提取失败: {e}")
|
||
return []
|
||
|
||
return images
|
||
|
||
|
||
def extract_images_batch(
|
||
pdf_dir: str,
|
||
output_dir: str,
|
||
**kwargs
|
||
) -> Dict[str, List[ImageInfo]]:
|
||
"""
|
||
批量提取 PDF 目录下所有文件的图片
|
||
|
||
Args:
|
||
pdf_dir: PDF 文件目录
|
||
output_dir: 图片输出目录
|
||
**kwargs: 传递给 extract_images_from_pdf 的参数
|
||
|
||
Returns:
|
||
{文件名: [ImageInfo, ...], ...}
|
||
"""
|
||
pdf_dir = Path(pdf_dir)
|
||
output_dir = Path(output_dir)
|
||
|
||
results = {}
|
||
|
||
for pdf_file in pdf_dir.glob("**/*.pdf"):
|
||
try:
|
||
# 为每个 PDF 创建子目录
|
||
pdf_output_dir = output_dir / pdf_file.stem
|
||
|
||
images = extract_images_from_pdf(
|
||
str(pdf_file),
|
||
str(pdf_output_dir),
|
||
**kwargs
|
||
)
|
||
|
||
if images:
|
||
results[pdf_file.name] = images
|
||
logger.info(f"{pdf_file.name}: 提取 {len(images)} 张图片")
|
||
|
||
except Exception as e:
|
||
logger.error(f"{pdf_file.name}: {e}")
|
||
|
||
return results
|
||
|
||
|
||
def get_images_base_path() -> str:
|
||
"""获取图片存储的基础路径"""
|
||
try:
|
||
from config import DOCUMENTS_PATH
|
||
return os.path.join(DOCUMENTS_PATH, "images")
|
||
except ImportError:
|
||
return "documents/images"
|
||
|
||
|
||
# ==================== 集成到现有解析器 ====================
|
||
|
||
def filter_noise_images(
|
||
images: List[Any],
|
||
min_size: int = 100,
|
||
max_size: int = 2000,
|
||
enable_hash_dedup: bool = True,
|
||
enable_content_check: bool = True,
|
||
max_aspect_ratio: float = 10.0
|
||
) -> List[Any]:
|
||
"""
|
||
三级噪音图片过滤管道
|
||
|
||
Level 1: 尺寸过滤（过滤图标、跨页底纹）
|
||
Level 2: Hash 去重（过滤重复图片）
|
||
Level 3: 内容检测（过滤纯色背景、装饰横幅）
|
||
|
||
Args:
|
||
images: 图片列表（ImageInfo 或 dict）
|
||
min_size: 最小尺寸阈值（像素），过滤图标
|
||
max_size: 最大尺寸阈值（像素），过滤跨页底纹
|
||
enable_hash_dedup: 启用 Hash 去重
|
||
enable_content_check: 启用内容相关性检测
|
||
max_aspect_ratio: 最大宽高比阈值，过滤装饰横幅
|
||
|
||
Returns:
|
||
过滤后的图片列表
|
||
"""
|
||
if not images:
|
||
return images
|
||
|
||
filtered = []
|
||
seen_hashes = set()
|
||
|
||
for img in images:
|
||
# 支持 dataclass 和 dict 两种格式
|
||
if hasattr(img, 'width'):
|
||
width, height = img.width, img.height
|
||
storage_path = getattr(img, 'storage_path', '')
|
||
else:
|
||
width = img.get('width', 0)
|
||
height = img.get('height', 0)
|
||
storage_path = img.get('storage_path', '')
|
||
|
||
# Level 1: 尺寸过滤
|
||
if width < min_size or height < min_size:
|
||
continue # 图标、装饰线条
|
||
if width > max_size or height > max_size:
|
||
continue # 跨页底纹、背景横幅
|
||
|
||
# Level 2: Hash 去重
|
||
if enable_hash_dedup and storage_path:
|
||
try:
|
||
img_hash = _compute_image_hash(storage_path)
|
||
if img_hash in seen_hashes:
|
||
continue # 重复图片
|
||
seen_hashes.add(img_hash)
|
||
except Exception:
|
||
pass # Hash 计算失败时跳过去重
|
||
|
||
# Level 3: 内容相关性检测
|
||
if enable_content_check:
|
||
aspect_ratio = max(width, height) / max(min(width, height), 1)
|
||
|
||
# 过滤极端宽高比（装饰横幅）
|
||
if aspect_ratio > max_aspect_ratio:
|
||
continue
|
||
|
||
# 过滤纯色/渐变背景
|
||
if storage_path and _is_solid_color_image(storage_path):
|
||
continue
|
||
|
||
filtered.append(img)
|
||
|
||
return filtered
|
||
|
||
|
||
def _compute_image_hash(image_path: str) -> str:
|
||
"""
|
||
计算图片文件的 Hash 值
|
||
|
||
Args:
|
||
image_path: 图片文件路径
|
||
|
||
Returns:
|
||
MD5 Hash 字符串
|
||
"""
|
||
hash_md5 = hashlib.md5()
|
||
with open(image_path, 'rb') as f:
|
||
for chunk in iter(lambda: f.read(4096), b''):
|
||
hash_md5.update(chunk)
|
||
return hash_md5.hexdigest()
|
||
|
||
|
||
def _is_solid_color_image(image_path: str, threshold: float = 0.95) -> bool:
|
||
"""
|
||
检测图片是否为纯色/渐变背景（装饰性横幅）
|
||
|
||
使用简单的颜色分布检测：
|
||
- 如果图片 95% 以上像素属于同一颜色范围，判定为纯色背景
|
||
|
||
Args:
|
||
image_path: 图片文件路径
|
||
threshold: 纯色判定阈值
|
||
|
||
Returns:
|
||
True 表示是纯色背景图片
|
||
"""
|
||
try:
|
||
from PIL import Image
|
||
import numpy as np
|
||
except ImportError:
|
||
# PIL/numpy 未安装，跳过检测
|
||
return False
|
||
|
||
try:
|
||
img = Image.open(image_path)
|
||
# 缩小图片加速处理
|
||
img.thumbnail((100, 100))
|
||
img_array = np.array(img)
|
||
|
||
if img_array.ndim == 2:
|
||
# 灰度图
|
||
unique, counts = np.unique(img_array, return_counts=True)
|
||
elif img_array.ndim == 3:
|
||
# 彩色图，计算颜色直方图
|
||
pixels = img_array.reshape(-1, img_array.shape[-1])
|
||
# 量化颜色（减少颜色数量）
|
||
quantized = (pixels // 32) * 32
|
||
unique, counts = np.unique(quantized, axis=0, return_counts=True)
|
||
else:
|
||
return False
|
||
|
||
# 如果主颜色占比超过阈值，判定为纯色背景
|
||
max_color_ratio = max(counts) / sum(counts)
|
||
return max_color_ratio > threshold
|
||
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def enrich_chunks_with_images(
|
||
chunks: List[Any],
|
||
images: List[ImageInfo],
|
||
source_file: str
|
||
) -> List[Any]:
|
||
"""
|
||
为分块添加图片信息
|
||
|
||
根据页码将图片关联到对应的分块
|
||
|
||
Args:
|
||
chunks: 分块列表（ChunkMetadata 或 dict）
|
||
images: 图片信息列表
|
||
source_file: 源文件名
|
||
|
||
Returns:
|
||
添加了图片信息的分块列表
|
||
"""
|
||
if not images:
|
||
return chunks
|
||
|
||
# 按页码分组图片
|
||
page_to_images = {}
|
||
for img in images:
|
||
page = img.page
|
||
if page not in page_to_images:
|
||
page_to_images[page] = []
|
||
page_to_images[page].append({
|
||
"id": img.image_id,
|
||
"caption": img.caption,
|
||
"page": img.page,
|
||
"width": img.width,
|
||
"height": img.height
|
||
})
|
||
|
||
# 为每个分块添加图片信息
|
||
for chunk in chunks:
|
||
# 支持 dataclass 和 dict 两种格式
|
||
if hasattr(chunk, 'page_start'):
|
||
page_start = chunk.page_start
|
||
page_end = getattr(chunk, 'page_end', page_start)
|
||
else:
|
||
page_start = chunk.get('page_start', 1)
|
||
page_end = chunk.get('page_end', page_start)
|
||
|
||
# 单页绑定：仅在切片不跨页时绑定图片
|
||
chunk_images = []
|
||
if page_start == page_end and page_start in page_to_images:
|
||
chunk_images = page_to_images[page_start]
|
||
|
||
# 应用噪音过滤
|
||
chunk_images = filter_noise_images(chunk_images)
|
||
|
||
# 限制每切片最多 3 张图片
|
||
chunk_images = chunk_images[:3]
|
||
|
||
# 添加到分块
|
||
if chunk_images:
|
||
if hasattr(chunk, '__dict__'):
|
||
# dataclass
|
||
chunk.images = chunk_images
|
||
else:
|
||
# dict
|
||
chunk['images'] = chunk_images
|
||
|
||
return chunks
|
||
|
||
|
||
# ==================== 测试 ====================
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
|
||
if sys.platform == 'win32':
|
||
sys.stdout.reconfigure(encoding='utf-8')
|
||
|
||
print("=" * 60)
|
||
print("PDF 图片提取模块测试")
|
||
print("=" * 60)
|
||
|
||
# 检查依赖
|
||
try:
|
||
import fitz
|
||
print("[OK] PyMuPDF 已安装")
|
||
except ImportError:
|
||
print("[错误] PyMuPDF 未安装，请运行: pip install PyMuPDF")
|
||
sys.exit(1)
|
||
|
||
# 测试提取
|
||
if len(sys.argv) >= 2:
|
||
pdf_path = sys.argv[1]
|
||
output_dir = sys.argv[2] if len(sys.argv) >= 3 else "documents/images"
|
||
|
||
print(f"\n提取图片: {pdf_path}")
|
||
print(f"输出目录: {output_dir}")
|
||
|
||
images = extract_images_from_pdf(pdf_path, output_dir)
|
||
|
||
print(f"\n提取结果: {len(images)} 张图片")
|
||
for img in images[:10]:
|
||
print(f" - {img.image_id}: {img.width}x{img.height}, {img.size_bytes} bytes, 页码 {img.page}")
|
||
else:
|
||
print("\n用法: python image_extractor.py <pdf_path> [output_dir]")
|
||
print("\n功能演示: 创建模拟图片信息")
|
||
|
||
# 创建模拟数据演示功能
|
||
mock_images = [
|
||
ImageInfo(
|
||
image_id="test_p1_img1",
|
||
original_name="test.pdf",
|
||
storage_path="images/test_p1_img1.png",
|
||
page=1,
|
||
width=800,
|
||
height=600,
|
||
format="png",
|
||
size_bytes=45000,
|
||
caption="流程图"
|
||
),
|
||
ImageInfo(
|
||
image_id="test_p3_img1",
|
||
original_name="test.pdf",
|
||
storage_path="images/test_p3_img1.jpg",
|
||
page=3,
|
||
width=1200,
|
||
height=900,
|
||
format="jpg",
|
||
size_bytes=120000,
|
||
caption="组织架构图"
|
||
)
|
||
]
|
||
|
||
print("\n模拟图片信息:")
|
||
for img in mock_images:
|
||
print(f" ID: {img.image_id}")
|
||
print(f" 页码: {img.page}")
|
||
print(f" 尺寸: {img.width}x{img.height}")
|
||
print(f" 格式: {img.format}")
|
||
print(f" 大小: {img.size_bytes} bytes")
|
||
print()
|