用YOLO和LLM增强的OCR

news2025/4/5 16:21:29

虽然最近我花了很多时间在大型语言模型 (LLM) 上进行实验，但我对计算机视觉的热情始终未减。因此，当我有机会将两者融合在一起时，我迫不及待地想要立即开始。在 Goodreads 上扫描书籍封面并将其标记为已读一直感觉有点神奇，我很兴奋自己尝试一下。

将自定义训练的 YOLOv10 模型与 OCR 技术相结合可显著提高准确率，但真正的转变发生在集成 LLM（如 Llama 3.1）时——它将杂乱的 OCR 输出转换为可用于实际应用的精致文本。

NSDT工具推荐： Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - AI模型在线查看 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割

1、为什么需要 YOLO 、Ollama 和OCR？

传统的 OCR（光学字符识别）方法非常适合从简单图像中提取文本，但当文本与其他视觉元素交织在一起时，通常会遇到困难。通过首先使用自定义 YOLO 模型检测文本区域等对象，我们可以隔离这些区域以进行 OCR，从而显著减少噪音并提高准确率。

让我们通过在没有 YOLO 的图像上运行基本的 OCR 示例来证明这一点，以强调单独使用 OCR 的挑战：

import easyocr
import cv2
# Initialize EasyOCR
reader = easyocr.Reader(['en'])
# Load the image
image = cv2.imread('book.jpg')
# Run OCR directly
results = reader.readtext(image)
# Display results
for (bbox, text, prob) in results:
    print(f"Detected Text: {text} (Probability: {prob})")

输出结果如下：

THE 0 R |G |NAL B E STSELLE R THE SECRET HISTORY DONNA TARTT Haunting, compelling and brilliant The Times

虽然这种方法适用于较简单的图像，但当存在噪声或复杂的视觉模式时，你会注意到错误会增加。这时 YOLO 模型可以发挥巨大作用。

2、使用自定义数据集训练 Yolov10

使用对象检测增强 OCR 的第一步是在数据集上训练自定义 YOLO 模型。

YOLO（You Only Look Once）是一种功能强大的实时对象检测模型，它将图像划分为网格，使其能够在一次前向传递中识别多个对象。这种方法非常适合检测图像中的文本，尤其是当你想要通过隔离特定区域来改善 OCR 结果时。

我们将使用此处链接的预标注书籍封面数据集，并在其上训练 YOLOv10 模型。YOLOv10 针对较小的物体进行了优化，使其非常适合在视频或扫描文档等具有挑战性的环境中检测文本。

from ultralytics import YOLO

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

就我而言，在 Google Colab 上训练此模型大约需要 6 个小时，共 50 个 epoch。你可以调整 epoch 数量、数据集大小或超参数等参数来提高模型的性能和准确性。

2、在视频上运行自定义模型以获取边界框

YOLO 模型训练完成后，您可以将其应用于视频以检测文本区域周围的边界框。这些边界框隔离了感兴趣的区域，确保 OCR 过程更加清晰：

import cv2
# Open video file
video_path = 'books.mov'
cap = cv2.VideoCapture(video_path)
# Load YOLO model
model = YOLO('model.pt')
# Function for object detection and drawing bounding boxes
def predict_and_detect(model, frame, conf=0.5):
    results = model.predict(frame, conf=conf)
    for result in results:
        for box in result.boxes:
            # Draw bounding box
            x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
    return frame, results
# Process video frames
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    # Run object detection
    processed_frame, results = predict_and_detect(model, frame)
    # Show video with bounding boxes
    cv2.imshow('YOLO + OCR Detection', processed_frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
# Release video
cap.release()
cv2.destroyAllWindows()

该代码实时处理视频，在检测到的文本周围绘制边界框，并为下一步 OCR 准备这些区域。

3、在边界框上运行 OCR

现在我们已经使用 YOLO 隔离了文本区域，我们可以在这些特定区域内应用 OCR，与在整个图像上运行 OCR 相比，这大大提高了准确性：

import easyocr
# Initialize EasyOCR
reader = easyocr.Reader(['en'])
# Function to crop frames and perform OCR
def run_ocr_on_boxes(frame, boxes):
    ocr_results = []
    for box in boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
        cropped_frame = frame[y1:y2, x1:x2]
        ocr_result = reader.readtext(cropped_frame)
        ocr_results.append(ocr_result)
    return ocr_results
# Perform OCR on detected bounding boxes
for result in results:
    ocr_results = run_ocr_on_boxes(frame, result.boxes)
    # Extract and display the text from OCR results
    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'

结果明显改善，因为 OCR 引擎现在只处理明确标识为包含文本的区域，从而降低了因不相关图像元素而产生误解的风险。

4、使用 Ollama 改进文本

使用 easyocr 提取文本后，Ollama 的 Llama 3.1 可以进一步完善通常不完美且混乱的结果。OCR 功能强大，但它仍然可能误解文本或无序返回数据，尤其是书名或作者姓名。

Ollama 的 Llama 3.1 介入清理输出，从原始 OCR 结果中提供结构化、连贯的文本。通过向 Llama 3.1 提供识别和组织文本的具体说明，我们可以将不完美的 OCR 输出转换为格式整齐的书名和作者姓名。

import ollama
# Construct a prompt to clean up the OCR output
prompt = f"""
- Below is a text extracted from an OCR. The text contains mentions of famous books and their corresponding authors.
- Some words may be slightly misspelled or out of order.
- Your task is to identify the book titles and corresponding authors from the text.
- Output the text in the format: '<Name of the book> : <Name of the author>'.
- Do not generate any other text except the book title and the author.
TEXT:
{output_text}
"""
# Use Ollama to clean and structure the OCR output
response = ollama.chat(
    model="llama3",
    messages=[{"role": "user", "content": prompt}]
)
# Extract cleaned text
cleaned_text = response['message']['content'].strip()
print(cleaned_text)

输出结果如下：