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/NaviDC-OCR", trust_remote_code=True, use_fast=True)
model = AutoModel.from_pretrained(
"StarDoc-AI/NaviDC-OCR",
trust_remote_code=True,
torch_dtype=torch.bfloat16,
).cuda().eval()
image=Image.open("./assets/text.png").convert("RGB")
raw_text = infer(image, "Please output the text content from the image.")
print(raw_text.strip())
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))
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)
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())
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())
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())
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))