最近，我大部分時(shí)間都在玩大型語(yǔ)言模型（LLMs），但我對(duì)計(jì)算機(jī)視覺(jué)的熱愛(ài)從未真正消退。因此，當(dāng)有機(jī)會(huì)將兩者結(jié)合起來(lái)時(shí)，我迫不及待地想要深入研究。在Goodreads上掃描書的封面并將其標(biāo)記為“已讀”總是感覺(jué)像一種魔法，我忍不住想要為自己重現(xiàn)這種體驗(yàn)。

通過(guò)結(jié)合自定義訓(xùn)練的YOLOv10模型和OCR技術(shù)，你可以獲得巨大的準(zhǔn)確性提升。但當(dāng)你加入一個(gè)LLM（Llama 3）時(shí)，真正的魔法就發(fā)生了——那些混亂的OCR輸出突然變成了干凈、可用的文本，非常適合實(shí)際應(yīng)用。

為什么我們需要在OCR中使用YOLO和Ollama？

傳統(tǒng)的OCR（光學(xué)字符識(shí)別）方法在從簡(jiǎn)單圖像中提取文本方面做得很好，但當(dāng)文本與其他視覺(jué)元素交織在一起時(shí)，往往難以應(yīng)對(duì)。通過(guò)使用自定義的YOLO模型首先檢測(cè)文本區(qū)域等對(duì)象，我們可以為OCR隔離這些區(qū)域，顯著減少噪聲并提高準(zhǔn)確性。讓我們通過(guò)在沒(méi)有YOLO的圖像上運(yùn)行一個(gè)基本的OCR示例來(lái)演示這一點(diǎn)，以突出單獨(dú)使用OCR的挑戰(zhàn)：

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

雖然它可以很好地處理更簡(jiǎn)單的圖像，但當(dāng)存在噪聲或復(fù)雜的視覺(jué)模式時(shí)，錯(cuò)誤就開始堆積。這時(shí)，YOLO模型介入并真正發(fā)揮作用。

1. 訓(xùn)練自定義Yolov10數(shù)據(jù)集

用對(duì)象檢測(cè)增強(qiáng)OCR的第一步是在你數(shù)據(jù)集上訓(xùn)練一個(gè)自定義的YOLO模型。YOLO（You Only Look Once）是一個(gè)強(qiáng)大的實(shí)時(shí)對(duì)象檢測(cè)模型，它將圖像劃分為網(wǎng)格，允許它在單次前向傳遞中識(shí)別多個(gè)對(duì)象。這種方法非常適合檢測(cè)圖像中的文本，特別是當(dāng)你想要通過(guò)隔離特定區(qū)域來(lái)提高OCR結(jié)果時(shí)。

書籍封面數(shù)據(jù)集

我們將使用這里鏈接的預(yù)注釋書籍封面數(shù)據(jù)集，并在它上面訓(xùn)練一個(gè)YOLOv10模型。YOLOv10針對(duì)較小的對(duì)象進(jìn)行了優(yōu)化，使其非常適合在視頻或掃描文檔等具有挑戰(zhàn)性的環(huán)境中檢測(cè)文本。

from ultralytics import YOLO

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

你可以調(diào)整周期數(shù)量、數(shù)據(jù)集大小等參數(shù)，或者嘗試調(diào)整超參數(shù)以提高模型的性能和準(zhǔn)確性。

YOLOv10自定義數(shù)據(jù)集訓(xùn)練的關(guān)鍵指標(biāo)

2. 在視頻上運(yùn)行自定義模型以獲取邊界框

一旦你的YOLO模型訓(xùn)練完成，你可以將其應(yīng)用于視頻以檢測(cè)文本區(qū)域周圍的邊界框。這些邊界框隔離了感興趣的區(qū)域，確保了更干凈的OCR過(guò)程：

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()

這段代碼實(shí)時(shí)處理視頻，繪制檢測(cè)到的文本周圍的邊界框，并隔離這些區(qū)域，為下一步——OCR——做好了完美的準(zhǔn)備。

3. 在邊界框上運(yùn)行OCR

現(xiàn)在我們已經(jīng)用YOLO隔離了文本區(qū)域，我們可以在這些特定區(qū)域內(nèi)應(yīng)用OCR，與在整個(gè)圖像上運(yùn)行OCR相比，大大提高了準(zhǔn)確性：

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'

結(jié)果明顯改善，因?yàn)镺CR引擎現(xiàn)在只處理被特別識(shí)別為包含文本的區(qū)域，減少了從無(wú)關(guān)圖像元素中誤解的風(fēng)險(xiǎn)。

4. 使用Ollama改進(jìn)文本

使用easyocr提取文本后，Llama 3可以進(jìn)一步通過(guò)完善通常不完美和混亂的結(jié)果。OCR功能強(qiáng)大，但它仍然可能誤解文本或返回?zé)o序的數(shù)據(jù)，特別是書籍標(biāo)題或作者名稱。LLM介入整理輸出，將原始OCR結(jié)果轉(zhuǎn)化為結(jié)構(gòu)化、連貫的文本。通過(guò)用特定提示引導(dǎo)Llama 3來(lái)識(shí)別和組織內(nèi)容，我們可以將不完美的OCR數(shù)據(jù)完善為整潔格式化的書籍標(biāo)題和作者名稱。最好的部分？你可以使用Ollama在本地運(yùn)行它！

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

一旦LLM清理了文本，拋光后的輸出可以存儲(chǔ)在數(shù)據(jù)庫(kù)中，或用于各種實(shí)際應(yīng)用，例如：

• 數(shù)字圖書館或書店：自動(dòng)分類并顯示書籍標(biāo)題及其作者。
• 檔案系統(tǒng)：將掃描的書籍封面或文件轉(zhuǎn)換為可搜索的數(shù)字記錄。
• 自動(dòng)元數(shù)據(jù)生成：根據(jù)提取的信息為圖像、PDF或其他數(shù)字資產(chǎn)生成元數(shù)據(jù)。
• 數(shù)據(jù)庫(kù)輸入：將清理后的文本直接插入數(shù)據(jù)庫(kù)，確保更大系統(tǒng)中的數(shù)據(jù)結(jié)構(gòu)化和一致性。

通過(guò)結(jié)合對(duì)象檢測(cè)、OCR和LLMs，你解鎖了一個(gè)強(qiáng)大的管道，用于更結(jié)構(gòu)化的數(shù)據(jù)處理，非常適合需要高精度的應(yīng)用。

結(jié)論

你可以通過(guò)結(jié)合自定義訓(xùn)練的YOLOv10模型和EasyOCR，并使用LLM增強(qiáng)結(jié)果，顯著改進(jìn)文本識(shí)別工作流程。無(wú)論你是在處理棘手的圖像或視頻中的文本，清理OCR混亂，還是使一切超級(jí)拋光，這個(gè)管道都能為你提供實(shí)時(shí)、激光精確的文本提取和細(xì)化。

完整代碼：https://github.com/tapanBabbar9/yolov10/tree/main/book-cover