preprocessor.py

import cv2
import numpy as np

class Preprocessor:
    @staticmethod
    def to_grayscale(img):
        if len(img.shape) == 3:
            return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return img.copy()

    @staticmethod
    def denoise_bilateral(img, d=9, sigma_color=75, sigma_space=75):
        """
        Bilateral filtering smooths scan smudges and noise while preserving crisp line boundaries.
        """
        return cv2.bilateralFilter(img, d, sigma_color, sigma_space)

    @staticmethod
    def adaptive_threshold(img, block_size=11, c=2):
        """
        Adapts to variable scan brightness, converting background to black (0) and lines to white (255).
        """
        return cv2.adaptiveThreshold(
            img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, block_size, c
        )

    @staticmethod
    def canny_edges(img, low=50, high=150):
        """
        Canny edge detection to extract thin outlines.
        """
        return cv2.Canny(img, low, high)

    @staticmethod
    def morphological_close(img, kernel_size=3):
        """
        Applies morphological closing to bridge micro-cracks and disconnections in scanned drawings.
        """
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
        return cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)

    @staticmethod
    def extract_and_mask_wires(binary_img, line_len=50):
        """
        Detects and isolates long horizontal and vertical connecting wires in schematics
        using morphological structural elements. Masking them prevents wires from distorting
        symbol matching scores.
        """
        # Create structural elements
        horiz_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (line_len, 1))
        vert_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, line_len))
        
        # Extract wires
        horiz_wires = cv2.morphologyEx(binary_img, cv2.MORPH_OPEN, horiz_kernel)
        vert_wires = cv2.morphologyEx(binary_img, cv2.MORPH_OPEN, vert_kernel)
        
        # Combine wires
        all_wires = cv2.bitwise_or(horiz_wires, vert_wires)
        
        # Subtract wires from main image to isolate symbols, preserving junction vertices
        isolated_symbols = cv2.subtract(binary_img, all_wires)
        return isolated_symbols, all_wires

    @staticmethod
    def skeletonize(img):
        """
        Mathematical skeletonization loop. Reduces thick/fuzzy drawn lines to a 1-pixel thin skeleton.
        This normalizes all line drawings, ensuring perfect scale and thickness invariance.
        """
        skel = np.zeros(img.shape, np.uint8)
        element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
        temp = img.copy()
        
        # Limit iterations to prevent potential infinite loops on complex noise
        max_iters = 100
        for _ in range(max_iters):
            eroded = cv2.erode(temp, element)
            temp_open = cv2.dilate(eroded, element)
            temp_sub = cv2.subtract(temp, temp_open)
            skel = cv2.bitwise_or(skel, temp_sub)
            temp = eroded.copy()
            if cv2.countNonZero(temp) == 0:
                break
                
        return skel

    def process(
        self, img, 
        method='canny', 
        canny_low=50, 
        canny_high=150, 
        denoise=False, 
        morph_kernel=3, 
        mask_wires=False, 
        do_skeletonize=False,
        dilate_kernel=0
    ):
        # 1. Convert to grayscale
        gray = self.to_grayscale(img)
        
        # 2. Denoise if enabled
        if denoise:
            gray = self.denoise_bilateral(gray)
            
        # 3. Base thresholding/edge mapping
        if method == 'Canny Edge Detection':
            proc = self.canny_edges(gray, canny_low, canny_high)
        elif method == 'Adaptive Thresholding':
            proc = self.adaptive_threshold(gray)
        else: # Inverted Thresholding
            _, proc = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
            
        # 4. Bridge scanned broken lines
        if morph_kernel > 1:
            proc = self.morphological_close(proc, morph_kernel)
            
        # 5. Suppress long connecting wires if enabled
        wires = None
        if mask_wires:
            proc, wires = self.extract_and_mask_wires(proc)
            
        # 6. Normalize line widths if enabled
        if do_skeletonize:
            proc = self.skeletonize(proc)
            
        # 7. Apply dilation if requested (thicken lines for spatial tolerance)
        if dilate_kernel > 0:
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (dilate_kernel, dilate_kernel))
            proc = cv2.dilate(proc, kernel)
            
        return proc, wires