最近,我花了大部分时间玩大型语言模型(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自定义数据集训练的关键指标数据
一旦你的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处理完文本后,精炼后的输出可以存入数据库或用于各种实际应用,例如:
通过结合目标检测、OCR 和 LLMs,你将释放一个强大的处理管道,用于结构化数据处理,非常适合需要高精度的应用。
你可以通过将自定义训练的 YOLOv10 模型与 EasyOCR 结合,并利用 LLM 来增强结果,显著提升文本识别工作流程的效率。无论你是在处理复杂图像或视频中的文字,清理 OCR 产生的混乱,还是让所有内容变得超级干净整洁,这个流程都能为你提供实时、精确的文本提取和优化。
完整的源代码和Jupyter Notebook可在我们的GitHub仓库(点击这里)中找到。如果您有关于优化指令或其他方面的改进的想法,随时欢迎提问或提供反馈。