foxygit / RPA-Remisser Log in
commits tags

/rpa/ocr.py · 7.3 KB

raw
"""OCR för ocr_read-steget (rpa/flow_engine.py).

Läser text ur en inskannad PDF, ett fönster, ett skärmområde eller en bildfil.
Windows-only: winocr (Windows.Media.Ocr) + pymupdf för PDF-rendering. Importeras
bakom try/except så flow_engine kan importeras överallt – anropas ocr_read utan
paketen kastas OcrUnavailable i stället för ImportError.
"""

from __future__ import annotations

import re
from pathlib import Path

from . import config

try:  # Windows-only
    import winocr
except Exception:  # noqa: BLE001
    winocr = None

try:
    import pymupdf
except Exception:  # noqa: BLE001
    pymupdf = None

try:
    from PIL import Image, ImageGrab
except Exception:  # noqa: BLE001
    Image = None
    ImageGrab = None

DEFAULT_LANG = "sv-SE"

# Om textlagret i en PDF ger minst så här många tecken hoppar vi OCR helt.
_TEXT_LAYER_MIN = 40


class OcrUnavailable(RuntimeError):
    """winocr/pymupdf/PIL saknas eller OCR-språket är inte installerat i Windows."""


class OcrError(RuntimeError):
    """Filen/fönstret/området kunde inte OCR:as."""


def _need_winocr() -> None:
    if winocr is None:
        raise OcrUnavailable(
            "winocr kunde inte laddas – OCR fungerar bara på Windows. "
            "Installera med: pip install -r requirements.txt"
        )


def _need_pil() -> None:
    if Image is None:
        raise OcrUnavailable("Pillow kunde inte laddas – pip install -r requirements.txt")


# --------------------------------------------------------------------------
# Låg nivå: en PIL-bild -> text + ordboxar
# --------------------------------------------------------------------------

def image_to_text(pil_image, lang: str = DEFAULT_LANG) -> tuple[str, list[dict]]:
    """(text, words). words = [{text, x, y, w, h}] i bildpixlar."""
    _need_winocr()
    try:
        result = winocr.recognize_pil_sync(pil_image, lang)
    except AssertionError as exc:
        # winocr assertar med exakt DISM-kommandot när språket saknas.
        raise OcrUnavailable(
            f"OCR-språket {lang!r} är inte installerat i Windows. Installera det "
            "(admin-terminal):\n    " + str(exc) + "\neller via Inställningar → Tid "
            "och språk → Språk → lägg till språket → Språkalternativ → 'Optisk "
            "teckenläsning'."
        ) from exc
    except Exception as exc:  # noqa: BLE001
        raise OcrError(f"OCR misslyckades: {exc}") from exc

    lines = result.get("lines", []) or []
    words: list[dict] = []
    for line in lines:
        for w in line.get("words", []) or []:
            r = w.get("bounding_rect") or {}
            words.append({
                "text": w.get("text", ""),
                "x": r.get("x", 0), "y": r.get("y", 0),
                "w": r.get("width", 0), "h": r.get("height", 0),
            })
    # Bygg texten radvis (winocrs "text" slår ihop rader med mellanslag – sämre
    # för radbaserad extraktion i value_near).
    text = "\n".join(l.get("text", "") for l in lines) or (result.get("text") or "")
    return text, words


# --------------------------------------------------------------------------
# Källor
# --------------------------------------------------------------------------

def pdf_to_text(path: str | Path, lang: str = DEFAULT_LANG) -> str:
    """pypdf-textlager först; är det tomt renderas sidorna och OCR:as."""
    p = Path(path)
    if not p.exists():
        raise OcrError(f"PDF:en finns inte: {p}")

    layer = _pdf_text_layer(p)
    if len(layer.strip()) >= _TEXT_LAYER_MIN:
        return layer

    if pymupdf is None:
        raise OcrUnavailable(
            "pymupdf kunde inte laddas – behövs för att OCR:a inskannade PDF:er. "
            "pip install -r requirements.txt"
        )
    _need_pil()
    parts: list[str] = []
    doc = pymupdf.open(str(p))
    try:
        for page in doc:
            pix = page.get_pixmap(dpi=300)
            mode = "RGBA" if pix.alpha else "RGB"
            img = Image.frombytes(mode, (pix.width, pix.height), pix.samples)
            parts.append(image_to_text(img.convert("RGB"), lang)[0])
    finally:
        doc.close()
    return "\n\n".join(parts).strip()


def _pdf_text_layer(p: Path) -> str:
    try:
        from pypdf import PdfReader

        reader = PdfReader(str(p))
        return "\n".join((page.extract_text() or "") for page in reader.pages)
    except Exception:  # noqa: BLE001
        return ""


def window_to_text(session, title: str | None, lang: str = DEFAULT_LANG) -> tuple[str, list[dict]]:
    """session = rpa.desktop.DesktopSession. Fångar fönstret och OCR:ar."""
    _need_pil()
    win = session._window(title or None)  # noqa: SLF001 – avsiktlig återanvändning
    img = win.capture_as_image()
    return image_to_text(img, lang)


def region_to_text(bbox: tuple[int, int, int, int], lang: str = DEFAULT_LANG) -> tuple[str, list[dict]]:
    """bbox = (x, y, w, h) i skärmkoordinater."""
    _need_pil()
    if ImageGrab is None:
        raise OcrUnavailable("PIL.ImageGrab saknas.")
    x, y, w, h = bbox
    img = ImageGrab.grab(bbox=(x, y, x + w, y + h), all_screens=True)
    return image_to_text(img, lang)


def image_file_to_text(path: str | Path, lang: str = DEFAULT_LANG) -> tuple[str, list[dict]]:
    _need_pil()
    p = Path(path)
    if not p.is_absolute():
        p = config.DESKTOP_IMAGE_DIR / p
    if not p.exists():
        raise OcrError(f"Bildfilen finns inte: {p}")
    return image_to_text(Image.open(str(p)), lang)


# --------------------------------------------------------------------------
# Extraktion ur OCR-text
# --------------------------------------------------------------------------

def apply_pattern(text: str, pattern: str) -> str:
    """re.search mot texten. Har mönstret en namngiven grupp `(?P<value>…)` returneras
    den, annars hela träffen. Ingen träff -> ''."""
    m = re.search(pattern, text, re.IGNORECASE | re.MULTILINE)
    if not m:
        return ""
    if "value" in (m.groupdict() or {}):
        return m.group("value") or ""
    return m.group(0)


def value_near(text: str, words: list[dict], label: str, mode: str = "same_line") -> str:
    """Värdet bredvid/under en etikett.

    same_line: resten av den rad som innehåller `label` (efter label + ev. ':').
    below: närmaste rad UNDER etiketten, ungefär i samma kolumn (kräver ordboxar).
    """
    lines = [ln for ln in text.splitlines()]
    if mode == "same_line":
        for ln in lines:
            if label.lower() in ln.lower():
                idx = ln.lower().index(label.lower()) + len(label)
                rest = ln[idx:].strip(" :\t-")
                if rest:
                    return rest
        return ""

    # below: hitta etikettordens y/x, ta närmaste ordrad under
    lbl_words = [w for w in words if label.lower() in w["text"].lower()]
    if not lbl_words and " " in label:
        first = label.split()[0].lower()
        lbl_words = [w for w in words if w["text"].lower() == first]
    if not lbl_words:
        return ""
    ly = min(w["y"] for w in lbl_words)
    lx = min(w["x"] for w in lbl_words)
    lh = max(w["h"] for w in lbl_words)
    below = [w for w in words if w["y"] > ly + lh - 5 and abs(w["x"] - lx) < 300]
    if not below:
        return ""
    row_y = min(w["y"] for w in below)
    row = [w for w in below if abs(w["y"] - row_y) < 15]
    row.sort(key=lambda w: w["x"])
    return " ".join(w["text"] for w in row).strip()