【WSL Container】Unlimited OCR 2 を使ってみる

環境構築

Dockerfile pyproject.toml app.py の3つのファイルを同じフォルダ内に用意して以下を実行するだけです。

ただし、モデルのダウンロード先としてホスト側（Windows側）にあらかじめ D:\HF_HOME というディレクトリを作ってあります。

wslc build -t unlimited-ocr .

wslc run -d `
   --gpus all `
   -p 7860:7860 `
   -v D:\HF_HOME:/home/ubuntu/HF_HOME `
   --name unlimited-ocr `
   unlimited-ocr

Dockerfile

FROM nvidia/cuda:13.2.1-runtime-ubuntu24.04

# 対話型プロンプトの抑止と環境変数の設定
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PATH="/root/.local/bin:${PATH}"
ENV HF_HOME="/home/ubuntu/HF_HOME"

# build-essential のインストール
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
 && rm -rf /var/lib/apt/lists/*

# uv のインストール
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/

# ocr ディレクトリを作成してワーキングディレクトリに設定
WORKDIR /home/ubuntu/ocr

# app.py と pyproject.toml をコピー
COPY app.py pyproject.toml ./

# 仮想環境 (.venv) を作成
RUN uv sync

# アプリの起動
CMD ["uv", "run", "app.py"]

pyproject.toml

[project]
name = "ocr"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13, <3.14"

dependencies = [
    "addict==2.4.0",
    "easydict==1.13",
    "einops==0.8.2",
    "gradio==6.17.3",
    "huggingface-hub==0.36.2",
    "pillow==12.3.0",
    "tokenizers==0.22.2",
    "torch==2.14.1+cu132",
    "torchvision==0.29.1+cu132",
    "transformers==4.57.6",
]

[[tool.uv.index]]
name = "torch-cuda"
url = "https://download.pytorch.org/whl/cu132"
explicit = true

[tool.uv.sources]
torch = [{ index = "torch-cuda" }]
torchvision = [{ index = "torch-cuda" }]

app.py

import os
import glob
import torch
import gradio as gr
from transformers import AutoModel, AutoTokenizer

# ── 1. モデルと Tokenizer のロード (起動時に1回実行) ──
model_name = 'baidu/Unlimited-OCR-2'

print("Loading tokenizer and model...")
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=torch.bfloat16,
)
model = model.eval().cuda()
print("Model loaded successfully!")

# ── 2. 推論関数の定義 ──
def process_ocr(image_path):
    if image_path is None:
        return "画像をアップロードしてください。", ""

    output_dir = "./ocr_results"
    os.makedirs(output_dir, exist_ok=True)

    try:
        # 単ページ推論 (デフォルトで動的タイリング = gundam モード)
        text = model.infer(
            tokenizer,
            image_file=image_path,
            output_path=output_dir,
            save_results=True,
        )

        result_content = text if isinstance(text, str) and text else None

        # 戻り値が得られない場合は、保存された最新の結果ファイルを読み込む
        if result_content is None:
            result_content = "結果ファイルが見つかりませんでした。"
            saved_files = glob.glob(os.path.join(output_dir, "*.md")) + glob.glob(os.path.join(output_dir, "*.txt"))
            if saved_files:
                latest_file = max(saved_files, key=os.path.getmtime)
                with open(latest_file, "r", encoding="utf-8") as f:
                    result_content = f.read()

        # レンダリング用とRawソース用の両方に返却
        return result_content, result_content

    except Exception as e:
        error_msg = f"エラーが発生しました: {str(e)}"
        return error_msg, error_msg

# ── 3. カスタムCSS（テーブルの崩れ防止など） ──
custom_css = """
.markdown-body table {
    display: block;
    overflow-x: auto;
    white-space: nowrap;
}
"""

# ── 4. Gradio UI の構築 ──
with gr.Blocks(title="Unlimited-OCR-2 Demo") as demo:
    gr.Markdown("# Unlimited-OCR-2 Web UI")

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="filepath", label="画像をアップロード")
            submit_btn = gr.Button("OCR実行", variant="primary")

        with gr.Column(scale=1):
            # 表示形式を「プレビュー」と「Rawソース（コピー用）」で切り替え可能に設定
            with gr.Tabs():
                with gr.Tab("プレビュー"):
                    output_markdown = gr.Markdown(label="解析結果", sanitize_html=False)
                with gr.Tab("Markdownソース"):
                    output_raw = gr.Textbox(
                        label="Rawテキスト",
                        lines=15,
                        buttons=["copy"]
                    )

    submit_btn.click(
        fn=process_ocr,
        inputs=[input_image],
        outputs=[output_markdown, output_raw]
    )

# ── 5. アプリの起動 ──
if __name__ == "__main__":
    print("\n>>> ブラウザで http://localhost:7860 にアクセスしてください <<<\n", flush=True)
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        css=custom_css
    )

---