使用自定义的 Yolov10 和 Ollama (Llama 3) 来提升 OCR 能力

最近，我花了大部分时间玩大型语言模型（LLMs），但我对计算机视觉的热爱从未真正消逝。因此，当有机会将两者结合起来时，我迫不及待地一头扎进去。在Goodreads上拍下一本书的封面照片并将其标记为“已读”总是感觉像是有点魔法，忍不住想自己也试试这种感觉。

结合自定义训练的 YOLOv10 模型与OCR技术显著提升了准确性。但真正厉害的地方在于当你引入一个LLM（Llama 3）时——突然，那些杂乱无章的OCR输出变成了整洁、可用的文本，非常适合实际使用场景。

传统OCR（光学字符识别）方法在处理简单图像中的文本提取方面表现良好，但当处理与其它视觉元素交织在一起的文本时，往往表现不佳。通过使用自定义的YOLO模型先识别出文本区域等对象，我们可以将这些区域隔离开来进行OCR处理，从而大大减少噪声并提高识别准确度。

我们来看一个没有使用YOLO的图像上的基本OCR示例，来突出单独使用OCR挑战。

    import easyocr  
    import cv2  
    # 初始化EasyOCR引擎  
    reader = easyocr.Reader(['en'])  
    # 加载图像  
    image = cv2.imread('book.jpg')  
    # 直接执行OCR  
    results = reader.readtext(image)  
    # 显示结果如下  
    for (bbox, text, prob) in results:  
        print(f"检测到的文本: {text} (识别概率: {prob})")

    原版畅销书《秘密历史》一书，作者唐娜·塔特 扣人心弦，引人入胜且才华横溢 泰晤士报

这可不像你想要的效果吧？虽然它处理简单图像没问题，但是当图像包含噪声或复杂视觉图案时，错误就开始累积了。这时一个 YOLO 模型就能真正发挥作用了。

提高OCR性能的第一步结合对象检测是，在您的数据集上训练一个定制的YOLO模型。YOLO（You Only Look Once，简称YOLO）是一个强大的实时对象检测模型，它将图像分割成网格，从而能够在一次传递中识别多个对象。这种方法非常适合检测图像中的文本，特别是在您希望隔离特定区域以提升OCR效果时。

书籍封面数据集资料由kernst制作

我们将使用在该链接中的预注释的图书封面数据集训练YOLOv10模型：https://universe.roboflow.com/kernst/book-covers-2/dataset/1。YOLOv10优化了对较小对象的检测，因此非常适合用来检测视频或扫描文档等具有挑战性环境中的文本，例如。

from ultralytics 导入 YOLO 模型  

model = YOLO("yolov10n.pt")  
# 开始训练模型  
model.train(data="datasets/data.yaml", epochs=50, imgsz=640)

在我的情况下，我在Google Colab上训练这个模型大约花了六个小时（呼~！），用了50个训练周期。你可以调整一些参数，比如迭代次数和数据集大小，或者通过调整超参数来优化模型的性能和精确度。

来自YOLOv10自定义数据集训练的关键指标数据

1. 在视频上运行自定义模型以生成边界框

一旦你的YOLO模型训练完成，你可以将其应用于视频，来检测视频中文字区域的边界框。这些边界框以便隔离出感兴趣的区域，从而让OCR处理更加干净。

    import cv2  
    # 打开视频文件  
    video_path = 'books.mov'  
    cap = cv2.VideoCapture(video_path)  
    # 加载YOLO模型文件  
    model = YOLO('model.pt')  
    # 用于检测并绘制边界框的函数  
    def 检测并绘制边界框(model, frame, conf=0.5):  
        results = model.predict(frame, conf=conf)  
        for result in results:  
            for box in result.boxes:  
                # 绘制边界框  
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())  
                cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)  
        return frame, results  
    # 处理视频帧  
    while cap.isOpened():  
        ret, frame = cap.read()  
        if not ret:  
            break  
        # 运行对象检测  
        processed_frame, results = 检测并绘制边界框(model, frame)  
        # 显示视频中的边界框  
        cv2.imshow('YOLO+OCR检测', processed_frame)  
        if cv2.waitKey(1) & 0xFF == ord('q'):  
            break  
    # 释放视频资源  
    cap.release()  
    cv2.destroyAllWindows()

这段代码实时处理视频中的内容，在检测到的文字周围画出边界框，并隔离这些区域，为OCR做好准备。（OCR：光学字符识别）

现在我们已经利用YOLO隔离了文本区域，我们可以在这些特定区域里应用OCR，这样做相比于在整个图像上运行OCR，可以大大提高准确性。

    import easyocr  
    # 初始化EasyOCR识别器  
    reader = easyocr.Reader(['en'])  
    # 此函数用于裁剪帧并执行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  
    # 在检测到的边界框上执行OCR识别  
    for result in results:  
        ocr_results = run_ocr_on_boxes(frame, result.boxes)  
        # 从OCR结果中提取并打印文本  
        extracted_text = [detection[1] for ocr in ocr_results for detection in ocr[1]]  
        print(f"提取的文本: {', '.join(extracted_text)}")

'THE, 秘密的, 历史的, 唐娜, 塔特'

结果明显改善了，因为OCR（光学字符识别）引擎现在只处理特别标记为包含文本的区域，减少了由于无关图像元素引起的误读风险。

使用 easyocr 提取文本后，Llama 3 可以进一步优化这些结果，以改进 OCR 经常产生的不准确且混乱的结果。OCR 很强大，但仍可能误读文本或返回乱序数据，特别是在处理书名和作者名时。

LLM 会整理输出，将原始 OCR 数据转化为结构化、连贯的文本。通过给 Llama 3 提供特定提示，使其识别和组织内容，我们可以将不连贯的原始 OCR 数据转变为结构化、连贯的文本，从而将不完美的 OCR 数据精炼成整洁排版的书名和作者名。最好的一点是，你还可以在本地通过 Ollama 运行它！

    导入库 ollama  
    # 构建一个提示，用于清理 OCR 输出  
    prompt = f"""  
    - 下面是从 OCR 中提取的文本。文中提到了一些著名书籍及其对应的作者。  
    - 文中的一些单词可能拼写有轻微错误或顺序混乱。  
    - 你需要从文本中识别书籍名称及其对应的作者。  
    - 输出格式为：'<书名> : <作者名>'。  
    - 除了书名和作者名外，不要再生成其他任何内容。  
    TEXT:  
    {output_text}  
    """  
    # 利用 Ollama 清理并结构化 OCR 输出  
    response = ollama.chat(  
        model="llama3",  
        messages=[{"角色": "用户", "内容": prompt}]  
    )  
    # 提取清理后的内容  
    cleaned_text = response['message']['content'].strip()  
    print(cleaned_text)

秘密史：唐娜·塔特

没错，就是这样！一旦LLM处理完文本后，精炼后的输出可以存入数据库或用于各种实际应用，例如：

• 数字图书馆和书店：自动分类并展示书名及其作者。
• 档案系统：将扫描的书籍封面或文档转换为可搜索的数字档案。
• 自动化元数据生成：根据提取的信息自动生成图片、PDF或其他数字资产的元数据。
• 数据库录入：将清理过的文本直接录入数据库，确保在大型系统中结构化和一致的数据。

通过结合目标检测、OCR 和 LLMs，你将释放一个强大的处理管道，用于结构化数据处理，非常适合需要高精度的应用。

你可以通过将自定义训练的 YOLOv10 模型与 EasyOCR 结合，并利用 LLM 来增强结果，显著提升文本识别工作流程的效率。无论你是在处理复杂图像或视频中的文字，清理 OCR 产生的混乱，还是让所有内容变得超级干净整洁，这个流程都能为你提供实时、精确的文本提取和优化。

完整的源代码和Jupyter Notebook可在我们的GitHub仓库（点击这里）中找到。如果您有关于优化指令或其他方面的改进的想法，随时欢迎提问或提供反馈。