ocr_studio/preprocess.py

from __future__ import annotations

from PIL import Image

from ocr_studio.config import DESKEW_MAX_ANGLE


def _rotate(image: Image.Image, angle: float) -> Image.Image:
    if abs(angle) < 0.25:
        return image
    fill = (255, 255, 255)
    return image.convert("RGB").rotate(angle, expand=True, fillcolor=fill, resample=Image.Resampling.BICUBIC)


def enhance_scan(image: Image.Image) -> Image.Image:
    try:
        import cv2
        import numpy as np
    except Exception:
        return image.convert("RGB")

    rgb = np.array(image.convert("RGB"))
    lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
    lightness, axis_a, axis_b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    lightness = clahe.apply(lightness)
    merged = cv2.merge((lightness, axis_a, axis_b))
    enhanced = cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)
    return Image.fromarray(enhanced)


def estimate_deskew_angle(image: Image.Image) -> float:
    try:
        import cv2
        import numpy as np
    except Exception:
        return 0.0

    gray = np.array(image.convert("L"))
    if gray.size < 400:
        return 0.0
    inverted = cv2.bitwise_not(gray)
    _, binary = cv2.threshold(inverted, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    coords = np.column_stack(np.where(binary > 0))
    if coords.shape[0] < 80:
        return 0.0
    _rect, _size, raw_angle = cv2.minAreaRect(coords)
    angle = raw_angle
    if angle < -45:
        angle = 90.0 + angle
    if abs(angle) > DESKEW_MAX_ANGLE:
        return 0.0
    return float(-angle)


def deskew_image(image: Image.Image) -> tuple[Image.Image, float]:
    rgb = image.convert("RGB")
    angle = estimate_deskew_angle(rgb)
    straightened = _rotate(rgb, angle)
    return enhance_scan(straightened), angle


def prepare_page(image: Image.Image, deskew: bool) -> tuple[Image.Image, float]:
    if not deskew:
        return enhance_scan(image.convert("RGB")), 0.0
    return deskew_image(image)