"""OCR text cleaning and extraction utilities. This module provides functions to clean and process OCR output from Surya, including handling common OCR artifacts and extracting text for specific regions. """ from __future__ import annotations import logging import re import unicodedata from typing import Any from pdf2zh.parser.utils.bbox import bbox_area, bbox_intersection, polygon_to_bbox logger = logging.getLogger(__name__) def adjust_cell_bbox( matching_cell_lines: list[Any], cell_bbox_pdf: list[float], cell_bbox_image: list[float], padding: float = 0.0, ) -> list[float]: """ Co nhỏ cell_bbox_pdf lại để ôm sát vào phân vùng chứa textlines thực tế, sau đó bổ sung thêm một lượng padding. Args: matching_cell_lines: Danh sách các dòng OCR tìm thấy trong ô cell_bbox_pdf: Bounding box của ô ở hệ PDF [x0, y0, x1, y1] cell_bbox_image: Bounding box của ô ở hệ ảnh [x0, y0, x1, y1] padding: Khoảng cách đệm thêm vào các cạnh (đơn vị: points) Returns: Bounding box mới hệ PDF [x0, y0, x1, y1] đã được điều chỉnh ôm sát text """ if not matching_cell_lines: return cell_bbox_pdf text_x0 = float("inf") text_y0 = float("inf") text_x1 = float("-inf") text_y1 = float("-inf") for line in matching_cell_lines: line_bbox = _get_ocr_bbox(line) if line_bbox is None: continue text_x0 = min(text_x0, line_bbox[0]) text_y0 = min(text_y0, line_bbox[1]) text_x1 = max(text_x1, line_bbox[2]) text_y1 = max(text_y1, line_bbox[3]) if text_x0 == float("inf"): return cell_bbox_pdf img_w = cell_bbox_image[2] - cell_bbox_image[0] img_h = cell_bbox_image[3] - cell_bbox_image[1] pdf_w = cell_bbox_pdf[2] - cell_bbox_pdf[0] pdf_h = cell_bbox_pdf[3] - cell_bbox_pdf[1] scale_x = pdf_w / img_w if img_w > 0 else 1.0 scale_y = pdf_h / img_h if img_h > 0 else 1.0 dx0 = max(0.0, text_x0 - cell_bbox_image[0]) dy0 = max(0.0, text_y0 - cell_bbox_image[1]) dx1 = max(0.0, cell_bbox_image[2] - text_x1) dy1 = max(0.0, cell_bbox_image[3] - text_y1) new_pdf_x0 = cell_bbox_pdf[0] + max(0.0, dx0 * scale_x - padding) new_pdf_y0 = cell_bbox_pdf[1] + max(0.0, dy0 * scale_y - padding) new_pdf_x1 = cell_bbox_pdf[2] - max(0.0, dx1 * scale_x - padding) new_pdf_y1 = cell_bbox_pdf[3] - max(0.0, dy1 * scale_y - padding) final_x0 = max(cell_bbox_pdf[0], min(new_pdf_x0, cell_bbox_pdf[2])) final_y0 = max(cell_bbox_pdf[1], min(new_pdf_y0, cell_bbox_pdf[3])) final_x1 = max(final_x0, min(new_pdf_x1, cell_bbox_pdf[2])) final_y1 = max(final_y0, min(new_pdf_y1, cell_bbox_pdf[3])) return [final_x0, final_y0, final_x1, final_y1] def clean_ocr_text(text: str) -> str: """Clean OCR text by removing artifacts and normalizing whitespace. Processing steps: 1. Normalize Unicode (NFC form) 2. Remove control characters except newlines and tabs 3. Fix common OCR artifacts (ligatures, smart quotes, etc.) 4. Normalize whitespace (collapse multiple spaces, trim lines) 5. Remove empty lines at start/end Args: text: Raw OCR text from Surya Returns: Cleaned text string """ if not text: return "" # Step 1: Unicode normalization text = unicodedata.normalize("NFC", text) # Step 2: Remove control characters except newlines and tabs cleaned_chars = [] for char in text: if char in ("\n", "\t"): cleaned_chars.append(char) elif unicodedata.category(char)[0] != "C": cleaned_chars.append(char) text = "".join(cleaned_chars) # # Step 3: Fix common OCR artifacts # # Ligatures # text = text.replace("\ufb01", "fi") # text = text.replace("\ufb02", "fl") # text = text.replace("\ufb00", "ff") # text = text.replace("\ufb03", "ffi") # text = text.replace("\ufb04", "ffl") # # Smart quotes to straight quotes # text = text.replace("\u2018", "'") # Left single quote # text = text.replace("\u2019", "'") # Right single quote # text = text.replace("\u201c", '"') # Left double quote # text = text.replace("\u201d", '"') # Right double quote # # Dashes # text = text.replace("\u2013", "-") # En dash # text = text.replace("\u2014", "-") # Em dash # text = text.replace("\u2212", "-") # Minus sign # # Other common artifacts # text = text.replace("\u00a0", " ") # Non-breaking space # text = text.replace("\u2026", "...") # Ellipsis text = text.replace("
", "\n") # Line break tags # Step 4: Normalize whitespace # Collapse multiple spaces into one text = re.sub(r" +", " ", text) # Trim each line lines = text.split("\n") lines = [line.strip() for line in lines] # Step 5: Remove empty lines at start and end while lines and not lines[0]: lines.pop(0) while lines and not lines[-1]: lines.pop() return "\n".join(lines) def collect_ocr_text(ocr_result: Any) -> str: """Collect all text lines from an OCR result into a single string. Used after crop-then-OCR: the entire OCR result belongs to one layout region, so we simply concatenate all detected text lines. Args: ocr_result: Surya OCR result with ``text_lines`` attribute Returns: Cleaned concatenated text """ if not hasattr(ocr_result, "text_lines"): return "" lines = [] for line in ocr_result.text_lines: if hasattr(line, "text") and line.text: lines.append(line.text) return clean_ocr_text(" ".join(lines)) def smart_join_text_lines(lines: list[Any]) -> str: if not lines: return "" result = [] last_valid_text = "" # Lưu lại văn bản của dòng có chữ gần nhất for line in lines: current_text = getattr(line, "text", "").strip() if not current_text: continue # Nếu là dòng chứa chữ đầu tiên, chỉ cần thêm vào kết quả if not result: result.append(current_text) last_valid_text = current_text continue ends_with_punctuation = last_valid_text[-1] in {".", "!", "?"} starts_with_upper = current_text[0].isupper() ends_with_hyphen = last_valid_text.endswith("-") if ends_with_hyphen: result.append(current_text) elif not ends_with_punctuation and starts_with_upper: result.append("\n" + current_text) else: result.append(" " + current_text) last_valid_text = current_text return clean_ocr_text("".join(result)) def sort_text_lines(lines: list[Any]) -> list[Any]: """ Sort OCR text lines in reading order (top-to-bottom, left-to-right). """ if not lines: return [] first_line = lines[0] if hasattr(first_line, "bbox") and first_line.bbox: def get_full_bbox(line): b = line.bbox return b[0], b[1], b[2], b[3] elif hasattr(first_line, "polygon"): def get_full_bbox(line): poly = line.polygon xs = [p[0] for p in poly] ys = [p[1] for p in poly] return min(xs), min(ys), max(xs), max(ys) else: return lines boxes = [] for line in lines: x_min, y_min, x_max, y_max = get_full_bbox(line) y_center = (y_min + y_max) / 2.0 boxes.append((y_min, y_center, x_min, y_max, line)) boxes.sort() rows = [] current_row = [] anchor_y_center = None for box in boxes: y_min, y_center, x_min, y_max, line = box if not current_row: current_row.append((x_min, line)) anchor_y_center = y_center else: if y_min <= anchor_y_center <= y_max: current_row.append((x_min, line)) else: rows.append(current_row) current_row = [(x_min, line)] anchor_y_center = y_center if current_row: rows.append(current_row) sorted_lines = [] for row in rows: row.sort() for _, line in row: sorted_lines.append(line) return sorted_lines def extract_text_for_region( ocr_result: Any, region_bbox: list[float], overlap_threshold: float = 0.5, ) -> list[Any]: """Extract OCR text that falls within a region. Finds all text lines from the OCR result that overlap significantly with the given region and concatenates them. Args: ocr_result: Surya OCR result object with text_lines attribute region_bbox: [x0, y0, x1, y1] in image pixels image_width: Image width for coordinate validation image_height: Image height for coordinate validation overlap_threshold: Minimum overlap ratio to include a line Returns: Concatenated text from overlapping lines and estimated font size """ if not hasattr(ocr_result, "text_lines"): return [] matching_lines = _collect_region_matches( getattr(ocr_result, "text_lines", []), region_bbox, overlap_threshold, ) return sort_text_lines(matching_lines) def _collect_region_matches( items: list[Any], region_bbox: list[float], overlap_threshold: float, ) -> list[Any]: matching_items: list[Any] = [] for item in items: if not hasattr(item, "text"): continue item_bbox = _get_ocr_bbox(item) if item_bbox is None: continue intersection = bbox_intersection(region_bbox, item_bbox) if intersection is None: continue item_area = max(1.0, bbox_area(item_bbox)) overlap_ratio = bbox_area(intersection) / item_area if overlap_ratio >= overlap_threshold: matching_items.append(item) return matching_items def _get_ocr_bbox(item: Any) -> list[float] | None: item_bbox = getattr(item, "bbox", None) if item_bbox is not None: return list(item_bbox) if hasattr(item, "polygon"): return polygon_to_bbox(item.polygon) return None def log_toc_hints(elements: list[Any], page_index: int) -> None: """Log potential Table of Contents entries for debugging. Looks for Section-header elements that might indicate chapter structure and logs them for manual review during development. Args: elements: List of ElementData objects from the page page_index: 0-based page number for logging context """ toc_hints = [] for elem in elements: label = getattr(elem, "label", "") # Look for section headers and TOC elements if label in ("Section-header", "Table-of-contents"): text = getattr(elem, "source_text", "") if text: # Truncate long text display_text = text[:80] + "..." if len(text) > 80 else text toc_hints.append(f" [{label}] {display_text}") if toc_hints: logger.debug(f"Page {page_index} TOC hints:\n" + "\n".join(toc_hints)) def join_raw_text(elements: list[Any]) -> str: """Concatenate ``source_text`` from translatable layout elements. Collects source text from every :class:`~pdf2zh.scanned.enums.ElementCategory` that carries translatable content (``FLOWING_TEXT`` and ``IN_PLACE``) and joins them with newlines to form the ``raw_text`` field of :class:`~pdf2zh.scanned.models.PageData`. BYPASS, TABLE, and EQUATION categories are intentionally excluded: BYPASS has no text; TABLE text is stored per-cell; EQUATION text is a placeholder handled separately. Args: elements: Ordered list of :class:`~pdf2zh.scanned.models.ElementData` objects (or any object with ``category`` and ``source_text`` attributes) for a single page. Returns: Single string with element texts joined by ``"\n"``, or an empty string if no translatable elements are present. """ from pdf2zh.parser.enums import ElementCategory text_parts = [] for elem in elements: category = getattr(elem, "category", None) # Only include FLOWING_TEXT and IN_PLACE categories if category in (ElementCategory.FLOWING_TEXT, ElementCategory.IN_PLACE): source_text = getattr(elem, "source_text", "") if source_text: text_parts.append(source_text) return "\n".join(text_parts)