使用自定义 YOLOv10 和 Ollama (Llama 3.1) 提升 OCR 效果

YOLO 与 Ollama 在 OCR 中的应用价值

传统的 OCR 方法在提取简单图像中的文本时效果较好，但在处理包含其他视觉元素的图像时常遇到困难。通过使用自定义 YOLO 模型先检测文本区域，可以在 OCR 阶段仅对所需区域进行识别，显著降低干扰并提升准确率。

1. 训练自定义 YOLOv10 数据集

提升 OCR 精度的第一步是为数据集训练自定义 YOLO 模型。YOLO 通过将图像划分为网格进行物体检测，是检测图像中文本的理想选择。我们可以使用已标注的图书封面数据集，并在其上训练 YOLOv10 模型。YOLOv10 尤其适合检测视频或扫描文档中的小物体。

from ultralytics import YOLO

model = YOLO("yolov10n.pt")
model.train(data="datasets/data.yaml", epochs=50, imgsz=640)

2. 视频中的文本检测与边界框绘制

YOLO 模型训练完成后，可应用于视频中以检测文本区域的边界框。这些边界框可将目标区域隔离出来，为下一步 OCR 打下基础：

pythonCopy codeimport cv2
# 打开视频文件
video_path = 'books.mov'
cap = cv2.VideoCapture(video_path)
# 加载 YOLO 模型
model = YOLO('model.pt')
# 对视频帧进行目标检测
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    # 绘制边界框
    processed_frame, results = predict_and_detect(model, frame)
    cv2.imshow('YOLO + OCR Detection', processed_frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()

3. 在边界框区域内运行 OCR

使用 YOLO 隔离文本区域后，可在这些特定区域内运行 OCR，精度较直接处理整张图片大幅提升：

import easyocr
# 初始化 EasyOCR
reader = easyocr.Reader(['en'])
# OCR 运行在边界框内
for result in results:
    ocr_results = run_ocr_on_boxes(frame, result.boxes)
    extracted_text = [detection[1] for ocr in ocr_results for detection in ocr]
    print(f"Extracted Text: {', '.join(extracted_text)}")

#输出：'THE, SECRET, HISTORY, DONNA, TARTT'

4. 使用 Ollama 优化文本

提取的文本经常包含 OCR 误差。此时，可以引入 Llama 3 来优化，使用提示词指导 LLM 辨识并整理内容，将杂乱无章的文本转化为结构化的数据。

import ollama
prompt = f"""
- 以下是 OCR 提取的文本，包含著名书籍及其作者的名称。
- 请识别书名和对应的作者，格式为：'<书名> : <作者名>'
文本:
{output_text}
"""
response = ollama.chat(
    model="llama3",
    messages=[{"role": "user", "content": prompt}]
)
cleaned_text = response['message']['content'].strip()
print(cleaned_text)
#输出 The Secret History : Donna Tartt

总结
通过自定义 YOLOv10、EasyOCR 和 LLM 的结合，大幅提升文本识别效果。这一流程适用于文本识别精度要求较高的应用场景，如数字图书馆、图书封面识别、档案系统等。