preprocess_for_ocr.py

#!/usr/bin/env python3
"""
Preprocess handwritten-text images for OCR.

For every image in the current folder, produces a cleaned version in ./processed/
Pipeline: grayscale → denoise → normalize lighting → contrast enhance → upscale.

Modern OCR engines perform best on clean grayscale (not harsh binarization),
so this pipeline keeps the output in grayscale and saves as lossless PNG.

Requirements:  pip install opencv-python numpy
"""

import cv2
import numpy as np
from pathlib import Path

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp"}
OUTPUT_DIR = Path(__file__).parent / "processed"
UPSCALE_FACTOR = 2.0


def to_grayscale(img: np.ndarray) -> np.ndarray:
    if len(img.shape) == 3:
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img


def denoise(gray: np.ndarray) -> np.ndarray:
    """Non-local means denoising — best at preserving edges in text."""
    return cv2.fastNlMeansDenoising(gray, h=12, templateWindowSize=7, searchWindowSize=21)


def normalize_lighting(gray: np.ndarray) -> np.ndarray:
    """Divide by a blurred version of itself to flatten uneven illumination."""
    bg = cv2.GaussianBlur(gray, (0, 0), sigmaX=51)
    normalized = cv2.divide(gray, bg, scale=255)
    return normalized


def enhance_contrast(gray: np.ndarray) -> np.ndarray:
    """CLAHE for local contrast enhancement — makes faint ink readable."""
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    return clahe.apply(gray)


def sharpen(gray: np.ndarray) -> np.ndarray:
    """Gentle unsharp-mask to crisp up stroke edges."""
    blurred = cv2.GaussianBlur(gray, (0, 0), sigmaX=2)
    return cv2.addWeighted(gray, 1.5, blurred, -0.5, 0)


def deskew(gray: np.ndarray) -> np.ndarray:
    """Straighten slightly rotated scans using Otsu-detected text pixels."""
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    coords = np.column_stack(np.where(thresh > 0))
    if len(coords) < 100:
        return gray
    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = 90 + angle
    if abs(angle) > 10:
        return gray
    h, w = gray.shape
    M = cv2.getRotationMatrix2D((w // 2, h // 2), angle, 1.0)
    return cv2.warpAffine(gray, M, (w, h), flags=cv2.INTER_CUBIC, borderValue=255)


def upscale(img: np.ndarray, factor: float = UPSCALE_FACTOR) -> np.ndarray:
    """Upscale for better OCR accuracy on small handwriting."""
    if factor <= 1.0:
        return img
    return cv2.resize(img, None, fx=factor, fy=factor, interpolation=cv2.INTER_CUBIC)


def preprocess(path: Path) -> np.ndarray:
    img = cv2.imread(str(path), cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"Cannot read image: {path}")

    gray = to_grayscale(img)
    gray = denoise(gray)
    gray = normalize_lighting(gray)
    gray = enhance_contrast(gray)
    gray = sharpen(gray)
    gray = deskew(gray)
    gray = upscale(gray)
    return gray


def main() -> None:
    src_dir = Path(__file__).parent
    images = sorted(
        p for p in src_dir.iterdir()
        if p.suffix.lower() in IMAGE_EXTENSIONS and p.is_file()
    )

    if not images:
        print("No images found in", src_dir)
        return

    OUTPUT_DIR.mkdir(exist_ok=True)
    print(f"Found {len(images)} image(s). Processing …\n")

    for path in images:
        try:
            result = preprocess(path)
            out_path = OUTPUT_DIR / f"{path.stem}_ocr.png"  # lossless PNG
            cv2.imwrite(str(out_path), result)
            print(f"  ✔ {path.name}  →  {out_path.relative_to(src_dir)}")
        except Exception as e:
            print(f"  ✘ {path.name}  — {e}")

    print(f"\nDone. Processed files are in ./{OUTPUT_DIR.name}/")


if __name__ == "__main__":
    main()