Spaces:

lylron
/

wonq-backend

Sleeping

wonq-backend / processors /paddleocr_processor.py

Aaron

Sync backend with GitHub main (f6aa117): OCR word-coherence filter, full_text consistency fix

e87916b 4 months ago

6.3 kB

"""
PaddleOCR Processor
===================
Alternative to Tesseract (see ocr_processor.py) backed by PaddleOCR
(Apache 2.0, https://github.com/PaddlePaddle/PaddleOCR).
Why this is a separate, simpler class rather than another Tesseract config:
PaddleOCR's detector locates text at the phrase/line level directly, so none
of ocr_processor.py's word->phrase grouping, gap-splitting, or tick-mark
height filtering is needed here — those all exist specifically to work
around Tesseract returning individual words with no semantic grouping.
Classification (dimensions/room_names/numbers/labels) is shared via
ocr_classification.py so both engines bucket text identically.
Status: NOT YET BENCHMARKED. Model weights could not be downloaded in the
sandboxed dev environment this was written in — see CLAUDE.md for why
(corporate TLS-inspecting proxy whose root CA isn't in Python's default
trust store). The API calls below are verified against the installed
paddleocr==3.7.0 / paddlex package source directly (constructor params,
and the exact result schema: `rec_texts`/`rec_scores`/`rec_boxes`, with
`rec_boxes[i] == [left, top, right, bottom]` — see
venv/lib/python3.12/site-packages/paddlex/inference/pipelines/ocr/pipeline.py
and .../components/common/convert_points_and_boxes.py), but real accuracy
on this project's sample images has not been measured. Run
`benchmark/run_benchmark.py` before relying on this over Tesseract.
"""

from __future__ import annotations
import logging
import os
from typing import Any, Dict, List, Optional

import numpy as np

from .ocr_classification import classify, clean_elements, empty_result

logger = logging.getLogger(__name__)

# PP-OCRv6's detection/recognition models crash under paddlex's default CPU
# MKLDNN/PIR execution path on this Windows/CPU combo:
#   NotImplementedError: (Unimplemented) ConvertPirAttribute2RuntimeAttribute
#   not support [pir::ArrayAttribute<pir::DoubleAttribute>]
# Neither model is in paddlex's MKLDNN_BLOCKLIST (an upstream gap), so the
# workaround is disabling MKLDNN-by-default before paddlex's flags module
# is imported. Must happen before `import paddleocr`, not after.
os.environ.setdefault("PADDLE_PDX_ENABLE_MKLDNN_BYDEFAULT", "False")

try:
    from paddleocr import PaddleOCR as _PaddleOCREngine
    _PADDLEOCR_AVAILABLE = True
    _PADDLEOCR_ERROR: Optional[str] = None
except ImportError as _e:
    _PADDLEOCR_AVAILABLE = False
    _PADDLEOCR_ERROR = f"paddleocr / paddlepaddle not installed: {_e}"


class PaddleOCRProcessor:
    """
    Drop-in alternative to OCRProcessor — same extract_text() contract
    (full_text / element_count / elements / classified), so main.py can
    swap between engines without any other code changing.
    """

    def __init__(self, lang: str = "en") -> None:
        # PaddleOCR loads one recognition model per language (unlike
        # Tesseract's combined "jpn+eng"); "en" covers this project's
        # sample images, pass lang="japan" for Japanese-only drawings.
        self._lang = lang
        self._engine = None  # lazy-loaded on first use — downloads weights

    @property
    def is_available(self) -> bool:
        return _PADDLEOCR_AVAILABLE

    @property
    def unavailable_reason(self) -> Optional[str]:
        return _PADDLEOCR_ERROR

    def extract_text(self, image: np.ndarray) -> Dict[str, Any]:
        """
        Parameters
        ----------
        image : enhanced binary or grayscale (or BGR) image
        Returns
        -------
        Same schema as OCRProcessor.extract_text() — see that module's
        docstring.
        """
        if not _PADDLEOCR_AVAILABLE:
            logger.warning("paddleocr not installed – OCR skipped")
            return empty_result("paddleocr not installed")

        try:
            engine = self._get_engine()
            results = engine.predict(image)
        except Exception as exc:
            logger.exception("PaddleOCR error")
            return empty_result(str(exc))

        if not results:
            return empty_result("no text regions detected")

        result = results[0]
        raw_elements = self._to_elements(
            result.get("rec_texts", []),
            result.get("rec_scores", []),
            result.get("rec_boxes", []),
        )
        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 _get_engine(self):
        if self._engine is not None:
            return self._engine
        logger.info("Loading PaddleOCR engine (lang=%s) — downloads weights on first use", self._lang)
        self._engine = _PaddleOCREngine(
            lang=self._lang,
            # Our own enhancement pipeline already deskews/upscales the
            # image, so skip PaddleOCR's own document-orientation
            # classification and unwarping sub-pipelines — one less model
            # to download and one less inference pass per image.
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
        )
        return self._engine

    @staticmethod
    def _to_elements(
        texts: List[str], scores: List[float], boxes: Any
    ) -> List[Dict[str, Any]]:
        elements: List[Dict[str, Any]] = []
        for text, score, box in zip(texts, scores, boxes):
            text = str(text).strip()
            if not text:
                continue
            x1, y1, x2, y2 = (float(v) for v in box)
            elements.append({
                "text":       text,
                "confidence": round(float(score) * 100),  # match Tesseract's 0-100 scale
                "x":          int(round(x1)),
                "y":          int(round(y1)),
                "width":      int(round(x2 - x1)),
                "height":     int(round(y2 - y1)),
            })
        return elements

	"""
	PaddleOCR Processor
	===================
	Alternative to Tesseract (see ocr_processor.py) backed by PaddleOCR
	(Apache 2.0, https://github.com/PaddlePaddle/PaddleOCR).

	Why this is a separate, simpler class rather than another Tesseract config:
	PaddleOCR's detector locates text at the phrase/line level directly, so none
	of ocr_processor.py's word->phrase grouping, gap-splitting, or tick-mark
	height filtering is needed here — those all exist specifically to work
	around Tesseract returning individual words with no semantic grouping.
	Classification (dimensions/room_names/numbers/labels) is shared via
	ocr_classification.py so both engines bucket text identically.

	Status: NOT YET BENCHMARKED. Model weights could not be downloaded in the
	sandboxed dev environment this was written in — see CLAUDE.md for why
	(corporate TLS-inspecting proxy whose root CA isn't in Python's default
	trust store). The API calls below are verified against the installed
	paddleocr==3.7.0 / paddlex package source directly (constructor params,
	and the exact result schema: `rec_texts`/`rec_scores`/`rec_boxes`, with
	`rec_boxes[i] == [left, top, right, bottom]` — see
	venv/lib/python3.12/site-packages/paddlex/inference/pipelines/ocr/pipeline.py
	and .../components/common/convert_points_and_boxes.py), but real accuracy
	on this project's sample images has not been measured. Run
	`benchmark/run_benchmark.py` before relying on this over Tesseract.
	"""

	from __future__ import annotations
	import logging
	import os
	from typing import Any, Dict, List, Optional

	import numpy as np

	from .ocr_classification import classify, clean_elements, empty_result

	logger = logging.getLogger(__name__)

	# PP-OCRv6's detection/recognition models crash under paddlex's default CPU
	# MKLDNN/PIR execution path on this Windows/CPU combo:
	# NotImplementedError: (Unimplemented) ConvertPirAttribute2RuntimeAttribute
	# not support [pir::ArrayAttribute<pir::DoubleAttribute>]
	# Neither model is in paddlex's MKLDNN_BLOCKLIST (an upstream gap), so the
	# workaround is disabling MKLDNN-by-default before paddlex's flags module
	# is imported. Must happen before `import paddleocr`, not after.
	os.environ.setdefault("PADDLE_PDX_ENABLE_MKLDNN_BYDEFAULT", "False")

	try:
	from paddleocr import PaddleOCR as _PaddleOCREngine
	_PADDLEOCR_AVAILABLE = True
	_PADDLEOCR_ERROR: Optional[str] = None
	except ImportError as _e:
	_PADDLEOCR_AVAILABLE = False
	_PADDLEOCR_ERROR = f"paddleocr / paddlepaddle not installed: {_e}"


	class PaddleOCRProcessor:
	"""
	Drop-in alternative to OCRProcessor — same extract_text() contract
	(full_text / element_count / elements / classified), so main.py can
	swap between engines without any other code changing.
	"""

	def __init__(self, lang: str = "en") -> None:
	# PaddleOCR loads one recognition model per language (unlike
	# Tesseract's combined "jpn+eng"); "en" covers this project's
	# sample images, pass lang="japan" for Japanese-only drawings.
	self._lang = lang
	self._engine = None # lazy-loaded on first use — downloads weights

	@property
	def is_available(self) -> bool:
	return _PADDLEOCR_AVAILABLE

	@property
	def unavailable_reason(self) -> Optional[str]:
	return _PADDLEOCR_ERROR

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

	Returns
	-------
	Same schema as OCRProcessor.extract_text() — see that module's
	docstring.
	"""
	if not _PADDLEOCR_AVAILABLE:
	logger.warning("paddleocr not installed – OCR skipped")
	return empty_result("paddleocr not installed")

	try:
	engine = self._get_engine()
	results = engine.predict(image)
	except Exception as exc:
	logger.exception("PaddleOCR error")
	return empty_result(str(exc))

	if not results:
	return empty_result("no text regions detected")

	result = results[0]
	raw_elements = self._to_elements(
	result.get("rec_texts", []),
	result.get("rec_scores", []),
	result.get("rec_boxes", []),
	)
	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 _get_engine(self):
	if self._engine is not None:
	return self._engine
	logger.info("Loading PaddleOCR engine (lang=%s) — downloads weights on first use", self._lang)
	self._engine = _PaddleOCREngine(
	lang=self._lang,
	# Our own enhancement pipeline already deskews/upscales the
	# image, so skip PaddleOCR's own document-orientation
	# classification and unwarping sub-pipelines — one less model
	# to download and one less inference pass per image.
	use_doc_orientation_classify=False,
	use_doc_unwarping=False,
	use_textline_orientation=False,
	)
	return self._engine

	@staticmethod
	def _to_elements(
	texts: List[str], scores: List[float], boxes: Any
	) -> List[Dict[str, Any]]:
	elements: List[Dict[str, Any]] = []
	for text, score, box in zip(texts, scores, boxes):
	text = str(text).strip()
	if not text:
	continue
	x1, y1, x2, y2 = (float(v) for v in box)
	elements.append({
	"text": text,
	"confidence": round(float(score) * 100), # match Tesseract's 0-100 scale
	"x": int(round(x1)),
	"y": int(round(y1)),
	"width": int(round(x2 - x1)),
	"height": int(round(y2 - y1)),
	})
	return elements