dsh-ab-ocr 0.1.0

package/python/layout.py

@@ -0,0 +1,403 @@
+"""Box geometry, line assembly, column detection, and heading inference.
+
+The OCR engine returns one bounding box per detected text region, in no
+guaranteed order. Everything that turns those boxes into ordered lines and a
+document outline lives here, with no engine, file, or process dependency, so the
+layout rules can be driven directly by tests.
+"""
+
+from __future__ import annotations
+
+import re
+import unicodedata
+from dataclasses import dataclass
+from statistics import median
+from typing import Sequence
+
+#: Two boxes on one row further apart than this multiple of the row height are
+#: treated as separate fields rather than as words of one line.
+WIDE_GAP_FACTOR = 1.6
+
+#: Vertical overlap, relative to the shorter box, that puts two boxes on one row.
+ROW_OVERLAP = 0.5
+
+#: Horizontal run of free space, as a fraction of page width, a gutter must span.
+GUTTER_MIN_WIDTH = 0.02
+
+#: Fraction of rows that must be free across a gutter for it to be a column split.
+GUTTER_ROW_RATIO = 0.60
+
+#: Fewest rows a page needs before column detection is attempted at all.
+GUTTER_MIN_ROWS = 8
+
+
+@dataclass
+class Line:
+    """One detected text region, with the geometry its layout depends on."""
+
+    text: str
+    score: float
+    x0: float
+    y0: float
+    x1: float
+    y1: float
+    column: int = 0
+
+    @property
+    def height(self) -> float:
+        """Vertical extent of the region, which approximates the glyph size."""
+        return self.y1 - self.y0
+
+    @property
+    def width(self) -> float:
+        """Horizontal extent of the region."""
+        return self.x1 - self.x0
+
+    @property
+    def center_y(self) -> float:
+        """Vertical midpoint, used to order rows."""
+        return (self.y0 + self.y1) / 2.0
+
+
+def box_geometry(box: Sequence[Sequence[float]]) -> tuple[float, float, float, float]:
+    """Reduce one four-point OCR box to its upright bounding rectangle.
+
+    @param box: four (x, y) corners, in any rotation order.
+    @returns: (x0, y0, x1, y1) of the enclosing rectangle.
+    """
+    xs = [float(point[0]) for point in box]
+    ys = [float(point[1]) for point in box]
+    return min(xs), min(ys), max(xs), max(ys)
+
+
+def is_cjk(character: str) -> bool:
+    """Report whether a character is written without inter-word spacing.
+
+    @param character: one character.
+    @returns: True for CJK ideographs, kana, and full-width punctuation.
+    """
+    code = ord(character)
+    return (
+        0x3000 <= code <= 0x303F
+        or 0x3040 <= code <= 0x30FF
+        or 0x3400 <= code <= 0x4DBF
+        or 0x4E00 <= code <= 0x9FFF
+        or 0xF900 <= code <= 0xFAFF
+        or 0xFF00 <= code <= 0xFFEF
+        or 0x20000 <= code <= 0x2FA1F
+    )
+
+
+def collapse_space(text: str) -> str:
+    """Collapse whitespace runs and drop the invisible characters OCR emits.
+
+    @param text: raw recognized text.
+    @returns: the text with single spaces and no zero-width or control characters.
+    """
+    cleaned = text.replace('\u3000', ' ')
+    cleaned = ''.join(
+        character
+        for character in cleaned
+        if character.isspace() or unicodedata.category(character)[0] != 'C'
+    )
+    return re.sub(r'[ \t\r\n]+', ' ', cleaned).strip()
+
+
+def _rows(items: Sequence[Line]) -> list[list[Line]]:
+    """Group same-row boxes by vertical overlap.
+
+    @param items: boxes of one column, in any order.
+    @returns: rows, each a list of boxes, ordered top to bottom.
+    """
+    rows: list[list[Line]] = []
+    for item in sorted(items, key=lambda line: line.y0):
+        target: list[Line] | None = None
+        for row in rows:
+            reference = row[0]
+            overlap = min(item.y1, reference.y1) - max(item.y0, reference.y0)
+            if overlap >= ROW_OVERLAP * min(item.height, reference.height):
+                target = row
+                break
+        if target is None:
+            rows.append([item])
+        else:
+            target.append(item)
+    return rows
+
+
+def merge_row(row: Sequence[Line], column: int) -> Line:
+    """Flatten one row of boxes into a single line in left-to-right order.
+
+    Adjacent CJK fragments are joined with no space; a Latin fragment boundary
+    keeps one, because there the space is part of the text.
+    @param row: the row's boxes.
+    @param column: column index the row belongs to.
+    @returns: the merged line.
+    """
+    ordered = sorted(row, key=lambda line: line.x0)
+    pieces: list[tuple[str, float]] = []
+    for index, item in enumerate(ordered):
+        text = collapse_space(item.text)
+        if text == '':
+            continue
+        gap = 0.0 if not pieces else item.x0 - ordered[index - 1].x1
+        pieces.append((text, gap))
+    joined = ''
+    for index, (piece, gap) in enumerate(pieces):
+        if index > 0 and joined != '':
+            wide = gap > WIDE_GAP_FACTOR * max(ordered[0].height, 1.0)
+            if not (is_cjk(joined[-1]) and is_cjk(piece[0])):
+                joined += '  ' if wide else ' '
+        joined += piece
+    return Line(
+        text=joined.strip(),
+        score=min(line.score for line in ordered),
+        x0=min(line.x0 for line in ordered),
+        y0=min(line.y0 for line in ordered),
+        x1=max(line.x1 for line in ordered),
+        y1=max(line.y1 for line in ordered),
+        column=column,
+    )
+
+
+def free_runs(row: Sequence[Line], width: float) -> list[tuple[float, float]]:
+    """Compute the horizontal spans of one row that carry no box.
+
+    @param row: the row's boxes.
+    @param width: page width in pixels.
+    @returns: free spans as (start, end) pairs, left to right.
+    """
+    spans = sorted((item.x0, item.x1) for item in row)
+    runs: list[tuple[float, float]] = []
+    cursor = 0.0
+    for start, end in spans:
+        if start - cursor > 0.0:
+            runs.append((cursor, start))
+        cursor = max(cursor, end)
+    if width - cursor > 0.0:
+        runs.append((cursor, width))
+    return runs
+
+
+def find_gutter(rows: Sequence[Sequence[Line]], width: float) -> float | None:
+    """Locate a vertical gutter that splits the page into reading columns.
+
+    A gutter is a horizontal position that falls inside a free span of most
+    rows, so text left of it and text right of it are separate reading columns
+    rather than two fragments of one line.
+    @param rows: every row of the page.
+    @param width: page width in pixels.
+    @returns: the gutter's x position, or None when the page is single-column.
+    """
+    if len(rows) < GUTTER_MIN_ROWS or width <= 0:
+        return None
+    steps = 50
+    best_coverage = 0
+    best_positions: list[float] = []
+    for step in range(1, steps):
+        position = width * step / steps
+        coverage = 0
+        for row in rows:
+            for start, end in free_runs(row, width):
+                # A page margin is free on every row, so only an interior run
+                # can be the gutter between two columns.
+                if start <= 0.0 or end >= width:
+                    continue
+                if start < position < end and (end - start) >= GUTTER_MIN_WIDTH * width:
+                    coverage += 1
+                    break
+        if coverage > best_coverage:
+            best_coverage = coverage
+            best_positions = [position]
+        elif coverage == best_coverage and coverage > 0:
+            best_positions.append(position)
+    if best_coverage < GUTTER_ROW_RATIO * len(rows) or not best_positions:
+        return None
+    position = median(best_positions)
+    left = sum(1 for row in rows for item in row if item.x0 + item.x1 <= 2.0 * position)
+    right = sum(1 for row in rows for item in row if item.x0 + item.x1 > 2.0 * position)
+    if left < 3 or right < 3:
+        return None
+    return position
+
+
+def assemble_lines(
+    boxes: Sequence[Sequence[Sequence[float]]],
+    texts: Sequence[str],
+    scores: Sequence[float],
+    width: float,
+    min_score: float,
+    detect_columns: bool,
+) -> list[Line]:
+    """Turn engine output into reading-order lines, splitting columns when present.
+
+    @param boxes: one four-point box per recognized region.
+    @param texts: recognized text per region.
+    @param scores: recognition confidence per region.
+    @param width: page width in pixels.
+    @param min_score: regions below this confidence are dropped.
+    @param detect_columns: whether a multi-column page is split into columns.
+    @returns: the page's lines, in the order a reader reads them.
+    """
+    items = [
+        Line(text=text, score=float(score), **dict(zip(('x0', 'y0', 'x1', 'y1'), box_geometry(box))))
+        for box, text, score in zip(boxes, texts, scores)
+        if collapse_space(text) != '' and float(score) >= min_score
+    ]
+    if not items:
+        return []
+    rows = _rows(items)
+    gutter = find_gutter(rows, width) if detect_columns else None
+    if gutter is None:
+        return [merge_row(row, 0) for row in sorted(rows, key=lambda row: row[0].y0)]
+    left: list[Line] = []
+    right: list[Line] = []
+    for row in rows:
+        for item in row:
+            (left if item.x0 + item.x1 <= 2.0 * gutter else right).append(item)
+    return [merge_row([line], 0) for line in sorted(left, key=lambda line: (line.y0, line.x0))] + [
+        merge_row([line], 1) for line in sorted(right, key=lambda line: (line.y0, line.x0))
+    ]
+
+
+def body_height(lines: Sequence[Line]) -> float:
+    """Estimate the document's body glyph height from its line heights.
+
+    Each line votes with its character count, so a few oversized titles cannot
+    drag the body estimate up and short captions cannot drag it down.
+    @param lines: every kept line of the document.
+    @returns: the weighted median line height, or 0 when there are no lines.
+    """
+    weighted: list[float] = []
+    for line in lines:
+        weighted.extend([line.height] * max(1, min(len(line.text), 60)))
+    return median(weighted) if weighted else 0.0
+
+
+#: Numbering and section words that fix a heading's level independent of size.
+_CHAPTER = re.compile(r'^第\s*[0-9一二三四五六七八九十百零〇]+\s*[章篇部节]')
+_PART = re.compile(r'^第\s*[0-9一二三四五六七八九十百零〇]+\s*部分')
+#: A multi-level section number such as 1.1 or 1.1.2. OCR often drops the thin
+#: space that follows it, so the number alone is enough to recognize the level.
+_SECTION_NUMBER = re.compile(r'^(\d{1,3}(?:\.\d{1,3})+)')
+#: A single-level outline number. The ASCII full stop is deliberately excluded:
+#: a line opening with an ASCII full stop is a numbered list item, not a heading.
+_SIMPLE_NUMBER = re.compile(r'^(\d{1,2})\s*[、．:：]\s*(?=\S)|^(\d{1,2})\s+(?=\S)')
+_CJK_ORDINAL = re.compile(r'^[一二三四五六七八九十百零〇]+\s*[、.．]\s*')
+_NAMED = re.compile(
+    r'^(?:摘\s*要|关\s*键\s*词|目\s*录|前\s*言|引\s*言|绪\s*论|结\s*论|结\s*语|'
+    r'参\s*考\s*文\s*献|致\s*谢|附\s*录|附\s*件|正\s*文|简\s*介|背\s*景|'
+    r'abstract|introduction|conclusion|conclusions|references|appendix|'
+    r'contents|acknowledg\w*|keywords?)\b',
+    re.IGNORECASE,
+)
+_LIST_ITEM = re.compile(r'^(?:[-•·▪◦*]|\(\d+\)|（\d+）|\d+[.)、])\s*')
+#: Punctuation that ends a sentence, so a line carrying it is prose, not a heading.
+_SENTENCE_END = '。！？；!?;：'
+#: The number prefix a heading carries, used to put back the space after it.
+_HEADING_PREFIX = re.compile(
+    r'^(?:第\s*[\d一二三四五六七八九十百零〇]+\s*[章篇部节]'
+    r'|第\s*[\d一二三四五六七八九十百零〇]+\s*部分'
+    r'|\d{1,3}(?:\.\d{1,3})+'
+    r'|\d{1,2}(?=[、．:：\s]))'
+)
+
+
+def heading_level(text: str, height: float, base: float, min_ratio: float) -> int | None:
+    """Decide whether a line is a heading, and at which Markdown level.
+
+    Explicit numbering and section words win over glyph size, because a numbered
+    heading keeps its place in the outline even when it is set in body type.
+    @param text: the line's text.
+    @param height: the line's glyph height.
+    @param base: the document's body glyph height.
+    @param min_ratio: how much taller than the body a line must be to be a heading.
+    @returns: a Markdown level from 1 to 6, or None when the line is body text.
+    """
+    stripped = text.strip()
+    if stripped == '' or len(stripped) > 60:
+        return None
+    numbered = section_number(stripped)
+    if _LIST_ITEM.match(stripped) and numbered is None:
+        return None
+    if _CHAPTER.match(stripped) or _PART.match(stripped) or is_named_section(stripped):
+        return 1
+    if numbered is not None:
+        depth, offset = numbered
+        tail = stripped[offset:].strip()
+        if tail != '' and len(tail) <= 40 and not tail.endswith(tuple(_SENTENCE_END)):
+            return depth
+    if _CJK_ORDINAL.match(stripped) and len(stripped) <= 30:
+        return 2
+    if base > 0:
+        ratio = height / base
+        if ratio >= min_ratio and not stripped.endswith(tuple(_SENTENCE_END)):
+            if ratio >= min_ratio * 1.55:
+                return 1
+            if ratio >= min_ratio * 1.28:
+                return 2
+            return 3
+    return None
+
+
+def is_named_section(text: str) -> bool:
+    """Report whether a line opens with one of the document's named sections.
+
+    @param text: the line's text.
+    @returns: True for a named part such as an abstract, a contents list, or a
+        reference list, whatever its glyph size.
+    """
+    return _NAMED.match(text.strip()) is not None
+
+
+def section_number(text: str) -> tuple[int, int] | None:
+    """Read a leading outline number and the offset it ends at.
+
+    @param text: the line's text.
+    @returns: the outline depth and the offset after the number, or None when
+        the line does not open with an outline number.
+    """
+    match = _SECTION_NUMBER.match(text)
+    if match is not None:
+        return min(match.group(1).count('.') + 1, 6), match.end()
+    match = _SIMPLE_NUMBER.match(text)
+    if match is not None:
+        return 1, match.end()
+    return None
+
+
+def normalize_heading(text: str) -> str:
+    """Restore the space a section number and its title are set with.
+
+    Recognition frequently drops the thin space after an outline number, so a
+    heading arrives as one run of characters; the outline reads better and still
+    matches its source when the space is put back.
+    @param text: the heading's text.
+    @returns: the heading with the space after its number restored.
+    """
+    stripped = text.strip()
+    match = _HEADING_PREFIX.match(stripped)
+    if match is None:
+        return stripped
+    offset = match.end()
+    if offset == 0 or offset >= len(stripped) or stripped[offset].isspace():
+        return stripped
+    return stripped[:offset] + ' ' + stripped[offset:]
+
+
+def is_list_item(text: str) -> bool:
+    """Report whether a line already carries a list marker.
+
+    @param text: the line's text.
+    @returns: True when the line opens with a bullet or an enumerated marker.
+    """
+    return _LIST_ITEM.match(text.strip()) is not None
+
+
+def merge_list_markers(text: str) -> str:
+    """Normalize list markers to Markdown bullets.
+
+    @param text: a line recognized as a list item.
+    @returns: the line with its marker replaced by a Markdown bullet.
+    """
+    return '- ' + _LIST_ITEM.sub('', text.strip(), count=1).strip()