processors/ocr_processor.py

"""
OCR Processor
=============
Extracts and classifies text from enhanced architectural drawings using
Tesseract 5 (LSTM engine via pytesseract).

Words are grouped into phrases (see `_group_phrases`) before classification
so that multi-word labels such as "Living Room" or "3,640 x 3,640" are
treated as single tokens rather than being split apart and losing meaning.

Classification buckets:
  - dimensions  : numeric strings WITH a recognised unit suffix (mm/cm/m/ft/in/'/")
                  or a "WIDTH x HEIGHT" pair (with or without a unit suffix)
  - room_names  : multi-character capitalised phrases (title-block keywords excluded)
  - numbers     : bare numeric tokens (no unit suffix)
  - labels      : everything else with sufficient confidence
"""

import os
import re
import shutil
import cv2
import numpy as np
import logging
from typing import Any, Dict, List

from .ocr_classification import _CJK_RE, classify, clean_elements, empty_result

try:
    import pytesseract
    _TESSERACT_AVAILABLE = True
    # On Windows the UB-Mannheim installer doesn't add tesseract to PATH —
    # fall back to its default install location so pytesseract can spawn it.
    if os.name == "nt" and shutil.which("tesseract") is None:
        _win_default = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
        if os.path.exists(_win_default):
            pytesseract.pytesseract.tesseract_cmd = _win_default
except ImportError:
    _TESSERACT_AVAILABLE = False

logger = logging.getLogger(__name__)

# Tesseract page-segmentation-mode:
#   3 = fully automatic page segmentation (default) — floor plans scatter
#       room labels and dimension callouts across a 2D grid rather than one
#       paragraph, so psm 6 ("assume a uniform block of text") was forcing
#       a single reading order across unrelated regions: it merged distant
#       labels onto one "line" and regularly dropped isolated words (e.g.
#       a lone "Bedroom" label with nothing beside it) that automatic
#       segmentation finds correctly.
# OEM 3 = LSTM + legacy combined (most accurate)
# lang: jpn+eng handles Japanese architectural drawings; Tesseract falls back
# gracefully to eng-only if the jpn data file is not installed.
_TESS_CONFIG      = "--psm 3 --oem 3"
# Supplementary pass: psm 11 ("sparse text, no particular order") recovers
# isolated labels that automatic segmentation sometimes drops entirely —
# e.g. a single room name with nothing else nearby. Only phrases that don't
# overlap the primary pass are kept (see the merge in extract_text).
_TESS_CONFIG_SPARSE = "--psm 11 --oem 3"
_TESS_LANG      = "jpn+eng"

# Minimum confidence (0-100) to include a token
_MIN_CONFIDENCE = 30

# Adaptive resolution boost — see _estimate_upscale_factor().
#
# A blind pixel-count floor doesn't work here: the project's own clean,
# computer-rendered benchmark images are already well above any reasonable
# floor (e.g. 1400x1000), so a floor keyed to *overall image size* never
# fires for them — but a naive fix that instead forces a large fixed
# upscale on *every* image regressed the clean-sample benchmark hard
# (OCR Token F1 96.3% -> 68.5%): interpolation blur/re-splitting hurts
# text that was already legible. The real fix has to be conditional on
# whether the text is *actually* too small, not on the image's pixel count.
#
# Measured directly: at native resolution, Tesseract found
#   - floor_plan_level1.jpg   (clean render):  34 word boxes, median height 21px
#   - synthetic_floor_plan.png(clean render):   5 word boxes, median height 36px
#   - real_faded_1.jpg   (real 1970s drawing):  4 word boxes, median height  9-11px
# Tesseract's LSTM models are trained on ~25-30px glyph height — the clean
# samples already sit in that range and must be left alone; the real
# drawing is running at roughly a third of it, which is why OCR was
# effectively finding almost nothing (not primarily a "handwriting is
# hard" problem — a resolution problem). Upscaling the real drawing ~3.2x
# brought its median glyph height to ~32px and raw box count from 4 to 17+,
# recovering real dimension digits and label fragments that were
# completely invisible before.
_OCR_TARGET_GLYPH_HEIGHT_PX = 28
_OCR_MIN_HEALTHY_HEIGHT_PX  = 18   # skip upscaling if already at least this tall
_OCR_MAX_UPSCALE            = 4.0  # safety cap against a single stray tiny box

# Phrases shorter than this (pixels, post _prepare()'s upscale-to-≥800px-side
# floor) are almost never real text — they're dimension tick-marks, arrowheads,
# and extension-line stubs that Tesseract occasionally misreads as a stray
# digit or punctuation mark. Real content in practice starts around 17px at
# this preprocessing scale (the smallest dimension-callout digits); the
# smallest observed misread artifacts sit at 10-11px, so 14px cleanly
# separates the two with margin on both sides. Applied to the *merged
# phrase's* height (not each word's) so a short lowercase separator like the
# "x" in "3,640 x 3,640" — legitimately ~11px tall on its own — doesn't get
# dropped: it survives because the phrase's bounding box is dominated by its
# taller digit neighbours.
_MIN_PHRASE_HEIGHT_PX = 14

# Words on the same Tesseract line are only merged into one phrase if the
# horizontal gap between them is below this many pixels — otherwise a
# single wide title-block row (e.g. "SCALE 1:100    FLOOR PLAN    6,400mm")
# would collapse into one meaningless blob. Scaled by the line's own text
# height so it works across resolutions.
_MAX_WORD_GAP_RATIO = 2.5

# Tesseract occasionally misreads the comma in a thousands-grouped mm value
# as a period (e.g. "6,400mm" -> "6.400mm"). Only correct this for "mm"
# specifically: a real mm measurement is never expressed to 3 decimal
# places (that would be sub-micrometre precision), so "\d.\d\d\dmm" is
# always a misread comma, not a genuine decimal — unlike metres, where
# 3-decimal precision (e.g. "12.345m") is a completely normal way to state
# a millimetre-accurate length. Scoped to mm only so real decimals survive.
_FIX_MM_THOUSANDS_RE = re.compile(r"(\d)\.(\d{3})(?=\s*mm\b)", re.IGNORECASE)

# Tokens made up entirely of dash/line/box-drawing glyphs. Tesseract
# regularly misreads isolated line strokes (door-swing arcs, leader lines,
# dashed hidden-line marks) as these characters — most commonly the
# katakana chōon mark (ー) when jpn+eng is active. They are dropped
# unconditionally: real text never consists solely of these glyphs, and
# they otherwise survive confidence filtering with deceptively high scores.
_NOISE_RE = re.compile(r"^[\-_|.·‐-―─-╿ー]+$")

# A phrase mixing CJK script with Latin letters/digits (e.g. a hatching
# pattern or leader-line fragment misread as "5R5R5t5や5やつべ") never
# occurs in real architectural drawing text — legitimate content is either
# pure ASCII (dimensions, room codes) or pure Japanese. Combined jpn+eng
# recognition hallucinates this kind of mixed gibberish from fine line
# art (hatching, tick marks, arc endpoints), so it is dropped outright.
# Shared definition lives in ocr_classification.py (also used by the
# word-coherence filter); re-imported under this name for compatibility.
_LATIN_DIGIT_RE  = re.compile(r"[A-Za-z0-9]")


class OCRProcessor:
    """Extracts text elements and their positions from architectural drawings."""

    def extract_text(self, image: np.ndarray) -> Dict[str, Any]:
        """
        Parameters
        ----------
        image : enhanced binary or grayscale image

        Returns
        -------
        dict with keys:
          full_text    – cleaned phrases joined by newline (see module
                         docstring and ocr_classification.clean_elements —
                         *not* Tesseract's raw image_to_string dump, which
                         skips every quality filter elements/classified
                         already go through and reads as gibberish on
                         degraded scans)
          element_count
          elements     – list of {text, x, y, width, height, confidence}
                         (phrase-grouped — see module docstring)
          classified   – dict of {dimensions, room_names, numbers, labels}
        """
        if not _TESSERACT_AVAILABLE:
            logger.warning("pytesseract not installed – OCR skipped")
            return empty_result("pytesseract not installed")

        prepared = self._prepare(image)

        try:
            primary_phrases = self._group_phrases(
                self._parse_words(self._run_tesseract_data(prepared, _TESS_CONFIG))
            )
            sparse_phrases = self._group_phrases(
                self._parse_words(self._run_tesseract_data(prepared, _TESS_CONFIG_SPARSE))
            )
        except Exception as exc:
            logger.exception("Tesseract error")
            return empty_result(str(exc))

        # Keep every primary-pass phrase, plus any sparse-pass phrase that
        # doesn't spatially overlap something the primary pass already found
        # — this recovers isolated labels automatic segmentation drops
        # (see _TESS_CONFIG_SPARSE) without duplicating anything.
        raw_elements = primary_phrases + [
            p for p in sparse_phrases if not self._overlaps_any(p, primary_phrases)
        ]
        elements = clean_elements(raw_elements)
        classified = classify(elements)

        return {
            "full_text":     "\n".join(e["text"] for e in elements),
            "element_count": len(elements),
            "elements":      elements,
            "classified":    classified,
        }

    # ------------------------------------------------------------------ #
    # Private helpers                                                       #
    # ------------------------------------------------------------------ #

    def _run_tesseract_data(self, prepared: np.ndarray, config: str) -> Dict[str, List]:
        return pytesseract.image_to_data(
            prepared, config=config, lang=_TESS_LANG, output_type=pytesseract.Output.DICT,
        )

    @staticmethod
    def _overlaps_any(phrase: Dict[str, Any], others: List[Dict[str, Any]]) -> bool:
        """True if `phrase`'s box intersects any box in `others` at all."""
        px0, py0 = phrase["x"], phrase["y"]
        px1, py1 = px0 + phrase["width"], py0 + phrase["height"]
        for other in others:
            ox0, oy0 = other["x"], other["y"]
            ox1, oy1 = ox0 + other["width"], oy0 + other["height"]
            if px0 < ox1 and px1 > ox0 and py0 < oy1 and py1 > oy0:
                return True
        return False

    def _prepare(self, image: np.ndarray) -> np.ndarray:
        """
        Upscale genuinely-too-small text and ensure white-background /
        black-text convention that Tesseract expects.
        """
        # Grayscale
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()

        # Absolute floor for pathologically tiny images (unchanged from
        # before) — resolution-adaptive upscaling below still needs enough
        # pixels for its own probe pass to find anything at all.
        h, w = gray.shape
        if min(h, w) < 800:
            scale = 800 / min(h, w)
            gray = cv2.resize(
                gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC
            )

        # Tesseract expects dark-text-on-light-background
        if np.mean(gray) < 127:
            gray = cv2.bitwise_not(gray)

        scale = self._estimate_upscale_factor(gray)
        if scale > 1.0:
            gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)

        return gray

    def _estimate_upscale_factor(self, gray: np.ndarray) -> float:
        """
        Probe Tesseract's own detected glyph heights at current resolution
        to decide whether upscaling is actually needed — see the
        module-level comment above _OCR_TARGET_GLYPH_HEIGHT_PX for why a
        blind pixel-count floor doesn't work (it regressed the clean-sample
        benchmark) and how these specific thresholds were measured.
        """
        try:
            data = pytesseract.image_to_data(
                gray, config=_TESS_CONFIG, lang=_TESS_LANG, output_type=pytesseract.Output.DICT,
            )
        except Exception:
            return 1.0

        heights = [
            data["height"][i] for i in range(len(data["level"]))
            if (data["text"][i] or "").strip() and int(data["conf"][i]) > 0
        ]
        if heights and np.median(heights) >= _OCR_MIN_HEALTHY_HEIGHT_PX:
            return 1.0  # already legible at native resolution — leave it alone

        sample_height = float(np.median(heights)) if heights else 10.0
        scale = _OCR_TARGET_GLYPH_HEIGHT_PX / max(sample_height, 1.0)
        return float(min(scale, _OCR_MAX_UPSCALE))

    def _parse_words(self, data: Dict[str, List]) -> List[Dict[str, Any]]:
        """Filter and normalise Tesseract word-level output."""
        words: List[Dict[str, Any]] = []
        n = len(data["level"])

        for i in range(n):
            text = (data["text"][i] or "").strip()
            conf = int(data["conf"][i])

            if conf < _MIN_CONFIDENCE or not text:
                continue
            if _NOISE_RE.match(text):
                continue

            words.append(
                {
                    "text":       text,
                    "confidence": conf,
                    "x":          int(data["left"][i]),
                    "y":          int(data["top"][i]),
                    "width":      int(data["width"][i]),
                    "height":     int(data["height"][i]),
                    "block_num":  data["block_num"][i],
                    "par_num":    data["par_num"][i],
                    "line_num":   data["line_num"][i],
                    "word_num":   data["word_num"][i],
                }
            )

        return words

    def _group_phrases(self, words: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Merge words on the same Tesseract line into phrases, splitting on
        large horizontal gaps so unrelated title-block cells that happen to
        share a line don't collapse into one token (see _MAX_WORD_GAP_RATIO).
        """
        phrases: List[Dict[str, Any]] = []
        current: List[Dict[str, Any]] = []

        def flush() -> None:
            if not current:
                return
            text = " ".join(w["text"] for w in current)
            text = _FIX_MM_THOUSANDS_RE.sub(r"\1,\2", text)
            # Stray quote/backtick marks are a common misread of a tick-mark
            # or witness-line endpoint fused onto the adjacent character by
            # Tesseract (e.g. '"SCALE'); real drawing text never legitimately
            # starts or ends with one, so it's safe to trim unconditionally.
            text = text.strip("\"'`")
            if not text:
                return
            if _CJK_RE.search(text) and _LATIN_DIGIT_RE.search(text):
                return  # mixed CJK/Latin gibberish — see _CJK_RE comment
            xs0 = min(w["x"] for w in current)
            ys0 = min(w["y"] for w in current)
            xs1 = max(w["x"] + w["width"] for w in current)
            ys1 = max(w["y"] + w["height"] for w in current)
            if ys1 - ys0 < _MIN_PHRASE_HEIGHT_PX:
                return  # tick-mark/arrowhead misread as a character — see _MIN_PHRASE_HEIGHT_PX
            phrases.append({
                "text":       text,
                "confidence": round(sum(w["confidence"] for w in current) / len(current)),
                "x":          xs0,
                "y":          ys0,
                "width":      xs1 - xs0,
                "height":     ys1 - ys0,
            })

        prev = None
        for word in words:
            same_line = (
                prev is not None
                and word["block_num"] == prev["block_num"]
                and word["par_num"]   == prev["par_num"]
                and word["line_num"]  == prev["line_num"]
            )
            gap_ok = True
            if same_line:
                gap = word["x"] - (prev["x"] + prev["width"])
                max_gap = max(30, prev["height"] * _MAX_WORD_GAP_RATIO)
                gap_ok = gap <= max_gap

            if same_line and gap_ok:
                current.append(word)
            else:
                flush()
                current = [word]
            prev = word

        flush()
        return phrases

    def _classify(self, elements: List[Dict[str, Any]]) -> Dict[str, List]:
        """Kept as an instance method for backward compatibility with
        existing callers/tests; delegates to the provider-shared rules in
        ocr_classification.py so Tesseract and PaddleOCR bucket text
        identically."""
        return classify(elements)