src/data/preprocessing.py

"""
Signature preprocessing module for image normalization and preparation.
"""

import cv2
import numpy as np
import torch
from PIL import Image
from typing import Tuple, Union, Optional
import albumentations as A
from albumentations.pytorch import ToTensorV2


class SignaturePreprocessor:
    """
    Handles preprocessing of signature images for the verification model.
    """
    
    def __init__(self, target_size: Tuple[int, int] = (224, 224)):
        """
        Initialize the preprocessor.
        
        Args:
            target_size: Target size for signature images (height, width)
        """
        self.target_size = target_size
        self.transform = A.Compose([
            A.Resize(target_size[0], target_size[1]),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
    
    def load_image(self, image_path: str) -> np.ndarray:
        """
        Load image from file path.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Loaded image as numpy array
        """
        try:
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"Could not load image from {image_path}")
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            return image
        except Exception as e:
            raise ValueError(f"Error loading image {image_path}: {str(e)}")
    
    def preprocess_image(self, image: Union[str, np.ndarray, Image.Image]) -> torch.Tensor:
        """
        Preprocess a signature image for model input.
        
        Args:
            image: Image as file path, numpy array, or PIL Image
            
        Returns:
            Preprocessed image as torch tensor
        """
        # Convert to numpy array if needed
        if isinstance(image, str):
            image = self.load_image(image)
        elif isinstance(image, Image.Image):
            image = np.array(image)
        elif isinstance(image, torch.Tensor):
            image = image.numpy()
        
        # Ensure image is in RGB format
        if len(image.shape) == 3 and image.shape[2] == 3:
            pass  # Already RGB
        elif len(image.shape) == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        else:
            raise ValueError(f"Unsupported image format with shape: {image.shape}")
        
        # Apply transformations
        transformed = self.transform(image=image)
        return transformed['image']
    
    def enhance_signature(self, image: np.ndarray) -> np.ndarray:
        """
        Enhance signature image quality.
        
        Args:
            image: Input signature image
            
        Returns:
            Enhanced signature image
        """
        # Convert to grayscale for processing
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image.copy()
        
        # Apply adaptive thresholding to get binary image
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
        )
        
        # Morphological operations to clean up the signature
        kernel = np.ones((2, 2), np.uint8)
        cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel)
        
        # Convert back to RGB
        if len(image.shape) == 3:
            enhanced = cv2.cvtColor(cleaned, cv2.COLOR_GRAY2RGB)
        else:
            enhanced = cleaned
        
        return enhanced
    
    def normalize_signature(self, image: np.ndarray) -> np.ndarray:
        """
        Normalize signature image for consistent processing.
        
        Args:
            image: Input signature image
            
        Returns:
            Normalized signature image
        """
        # Convert to grayscale
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        else:
            gray = image.copy()
        
        # Find signature contours
        contours, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if not contours:
            return image
        
        # Get bounding box of the signature
        x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
        
        # Crop to signature area with some padding
        padding = 10
        x1 = max(0, x - padding)
        y1 = max(0, y - padding)
        x2 = min(image.shape[1], x + w + padding)
        y2 = min(image.shape[0], y + h + padding)
        
        cropped = image[y1:y2, x1:x2]
        
        # Resize to target size while maintaining aspect ratio
        h_orig, w_orig = cropped.shape[:2]
        aspect_ratio = w_orig / h_orig
        
        if aspect_ratio > 1:
            new_w = self.target_size[1]
            new_h = int(new_w / aspect_ratio)
        else:
            new_h = self.target_size[0]
            new_w = int(new_h * aspect_ratio)
        
        resized = cv2.resize(cropped, (new_w, new_h))
        
        # Create canvas with target size
        canvas = np.ones((self.target_size[0], self.target_size[1], 3), dtype=np.uint8) * 255
        
        # Center the signature on the canvas
        y_offset = (self.target_size[0] - new_h) // 2
        x_offset = (self.target_size[1] - new_w) // 2
        
        canvas[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized
        
        return canvas
    
    def preprocess_batch(self, images: list) -> torch.Tensor:
        """
        Preprocess a batch of signature images.
        
        Args:
            images: List of images to preprocess
            
        Returns:
            Batch of preprocessed images as torch tensor
        """
        processed_images = []
        for image in images:
            processed = self.preprocess_image(image)
            processed_images.append(processed)
        
        return torch.stack(processed_images)


class SignatureAugmentation:
    """
    Data augmentation for signature images during training.
    """
    
    def __init__(self, target_size: Tuple[int, int] = (224, 224)):
        """
        Initialize augmentation pipeline.
        
        Args:
            target_size: Target size for signature images
        """
        self.target_size = target_size
        
        # Training augmentations
        self.train_transform = A.Compose([
            A.Resize(target_size[0], target_size[1]),
            A.HorizontalFlip(p=0.3),
            A.Rotate(limit=15, p=0.5),
            A.RandomBrightnessContrast(
                brightness_limit=0.2, 
                contrast_limit=0.2, 
                p=0.5
            ),
            A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),
            A.ElasticTransform(
                alpha=1, 
                sigma=50, 
                alpha_affine=50, 
                p=0.3
            ),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
        
        # Validation augmentations (minimal)
        self.val_transform = A.Compose([
            A.Resize(target_size[0], target_size[1]),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
    
    def augment(self, image: np.ndarray, is_training: bool = True) -> torch.Tensor:
        """
        Apply augmentation to signature image.
        
        Args:
            image: Input signature image
            is_training: Whether to apply training augmentations
            
        Returns:
            Augmented image as torch tensor
        """
        transform = self.train_transform if is_training else self.val_transform
        transformed = transform(image=image)
        return transformed['image']
    
    def augment_batch(self, images: list, is_training: bool = True) -> torch.Tensor:
        """
        Apply augmentation to a batch of signature images.
        
        Args:
            images: List of images to augment
            is_training: Whether to apply training augmentations
            
        Returns:
            Batch of augmented images as torch tensor
        """
        augmented_images = []
        for image in images:
            augmented = self.augment(image, is_training)
            augmented_images.append(augmented)
        
        return torch.stack(augmented_images)