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

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

解析PDF中图像的原理

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

  1. 提取图像:从PDF文件中提取嵌入的图像。
  2. 图像处理:将提取的图像转换为可分析的格式(如JPEG/PNG)。
  3. 图像内容分析:使用OCR(光学字符识别)提取图像中的文本,或使用图像识别技术分析图像内容。
  4. 整合到RAG:将提取的内容作为上下文输入到RAG的检索和生成流程中。
AI教我做事之RAG开发-3 如何解析PDF中的图像

完整可用代码

#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中的图像

RA/SD 衍生者AI训练营。发布者:稻草人,转载请注明出处:https://www.shxcj.com/ai%e6%95%99%e6%88%91%e5%81%9a%e4%ba%8b%e4%b9%8brag%e5%bc%80%e5%8f%91-3-%e5%a6%82%e4%bd%95%e8%a7%a3%e6%9e%90pdf%e4%b8%ad%e7%9a%84%e5%9b%be%e5%83%8f/

(0)
上一篇 2天前
下一篇 1天前

相关推荐

发表回复

登录后才能评论
本文授权以下站点有原版访问授权 https://www.shxcj.com https://www.2img.ai https://www.2video.cn