ocr-postprocess 0.1.0__py3-none-any.whl

ocr_postprocess/engine/reconciler.py

@@ -0,0 +1,170 @@
+"""Stage 7 — Reconciler: merge multi-source candidates, cross-validate."""
+
+from __future__ import annotations
+
+import logging
+import unicodedata
+from typing import Any
+
+from ocr_postprocess.models import Candidate, CrossCheck, PipelineContext
+
+logger = logging.getLogger(__name__)
+
+# Priority order for conflict resolution (higher index = higher trust)
+_SOURCE_PRIORITY = [
+    "label_anchor",
+    "pattern",
+    "pattern_with_checksum",
+    "structured",
+    "constant",
+    "computed",
+]
+
+
+def _source_priority(sources: list[str]) -> int:
+    best = 0
+    for src in sources:
+        for i, key in enumerate(_SOURCE_PRIORITY):
+            if key in src:
+                best = max(best, i)
+    return best
+
+
+def _normalize_for_compare(value: Any) -> str:
+    """Normalise value to a comparable string."""
+    if value is None:
+        return ""
+    s = str(value).strip().lower()
+    # Remove diacritics
+    nfkd = unicodedata.normalize("NFKD", s)
+    return "".join(c for c in nfkd if not unicodedata.combining(c))
+
+
+def reconcile_stage(ctx: PipelineContext) -> None:
+    """Pipeline stage 7: reconcile and merge candidates."""
+    profile = ctx.profile
+
+    # Group by key
+    grouped: dict[str, list[Candidate]] = {}
+    for cand in ctx.candidates:
+        grouped.setdefault(cand.key, []).append(cand)
+
+    merged: list[Candidate] = []
+    cross_checks: list[CrossCheck] = []
+
+    for key, candidates in grouped.items():
+        if not candidates:
+            continue
+
+        if len(candidates) == 1:
+            merged.append(candidates[0])
+            continue
+
+        # Normalize values for comparison
+        norm_values = [_normalize_for_compare(c.value) for c in candidates]
+        unique_values = set(v for v in norm_values if v)
+
+        all_sources = []
+        for c in candidates:
+            all_sources.extend(c.sources)
+
+        if len(unique_values) <= 1:
+            # All agree — boost confidence
+            best = max(candidates, key=lambda c: c.confidence)
+            boosted_conf = min(1.0, 1.0 - (1.0 - best.confidence) * 0.7)
+            merged.append(
+                best.model_copy(
+                    update={
+                        "confidence": boosted_conf,
+                        "sources": list(dict.fromkeys(all_sources)),
+                    }
+                )
+            )
+        else:
+            # Conflict — choose by source priority, flag conflict
+            best = max(candidates, key=lambda c: (_source_priority(c.sources), c.confidence))
+            # Internal detail (for logs): full extractor=value pairs
+            detail = " vs ".join(f"{c.extractor}='{c.value}'" for c in candidates)
+            # Human-readable: unique values only
+            seen_vals: list[str] = []
+            for c in candidates:
+                v = str(c.value)
+                if v not in seen_vals:
+                    seen_vals.append(v)
+            field_label = key
+            if ctx.profile:
+                fdef = next((f for f in ctx.profile.fields if f.key == key), None)
+                if fdef and fdef.aliases:
+                    field_label = fdef.aliases[0]
+            readable_vals = ", ".join(f'"{v}"' for v in seen_vals)
+            warning_msg = f"Conflict: {field_label} — nhiều giá trị khác nhau: {readable_vals}"
+            logger.warning("Conflict for field '%s': %s", key, detail)
+            ctx.warnings.append(warning_msg)
+            merged.append(
+                best.model_copy(
+                    update={
+                        "conflict": True,
+                        "sources": list(dict.fromkeys(all_sources)),
+                        "notes": best.notes + [f"Conflict: {detail}"],
+                    }
+                )
+            )
+
+        # Cross-check: compare candidates from different extractors
+        if len(candidates) >= 2:
+            all_agree = len(unique_values) <= 1
+            # Collect unique original (non-normalized) values for human-readable display
+            seen: list[str] = []
+            for c in candidates:
+                v = str(c.value)
+                if v not in seen:
+                    seen.append(v)
+            # Build value_sources: for each unique value, pick best-confidence candidate as source
+            vs_map: dict[str, dict] = {}
+            for c in candidates:
+                v = str(c.value)
+                if v not in vs_map or c.confidence > vs_map[v]["confidence"]:
+                    vs_map[v] = {
+                        "value": v,
+                        "extractor": c.extractor,
+                        "confidence": c.confidence,
+                        "raw": c.raw or "",
+                        "line_index": c.line_index,
+                    }
+            cross_checks.append(
+                CrossCheck(
+                    field_key=key,
+                    sources=list(dict.fromkeys(all_sources)),
+                    matched=all_agree,
+                    detail=f"{norm_values[0]!r}" if all_agree else detail,
+                    values=seen,
+                    value_sources=list(vs_map.values()),
+                )
+            )
+
+    # Inject constant fields from profile
+    if profile:
+        existing_keys = {c.key for c in merged}
+        for field in profile.fields:
+            if field.constant is not None and field.key not in existing_keys:
+                merged.append(
+                    Candidate(
+                        key=field.key,
+                        value=field.constant,
+                        raw=str(field.constant),
+                        extractor="constant",
+                        sources=["constant"],
+                        confidence=1.0,
+                    )
+                )
+
+        # Warn on missing required fields
+        for field in profile.fields:
+            if field.required and not any(c.key == field.key for c in merged):
+                msg = f"Required field missing: {field.key}"
+                ctx.warnings.append(msg)
+                logger.warning(msg)
+
+    ctx.candidates = merged
+    ctx.cross_checks = cross_checks
+    logger.debug("Reconciler: %d candidates, %d cross-checks", len(merged), len(cross_checks))

ocr_postprocess/engine/reconstructor.py

@@ -0,0 +1,469 @@
+"""Stage 4 — LineReconstructor: 6 sub-steps to rebuild document structure."""
+
+from __future__ import annotations
+
+import logging
+import unicodedata
+from typing import Any
+
+import regex as re
+from rapidfuzz import fuzz
+
+from ocr_postprocess.models import LabelHit, Line, PipelineContext, Section
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _normalize_label(text: str) -> str:
+    """Lowercase + strip diacritics for fuzzy label comparison."""
+    nfkd = unicodedata.normalize("NFKD", text)
+    return "".join(c for c in nfkd if not unicodedata.combining(c)).lower().strip()
+
+
+def _all_aliases(profile: Any) -> list[str]:
+    """Return all field aliases from profile (flat list)."""
+    aliases = []
+    if not profile:
+        return aliases
+    for field in profile.fields:
+        aliases.extend(field.aliases)
+    return aliases
+
+
+# ---------------------------------------------------------------------------
+# Sub-step (a) — Section splitter
+# ---------------------------------------------------------------------------
+
+_ALL_CAPS_LINE = re.compile(r"^[A-ZÀÁÂÃÈÉÊÌÍÒÓÔÕÙÚĂĐĨŨƠƯẠ-Ỵ0-9 \-/()]{8,}$")
+_NUMBERED_HEADING = re.compile(r"^\d+\.\s+[A-ZÀÁÂÃÈÉÊÌÍÒÓÔÕÙÚĂĐĨŨƠƯẠ-Ỵ]")
+
+
+def _step_a_section_split(
+    text: str, section_defs: list[Any]
+) -> list[tuple[str, str | None, list[str]]]:
+    """Split text into sections.
+
+    Returns list of (section_id, title, lines).
+    """
+    lines = text.split("\n")
+
+    # Build compiled markers from profile
+    markers: list[tuple[str, str, bool]] = []  # (id, pattern, is_regex)
+    for sdef in section_defs:
+        for start in sdef.start:
+            markers.append((sdef.id, start, sdef.is_regex))
+
+    if not markers:
+        return [("_default", None, lines)]
+
+    sections: list[tuple[str, str | None, list[str]]] = []
+    current_id = "_default"
+    current_title: str | None = None
+    current_lines: list[str] = []
+
+    for line in lines:
+        matched_id = None
+        stripped = line.strip()
+
+        for sec_id, pattern, is_regex in markers:
+            if is_regex:
+                if re.search(pattern, stripped):
+                    matched_id = sec_id
+                    break
+            else:
+                if pattern.lower() in stripped.lower():
+                    matched_id = sec_id
+                    break
+
+        # Fallback heuristic
+        if matched_id is None and len(stripped) >= 8:
+            if _ALL_CAPS_LINE.match(stripped) or _NUMBERED_HEADING.match(stripped):
+                matched_id = re.sub(r"\s+", "_", stripped[:20]).upper()
+
+        if matched_id is not None and matched_id != current_id:
+            if current_lines:
+                sections.append((current_id, current_title, current_lines))
+            current_id = matched_id
+            current_title = stripped
+            current_lines = []
+        else:
+            current_lines.append(line)
+
+    if current_lines or not sections:
+        sections.append((current_id, current_title, current_lines))
+
+    return sections
+
+
+# ---------------------------------------------------------------------------
+# Sub-step (b) — Bilingual label normalizer
+# ---------------------------------------------------------------------------
+
+_BILINGUAL_SLASH = re.compile(r"(.+?)\s*/\s*([A-Za-z].+?)(\s*:)?$")
+_BILINGUAL_PAREN = re.compile(r"(.+?)\s*\(([A-Za-z][^)]+)\)(\s*:)?$")
+
+
+def _step_b_build_bilingual_map(
+    lines: list[str],
+    bilingual_pairs: list[list[str]],
+) -> dict[str, str]:
+    """Build VN↔EN alias map from explicit YAML pairs and auto-detected patterns."""
+    alias_map: dict[str, str] = {}
+
+    # Explicit pairs from YAML
+    for pair in bilingual_pairs:
+        if len(pair) == 2:
+            alias_map[_normalize_label(pair[0])] = pair[1]
+            alias_map[_normalize_label(pair[1])] = pair[0]
+
+    # Auto-detect from text
+    for line in lines:
+        stripped = line.strip().rstrip(":")
+        for pattern in (_BILINGUAL_SLASH, _BILINGUAL_PAREN):
+            m = pattern.match(stripped)
+            if m:
+                vn, en = m.group(1).strip(), m.group(2).strip()
+                alias_map[_normalize_label(vn)] = en
+                alias_map[_normalize_label(en)] = vn
+                break
+
+    return alias_map
+
+
+# ---------------------------------------------------------------------------
+# Sub-step (c) — Multi-label-line splitter
+# ---------------------------------------------------------------------------
+
+
+def _step_c_split_multi_label_lines(
+    lines: list[str],
+    known_labels: list[str],
+    threshold: float = 0.8,
+    min_count: int = 2,
+    alias_map: dict[str, str] | None = None,
+) -> list[str]:
+    """Split lines containing multiple labels into separate lines.
+
+    Bilingual label lines (e.g. "Họ và tên / Full name:") are NOT split because
+    both labels refer to the same semantic field — splitting them causes downstream
+    extractors to see the English translation as the field value.
+    """
+    if not known_labels:
+        return lines
+
+    _alias_map = alias_map or {}
+
+    def _is_bilingual_pair(labels: list[str]) -> bool:
+        """Return True if every label in *labels* is a bilingual alias of another label."""
+        if len(labels) < 2:
+            return False
+        norm = [_normalize_label(la) for la in labels]
+        for a in norm:
+            for b in norm:
+                if a != b and _alias_map.get(a) and _normalize_label(_alias_map[a]) == b:
+                    return True
+        return False
+
+    result: list[str] = []
+
+    for line in lines:
+        # Find positions of known labels in this line
+        positions: list[tuple[int, int, str]] = []  # (start, end, label)
+        line_lower = line.lower()
+
+        for label in known_labels:
+            label_lower = label.lower()
+            for m in re.finditer(re.escape(label_lower), line_lower):
+                positions.append((m.start(), m.end(), label))
+
+        # Also fuzzy-match labels
+        if len(line) > 10:
+            words = line.split()
+            for i in range(len(words)):
+                for j in range(i + 1, min(i + 5, len(words) + 1)):
+                    fragment = " ".join(words[i:j])
+                    for label in known_labels:
+                        score = fuzz.ratio(fragment.lower(), label.lower()) / 100.0
+                        if score >= threshold and (fragment.lower() != label.lower()):
+                            start = line.lower().find(fragment.lower())
+                            if start >= 0:
+                                positions.append((start, start + len(fragment), label))
+
+        # Remove duplicates and overlaps, sort by position
+        positions.sort(key=lambda x: x[0])
+        filtered: list[tuple[int, int, str]] = []
+        last_end = -1
+        for start, end, label in positions:
+            if start >= last_end:
+                filtered.append((start, end, label))
+                last_end = end
+
+        found_labels = [la for _, _, la in filtered]
+
+        if len(filtered) >= min_count and not _is_bilingual_pair(found_labels):
+            # Split line at each label position
+            parts: list[str] = []
+            for idx, (start, end, label) in enumerate(filtered):
+                # Value = text between this label's end and next label's start
+                next_start = filtered[idx + 1][0] if idx + 1 < len(filtered) else len(line)
+                value = line[end:next_start].strip().lstrip(": ")
+                parts.append(f"{label}: {value}" if value else label)
+            result.extend(parts)
+        else:
+            result.append(line)
+
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Sub-step (d) — Wrap rejoin
+# ---------------------------------------------------------------------------
+
+
+def _step_d_wrap_rejoin(lines: list[str]) -> list[str]:
+    """Rejoin lines that were wrapped mid-sentence by OCR."""
+    result: list[str] = []
+    i = 0
+    while i < len(lines):
+        line = lines[i]
+        stripped = line.strip()
+
+        if not stripped:
+            result.append(line)
+            i += 1
+            continue
+
+        # Line ends with colon: next line is value
+        if stripped.endswith(":") and i + 1 < len(lines):
+            next_line = lines[i + 1].strip()
+            if next_line and not next_line.endswith(":"):
+                result.append(f"{stripped} {next_line}")
+                i += 2
+                continue
+
+        # Continuation: next line starts with lowercase (wrapped)
+        if i + 1 < len(lines) and not stripped.endswith((".", "!", "?")):
+            next_stripped = lines[i + 1].strip()
+            if next_stripped and next_stripped[0].islower():
+                result.append(f"{stripped} {next_stripped}")
+                i += 2
+                continue
+
+        result.append(line)
+        i += 1
+
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Sub-step (e) — EN-paren-rejoin
+# ---------------------------------------------------------------------------
+
+_EN_PAREN_LINE = re.compile(r"^\s*\(([A-Za-z][^)]*)\)\s*$")
+
+
+def _step_e_en_paren_rejoin(lines: list[str]) -> list[str]:
+    """Merge EN-only paren lines as aliases on the label above."""
+    result: list[str] = []
+    for i, line in enumerate(lines):
+        m = _EN_PAREN_LINE.match(line)
+        if m and result:
+            # Append as alias to last line
+            result[-1] = f"{result[-1].rstrip()} / {m.group(1)}"
+        else:
+            result.append(line)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Sub-step (f) — Label-boundary splitter
+# ---------------------------------------------------------------------------
+
+
+def _step_f_label_boundary_split(
+    lines: list[str], known_labels: list[str], max_iters: int = 5
+) -> list[str]:
+    """Split lines where a label is glued directly after a value."""
+    if not known_labels:
+        return lines
+
+    for _ in range(max_iters):
+        changed = False
+        new_lines: list[str] = []
+        for line in lines:
+            split_done = False
+            for label in known_labels:
+                # Pattern: value immediately followed by label (no space separator)
+                escaped = re.escape(label)
+                pattern = re.compile(rf"(\S)({escaped})", re.IGNORECASE)
+                m = pattern.search(line)
+                if m and m.start() > 0:
+                    before = line[: m.start() + 1].strip()
+                    after = line[m.start() + 1 :].strip()
+                    if before and after:
+                        new_lines.append(before)
+                        new_lines.append(after)
+                        split_done = True
+                        changed = True
+                        break
+            if not split_done:
+                new_lines.append(line)
+        lines = new_lines
+        if not changed:
+            break
+
+    return lines
+
+
+# ---------------------------------------------------------------------------
+# Label index builder
+# ---------------------------------------------------------------------------
+
+
+def _build_label_index(
+    sections: list[Section],
+    known_labels: list[str],
+    alias_map: dict[str, str],
+    threshold: float = 0.8,
+) -> dict[str, list[LabelHit]]:
+    """Build label_index: lowercased label → list of LabelHit."""
+    label_index: dict[str, list[LabelHit]] = {}
+
+    for section in sections:
+        for line in section.lines:
+            for label in known_labels:
+                label_lower = label.lower()
+                text_lower = line.text.lower()
+                start = text_lower.find(label_lower)
+                if start < 0:
+                    # Fuzzy search
+                    score = fuzz.partial_ratio(label_lower, text_lower) / 100.0
+                    if score < threshold:
+                        continue
+                    start = 0
+                    end = len(line.text)
+                else:
+                    end = start + len(label)
+                    score = 1.0
+
+                hit = LabelHit(
+                    label=label,
+                    aliases_matched=[label],
+                    section_id=section.id,
+                    line_index=line.index,
+                    char_start=start,
+                    char_end=end,
+                    fuzzy_score=score,
+                )
+                label_index.setdefault(label_lower, []).append(hit)
+
+                # Also index alias
+                if label_lower in alias_map:
+                    alias = alias_map[label_lower]
+                    label_index.setdefault(alias.lower(), []).append(hit)
+
+    return label_index
+
+
+# ---------------------------------------------------------------------------
+# Main reconstruct stage
+# ---------------------------------------------------------------------------
+
+
+def reconstruct_stage(ctx: PipelineContext) -> None:
+    """Pipeline stage 4: LineReconstructor."""
+    profile = ctx.profile
+    rcfg = profile.reconstruct if profile else None
+    section_defs = profile.sections if profile else []
+    bilingual_pairs = rcfg.bilingual_pairs if rcfg else []
+    fuzzy_threshold = rcfg.fuzzy_threshold if rcfg else 0.8
+    enabled = set(rcfg.enabled_steps if rcfg else ["a", "b", "c", "d", "e", "f"])
+
+    text = ctx.normalized_text or ctx.raw_text
+
+    # (a) Section split
+    raw_sections = (
+        _step_a_section_split(text, section_defs)
+        if "a" in enabled
+        else [("_default", None, text.split("\n"))]
+    )
+
+    sections: list[Section] = []
+    all_aliases = _all_aliases(profile)
+    alias_map: dict[str, str] = {}
+
+    for sec_id, sec_title, sec_lines in raw_sections:
+        lines = sec_lines
+
+        # (b) Bilingual map
+        if "b" in enabled:
+            try:
+                alias_map.update(_step_b_build_bilingual_map(lines, bilingual_pairs))
+            except Exception:
+                logger.exception("Reconstructor step-b failed (bilingual map); skipping")
+
+        # Build full label set: profile aliases + bilingual aliases
+        known_labels = list(set(all_aliases + list(alias_map.keys()) + list(alias_map.values())))
+        known_labels = [la for la in known_labels if la.strip()]
+
+        # (c) Multi-label-line split
+        if "c" in enabled:
+            try:
+                min_count = rcfg.multi_label_min_count if rcfg else 2
+                lines = _step_c_split_multi_label_lines(
+                    lines, known_labels, fuzzy_threshold, min_count, alias_map=alias_map
+                )
+            except Exception:
+                logger.exception("Reconstructor step-c failed (label split); skipping")
+
+        # (d) Wrap rejoin
+        if "d" in enabled:
+            try:
+                lines = _step_d_wrap_rejoin(lines)
+            except Exception:
+                logger.exception("Reconstructor step-d failed (wrap rejoin); skipping")
+
+        # (e) EN-paren-rejoin
+        if "e" in enabled:
+            try:
+                lines = _step_e_en_paren_rejoin(lines)
+            except Exception:
+                logger.exception("Reconstructor step-e failed (paren rejoin); skipping")
+
+        # (f) Label-boundary split
+        if "f" in enabled:
+            try:
+                lines = _step_f_label_boundary_split(lines, known_labels)
+            except Exception:
+                logger.exception("Reconstructor step-f failed (boundary split); skipping")
+
+        section = Section(
+            id=sec_id,
+            title=sec_title,
+            lines=[
+                Line(index=i, text=line)
+                for i, line in enumerate(lines)
+                if line.strip()  # skip blank lines
+            ],
+        )
+        sections.append(section)
+
+    ctx.sections = sections
+
+    # Build label_index using the full known_labels (profile aliases + auto-detected bilingual
+    # labels), so stop_labels and extractors can reference any label found in the text.
+    full_known_labels = list(set(all_aliases + list(alias_map.keys()) + list(alias_map.values())))
+    full_known_labels = [la for la in full_known_labels if la.strip()]
+    ctx.label_index = _build_label_index(sections, full_known_labels, alias_map, fuzzy_threshold)
+    total_lines = sum(len(s.lines) for s in sections)
+    logger.debug(
+        "Reconstructor: %d section(s), %d line(s), %d label entries",
+        len(sections),
+        total_lines,
+        len(ctx.label_index),
+    )