Aivanov_scan_ogc/pipeline/deskew.py

"""Détection d'angle de skew + redressement automatique des pages scannées.

Technique : Hough Transform sur les lignes détectées par Canny, puis moyenne
des angles des lignes « quasi horizontales » (±15° par rapport à l'horizontale).
Les fiches OGC ont énormément de traits de tableau → signal très fort.

Seuil : on ne corrige que si |angle| > `MIN_ANGLE_DEG` (0.3° par défaut) pour
éviter de toucher les scans déjà bien cadrés et introduire du bruit inutile.
"""
from __future__ import annotations

from pathlib import Path
from typing import Tuple

import numpy as np
from PIL import Image

try:
    import cv2  # type: ignore
    _HAS_CV2 = True
except ImportError:
    _HAS_CV2 = False


MIN_ANGLE_DEG = 0.3      # en-dessous, on ne corrige pas
MAX_ANGLE_DEG = 10.0     # au-dessus, c'est anormal → suspect, on ne corrige pas
NEAR_HORIZONTAL_BAND = 15.0  # degrés : bande autour de l'horizontale pour filtrer


def detect_skew_angle(img: Image.Image) -> float:
    """Retourne l'angle de skew en degrés (positif = tourné dans le sens
    des aiguilles d'une montre) à appliquer pour redresser l'image.

    Si aucune ligne horizontale n'est trouvée, retourne 0.0.
    Si l'angle détecté est hors [-MAX_ANGLE_DEG, +MAX_ANGLE_DEG], retourne 0.0
    (probablement une erreur de détection, on ne corrige pas).
    """
    if not _HAS_CV2:
        return 0.0
    gray = np.array(img.convert("L"))
    # Réduire l'image pour accélérer (max 1500 px de large)
    h, w = gray.shape
    if w > 1500:
        scale = 1500 / w
        gray = cv2.resize(gray, (1500, int(h * scale)), interpolation=cv2.INTER_AREA)

    # Canny edges — paramètres standards documents
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    # Hough Lines probabiliste : rapide et robuste
    lines = cv2.HoughLinesP(
        edges, rho=1, theta=np.pi / 180, threshold=200,
        minLineLength=gray.shape[1] // 4,  # au moins 25% de la largeur
        maxLineGap=20,
    )
    if lines is None or len(lines) == 0:
        return 0.0

    # Calculer l'angle de chaque ligne en degrés
    angles = []
    for line in lines:
        x1, y1, x2, y2 = line[0]
        if x2 == x1:
            continue  # ligne verticale, ignorée
        angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
        # On ne garde que les lignes proches de l'horizontale
        if abs(angle) < NEAR_HORIZONTAL_BAND:
            angles.append(angle)

    if not angles:
        return 0.0

    # Moyenne robuste : médiane plutôt que mean, moins sensible aux outliers
    angle = float(np.median(angles))
    if abs(angle) > MAX_ANGLE_DEG:
        return 0.0  # suspect → on ne corrige pas
    return angle


def deskew_image(img: Image.Image,
                 angle: float | None = None,
                 min_angle: float = MIN_ANGLE_DEG) -> Tuple[Image.Image, float]:
    """Redresse une image si le skew détecté dépasse `min_angle`.

    Retourne (image_eventuellement_rotee, angle_applique).
    Si |angle| < min_angle, retourne l'image inchangée et angle=0.0.
    """
    if angle is None:
        angle = detect_skew_angle(img)
    if abs(angle) < min_angle:
        return img, 0.0
    # PIL.Image.rotate : positive angle = counter-clockwise
    # detect_skew retourne positif = clockwise → on inverse pour PIL
    rotated = img.rotate(
        angle,
        resample=Image.Resampling.BICUBIC,
        expand=False,
        fillcolor="white",
    )
    return rotated, angle


def deskew_file(src: Path, dst: Path | None = None,
                min_angle: float = MIN_ANGLE_DEG) -> float:
    """Version fichier → fichier. Écrase `src` si `dst` est None.
    Retourne l'angle appliqué (0.0 si pas de rotation)."""
    img = Image.open(src)
    rotated, angle = deskew_image(img, min_angle=min_angle)
    out = dst or src
    rotated.save(out, "PNG", optimize=True)
    return angle


if __name__ == "__main__":
    import sys
    import glob
    paths = [Path(p) for p in (sys.argv[1:] or sorted(glob.glob(".cache/images/*/page_01.png")))]
    print(f"Deskew sur {len(paths)} images (seuil={MIN_ANGLE_DEG}°)...")
    total_corrected = 0
    for p in paths:
        angle = detect_skew_angle(Image.open(p))
        mark = "→" if abs(angle) >= MIN_ANGLE_DEG else "·"
        if abs(angle) >= MIN_ANGLE_DEG:
            total_corrected += 1
        print(f"  {mark} {p} : {angle:+.2f}°")
    print(f"\n{total_corrected}/{len(paths)} images auraient besoin d'un redressement.")