Blame for ocr_core.py

"""Lecture OCR et décompte de la 2e colonne d'un tableau photographié."""
import os
import re
import statistics
from difflib import get_close_matches

import cv2
import numpy as np
import pytesseract
from PIL import Image, ImageOps

COLS = ["PASF", "IPM", "PR", "MUT", "RESULTAT", "CONTROLE"]
ALIASES = {"ASS": "IPM",  "RESULTATS": "RESULTAT"}
KNOWN = COLS + list(ALIASES)
CONFUSIONS = str.maketrans({"J": "I", "1": "I", "L": "I", "|": "I", "!": "I", "0": "O", "5": "S", "$": "S"})


def load_bgr(source):
    """source : chemin de fichier ou objet fichier."""
    img = ImageOps.exif_transpose(Image.open(source)).convert("RGB")
    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)


def preprocess(bgr, size=2200):
    """Redresse la photo, neutralise les ombres, efface les traits de la grille."""
    h, w = bgr.shape[:2]

    s = size / max(h, w)
    bgr = cv2.resize(
        bgr, None, fx=s, fy=s,
        interpolation=cv2.INTER_CUBIC
    )

    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)

    bw = cv2.adaptiveThreshold(
        gray, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        41, 15
    )

    angle = 0.0

    segs = cv2.HoughLinesP(
        bw, 1, np.pi / 720,
        threshold=200,
        minLineLength=bw.shape[1] // 4,
        maxLineGap=20
    )

    if segs is not None:
        angs = []

        for segment in segs.reshape(-1, 4):
            x1, y1, x2, y2 = map(int, segment)

            a = np.degrees(
                np.arctan2(y2 - y1, x2 - x1)
            )

            if abs(a) < 20:
                angs.append(float(a))

        if angs:
            angle = float(np.median(angs))

    M = cv2.getRotationMatrix2D(
        (bw.shape[1] / 2, bw.shape[0] / 2),
        angle, 1
    )

    bw = cv2.warpAffine(
        bw, M,
        (bw.shape[1], bw.shape[0]),
        flags=cv2.INTER_LINEAR
    )

    hk = cv2.getStructuringElement(
        cv2.MORPH_RECT, (80, 1)
    )

    vk = cv2.getStructuringElement(
        cv2.MORPH_RECT, (1, 80)
    )

    grid = cv2.add(
        cv2.morphologyEx(bw, cv2.MORPH_OPEN, hk),
        cv2.morphologyEx(bw, cv2.MORPH_OPEN, vk)
    )

    bw = cv2.subtract(
        bw,
        cv2.dilate(grid, np.ones((3, 3), np.uint8))
    )

    return 255 - bw, angle


def read_words(bw, lang="eng", psm=11):
    d = pytesseract.image_to_data(bw, lang=lang, config=f"--psm {int(psm)}",
                                  output_type=pytesseract.Output.DICT)
    words = []
    for i, t in enumerate(d["text"]):
        if t.strip():
            words.append({"text": t.strip(), "left": d["left"][i], "top": d["top"][i],
                          "w": d["width"][i], "h": d["height"][i]})
    return words


def classify(token, fuzzy=False):
    """Renvoie le code canonique (PASF, IPM, PR, MUT, RESULTAT, CONTROLE) ou None."""
    u = re.sub(r"[^A-Z0-9]", "", token.upper())
    if len(u) < 2:
        return None
    cands = [u]
    if len(u) <= 4:                      # 1PM, |PM, lPM -> IPM
        cands.append(u.translate(CONFUSIONS))
    for c in cands:
        if c in KNOWN:
            return ALIASES.get(c, c)
    if fuzzy:
        for c in cands:
            m = get_close_matches(c, KNOWN, n=1, cutoff=0.75)
            if m:
                return ALIASES.get(m[0], m[0])
    return None


def count_column2(words, width, fuzzy=True):
    """Ne compte que les mots alignés sur la colonne des codes (la 2e colonne)."""
    anchors = [w for w in words if classify(w["text"])]
    counts = {c: 0 for c in COLS}
    notes = []
    if not anchors:
        return counts, ["Aucun code reconnu : image peu lisible ?"]
    center = statistics.median(w["left"] for w in anchors)
    cmid = statistics.median(w["left"] + w["w"] / 2 for w in anchors)
    tol = 0.04 * width

    def in_band(w):  # bord gauche OU centre (cellules centrées dans certains PDF)
        return abs(w["left"] - center) <= tol or abs(w["left"] + w["w"] / 2 - cmid) <= tol

    band = [w for w in words if in_band(w)]
    tops = sorted(w["top"] for w in anchors)
    pitch = statistics.median([y2 - y1 for y1, y2 in zip(tops, tops[1:]) if y2 - y1 > 2] or [40])
    first_y = min(w["top"] for w in anchors)
    unknown = []
    for w in band:
        code = classify(w["text"], fuzzy=fuzzy)
        if code:
            counts[code] += 1
        elif w["top"] >= first_y - 0.9 * pitch and re.fullmatch(r"[A-Za-z]{2,}", w["text"]):
            unknown.append(w["text"])
    outside = len(anchors) - sum(1 for w in anchors if in_band(w))
    if outside:
        notes.append(f"{outside} code(s) hors 2e colonne ignoré(s)")
    if unknown:
        notes.append("Non reconnu en 2e colonne : " + ", ".join(unknown[:6]))
    return counts, notes


def analyse(source, lang="eng", psm=11, fuzzy=True):
    bw, _ = preprocess(load_bgr(source))
    words = read_words(bw, lang, psm)
    return count_column2(words, bw.shape[1], fuzzy)


def _open_pdf(source):
    import pymupdf
    if isinstance(source, (str, os.PathLike)):
        return pymupdf.open(source)
    data = source.read() if hasattr(source, "read") else source
    return pymupdf.open(stream=data, filetype="pdf")


def analyse_pdf(source, lang="eng", psm=11, fuzzy=True, dpi=200):
    """Une page = un résultat (n° de page, comptes, remarques).
    Texte natif du PDF s'il existe, sinon OCR sur la page rendue en image."""
    doc = _open_pdf(source)
    results = []
    for i, page in enumerate(doc, 1):
        words = [{"text": w[4], "left": w[0], "top": w[1], "w": w[2] - w[0], "h": w[3] - w[1]}
                 for w in page.get_text("words")]
        counts, notes = count_column2(words, page.rect.width, fuzzy)
        method = "texte natif"
        if sum(counts.values()) == 0:
            pix = page.get_pixmap(dpi=dpi, colorspace=pymupdf_rgb(), alpha=False)
            rgb = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)
            bw, _ = preprocess(cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
            counts, notes = count_column2(read_words(bw, lang, psm), bw.shape[1], fuzzy)
            method = "OCR"
        results.append((i, counts, [f"Méthode : {method}"] + notes))
    return results


def pymupdf_rgb():
    import pymupdf
    return pymupdf.csRGB