百度360必应搜狗淘宝本站头条
当前位置:网站首页 > 技术分类 > 正文

用YOLO和LLM增强OCR(yolov4数据增强)

ztj100 2024-10-30 05:12 48 浏览 0 评论

虽然最近我花了很多时间在大型语言模型 (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)

输出结果如下:

The Secret History : Donna Tartt

一旦 Llama 3.1 清理了文本,经过润色的输出就可以存储在数据库中或用于各种实际应用,例如:

  • 数字图书馆或书店:自动对书名进行分类,并在作者旁边显示书名。
  • 档案系统:将扫描的书籍封面或文档转换为可搜索的数字记录。
  • 自动元数据生成:根据提取的信息为图像、PDF 或其他数字资产生成元数据。
  • 数据库输入:将清理后的文本直接插入数据库,确保大型系统的数据结构化和一致性。

通过结合对象检测、OCR 和 LLM,你可以解锁强大的管道,以实现更结构化的数据处理,非常适合需要高精度水平的应用程序。

5、结束语

通过将定制训练的 YOLOv10 模型与 EasyOCR 相结合并使用 Ollama 的 Llama 3.1 增强结果,你可以显著改善文本识别工作流程。无论是在检测复杂图像或视频中的文本、清理 OCR 结果还是使输出更易于使用,此管道都可以实现实时、高度准确的文本提取和细化。

完整的源代码和 Jupyter Notebook 可在 GitHub 存储库中找到。


原文链接:YOLO和LLM增强的OCR - BimAnt

相关推荐

其实TensorFlow真的很水无非就这30篇熬夜练

好的!以下是TensorFlow需要掌握的核心内容,用列表形式呈现,简洁清晰(含表情符号,<300字):1.基础概念与环境TensorFlow架构(计算图、会话->EagerE...

交叉验证和超参数调整:如何优化你的机器学习模型

准确预测Fitbit的睡眠得分在本文的前两部分中,我获取了Fitbit的睡眠数据并对其进行预处理,将这些数据分为训练集、验证集和测试集,除此之外,我还训练了三种不同的机器学习模型并比较了它们的性能。在...

机器学习交叉验证全指南:原理、类型与实战技巧

机器学习模型常常需要大量数据,但它们如何与实时新数据协同工作也同样关键。交叉验证是一种通过将数据集分成若干部分、在部分数据上训练模型、在其余数据上测试模型的方法,用来检验模型的表现。这有助于发现过拟合...

深度学习中的类别激活热图可视化

作者:ValentinaAlto编译:ronghuaiyang导读使用Keras实现图像分类中的激活热图的可视化,帮助更有针对性...

超强,必会的机器学习评估指标

大侠幸会,在下全网同名[算法金]0基础转AI上岸,多个算法赛Top[日更万日,让更多人享受智能乐趣]构建机器学习模型的关键步骤是检查其性能,这是通过使用验证指标来完成的。选择正确的验证指...

机器学习入门教程-第六课:监督学习与非监督学习

1.回顾与引入上节课我们谈到了机器学习的一些实战技巧,比如如何处理数据、选择模型以及调整参数。今天,我们将更深入地探讨机器学习的两大类:监督学习和非监督学习。2.监督学习监督学习就像是有老师的教学...

Python教程(三十八):机器学习基础

...

Python 模型部署不用愁!容器化实战,5 分钟搞定环境配置

你是不是也遇到过这种糟心事:花了好几天训练出的Python模型,在自己电脑上跑得顺顺当当,一放到服务器就各种报错。要么是Python版本不对,要么是依赖库冲突,折腾半天还是用不了。别再喊“我...

超全面讲透一个算法模型,高斯核!!

...

神经网络与传统统计方法的简单对比

传统的统计方法如...

AI 基础知识从0.1到0.2——用“房价预测”入门机器学习全流程

...

自回归滞后模型进行多变量时间序列预测

下图显示了关于不同类型葡萄酒销量的月度多元时间序列。每种葡萄酒类型都是时间序列中的一个变量。假设要预测其中一个变量。比如,sparklingwine。如何建立一个模型来进行预测呢?一种常见的方...

苹果AI策略:慢哲学——科技行业的“长期主义”试金石

苹果AI策略的深度原创分析,结合技术伦理、商业逻辑与行业博弈,揭示其“慢哲学”背后的战略智慧:一、反常之举:AI狂潮中的“逆行者”当科技巨头深陷AI军备竞赛,苹果的克制显得格格不入:功能延期:App...

时间序列预测全攻略,6大模型代码实操

如果你对数据分析感兴趣,希望学习更多的方法论,希望听听经验分享,欢迎移步宝藏公众号...

AI 基础知识从 0.4 到 0.5—— 计算机视觉之光 CNN

...

取消回复欢迎 发表评论: