Search

2026/01/25

CPUでOCR処理:LightOnOCR-2-1B

PythonでのOCR処理。EasyOCRでやってました。

【2025年12月版】Windows で EasyOCR と Poppler を利用して PDF を OCR する手順
https://nokoshitamono.blogspot.com/2025/12/202512windows-easyocr-poppler-pdf-ocr.html

LightOnOCR-2-1Bでもやってみました。

EasyOCRは高速処理。
LightOnOCR-2-1Bはそれなりに時間がかかるけど、段組みの認識は強い。
CPUでも動くので使えるかも。
WikipediaのページをOCR処理した結果







モデルはどっかにダウンロードする
git lfs install
git clone https://huggingface.co/lightonai/LightOnOCR-2-1B c:/llamamodels/LightOnOCR-2-1B

あとは各ライブラリをpipでインストール。
pip checkで過不足ない感じにすれば動くはず。
torchはCUDA対応にすればGPU使えますが、CPUだけでも動きます。


# python lightonocr_local.py test.jpg --out result.txt
# python lightonocr_local.py book.pdf --pages 1,2,5-12

import argparse
import io
import os
from pathlib import Path
from typing import List, Optional, Tuple

import torch
from PIL import Image

# PDF対応(任意)
try:
    import pypdfium2 as pdfium
except Exception:
    pdfium = None

from transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor

MODEL_PATH = r"c:/models/LightOnOCR-2-1B"

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}


def pick_device() -> Tuple[str, torch.dtype]:
    """
    LightOnOCR公式例に合わせて:
      - mps -> float32
      - cuda -> bfloat16(対応しないGPUなら float16 に切替推奨)
      - cpu -> float32
    """
    if torch.backends.mps.is_available():
        return "mps", torch.float32
    if torch.cuda.is_available():
        # bf16 が怪しい場合は float16 に変更してください
        return "cuda", torch.bfloat16
    return "cpu", torch.float32

def resolve_device_and_dtype(device_arg: str, dtype_arg: str, force_fp16: bool) -> Tuple[str, torch.dtype]:
    """
    CLI指定に基づいて device と dtype を安全に確定する。
    device_arg: auto|cpu|cuda|cuda:0|cuda:1|mps
    dtype_arg : auto|fp32|fp16|bf16
    """
    # --- device 決定 ---
    if device_arg == "auto":
        if torch.backends.mps.is_available():
            device = "mps"
        elif torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"
    else:
        device = device_arg

    # --- device 実在チェック ---
    if device.startswith("cuda"):
        if not torch.cuda.is_available():
            raise RuntimeError("CUDAが利用できません。--device cpu を指定してください。")
        if device != "cuda":
            # cuda:0 のような形式
            try:
                idx = int(device.split(":")[1])
            except Exception:
                raise RuntimeError("CUDA指定は cuda または cuda:0 の形式で指定してください。")

            if idx < 0 or idx >= torch.cuda.device_count():
                raise RuntimeError(f"{device} は存在しません。CUDAは {torch.cuda.device_count()}枚です。")

    if device == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError("MPS が利用できません。")

    # --- dtype 決定 ---
    if dtype_arg == "auto":
        if device.startswith("cuda"):
            dtype = torch.float16 if force_fp16 else torch.bfloat16
        else:
            dtype = torch.float32
    else:
        mp = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}
        if dtype_arg not in mp:
            raise ValueError("--dtype は auto|fp32|fp16|bf16 のいずれかです")
        dtype = mp[dtype_arg]

    # MPSはfp32が安全(fp16/bf16で不安定になりやすい)
    if device == "mps":
        dtype = torch.float32

    return device, dtype



def load_images_from_path(path: Path, page_spec: str = "") -> List[Image.Image]:
    """
    path が:
      - 画像ファイル: 1枚
      - フォルダ: 画像を全列挙
      - PDF: 各ページをレンダして画像化(pypdfium2必須)
    """
    if path.is_dir():
        imgs = []
        for p in sorted(path.rglob("*")):
            if p.suffix.lower() in IMAGE_EXTS:
                imgs.append(Image.open(p).convert("RGB"))
        return imgs

    if path.suffix.lower() in IMAGE_EXTS:
        return [Image.open(path).convert("RGB")]

    if path.suffix.lower() == ".pdf":
        if pdfium is None:
            raise RuntimeError("PDFを処理するには `pip install pypdfium2` が必要です。")

        pdf_data = path.read_bytes()
        pdf = pdfium.PdfDocument(pdf_data)

        page_indices = parse_page_range(page_spec, len(pdf))

        images = []
        for i in page_indices:
            page = pdf[i]
            pil_img = page.render(scale=2.77).to_pil()
            images.append(pil_img.convert("RGB"))

        return images

    raise ValueError(f"未対応の入力です: {path}")

def parse_page_range(spec: str, total_pages: int) -> List[int]:
    """
    ページ指定文字列を 0-based index のリストに変換
    """
    if not spec:
        return list(range(total_pages))

    pages = set()

    for part in spec.split(","):
        part = part.strip()

        if "-" in part:
            start, end = part.split("-", 1)

            start_i = int(start) - 1 if start else 0
            end_i = int(end) - 1 if end else total_pages - 1

            start_i = max(start_i, 0)
            end_i = min(end_i, total_pages - 1)

            for i in range(start_i, end_i + 1):
                pages.add(i)
        else:
            i = int(part) - 1
            if 0 <= i < total_pages:
                pages.add(i)

    return sorted(pages)



@torch.inference_mode()
def ocr_image(
    model: LightOnOcrForConditionalGeneration,
    processor: LightOnOcrProcessor,
    image: Image.Image,
    device: str,
    dtype: torch.dtype,
    max_new_tokens: int = 1024,
) -> str:
    """
    Transformersの公式モデルカードにある会話テンプレ方式でOCRする。 :contentReference[oaicite:2]{index=2}
    """
    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
            ],
        }
    ]

    inputs = processor.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    )

    # floatテンソルだけ dtype を当て、他は device のみ移動
    inputs = {
        k: (v.to(device=device, dtype=dtype) if v.is_floating_point() else v.to(device))
        for k, v in inputs.items()
    }

    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)
    generated_ids = output_ids[0, inputs["input_ids"].shape[1] :]
    text = processor.decode(generated_ids, skip_special_tokens=True)
    return text.strip()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("input", help="画像ファイル / 画像フォルダ / PDF のパス")
    ap.add_argument("--max_new_tokens", type=int, default=1024)
    ap.add_argument("--out", default="", help="出力テキスト保存先(未指定なら標準出力)")
    ap.add_argument("--page_sep", default="\n\n--- PAGE ---\n\n", help="複数ページの区切り")
    ap.add_argument("--force_fp16", action="store_true", help="CUDA時に fp16 を強制(bf16で不安定なら)")
    ap.add_argument("--pages", default="", help="PDFページ範囲(例: 1-5, 3-, -10, 1,3,5-7)※PDFのみ有効")
    ap.add_argument("--model", default="", help="HF model id or local path(未指定ならコード内の MODEL_PATH を使用)")

    # ★追加(GPU切替)
    ap.add_argument("--device", default="auto", help="実行デバイス: auto | cpu | cuda | cuda:0 | cuda:1 | mps")
    ap.add_argument("--dtype", default="auto", help="精度: auto | fp32 | fp16 | bf16(cuda以外はfp32推奨)")

    args = ap.parse_args()
    model_path = args.model if args.model else MODEL_PATH
    
    
    # ★追加(GPU切替)
    device, dtype = resolve_device_and_dtype(args.device, args.dtype, args.force_fp16)
    print(f"[INFO] device={device} dtype={dtype}")

    model = LightOnOcrForConditionalGeneration.from_pretrained(
        model_path,
        torch_dtype=dtype,
        local_files_only=True,
    ).to(device)

    processor = LightOnOcrProcessor.from_pretrained(
        model_path,
        local_files_only=True,
    )

    images = load_images_from_path(in_path, args.pages)
    if not images:
        raise RuntimeError("処理できる画像が見つかりませんでした。")

    texts: List[str] = []
    for idx, img in enumerate(images, start=1):
        print(f"[INFO] OCR {idx}/{len(images)} ...")
        txt = ocr_image(
            model=model,
            processor=processor,
            image=img,
            device=device,
            dtype=dtype,
            max_new_tokens=args.max_new_tokens,
        )
        texts.append(txt)

    final_text = args.page_sep.join(texts)

    if args.out:
        out_path = Path(args.out)
        out_path.write_text(final_text, encoding="utf-8")
        print(f"[INFO] saved: {out_path}")
    else:
        print(final_text)


if __name__ == "__main__":
    main()