Полное руководство по OCR с Baidu Unlimited-OCR: от установки до многостраничных PDF

В эпоху, когда объемы неструктурированных данных растут экспоненциально, задача извлечения текста из изображений и документов перестала быть просто технической необходимостью — она стала стратегическим преимуществом. Традиционные системы оптического распознавания символов (OCR) часто требуют сложной предварительной обработки, настройки параметров и интеграции нескольких разрозненных инструментов для анализа макета. Однако появление больших мультимодальных моделей, таких как Baidu Unlimited-OCR, меняет правила игры. Эта модель, основанная на архитектуре vision-language (VLM), способна не просто «видеть» текст, но и понимать структуру документа, распознавая заголовки, таблицы и абзацы в едином проходе декодирования.

В этой статье мы подробно разберем создание полного конвейера (pipeline) для работы с моделью Baidu Unlimited-OCR. Мы не просто запустим код, а поймем, как правильно настроить GPU-окружение, выбрать оптимальный режим инференса для разных типов документов (от плотных таблиц до чистых страниц) и обработать многостраничные PDF-файлы. Этот гайд предназначен для разработчиков, Data Scientist'ов и инженеров, которые хотят внедрить современное OCR-решение в свои проекты, будь то автоматизация документооборота, оцифровка архивов или анализ финансовых отчетов.

011. Подготовка окружения и установка зависимостей

Первый шаг к успешному запуску любой нейросетевой модели — это корректная настройка среды. Baidu Unlimited-OCR — это модель объемом 3 миллиарда параметров, что требует значительных вычислительных ресурсов. Для работы нам потребуется доступ к GPU с поддержкой CUDA. Если вы используете локальную машину, убедитесь, что драйверы NVIDIA и CUDA Toolkit установлены правильно. Для тех, кто предпочитает облачные решения, Google Colab предоставляет удобный доступ к GPU, что делает процесс тестирования максимально простым.

Ключевым моментом является выбор точной версии библиотеки transformers. В данном пайплайне используется версия 4.57.1, которая обеспечивает совместимость с новыми архитектурами моделей Hugging Face. Помимо основной библиотеки, нам понадобятся инструменты для работы с изображениями (Pillow), визуализации (matplotlib), обработки PDF (PyMuPDF) и управления памятью (accelerate, psutil).

Ниже приведен код для автоматической установки необходимых пакетов. Обратите внимание на флаг -q (quiet), который минимизирует вывод в консоль, делая процесс установки чище.

import subprocess
import sys

def pip_install(pkgs):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *pkgs])

print(">> Installing dependencies (1-2 min)...")
pip_install([
    "transformers==4.57.1",
    "Pillow",
    "matplotlib",
    "einops",
    "addict",
    "easydict",
    "pymupdf",
    "psutil",
    "accelerate"
])
print(">> Done.")

022. Загрузка модели и токенизатора

После установки зависимостей переходим к загрузке самой модели. Мы используем классы AutoModel и AutoTokenizer из библиотеки Hugging Face Transformers. Параметр trust_remote_code=True критически важен, так как архитектура Unlimited-OCR содержит кастомные компоненты, которые не входят в стандартный набор библиотек и требуют доверия к коду репозитория автора.

Модель автоматически определяет тип данных (dtype) на основе возможностей вашего GPU. Если поддерживается bfloat16, модель загружается в этом формате, что экономит память и часто ускоряет вычисления на современных картах. Если нет, используется float16. После загрузки модель переводится в режим оценки (eval()) и перемещается на GPU с помощью метода .cuda().

import os
import torch
from transformers import AutoModel, AutoTokenizer

# Проверка наличия GPU
assert torch.cuda.is_available(), "No GPU detected! In Colab: Runtime -> Change runtime type -> GPU."
gpu_name = torch.cuda.get_device_name(0)
print(f">> GPU: {gpu_name}")

# Выбор точности вычислений
use_bf16 = torch.cuda.is_bf16_supported()
DTYPE = torch.bfloat16 if use_bf16 else torch.float16
print(f">> Using dtype: {DTYPE}")

MODEL_NAME = "baidu/Unlimited-OCR"
print(">> Downloading model (~6 GB for 3B params in BF16). First run takes a while...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModel.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_safetensors=True,
    torch_dtype=DTYPE
)
model.eval().cuda()
print(">> Model loaded and moved to GPU.")

033. Генерация тестовых документов

Для объективной оценки качества работы OCR нам нужны тестовые данные, имитирующие реальные бизнес-документы. Мы создадим скрипт, который генерирует изображения страниц с помощью библиотеки PIL (Pillow). Эти страницы будут содержать заголовки, текстовые блоки, таблицы и подвалы с примечаниями — типичную структуру отчетов или финансовых документов.

Функция make_sample_page создает изображение размером 1240x1754 пикселей, что соответствует высокому разрешению, необходимому для точного распознавания мелкого текста. Мы используем шрифты DejaVu, так как они стандартно доступны в большинстве Linux-окружений (включая Colab). Генерация таблиц происходит программно: мы задаем координаты ячеек, границы и текст, чтобы создать реалистичную сетку.

from PIL import Image, ImageDraw, ImageFont
import textwrap
import os

def load_font(size):
    for path in ["/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
                 "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"]:
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    return ImageFont.load_default()

def make_sample_page(path, page_no):
    # Создание холста
    img = Image.new("RGB", (1240, 1754), "white")
    draw = ImageDraw.Draw(img)
    
    title_f = load_font(48)
    head_f = load_font(34)
    body_f = load_font(26)
    
    # Заголовок
    draw.text((80, 70), f"Quarterly Operations Report — Page {page_no}", fill="black", font=title_f)
    draw.line([(80, 145), (800, 145)], fill="black", width=2)
    
    # Основной текст
    body = ("This document demonstrates Unlimited-OCR's one-shot long-horizon "
            "parsing. The model reads an entire page — headings, paragraphs, "
            "and tables — and emits structured text in a single decoding pass. "
            "Unlike classic OCR pipelines, no separate layout-analysis stage "
            "is required.")
    y = 190
    for line in textwrap.wrap(body, width=72):
        draw.text((80, y), line, fill="black", font=body_f)
        y += 40
        
    # Таблица
    y += 30
    draw.text((80, y), f"Table {page_no}: Regional Revenue (USD, millions)", fill="black", font=head_f)
    y += 60
    
    rows = [
        ["Region", "Q1", "Q2", "Q3"],
        ["North", "12.4", "13.1", "15.0"],
        ["South", "9.8", "10.2", "11.7"],
        ["East", "14.3", "13.9", "16.2"],
        ["West", "11.1", "12.5", "12.9"]
    ]
    col_w, row_h = 56, 80
    x0 = 260
    
    for i, row in enumerate(rows):
        for j, cell in enumerate(row):
            x = x0 + j * col_w
            y_pos = y + i * row_h
            draw.rectangle([x, y_pos, x + col_w, y_pos + row_h], outline="black", width=1)
            draw.text((x + 14, y_pos + 12), cell, fill="black", font=body_f)
        y += row_h
        
    # Подвал
    y += 50
    footer = (f"Note {page_no}: Figures are illustrative. Multi-page mode stitches "
              "context across pages, so cross-page references remain coherent.")
    for line in textwrap.wrap(footer, width=72):
        draw.text((80, y), line, fill="black", font=body_f)
        y += 40
        
    img.save(path)
    return path

# Создание директорий
os.makedirs("inputs", exist_ok=True)
os.makedirs("outputs/single_gundam", exist_ok=True)
os.makedirs("outputs/single_base", exist_ok=True)
os.makedirs("outputs/multi_page", exist_ok=True)

# Генерация страниц
IMAGE_PATH = make_sample_page("inputs/sample_page_1.png", 1)
PAGE_2 = make_sample_page("inputs/sample_page_2.png", 2)
PAGE_3 = make_sample_page("inputs/sample_page_3.png", 3)
print(f">> Sample pages written: {IMAGE_PATH}, {PAGE_2}, {PAGE_3}")

044. Режим Gundam: Высокая детализация для сложных макетов

Одним из ключевых преимуществ Unlimited-OCR является наличие двух режимов инференса для одиночных изображений: Gundam и Base. Режим Gundam предназначен для документов со сложной структурой, мелким шрифтом или плотным расположением элементов. Он работает по принципу тайлинга (разбиения изображения на плитки).

В этом режиме модель сначала анализирует глобальный вид документа, а затем последовательно обрабатывает его фрагменты (tiles). Параметр crop_mode=True включает эту функцию. Мы устанавливаем base_size=1024 и image_size=640, что означает, что модель будет обрабатывать плитки размером 640x640 пикселей, сохраняя высокую детализацию текста. Это критически важно для распознавания таблиц и мелких подписей, которые могут «потеряться» при уменьшении масштаба.

Также мы настраиваем параметры генерации: max_length=32768 позволяет модели выводить очень длинные тексты, а no_repeat_ngram_size=35 предотвращает зацикливание (повторы) фраз, что часто случается при работе с длинными контекстами.

print("\n" + "="*76)
print("STEP 4: Single image — GUNDAM mode (tiled, high detail)")
print("="*76)

result_gundam = model.infer(
    tokenizer,
    prompt="document parsing.",
    image_file=IMAGE_PATH,
    output_path="outputs/single_gundam",
    base_size=1024,
    image_size=640,
    crop_mode=True,
    max_length=32768,
    no_repeat_ngram_size=35,
    ngram_window=128,
    save_results=True
)

055. Режим Base: Скорость для чистых документов

Если ваш документ имеет простой макет, крупный шрифт и не содержит сложных таблиц, режим Gundam может быть избыточным с точки зрения вычислительных затрат. Здесь вступает в игру режим Base. В этом режиме изображение обрабатывается целиком (single view) без разбиения на плитки.

Мы устанавливаем crop_mode=False и image_size=1024. Это позволяет модели увидеть всю страницу сразу, что значительно ускоряет инференс. Поскольку нет накладных расходов на обработку тайлов и их последующую сборку, скорость работы возрастает. Однако, для очень плотных текстовых блоков это может привести к потере мелких деталей. Поэтому выбор режима всегда зависит от типа входных данных.

print("\n" + "="*76)
print("STEP 5: Single image — BASE mode (single view, faster)")
print("="*76)

result_base = model.infer(
    tokenizer,
    prompt="document parsing.",
    image_file=IMAGE_PATH,
    output_path="outputs/single_base",
    base_size=1024,
    image_size=1024,
    crop_mode=False,
    max_length=32768,
    no_repeat_ngram_size=35,
    ngram_window=128,
    save_results=True
)

066. Обработка многостраничных PDF-документов

Реальная работа с документами редко ограничивается одиночными изображениями. Чаще всего у нас есть PDF-файлы, содержащие десятки или сотни страниц. Baidu Unlimited-OCR поддерживает обработку таких файлов через функцию infer_multi(). Однако, модель работает с изображениями, поэтому первый шаг — это растровизация (rasterization) PDF-страниц в изображения высокого разрешения.

Для этого мы используем библиотеку PyMuPDF (также известную как fitz). Функция pdf_to_images открывает PDF, преобразует каждую страницу в PNG с разрешением 300 DPI (что обеспечивает отличное качество текста) и сохраняет их во временную директорию. Затем мы передаем список путей к этим изображениям в infer_multi().

Ключевое отличие многостраничного режима — это параметр ngram_window=1024. Он расширяет окно контроля повторов, что помогает модели сохранять контекст и согласованность терминологии на протяжении всего документа, предотвращая деградацию качества на поздних страницах.

import tempfile
import fitz

def pdf_to_images(pdf_path, dpi=300):
    """Rasterize every PDF page to a PNG; return the list of image paths."""
    doc = fitz.open(pdf_path)
    tmp_dir = tempfile.mkdtemp(prefix="pdf_ocr_")
    mat = fitz.Matrix(dpi/72, dpi/72)
    paths = []
    
    for i, page in enumerate(doc):
        out = os.path.join(tmp_dir, f"page_{i + 1:04d}.png")
        page.get_pixmap(matrix=mat).save(out)
        paths.append(out)
    
    doc.close()
    return paths

# Создание тестового PDF из наших изображений
SAMPLE_PDF = "inputs/sample_doc.pdf"
pdf = fitz.open()
for img_path in [IMAGE_PATH, PAGE_2, PAGE_3]:
    img_doc = fitz.open(img_path)
    pdf_bytes = img_doc.convert_to_pdf()
    img_pdf = fitz.open("pdf", pdf_bytes)
    pdf.new_page(width=img_pdf[0].rect.width, height=img_pdf[0].rect.height)
    pdf[-1].show_pdf_page(pdf[-1].rect, img_pdf, 0)
pdf.save(SAMPLE_PDF)
pdf.close()

# Растеризация и обработка
page_images = pdf_to_images(SAMPLE_PDF, dpi=300)
print(f">> Rasterized {len(page_images)} pages")

result_multi = model.infer_multi(
    tokenizer,
    prompt="Multi page parsing.",
    image_files=page_images,
    output_path="outputs/multi_page",
    image_size=1024,
    max_length=32768,
    no_repeat_ngram_size=35,
    ngram_window=1024,
    save_results=True
)

077. Анализ результатов и вывод

После завершения инференса результаты сохраняются в указанные директории. Мы можем использовать скрипт для обхода этих директорий и вывода превью сгенерированных текстовых файлов (TXT, MD, JSON). Это позволяет быстро оценить качество распознавания, проверить, не потерялись ли данные из таблиц, и убедиться, что структура документа сохранена.

08Что это значит на практике

Внедрение Baidu Unlimited-OCR в ваш рабочий процесс означает переход от фрагментарного распознавания к целостному пониманию документа. Вам больше не нужно писать сложные эвристики для разделения текста и таблиц или использовать отдельные модели для детекции объектов. Одна модель справляется со всем.

Для бизнеса это означает ускорение обработки заявок, автоматизацию ввода данных из счетов-фактур и возможность мгновенной оцифровки бумажных архивов. Для разработчиков это готовый, воспроизводимый пайплайн, который можно легко адаптировать под свои нужды, просто заменив источник изображений. Главное — правильно выбрать режим инференса в зависимости от сложности входных данных, чтобы найти баланс между скоростью и точностью.

Источник: MarkTechPost ↗