Flame-OCR / README_zh.md
hopex-ai's picture
Update README_zh.md
acaf9d5 verified
|
Raw
History Blame Contribute Delete
14.8 kB

Flame-OCR

Logo

Static Badge
Static Badge
👉 Flame-OCR Github

English | 简体中文

Flame-OCR: A high-speed, high-precision OCR model inspired by the purifying power of flame.

灵感源于斗破苍穹中异火的净化之力

我们发布了 Flame-OCR,这是一款轻量级(0.8B 参数)的端到端文档页面解析模型。给定一份文档页面图像,能按自然阅读顺序生成 Markdown 表示,涵盖文本、公式、表格和视觉区域。

Flame-OCR 基于 OvisOCR2(Qwen3.5-0.8B) 模型,通过高质数据进行后训练,特别是扩充与增强训练,提升在发票,医疗,体检,金融文档,表格,公式,手写,古文的效果。该模型在保持较小体积并保持原有效果的同时,实现了强大的文档解析性能。

Flame-OCR 在 OmniDocBench v1.6 上取得了 94.46 的综合得分。

Performance

Long-horizon OCR demo

推理

  • 基于Transformers
pip install "transformers>=5.0.0" pillow
import time
import threading
import torch
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
from PIL import Image

MIN_PIXELS = 448 * 448
MAX_PIXELS = 2880 * 2880
MAX_NEW_TOKENS = 16384
STREAM_MIN_CHARS = 64
STREAM_MAX_INTERVAL = 0.25


model_path = '/path/to/model or model_id'

processor = AutoProcessor.from_pretrained(
    model_path,
    min_pixels=MIN_PIXELS,
    max_pixels=MAX_PIXELS,
)
model = Qwen3_5ForConditionalGeneration.from_pretrained(
    model_path,
    dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation=None # or Flash Attention
).eval()

prompt = (
    "\nExtract all readable content from the image in natural human reading order "
    "and output the result as a single Markdown document. For charts or images, "
    'represent them using an HTML image tag: <img src="/anyforge/Flame-OCR/resolve/main/images/bbox_%7Bleft%7D_%7Btop%7D_%7Bright%7D_%7Bbottom%7D.jpg" />, '
    "where left, top, right, bottom are bounding box coordinates scaled to [0, 1000). "
    "Format formulas as LaTeX. Format tables as HTML: <table>...</table>. "
    "Transcribe all other text as standard Markdown. Preserve the original text "
    "without translation or paraphrasing."
)

def clean_truncated_repeats(
    text: str,
    min_text_len: int = 8000,
    max_period: int = 200,
    min_period: int = 1,
    min_repeat_chars: int = 100,
    min_repeat_times: int = 5,
) -> str:
    """Remove a repeated suffix created when generation reaches its token ceiling."""
    n = len(text)
    if n < min_text_len:
        return text

    max_period = min(max_period, n - 1)
    for unit_len in range(min_period, max_period + 1):
        if text[n - 1] != text[n - 1 - unit_len]:
            continue
        match_len = 1
        idx = n - 2
        while idx >= unit_len and text[idx] == text[idx - unit_len]:
            match_len += 1
            idx -= 1
        total_len = match_len + unit_len
        repeat_times = total_len // unit_len
        tail_len = total_len % unit_len
        if repeat_times >= min_repeat_times and total_len >= min_repeat_chars:
            return text[: n - total_len + unit_len] + text[n - tail_len :]
    return text

def _model_inputs(page_image: Image.Image):
    if processor is None or model is None:
        raise RuntimeError("OvisOCR2 is not loaded.")

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": page_image},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    return processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        enable_thinking=False,
    ).to(model.device)


def generation_token_ids(active_processor) -> dict[str, int]:
    """Use tokenizer stop IDs; this checkpoint's config and tokenizer differ."""
    tokenizer = active_processor.tokenizer
    return {
        "eos_token_id": int(tokenizer.eos_token_id),
        "pad_token_id": int(tokenizer.pad_token_id),
    }

def infer_stream(page_image: Image.Image):
    from transformers import TextIteratorStreamer

    inputs = _model_inputs(page_image)
    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    errors: list[BaseException] = []

    def generate() -> None:
        try:
            with torch.inference_mode():
                model.generate(
                    **inputs,
                    streamer=streamer,
                    max_new_tokens=MAX_NEW_TOKENS,
                    do_sample=False,
                    temperature=None,
                    top_p=None,
                    top_k=None,
                    **generation_token_ids(processor),
                )
        except BaseException as error:
            errors.append(error)
            streamer.on_finalized_text("", stream_end=True)

    worker = threading.Thread(target=generate, name="ovisocr2-generate", daemon=True)
    worker.start()
    text = ""
    last_yielded = ""
    last_yield_time = time.monotonic()
    for fragment in streamer:
        text += fragment
        now = time.monotonic()
        if (
            len(text) - len(last_yielded) >= STREAM_MIN_CHARS
            or now - last_yield_time >= STREAM_MAX_INTERVAL
        ):
            yield text
            last_yielded = text
            last_yield_time = now

    worker.join()
    if errors:
        raise RuntimeError("Model generation failed.") from errors[0]
    final_text = clean_truncated_repeats(text.strip())
    if final_text and final_text != last_yielded:
        yield final_text

def infer(page_image: Image.Image):
    inputs = _model_inputs(page_image)
    outputs = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS)
    result = processor.decode(outputs[0][inputs["input_ids"].shape[-1]:-1])
    return result
    

image = Image.open("/path/to/image")

## Inference
result = infer(image)
print(result)

## Stream Inference
for line in infer_stream(image):
    if line:
        print(line)
  • 基于vllm
pip install "vllm==0.22.1" pillow
from PIL import Image
from vllm import LLM, SamplingParams


class CustomOCRParser:
    def __init__(self, model_name_or_path: str):
        self.model = LLM(
            model=model_name_or_path,
            tensor_parallel_size=1,
            gpu_memory_utilization=0.8,
            gdn_prefill_backend="triton"
        )

        prompt = '\nExtract all readable content from the image in natural human reading order and output the result as a single Markdown document. For charts or images, represent them using an HTML image tag: <' + 'img src="images/bbox_{left}_{top}_{right}_{bottom}.jpg" />, where left, top, right, bottom are bounding box coordinates scaled to [0, 1000). Format formulas as LaTeX. Format tables as HTML: <table>...</table>. Transcribe all other text as standard Markdown. Preserve the original text without translation or paraphrasing.'
        self.prompt = self.model.get_tokenizer().apply_chat_template(
            [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}],
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False
        )

        self.sampling_params = SamplingParams(
            max_tokens=16384,
            temperature=0.0
        )

    def _clean_truncated_repeats(
        self,
        text: str,
        min_text_len: int = 8000,
        max_period: int = 200,
        min_period: int = 1,
        min_repeat_chars: int = 100,
        min_repeat_times: int = 5
    ) -> str:
        n = len(text)
        if n < min_text_len:
            return text

        max_period = min(max_period, n - 1)
        for unit_len in range(min_period, max_period + 1):
            if text[n - 1] != text[n - 1 - unit_len]:
                continue

            match_len = 1
            idx = n - 2
            while idx >= unit_len and text[idx] == text[idx - unit_len]:
                match_len += 1
                idx -= 1

            total_len = match_len + unit_len
            repeat_times = total_len // unit_len
            tail_len = total_len % unit_len

            if repeat_times >= min_repeat_times and total_len >= min_repeat_chars:
                return text[: n - total_len + unit_len] + text[n - tail_len:]

        return text

    def parse(self, images: list[Image.Image], filter_imgtags: bool = False) -> list[str]:
        vllm_inputs = [
            {
                "prompt": self.prompt,
                "multi_modal_data": {"image": image},
                "mm_processor_kwargs": {
                    "images_kwargs": {
                        "min_pixels": 448 * 448,
                        "max_pixels": 2880 * 2880
                    }
                }
            }
            for image in images
        ]

        outputs = self.model.generate(vllm_inputs, self.sampling_params)

        markdowns = []
        for output in outputs:
            text = output.outputs[0].text.strip()
            if filter_imgtags:
                text = "\n\n".join(
                    block
                    for block in text.split("\n\n")
                    if not block.strip().startswith('<img src="images/bbox_')
                )
            markdowns.append(self._clean_truncated_repeats(text))

        return markdowns


if __name__ == "__main__":
    parser = CustomOCRParser("/path/to/model")
    images = [Image.open("/path/to/image"), Image.open("/path/to/image")]
    markdowns = parser.parse(images)
    print(markdowns[0])
  • vllm cli

# use modelscope
export VLLM_USE_MODELSCOPE=true

# or model local path
vllm serve anyforge/Flame-OCR \
--served-model-name "flame-ocr" \
--tensor-parallel-size 1 \
--gpu-memory-utilization 0.8 \
--gdn-prefill-backend triton \
--max-model-len 16384 \
--max-num-seqs 1 \
--enable-chunked-prefill \
--max-num-batched-tokens 16384



import base64
from openai import OpenAI

def image_to_base64(image_path: str) -> str:
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="EMPTY"
)

image_path = "/path/to/test.png"
base64_image = image_to_base64(image_path)

prompt = '\nExtract all readable content from the image in natural human reading order and output the result as a single Markdown document. For charts or images, represent them using an HTML image tag: <' + 'img src="images/bbox_{left}_{top}_{right}_{bottom}.jpg" />, where left, top, right, bottom are bounding box coordinates scaled to [0, 1000). Format formulas as LaTeX. Format tables as HTML: <table>...</table>. Transcribe all other text as standard Markdown. Preserve the original text without translation or paraphrasing.'

response = client.chat.completions.create(
    model="flame-ocr",
    messages=[
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{base64_image}"
                    }
                },
                {
                    "type": "text",
                    "text": prompt
                }
            ]
        }
    ],
    max_tokens=8192,
    temperature=0.0,
    extra_body={
        "mm_processor_kwargs": {
            "images_kwargs": {
                "min_pixels": 448 * 448,
                "max_pixels": 2880 * 2880
            }
        }
    }
)

print(response.choices[0].message.content)
  • 提取视觉区域的 HTML 图像标签
import re
from pathlib import Path

from PIL import Image


BBOX_IMAGE_PATTERN = re.compile(
    r'<img src=' + r'"images/bbox_(\d+)_(\d+)_(\d+)_(\d+)\.jpg" />'
)


def save_renderable_markdown_with_visual_regions(
    markdown: str,
    page_image: Image.Image,
    output_dir: str,
) -> None:
    output_dir = Path(output_dir)
    images_dir = output_dir / "images"
    images_dir.mkdir(parents=True, exist_ok=True)

    width, height = page_image.size
    for left, top, right, bottom in BBOX_IMAGE_PATTERN.findall(markdown):
        x1 = max(0, min(width, round(int(left) * width / 1000)))
        y1 = max(0, min(height, round(int(top) * height / 1000)))
        x2 = max(0, min(width, round(int(right) * width / 1000)))
        y2 = max(0, min(height, round(int(bottom) * height / 1000)))
        if x2 <= x1 or y2 <= y1:
            continue

        crop_path = images_dir / f"bbox_{left}_{top}_{right}_{bottom}.jpg"
        page_image.crop((x1, y1, x2, y2)).convert("RGB").save(crop_path)

    (output_dir / "output.md").write_text(markdown, encoding="utf-8")


parser = CustomOCRParser("/path/to/model")
page_image = Image.open("test1.jpg")
markdown = parser.parse([page_image], filter_imgtags=False)[0]
save_renderable_markdown_with_visual_regions(markdown, page_image, "output")

Buy me a coffee

致谢

我们衷心感谢 OvisOCR2、Qwen 和 OmniDocBench 提供了宝贵的代码、模型权重和基准数据集。同时感谢所有人为这一开源项目做出的贡献!

免责声明

由于现实世界中文档的多样性和复杂性,Flame-OCR 仍可能产生错误或不完整的输出。在关键应用场景中,请务必人工核验结果。