import htmlimport itertoolsimport jsonimport refrom dataclasses import dataclass
from PIL import Imageimport torchfrom transformers import AutoProcessor, AutoModel
@dataclassclass ContentBlock: type: str bbox: list[float] angle: int | None = None content: str | None = None
@dataclassclass 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()
# textimage=Image.open("./assets/text.png").convert("RGB")raw_text = infer(image, "Please output the text content from the image.")print(raw_text.strip())
# tableimage=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].contentprint(formula)
#codeimage=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())
# layoutimage=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 layoutlayout_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 figureimage=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))
