AI教我做事之RAG开发-3 如何解析PDF中的图像

最新推荐文章于 2025-05-08 10:04:30 发布

拉达曼迪斯II

最新推荐文章于 2025-05-08 10:04:30 发布

阅读量742

点赞数 14

分类专栏： AIGC学习 AI创业文章标签：人工智能 pdf windows mid journey AIGC stable-diffusion RAG

本文链接：https://blog.csdn.net/ms44/article/details/146172019

版权

AI创业同时被 2 个专栏收录

503 篇文章

订阅专栏

AIGC学习

470 篇文章

订阅专栏

该行为的核心是能解析PDF文件，并且其中的图片内容可以通过OCR获取到文字，进行识别和解析。

解析PDF中图像的原理

在RAG系统中，解析PDF中的图像通常涉及以下步骤：

提取图像：从PDF文件中提取嵌入的图像。
图像处理：将提取的图像转换为可分析的格式（如JPEG/PNG）。
图像内容分析：使用OCR（光学字符识别）提取图像中的文本，或使用图像识别技术分析图像内容。
整合到RAG：将提取的内容作为上下文输入到RAG的检索和生成流程中。

完整可用代码

#from langchain.document_loaders import PyPDFLoader
from langchain_community.document_loaders import PyPDFLoader

from langchain.text_splitter import RecursiveCharacterTextSplitter

#from langchain_community.llms import Ollama
from langchain_ollama import OllamaLLM
#from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings
#from langchain.vectorstores import FAISS
from langchain_community.vectorstores import FAISS

#from langchain.llms import HuggingFacePipeline
from langchain_community.llms import HuggingFacePipeline

from langchain.chains import RetrievalQA
from transformers import pipeline
from langchain.prompts import PromptTemplate
from pathlib import Path
import fitz  # PyMuPDF，用于PDF处理
import pytesseract  # 用于OCR
from PIL import Image  # 图像处理
import io
import os
import torch

from langchain.docstore.document import Document

def extract_images_from_pdf(pdf_path, output_dir="extracted_images"):
    """
    从PDF中提取图像并保存到指定目录
    参数：
        pdf_path: PDF文件路径
        output_dir: 提取图像的保存目录
    返回：
        提取的图像文件路径列表
    """
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 打开PDF文件
    pdf_document = fitz.open(pdf_path)
    image_paths = []

    # 遍历每一页
    for page_num in range(len(pdf_document)):
        page = pdf_document[page_num]
        image_list = page.get_images(full=True)  # 获取页面中的所有图像

        # 遍历页面中的每张图像
        for img_index, img in enumerate(image_list):
            xref = img[0]  # 图像的xref引用
            base_image = pdf_document.extract_image(xref)  # 提取图像数据
            image_bytes = base_image["image"]  # 图像的字节数据
            image_ext = base_image["ext"]  # 图像扩展名（如jpeg, png）

            # 保存图像
            image_filename = f"{output_dir}/image_page{page_num+1}_{img_index}.{image_ext}"
            with open(image_filename, "wb") as image_file:
                image_file.write(image_bytes)
            image_paths.append(image_filename)

    pdf_document.close()
    return image_paths

def extract_text_from_image(image_path):
    """
    使用OCR从图像中提取文本
    参数：
        image_path: 图像文件路径
    返回：
        提取的文本内容
    """
    # 打开图像
    image = Image.open(image_path)
    # 使用pytesseract进行OCR
    text = pytesseract.image_to_string(image)
    return text

def process_pdf_images_for_rag(pdf_path):
    """
    处理PDF中的图像并提取内容，用于RAG
    参数：
        pdf_path: PDF文件路径
    返回：
        包含图像内容的字典列表
    """
    # 提取图像
    image_paths = extract_images_from_pdf(pdf_path)
    image_contents = []

    # 对每张图像进行OCR并提取文本
    for img_path in image_paths:
        text = extract_text_from_image(img_path)
        image_contents.append({
            "image_path": img_path,
            "text_content": text.strip()
        })

    return image_contents

# 1. 加载和处理文档
def load_documents(file_path):
    # 使用PyPDFLoader加载PDF文档
    loader = PyPDFLoader(file_path)
    documents = loader.load()
    return documents

def split_document(documents):
    # 分割文档为小块，便于检索,每个块约500字符，适合短文档；对于长文档可调整到1000
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = text_splitter.split_documents(documents)
    return docs

# 2. 创建向量存储
def create_vector_store(docs):
    # 使用HuggingFace的嵌入模型生成向量,all-MiniLM-L6-v2是一个轻量高效的嵌入模型，适合快速开发；生产环境可考虑更强的模型如all-mpnet-base-v2。
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    
    # 使用FAISS创建向量存储
    
    vector_store = FAISS.from_documents(docs, embeddings)
    return vector_store

# 3. 初始化语言模型
def initialize_llm():
    # 使用Hugging Face的生成模型（这里用distilgpt2作为示例）,Demo中使用distilgpt2（轻量但生成质量一般）；建议替换为gpt-neo-1.3B或LLaMA（需本地部署）。
    # text_generation_pipeline = pipeline("text-generation", model="distilgpt2", max_length=200)
    # llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
    myLLM = OllamaLLM(model="llama2")
    return myLLM

# 4. 创建RAG链
def create_rag_chain(vector_store, llm):
# 自定义提示模板，只要求简洁回答
    prompt_template = """根据以下上下文，直接用中文回答问题，不要添加多余内容。
上下文: {context}
问题: {question}
回答: """
    
    PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    
    # 创建RetrievalQA链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
        return_source_documents=False,  # 不返回来源文档
        chain_type_kwargs={"prompt": PROMPT}  # 使用自定义提示
    )
    return qa_chain

# 5. 主函数
def main():
    pytesseract.pytesseract.tesseract_cmd = 'tesseract.exe'
    pdf_file = Path('testdata') / 'RAG-TestImage-3.pdf'
    
    # 处理PDF中的图像
    results = process_pdf_images_for_rag(pdf_file)
    
    # 打印结果
    myStrings = []
    for result in results:
        print(f"Image: {result['image_path']}")
        print(f"Extracted Text: {result['text_content']}")
        myStrings.append(result['text_content'])
        print("-" * 50)

    # 加载并分割文档
    print("加载和分割文档...")
    # 将字符串内容认为是documents
    documents = [Document(page_content=s) for s in myStrings]
    # 分割文档（这是为了保持原来的逻辑和架构）
    docs=split_document(documents)

    # 创建向量存储
    print("创建向量存储...")
    vector_store = create_vector_store(docs)
    
    # 初始化语言模型
    print("初始化语言模型...")
    llm = initialize_llm()
    
    # 创建RAG链
    print("创建RAG链...")
    rag_chain = create_rag_chain(vector_store, llm)
    
    # 示例查询
    #query = "What is the main topic of the document?"
    query = "字形绘梦是哪个公司的产品?"
    print(f"查询: {query}")
    
    # 执行RAG查询,获取精准答案
    #result = rag_chain({"query": query})
    result = rag_chain.invoke({"query": query})
    print(f"回答: {result['result'].strip()}")

if __name__ == "__main__":
    main()

注意事项

Tesseract配置：Windows用户需指定tesseract.exe路径，Linux/Mac用户确保Tesseract在PATH中。 On Windows 将 C:\Program Files\Tesseract-OCR 加入环境变量PATH中。
二进制安装文件从这里下载

download binary from https://github.com/UB-Mannheim/tesseract/wiki. then add pytesseract.pytesseract.tesseract_cmd = 'C:\Program Files (x86)\Tesseract-OCR\tesseract.exe' to your script. 完整信息：https://github.com/UB-Mannheim/tesseract/wiki

语言支持：OCR默认使用英语，若需其他语言，需安装语言包并在image_to_string中指定（如lang=’chi_sim’表示简体中文）。
图像质量：低质量图像可能导致OCR结果不准确，可预处理图像（如增强对比度）。

. RAG中的应用

检索阶段：将text_content输入到向量数据库（如FAISS）中，构建索引。在使用 FAISS（Facebook AI Similarity Search）将字符串存入 embeddings 时，FAISS 本身并不直接处理字符串，而是操作数值型的向量（embeddings）。因此，你需要先将字符串转换为向量表示（通常使用预训练的嵌入模型，如 BERT、Word2Vec 或 SentenceTransformers），然后将这些向量存入 FAISS 索引中
生成阶段：将提取的文本作为上下文，提供给语言模型生成回答。

这个DEMO能解析图片。但是从图片读取的字符串效果不是很理想。