utils/preprocess.py

import cv2
import numpy as np


##################################################
# Convert to grayscale
##################################################

def to_grayscale(img):

    if len(img.shape) == 3:
        gray = cv2.cvtColor(
            img,
            cv2.COLOR_RGB2GRAY
        )
    else:
        gray = img.copy()

    return gray


##################################################
# Remove noise while preserving edges
##################################################

def denoise_image(gray):

    denoised = cv2.fastNlMeansDenoising(
        gray,
        None,
        h=10,
        templateWindowSize=7,
        searchWindowSize=21
    )

    return denoised


##################################################
# Increase contrast
##################################################

def enhance_contrast(gray):

    clahe = cv2.createCLAHE(
        clipLimit=2.0,
        tileGridSize=(8, 8)
    )

    enhanced = clahe.apply(gray)

    return enhanced


##################################################
# Adaptive thresholding
##################################################

def adaptive_binarize(gray):

    binary = cv2.adaptiveThreshold(
        gray,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        15,
        2
    )

    return binary


##################################################
# Close tiny gaps in boundaries
##################################################

def close_small_gaps(binary):

    kernel = np.ones((3, 3), np.uint8)

    closed = cv2.morphologyEx(
        binary,
        cv2.MORPH_CLOSE,
        kernel,
        iterations=1
    )

    return closed


##################################################
# Remove isolated noise pixels
##################################################

def remove_noise(binary):

    kernel = np.ones((2, 2), np.uint8)

    cleaned = cv2.morphologyEx(
        binary,
        cv2.MORPH_OPEN,
        kernel,
        iterations=1
    )

    return cleaned


##################################################
# Strengthen weak pencil lines
##################################################

def strengthen_lines(binary):

    kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT,
        (3, 3)
    )

    strengthened = cv2.dilate(
        binary,
        kernel,
        iterations=2
    )

    return strengthened


##################################################
# Complete preprocessing pipeline
##################################################

def preprocess_image(img):

    gray = to_grayscale(img)

    gray = denoise_image(gray)

    gray = enhance_contrast(gray)

    binary = adaptive_binarize(gray)

    binary = close_small_gaps(binary)

    binary = strengthen_lines(binary)

    return binary