AI Agent RAG Grok AI教我做事

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

Posted by Ramendeus on March 11, 2025

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

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

解析PDF中图像的原理

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

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

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

完整可用代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200

#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能解析图片。但是从图片读取的字符串效果不是很理想。

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

CATALOG
    FEATURED TAGS
    流式 Streaming WebRTC RAG AI教我做事 AI Agent Grok LLM AIGC Ollama 字形绘梦 微信小程序 GPT WebUI Lite 个人大模型 Deepseek HayStack 绘声 C# LlamaIndex SD
    FRIENDS