Xây dựng hệ thống AI OCR cục bộ với GOT-OCR2: Không Cloud, không thuê bao, chỉ cần Python

AI tutorial - IT technology blog
AI tutorial - IT technology blog

Sự bất tiện của OCR truyền thống

Việc trích xuất văn bản từ một tệp PDF sạch, độ phân giải cao là một bài toán đã có lời giải. Tuy nhiên, ngay khi bạn đưa vào một bức ảnh chụp bằng điện thoại bị mờ, một bảng biểu có các ô lồng nhau hoặc một công thức hóa học phức tạp, các công cụ như Tesseract thường sẽ “bó tay”. Tôi đã không còn nhớ nổi mình đã dành bao nhiêu giờ để chắp vá các tập lệnh regex và các pipeline tiền xử lý OpenCV chỉ để đạt được tỷ lệ chính xác 70%, con số đơn giản là không đủ cho các ứng dụng cấp độ production.

Nhiều nhà phát triển thường mặc định chọn các Cloud API như Google Vision hoặc AWS Textract để giải quyết vấn đề này. Mặc dù các dịch vụ này rất chính xác, nhưng chúng đi kèm với hóa đơn hàng tháng và những đánh đổi lớn về quyền riêng tư. Việc gửi các hồ sơ tài chính nhạy cảm hoặc tài liệu nhân sự nội bộ lên máy chủ của bên thứ ba là điều không thể chấp nhận đối với nhiều dự án doanh nghiệp. Cho đến gần đây, bạn buộc phải lựa chọn giữa “miễn phí nhưng hay lỗi” hoặc “chính xác nhưng đắt đỏ”.

Tại sao OCR thế hệ cũ gặp khó khăn với các cấu trúc phức tạp

Hầu hết các hệ thống OCR cũ coi việc phát hiện và nhận dạng văn bản là hai bước riêng biệt và cứng nhắc. Chúng tìm kiếm các bounding box (khung bao) trước, sau đó cố gắng đoán các ký tự bên trong đó. Cách tiếp cận này thất bại khi văn bản bị xoay, nằm trong biểu đồ hoặc được ép vào bố cục nhiều cột. Những mô hình này thiếu “ngữ cảnh hình ảnh” — chúng nhìn thấy các đường kẻ và điểm ảnh, nhưng không hiểu rằng một bảng biểu là một thực thể có cấu trúc thống nhất.

GOT-OCR2 (General OCR Theory) thay đổi hoàn toàn cuộc chơi bằng cách coi OCR là một tác vụ thị giác-ngôn ngữ (vision-language). Thay vì chỉ tìm kiếm các khung bao, nó sử dụng kiến trúc transformer thống nhất để “đọc” toàn bộ ngữ cảnh hình ảnh. Không quan trọng nếu văn bản nằm trong một bảng biểu răng cưa hay một bức ảnh thiếu sáng. Mô hình xử lý đầu vào hình ảnh và tạo ra đầu ra văn bản có cấu trúc chỉ trong một lần duy nhất.

Khởi động nhanh: Chạy GOT-OCR2 trong 5 phút

Bạn sẽ cần một máy tính có GPU để đạt được hiệu năng ổn định. Một card đồ họa NVIDIA với ít nhất 8GB VRAM (như RTX 3060 hoặc 4060) là lựa chọn lý tưởng. Mặc dù bạn có thể chạy trên CPU, nhưng thời gian suy luận (inference) có thể tăng từ 1-2 giây lên hơn 30 giây cho mỗi trang.

1. Thiết lập môi trường

Bắt đầu bằng cách cô lập các phụ thuộc (dependencies). Điều này giúp ngăn chặn xung đột phiên bản với các thư viện AI khác mà bạn có thể đã cài đặt.

# Tạo môi trường ảo
python -m venv got_ocr_env
source got_ocr_env/bin/activate  # Trên Windows: got_ocr_env\Scripts\activate

# Cài đặt các thư viện cốt lõi
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers tiktoken verovio accelerate

2. Tập lệnh thực thi

Lưu tệp này dưới tên ocr_engine.py. Tập lệnh này sẽ tải trực tiếp các trọng số (weights) đã được huấn luyện từ Hugging Face để xử lý các tác vụ nặng.

from transformers import AutoModel, AutoTokenizer
import torch

# Khởi tạo mô hình và tokenizer
model_name = "ucasyzx/GOT-OCR2_0"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id).eval().cuda()

def extract_text(image_path):
    # Loại 'ocr' là chế độ trích xuất văn bản tiêu chuẩn
    res = model.chat(tokenizer, image_path, ocr_type='ocr')
    return res

if __name__ == "__main__":
    image_path = "test_screenshot.png"
    print("--- Văn bản đã trích xuất ---")
    print(extract_text(image_path))

Sức mạnh thực sự: Xử lý dữ liệu cấu trúc

Sự linh hoạt của GOT-OCR2 là điểm cộng lớn nhất. Nó không chỉ là một trình trích xuất văn bản cơ bản; nó là một trình thông dịch cấu trúc. Trong các thử nghiệm của tôi, mô hình này vượt trội hơn các engine tiêu chuẩn, đặc biệt là khi bố cục không theo dạng tuyến tính.

Định dạng tài liệu quét

OCR tiêu chuẩn thường gộp văn bản từ các cột liền kề, tạo ra một mớ hỗn độn. Bằng cách đặt ocr_type thành ‘format’, mô hình sẽ tôn trọng luồng trang gốc. Đây là cứu cánh cho việc xử lý các bài báo khoa học hoặc bản tin nhiều cột.

Chuyển đổi bảng biểu sang Markdown

Đây là nơi mô hình thực sự tỏa sáng. Thay vì nhận được một chuỗi từ phẳng, bạn có thể yêu cầu mô hình xuất ra định dạng Markdown. Tôi đã sử dụng tính năng này để chuyển đổi các bảng tính tài chính phức tạp thành dữ liệu có cấu trúc mà các LLM có thể tiếp nhận mà không làm mất đi mối quan hệ giữa các hàng và cột.

# Lệnh này sẽ xuất ra một bảng Markdown sạch sẽ
res = model.chat(tokenizer, "table_image.png", ocr_type='format', render=True)

Kỹ thuật nâng cao: Cắt vùng chính xác

Bạn có thể gặp phải các trường hợp khó với các bản vẽ kỹ thuật khổng lồ hoặc các biểu mẫu mật độ cao. Khi việc quét toàn bộ (global scan) bỏ lỡ các chi tiết nhỏ, GOT-OCR2 cho phép “crop OCR”. Bạn có thể nhắm mục tiêu vào các tọa độ cụ thể trên ảnh độ phân giải cao để duy trì độ chính xác gần như hoàn hảo.

# Nhắm mục tiêu vào một vùng cụ thể (x1, y1, x2, y2)
res = model.chat(tokenizer, image_path, ocr_type='ocr', ocr_box=[150, 150, 600, 600])

Bằng cách chia một hình ảnh 4K thành các vùng logic, bạn có thể đạt được độ chính xác 99% trên các văn bản vốn không thể đọc được nếu thu nhỏ toàn trang theo cách thông thường.

Danh sách kiểm tra khi triển khai thực tế

Việc chuyển từ một tập lệnh chạy cục bộ sang một API thực tế đòi hỏi một vài tối ưu hóa về phần cứng và phần mềm để giữ cho hệ thống ổn định.

  • Hiệu quả VRAM: Mô hình tiêu thụ khoảng 5GB VRAM trong quá trình suy luận. Nếu bạn bị giới hạn về bộ nhớ, hãy sử dụng torch.cuda.empty_cache() giữa các batch lớn để tránh lỗi “Out of Memory” đáng sợ.
  • Tải Singleton: Đừng bao giờ tải lại mô hình cho mỗi yêu cầu. Hãy tải nó một lần vào bộ nhớ bằng một framework như FastAPI và giữ nó sẵn sàng cho các tác vụ tiếp theo.
  • Tiền xử lý thông minh: Mặc dù mô hình rất mạnh mẽ, nhưng việc khử nghiêng (straightening) cơ bản có thể giảm thời gian suy luận khoảng 15% vì transformer không phải “vất vả” để căn chỉnh văn bản.
  • Lượng tử hóa: Đối với những người chạy trên phần cứng cũ, hãy sử dụng lượng tử hóa 4-bit thông qua thư viện bitsandbytes. Điều này có thể giảm dung lượng VRAM xuống gần một nửa với mức giảm độ chính xác không đáng kể.

Lời kết? Xây dựng một hệ thống OCR cục bộ không chỉ là để cắt giảm chi phí. Đó là về chủ quyền dữ liệu hoàn toàn. Bằng cách chạy GOT-OCR2 trên phần cứng của riêng bạn, bạn loại bỏ độ trễ, giữ cho dữ liệu riêng tư và xử lý được các tài liệu có thể làm khó các công cụ truyền thống. Đây là một sự nâng cấp đáng kể cho bất kỳ nhà phát triển nào đang xây dựng các pipeline RAG (Retrieval-Augmented Generation) hoặc quy trình làm việc tài liệu tự động.

Share: