"""Post-processing utilities for Nemotron Parse model output. Parses raw model output into structured elements (classes, bboxes, text) and applies format conversions for tables and plain text. """ import re from .latex2html import convert_html_tables_to_markdown, latex_table_to_html def extract_classes_bboxes( text: str, ) -> tuple[list[str], list[tuple[float, float, float, float]], list[str]]: """Parse raw model output into semantic classes, bounding boxes, and text. :param str text: raw decoded output containing ```` coordinate tokens, content, and ```` tags :raises AssertionError: if a ``Page-number`` class is detected (unsupported) :return tuple: ``(classes, bboxes, texts)`` — class names, normalised ``(x1, y1, x2, y2)`` tuples, and extracted text strings """ _re_extract_class_bbox = re.compile( r"(.*?)]+)>", re.DOTALL, ) classes = [] bboxes = [] texts = [] for m in _re_extract_class_bbox.finditer(text): x1, y1, text, x2, y2, cls = m.groups() classes.append(cls) bboxes.append((float(x1), float(y1), float(x2), float(y2))) texts.append(text) # TODO: Remove when fixed classes = ["Formula" if cls == "Inline-formula" else cls for cls in classes] assert "Page-number" not in classes return classes, bboxes, texts def transform_bbox_to_original( bbox: tuple[float, float, float, float], original_width: int, original_height: int, target_w: int = 1664, target_h: int = 2048, ) -> tuple[float, float, float, float]: """Map a normalised bounding box back to original image pixel coordinates. Replicates the resize-then-pad logic from the image preprocessor to invert the coordinate transform applied during inference. :param tuple bbox: normalised ``(x1, y1, x2, y2)`` in 0-1 range relative to the padded target image :param int original_width: width of the original input image in pixels :param int original_height: height of the original input image in pixels :param int target_w: padded image width used during preprocessing, defaults to 1664 :param int target_h: padded image height used during preprocessing, defaults to 2048 :return tuple: ``(left, top, right, bottom)`` in original image pixel coordinates """ # Replicate exact resize logic aspect_ratio = original_width / original_height new_height = original_height new_width = original_width if original_height > target_h: new_height = target_h new_width = int(new_height * aspect_ratio) if new_width > target_w: new_width = target_w new_height = int(new_width / aspect_ratio) resized_width = new_width resized_height = new_height # Calculate padding pad_left = (target_w - resized_width) // 2 pad_top = (target_h - resized_height) // 2 # # Transform: use the ACTUAL resized dimensions, not the scale # # X coords left = ((bbox[0] * target_w) - pad_left) * original_width / resized_width right = ((bbox[2] * target_w) - pad_left) * original_width / resized_width # # Y coords - using original_height / resized_height directly top = ((bbox[1] * target_h) - pad_top) * original_height / resized_height bottom = ((bbox[3] * target_h) - pad_top) * original_height / resized_height return left, top, right, bottom def postprocess_text( text: str, cls: str = "Text", text_format: str = "markdown", table_format: str = "latex", blank_text_in_figures: bool = False, ) -> str: """Apply format conversions to extracted text for a single element. For tables, converts LaTeX tabular markup to the requested format. For non-table elements, optionally strips markdown formatting. :param str text: raw text content from the model output :param str cls: semantic class (e.g. ``"Table"``, ``"Text"``, ``"Picture"``) :param str text_format: ``"markdown"`` or ``"plain"``, defaults to ``"markdown"`` :param str table_format: ``"latex"``, ``"HTML"``, or ``"markdown"``, defaults to ``"latex"`` :param bool blank_text_in_figures: if True, return empty string for ``"Picture"`` elements, defaults to False :raises AssertionError: if *text_format* or *table_format* is not a supported value :return str: post-processed text string """ assert text_format in ["markdown", "plain"], ( "Unknown text format. Supported: markdown | plain" ) assert table_format in ["latex", "HTML", "markdown"], ( "Unknown table format. Supported: latex | HTML | markdown" ) if cls != "Table": if text_format == "plain": text = convert_mmd_to_plain_text_ours(text) elif table_format == "HTML": text = latex_table_to_html(text) elif table_format == "markdown": text = convert_html_tables_to_markdown(latex_table_to_html(text)) if blank_text_in_figures and cls == "Picture": text = "" return text def remove_nemotron_formatting(text: str) -> str: """Strip Nemotron-specific control tokens (````, ``\\unknown``, etc.). :param str text: raw text that may contain model-specific tokens :return str: cleaned text with control tokens removed """ text = text.replace("", "") text = text.replace("\\<|unk|\\>", "") text = text.replace("\\unknown", "") return text def convert_mmd_to_plain_text_ours(mmd_text: str) -> str: """Convert markdown/MMD formatted text to plain text. Strips markdown headers, bold, italic, and emphasis markers while preserving superscript/subscript content in LaTeX-style notation. :param str mmd_text: markdown-formatted text from the model :return str: plain text with formatting markers removed """ mmd_text = re.sub(r"(.*?)", r"^{\\1}", mmd_text, flags=re.DOTALL) mmd_text = re.sub(r"(.*?)", r"_{\\1}", mmd_text, flags=re.DOTALL) mmd_text = mmd_text.replace("
", "\n") # Remove headers (e.g., ##) mmd_text = re.sub(r"#+\s", "", mmd_text) # Remove bold (e.g., **) mmd_text = re.sub(r"\*\*(.*?)\*\*", r"\1", mmd_text) # mmd_text = mmd_text.replace("**","") # Remove italic (e.g., *) mmd_text = re.sub(r"\*(.*?)\*", r"\1", mmd_text) # Remove emphasized text formatting (e.g., _) mmd_text = re.sub(r"(?