C/C++教程

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

本文主要是介绍使用自定义的 Yolov10 和 Ollama (Llama 3) 来提升 OCR 能力,对大家解决编程问题具有一定的参考价值,需要的程序猿们随着小编来一起学习吧!

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

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

我们为什么需要带有OCR的YOLO和Ollama?

传统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 模型就能真正发挥作用了。

1. 训练自定义YOLOv10模型的数据集

提高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:光学字符识别)

3. 在边框上进行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(光学字符识别)引擎现在只处理特别标记为包含文本的区域,减少了由于无关图像元素引起的误读风险。

4 改进文本:使用Ollama

使用 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仓库(点击这里)中找到。如果您有关于优化指令或其他方面的改进的想法,随时欢迎提问或提供反馈。

这篇关于使用自定义的 Yolov10 和 Ollama (Llama 3) 来提升 OCR 能力的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!