src/pageparse/preprocess.py

from __future__ import annotations

import cv2
import numpy as np


def preprocess(img: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    gray = _deskew(gray)
    thresh = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 31, 10,
    )
    denoised = cv2.fastNlMeansDenoising(thresh, h=30)
    return denoised


def preprocess_otsu(img: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return thresh


def preprocess_light(img: np.ndarray) -> np.ndarray:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(gray)
    _, thresh = cv2.threshold(enhanced, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return thresh


def adaptive_preprocess(img: np.ndarray) -> tuple[np.ndarray, dict]:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    h, w = gray.shape
    mean_brightness = gray.mean()
    std_brightness = gray.std()

    params = {
        "adaptive_block_size": 31,
        "adaptive_c": 10,
        "denoise_h": 30,
        "clahe_clip": 2.0,
        "clahe_grid": 8,
    }

    if mean_brightness < 50:
        params["clahe_clip"] = 3.0
        params["adaptive_c"] = 8
    elif mean_brightness > 200:
        params["adaptive_block_size"] = 21
        params["adaptive_c"] = 12

    if std_brightness < 30:
        params["clahe_clip"] = 4.0
        params["denoise_h"] = 20

    params["adaptive_block_size"] += (params["adaptive_block_size"] + 1) % 2

    gray = _deskew(gray)
    clahe = cv2.createCLAHE(clipLimit=params["clahe_clip"], tileGridSize=(params["clahe_grid"], params["clahe_grid"]))
    enhanced = clahe.apply(gray)
    thresh = cv2.adaptiveThreshold(
        enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,
        params["adaptive_block_size"], params["adaptive_c"],
    )
    denoised = cv2.fastNlMeansDenoising(thresh, h=params["denoise_h"])

    return denoised, params


def preprocess_with_scale(img: np.ndarray, scale: float = 2.0) -> np.ndarray:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    scaled = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    return scaled


def _deskew(img: np.ndarray) -> np.ndarray:
    coords = np.column_stack(np.where(img < 128))
    if len(coords) < 5:
        return img
    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = 90 + angle
    if abs(angle) < 0.5:
        return img
    h, w = img.shape
    matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0)
    return cv2.warpAffine(
        img, matrix, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE,
    )