StarDoc-AI/TeleOCR

pip install transformers torch pillow

import html
import itertools
import json
import re
from dataclasses import dataclass

from PIL import Image
import torch
from transformers import AutoProcessor, AutoModel

@dataclass
class ContentBlock:
    type: str
    bbox: list[float]
    angle: int | None = None
    content: str | None = None


@dataclass
class TableCell:
    text: str
    start_row_offset_idx: int
    end_row_offset_idx: int
    start_col_offset_idx: int
    end_col_offset_idx: int
    row_span: int = 1
    col_span: int = 1


OTSL_NL = "<nl>"
OTSL_FCEL = "<fcel>"
OTSL_ECEL = "<ecel>"
OTSL_LCEL = "<lcel>"
OTSL_UCEL = "<ucel>"
OTSL_XCEL = "<xcel>"
OTSL_TOKENS = [OTSL_NL, OTSL_FCEL, OTSL_ECEL, OTSL_LCEL, OTSL_UCEL, OTSL_XCEL]


def _otsl_extract_tokens_and_text(text: str):
    pattern = "(" + "|".join(map(re.escape, OTSL_TOKENS)) + ")"
    tokens = re.findall(pattern, text)
    parts = [part for part in re.split(pattern, text) if part.strip()]
    return tokens, parts


def _count_right(rows, row_idx, col_idx, tokens):
    span = 0
    while col_idx < len(rows[row_idx]) and rows[row_idx][col_idx] in tokens:
        span += 1
        col_idx += 1
    return span


def _count_down(rows, row_idx, col_idx, tokens):
    span = 0
    while row_idx < len(rows) and col_idx < len(rows[row_idx]) and rows[row_idx][col_idx] in tokens:
        span += 1
        row_idx += 1
    return span


def _otsl_parse_texts(parts, tokens):
    rows = [list(row) for is_nl, row in itertools.groupby(tokens, lambda token: token == OTSL_NL) if not is_nl]
    if not rows:
        return [], []

    max_cols = max(len(row) for row in rows)
    for row in rows:
        row.extend([OTSL_ECEL] * (max_cols - len(row)))

    cells = []
    row_idx = 0
    col_idx = 0
    for idx, part in enumerate(parts):
        if part in (OTSL_FCEL, OTSL_ECEL):
            cell_text = ""
            right_offset = 1
            if part != OTSL_ECEL and idx + 1 < len(parts) and parts[idx + 1] not in OTSL_TOKENS:
                cell_text = parts[idx + 1].strip()
                right_offset = 2

            next_right = parts[idx + right_offset] if idx + right_offset < len(parts) else ""
            next_bottom = rows[row_idx + 1][col_idx] if row_idx + 1 < len(rows) and col_idx < len(rows[row_idx + 1]) else ""
            col_span = 1 + (_count_right(rows, row_idx, col_idx + 1, {OTSL_LCEL, OTSL_XCEL}) if next_right in {OTSL_LCEL, OTSL_XCEL} else 0)
            row_span = 1 + (_count_down(rows, row_idx + 1, col_idx, {OTSL_UCEL, OTSL_XCEL}) if next_bottom in {OTSL_UCEL, OTSL_XCEL} else 0)
            cells.append(TableCell(
                text=cell_text,
                row_span=row_span,
                col_span=col_span,
                start_row_offset_idx=row_idx,
                end_row_offset_idx=row_idx + row_span,
                start_col_offset_idx=col_idx,
                end_col_offset_idx=col_idx + col_span,
            ))
        if part in (OTSL_FCEL, OTSL_ECEL, OTSL_LCEL, OTSL_UCEL, OTSL_XCEL):
            col_idx += 1
        elif part == OTSL_NL:
            row_idx += 1
            col_idx = 0
    return cells, rows


def convert_otsl_to_html(otsl_content: str) -> str:
    if otsl_content.startswith("<table") and otsl_content.endswith("</table>"):
        return otsl_content

    tokens, parts = _otsl_extract_tokens_and_text(otsl_content)
    cells, rows = _otsl_parse_texts(parts, tokens)
    if not cells or not rows:
        return ""

    grid = [[None for _ in range(len(rows[0]))] for _ in range(len(rows))]
    for cell in cells:
        for row_idx in range(cell.start_row_offset_idx, min(cell.end_row_offset_idx, len(rows))):
            for col_idx in range(cell.start_col_offset_idx, min(cell.end_col_offset_idx, len(rows[0]))):
                grid[row_idx][col_idx] = cell

    html_rows = []
    for row_idx, row in enumerate(grid):
        html_rows.append("<tr>")
        for col_idx, cell in enumerate(row):
            if cell is None or cell.start_row_offset_idx != row_idx or cell.start_col_offset_idx != col_idx:
                continue
            attrs = ""
            if cell.row_span > 1:
                attrs += f' rowspan="{cell.row_span}"'
            if cell.col_span > 1:
                attrs += f' colspan="{cell.col_span}"'
            html_rows.append(f"<td{attrs}>{html.escape(cell.text.strip())}</td>")
        html_rows.append("</tr>")
    return "<table>" + "".join(html_rows) + "</table>"


def post_process(blocks: list[ContentBlock]) -> list[ContentBlock]:
    for block in blocks:
        if block.type == "table" and block.content:
            block.content = convert_otsl_to_html(block.content)
        elif block.type == "equation" and block.content:
            content = block.content.strip()
            content = content.removeprefix("\\[").removesuffix("\\]").strip()
            if not (content.startswith("$") and content.endswith("$")):
                content = f"$${content}$$"
            block.content = content
    return [block for block in blocks if block.type != "equation_block"]

def infer(image: Image.Image, prompt: str) -> str:
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": [
            {"type": "image"},
            {"type": "text", "text": prompt},
        ]},
    ]
    chat_prompt = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = processor(
        text=[chat_prompt],
        images=[image.convert("RGB")],
        padding=True,
        return_tensors="pt",
    ).to(device=model.device, dtype=model.dtype)
    output_ids = model.generate(
        **inputs,
        use_cache=True,
        max_new_tokens=4096,
        do_sample=False,
    )
    output_ids = output_ids.cpu().tolist()[0][len(inputs.input_ids[0]):]
    return processor.batch_decode(
        [output_ids],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0].strip()



processor = AutoProcessor.from_pretrained("StarDoc-AI/TeleOCR", trust_remote_code=True, use_fast=True)
model = AutoModel.from_pretrained(
    "StarDoc-AI/TeleOCR",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).cuda().eval()

# text
image=Image.open("./assets/text.png").convert("RGB")
raw_text = infer(image, "Please output the text content from the image.")
print(raw_text.strip())

# table
image=Image.open("./assets/table.png").convert("RGB")
raw_otsl = infer(image, "This is the image of a table. Please output the table in OTSL format.")
print(convert_otsl_to_html(raw_otsl))

# formula 
image=Image.open("./assets/formula.png").convert("RGB")
raw_formula = infer(image, "Please write out the expression of the formula in the image using LaTeX format.")
formula_block = ContentBlock("equation", [0.0, 0.0, 1.0, 1.0], content=raw_formula)
formula = post_process([formula_block])[0].content
print(formula)

#code
image=Image.open("./assets/code.png").convert("RGB")
raw_code = infer(image,"The image contains a code snippet, please output the parsing result.")
print(raw_code.strip())

# layout
image=Image.open("./assets/layout.jpg").convert("RGB")
image = image.resize((1036, 1036), Image.Resampling.BICUBIC)
raw_layout = infer(image, "Analyze the image layout.")
print(raw_layout.strip())

# Distorted document layout
layout_image = Image.open("./assets/layout_distorted.jpg").convert("RGB")
layout_image = layout_image.resize((1036, 1036), Image.Resampling.BICUBIC)
raw_layout = infer(layout_image, "\nMulti-point Layout Segmentation Analysis.")
print(raw_layout.strip())

#scientific figure
image=Image.open("./assets/scientific_figure.png").convert("RGB")
raw_scientific_figure = infer(image, "This is a scientific figure. Please extract the table implied by this figure.")
print(convert_otsl_to_html(raw_scientific_figure))

@article{teleocr,
  title={TeleOCR: Navigating Document Parsing Across Digital and Camera-Captured Documents},
  author={Cai, Peng and Zou, Zhaofan and Liu, Shifa and Wang, Yikun and Tang, Jiawei and Yang, Kaicheng and Tong, Meng and He, Zhongjiang and Sun, Hao},
  journal={arXiv preprint arXiv:2608.12898},
  year={2026}
}