大模型应用：多模态图文精准识别：基于本地化OCR模型应用实践.78

# 核心依赖：transformers（多模态核心）、modelscope（模型下载）
pip install transformers modelscope -i https://pypi.tuna.tsinghua.edu.cn/simple
# 基础框架：torch（模型运行，CPU版本自动适配，GPU需装对应CUDA版本）
pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple
# 图片处理/可视化：Pillow（图片读取）、matplotlib（图片展示）
pip install pillow matplotlib -i https://pypi.tuna.tsinghua.edu.cn/simple
# 可选依赖：opencv-python（图片预处理，拓展功能用）
pip install opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple

# 导入核心库
# 多模态模型与专用处理器
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
# ModelScope模型下载工具
from modelscope import snapshot_download
# 图片处理库
from PIL import Image
# 图片可视化库
import matplotlib.pyplot as plt
# 文件路径与系统操作
import os
# 异常处理相关（可选，增强鲁棒性）
import traceback

# -------------------------- 1. 模型基础配置（仅需修改此处少量参数） --------------------------
# 模型名称（ModelScope官方仓库，无需手动找权重）
MODEL_NAME = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
# 模型本地缓存路径（建议绝对路径，避免中文/空格/特殊字符）
CACHE_DIR = "D:\\modelscope\\hub"
# 测试图片路径（替换为你的图片路径，相对/绝对均可，建议绝对路径）
TEST_IMG_PATH = "test_ocr.png"
# 生成最大新token数（控制识别结果长度，按需调整，默认1024足够大部分场景）
MAX_NEW_TOKENS = 1024
# 确保缓存目录存在，避免模型下载时创建目录失败
os.makedirs(CACHE_DIR, exist_ok=True)

# -------------------------- 2. 模型本地下载/缓存校验 --------------------------
print("="*60)
print("📥 正在下载/校验Qwen2-VL-OCR-2B模型缓存...")
print("提示：首次运行需下载模型（约数GB），后续直接使用本地缓存，无需重复下载！")
try:
    local_model_path = snapshot_download(
        MODEL_NAME,    # 模型名称
        cache_dir=CACHE_DIR,      # 本地缓存路径
    )
    print(f"✅ 模型下载/校验完成，本地缓存路径：\n{local_model_path}")
except Exception as e:
    print(f"❌ 模型下载失败，错误信息：{e}")
    print(f"📝 错误详情：{traceback.format_exc()}")
    exit()  # 模型下载失败，直接退出程序

# -------------------------- 3. 加载多模态处理器与模型 --------------------------
print("="*60)
print("⚙️  正在加载多模态处理器与Qwen2-VL-OCR-2B模型...")
try:
    # 加载多模态专用处理器：负责图片编码、文本分词、对话模板拼接
    processor = AutoProcessor.from_pretrained(
        local_model_path,
        trust_remote_code=True  # 适配Qwen2-VL的自定义代码，必须开启
    )
    # 加载Qwen2-VL多模态模型，用于OCR推理
    model = Qwen2VLForConditionalGeneration.from_pretrained(
        local_model_path,
        trust_remote_code=True,  # 适配自定义代码，必须开启
        device_map="auto",       # 自动分配设备（CPU/GPU/显存自动适配）
        torch_dtype="auto"       # 自动适配数据类型，节省显存/内存
    ).eval()  # 模型设为评估模式，关闭训练相关层，提升推理速度和稳定性
    print("✅ 多模态处理器与模型加载成功！")
    print(f"🔧 模型运行设备：{model.device}")  # 打印模型实际运行的设备（CPU/GPU）
except Exception as e:
    print(f"❌ 处理器/模型加载失败，错误信息：{e}")
    print(f"📝 错误详情：{traceback.format_exc()}")
    exit()

# -------------------------- 4. 定义OCR识别+图片可视化核心函数 --------------------------
def qwen2_vl_ocr_detect(img_path, prompt="识别这张图片中的所有文字，保留原始排版"):
    """
    Qwen2-VL-OCR-2B多模态OCR识别函数，支持自定义指令
    :param img_path: 待识别图片的路径（相对/绝对）
    :param prompt: 识别指令（自然语言，可自定义，如多语言识别、指定区域识别）
    :return: 模型识别的文字结果
    """
    # 第一步：校验图片路径是否存在
    if not os.path.exists(img_path):
        raise FileNotFoundError(f"图片路径不存在，请检查：{img_path}")
    # 第二步：读取图片并转为RGB格式（解决PNG透明通道、BMP灰度图等格式兼容问题）
    img = Image.open(img_path).convert("RGB")
    
    # 第三步：图片可视化输出（图文对照，直观查看待识别内容）
    print("="*60)
    print("🖼️  正在展示待识别图片...")
    # 解决Matplotlib中文显示乱码问题（核心配置，必加）
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 加载中文字体（黑体）
    plt.rcParams['axes.unicode_minus'] = False    # 解决负号显示异常问题
    # 设置图片展示尺寸，适配不同大小的图片
    plt.figure(figsize=(10, 8))
    plt.imshow(img)  # 加载图片
    plt.axis("off")  # 关闭坐标轴，让图片展示更整洁
    plt.title("待识别图片（Qwen2-VL-OCR-2B多模态识别）", fontsize=18, pad=20)
    plt.show()  # 弹出图片展示窗口（Jupyter中内嵌，本地脚本中弹出独立窗口）
    
    # 第四步：多模态模型推理（核心步骤，按模型要求构造输入）
    print("="*60)
    print(f"🔍 正在执行多模态OCR识别，指令：{prompt}")
    # 构造模型输入的对话格式（符合Qwen2-VL的指令式输入要求）
    messages = [
        {
            "role": "user",  # 角色为用户
            "content": [
                {"type": "image", "image": img},  # 输入图片
                {"type": "text", "text": prompt}  # 输入自然语言识别指令
            ]
        }
    ]
    # 应用模型的对话模板，生成标准输入文本（无需手动拼接，处理器自动完成）
    text_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    # 将文本指令和图片转为模型可识别的张量，并分配到模型运行的设备（CPU/GPU）
    inputs = processor(text=[text_prompt], images=[img], return_tensors="pt").to(model.device)
    # 模型生成识别结果：基于输入的图片和指令，生成文字
    generated_ids = model.generate(
        **inputs,
        max_new_tokens=MAX_NEW_TOKENS,  # 最大生成token数，控制结果长度
        do_sample=False,                # 关闭采样，保证识别结果的确定性
        temperature=0.0                 # 温度值设为0，避免生成随机内容
    )
    # 将模型生成的token序列解码为人类可读的文字，跳过特殊符号
    ocr_result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return ocr_result

# -------------------------- 5. 主程序运行入口 --------------------------
if __name__ == "__main__":
    try:
        # 调用OCR识别函数，可自定义识别指令（示例：默认中文识别/英文识别/多语言识别）
        # 中文识别指令（默认）：prompt="识别这张图片中的所有文字，保留原始排版"
        # 英文识别指令：prompt="Recognize all text in this image and keep the original format"
        # 中英混合识别：prompt="识别这张图片中的中英文文字，保留原始排版"
        ocr_result = qwen2_vl_ocr_detect(
            img_path=TEST_IMG_PATH,
            prompt="识别这张图片中的所有文字，严格保留原始的排版和换行"
        )
        # 输出最终识别结果
        print("="*60)
        print("🎉 Qwen2-VL-OCR-2B多模态OCR识别完成！")
        print("📜 识别结果如下：")
        print("-"*60)
        print(ocr_result)
        print("-"*60)
    except Exception as e:
        print(f"❌ OCR识别过程失败，错误信息：{e}")
        print(f"📝 错误详情：{traceback.format_exc()}")

from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from modelscope import snapshot_download
from PIL import Image
import matplotlib.pyplot as plt
import os
import traceback

# ====================== 配置不变 ======================
model_name = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
cache_dir = "D:\\modelscope\\hub"
TEST_IMG_PATH = "image_understand.png"   # 可以是任意图片：风景、截图、海报、人物、物品
MAX_NEW_TOKENS = 1024
os.makedirs(cache_dir, exist_ok=True)

# ====================== 下载模型（不变） ======================
print("="*60)
print("📥 下载/校验模型缓存...")
try:
    local_model_path = snapshot_download(model_name, cache_dir=cache_dir)
    print(f"✅ 模型路径：{local_model_path}")
except Exception as e:
    print(f"❌ 下载失败：{e}")
    exit()

# ====================== 加载模型/处理器（不变） ======================
print("="*60)
print("⚙️ 加载多模态模型...")
try:
    processor = AutoProcessor.from_pretrained(local_model_path, trust_remote_code=True)
    model = Qwen2VLForConditionalGeneration.from_pretrained(
        local_model_path, trust_remote_code=True, device_map="auto", torch_dtype="auto"
    ).eval()
    print(f"✅ 加载成功，设备：{model.device}")
except Exception as e:
    print(f"❌ 加载失败：{e}")
    exit()

# ====================== 核心：识图函数（仅prompt不同） ======================
def image_understand(img_path, prompt="请详细描述这张图片的全部内容"):
    if not os.path.exists(img_path):
        raise FileNotFoundError(f"图片不存在：{img_path}")
    
    img = Image.open(img_path).convert("RGB")

    # 图片可视化输出（保留你要的图片展示）
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.figure(figsize=(10, 6))
    plt.imshow(img)
    plt.axis("off")
    plt.title("待识图图片（Qwen2-VL 多模态理解）", fontsize=16)
    plt.show()

    # 识图推理（格式完全同OCR，只是prompt变了）
    print("🔍 正在识图理解...")
    messages = [
        {"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": prompt}]}
    ]
    text_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text_prompt], images=[img], return_tensors="pt").to(model.device)

    generated_ids = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, temperature=0.0)
    result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return result

# ====================== 主程序：识图 ======================
if __name__ == "__main__":
    try:
        # -------------- 这里切换不同识图指令 --------------
        # 指令1：通用详细描述（最强通用识图）
        prompt = "请详细描述这张图片：场景、主体、物体、颜色、布局、文字、氛围、用途"
        
        # 指令2：问答式识图
        # prompt = "这张图片是什么场景？主要表达什么？"
        
        # 指令3：只识别主体物体
        # prompt = "请识别图中的主要物体/人物，不要多余描述"

        result = image_understand(TEST_IMG_PATH, prompt=prompt)
        
        print("="*60)
        print("🎉 识图完成！模型理解结果：")
        print("-"*60)
        print(result)
        print("-"*60)
        
    except Exception as e:
        print(f"❌ 识图失败：{e}")