Parse Scanned and Phone-Photo Documents with TeleOCR (1.2B)

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

# The id used in the card's own code. The card itself is hosted at XingChen-AGI/TeleOCR.
MODEL_ID = "StarDoc-AI/TeleOCR"

processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=True)
model = AutoModel.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).cuda().eval()

PROMPT_TEXT = "Please output the text content from the image."
PROMPT_TABLE = "This is the image of a table. Please output the table in OTSL format."
PROMPT_FORMULA = "Please write out the expression of the formula in the image using LaTeX format."
PROMPT_CODE = "The image contains a code snippet, please output the parsing result."
PROMPT_LAYOUT = "Analyze the image layout."
PROMPT_LAYOUT_DISTORTED = "\nMulti-point Layout Segmentation Analysis."
PROMPT_FIGURE = "This is a scientific figure. Please extract the table implied by this figure."


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()


if __name__ == "__main__":
    import sys

    image = Image.open(sys.argv[1]).convert("RGB")
    print(infer(image, PROMPT_TEXT))

import sys
from pathlib import Path

from PIL import Image

from teleocr import infer, PROMPT_TEXT

src = Path(sys.argv[1])
dst = Path(sys.argv[2])
dst.mkdir(parents=True, exist_ok=True)

for path in sorted(src.iterdir()):
    if path.suffix.lower() not in {".png", ".jpg", ".jpeg"}:
        continue
    image = Image.open(path).convert("RGB")
    text = infer(image, PROMPT_TEXT)
    (dst / f"{path.stem}.txt").write_text(text, encoding="utf-8")
    print(f"{path.name}: {len(text)} chars")

import html
import itertools
import re
from dataclasses import dataclass


@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>"

import sys

from PIL import Image

from otsl import convert_otsl_to_html
from teleocr import infer, PROMPT_TABLE, PROMPT_FIGURE

PROMPTS = {"table": PROMPT_TABLE, "figure": PROMPT_FIGURE}

mode, path = sys.argv[1], sys.argv[2]
if mode not in PROMPTS:
    sys.exit(f"Unknown mode {mode!r}: use 'table' or 'figure'")
prompt = PROMPTS[mode]

image = Image.open(path).convert("RGB")
raw_otsl = infer(image, prompt)
table_html = convert_otsl_to_html(raw_otsl)

if not table_html:
    print("Could not parse OTSL output:\n" + raw_otsl)
else:
    print(table_html)

import sys
from dataclasses import dataclass

from PIL import Image

from otsl import convert_otsl_to_html
from teleocr import infer, PROMPT_FORMULA


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


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"]


if __name__ == "__main__":
    image = Image.open(sys.argv[1]).convert("RGB")
    raw_formula = infer(image, PROMPT_FORMULA)
    formula_block = ContentBlock("equation", [0.0, 0.0, 1.0, 1.0], content=raw_formula)
    print(post_process([formula_block])[0].content)

import sys

from PIL import Image

from teleocr import infer, PROMPT_LAYOUT, PROMPT_LAYOUT_DISTORTED

path = sys.argv[1]
distorted = len(sys.argv) > 2 and sys.argv[2] == "--distorted"

image = Image.open(path).convert("RGB")
image = image.resize((1036, 1036), Image.Resampling.BICUBIC)

prompt = PROMPT_LAYOUT_DISTORTED if distorted else PROMPT_LAYOUT
print(infer(image, prompt).strip())