Spaces:
Sleeping
Sleeping
Aaron
Sync backend with GitHub main (f6aa117): OCR word-coherence filter, full_text consistency fix
e87916b | """ | |
| 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, | |
| ) | |
| 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) | |