src/data/augmentation.py

"""
Data augmentation utilities for signature verification training.
"""

import torch
import numpy as np
from typing import Tuple, List, Union
import albumentations as A
from albumentations.pytorch import ToTensorV2


class SignatureAugmentationPipeline:
    """
    Comprehensive augmentation pipeline for signature verification.
    """
    
    def __init__(self, 
                 target_size: Tuple[int, int] = (224, 224),
                 augmentation_strength: str = 'medium'):
        """
        Initialize augmentation pipeline.
        
        Args:
            target_size: Target size for signature images
            augmentation_strength: 'light', 'medium', or 'heavy'
        """
        self.target_size = target_size
        self.strength = augmentation_strength
        
        # Define augmentation strategies based on strength
        self._setup_augmentations()
    
    def _setup_augmentations(self):
        """Setup augmentation transforms based on strength."""
        
        if self.strength == 'light':
            self.train_transform = A.Compose([
                A.Resize(self.target_size[0], self.target_size[1]),
                A.HorizontalFlip(p=0.2),
                A.Rotate(limit=5, p=0.3),
                A.RandomBrightnessContrast(
                    brightness_limit=0.1, 
                    contrast_limit=0.1, 
                    p=0.3
                ),
                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
                ToTensorV2()
            ])
            
        elif self.strength == 'medium':
            self.train_transform = A.Compose([
                A.Resize(self.target_size[0], self.target_size[1]),
                A.HorizontalFlip(p=0.3),
                A.Rotate(limit=10, p=0.4),
                A.RandomBrightnessContrast(
                    brightness_limit=0.15, 
                    contrast_limit=0.15, 
                    p=0.4
                ),
                A.GaussNoise(var_limit=(5.0, 25.0), p=0.2),
                A.ElasticTransform(
                    alpha=0.5, 
                    sigma=25, 
                    alpha_affine=25, 
                    p=0.2
                ),
                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
                ToTensorV2()
            ])
            
        else:  # heavy
            self.train_transform = A.Compose([
                A.Resize(self.target_size[0], self.target_size[1]),
                A.HorizontalFlip(p=0.4),
                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.Perspective(scale=(0.05, 0.1), p=0.2),
                A.ShiftScaleRotate(
                    shift_limit=0.05, 
                    scale_limit=0.1, 
                    rotate_limit=10, 
                    p=0.3
                ),
                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
                ToTensorV2()
            ])
        
        # Validation transform (minimal)
        self.val_transform = A.Compose([
            A.Resize(self.target_size[0], self.target_size[1]),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
    
    def augment_image(self, image: np.ndarray, is_training: bool = True) -> torch.Tensor:
        """
        Apply augmentation to a single 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[np.ndarray], is_training: bool = True) -> torch.Tensor:
        """
        Apply augmentation to a batch of 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(image, is_training)
            augmented_images.append(augmented)
        
        return torch.stack(augmented_images)


class PairAugmentation:
    """
    Specialized augmentation for signature pairs in Siamese networks.
    """
    
    def __init__(self, target_size: Tuple[int, int] = (224, 224)):
        """
        Initialize pair augmentation.
        
        Args:
            target_size: Target size for signature images
        """
        self.target_size = target_size
        
        # Shared augmentations for both signatures in a pair
        self.shared_transform = A.Compose([
            A.Resize(self.target_size[0], self.target_size[1]),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ])
        
        # Individual augmentations for each signature
        self.individual_transform = A.Compose([
            A.HorizontalFlip(p=0.3),
            A.Rotate(limit=10, p=0.4),
            A.RandomBrightnessContrast(
                brightness_limit=0.15, 
                contrast_limit=0.15, 
                p=0.4
            ),
            A.GaussNoise(var_limit=(5.0, 25.0), p=0.2),
        ])
    
    def augment_pair(self, 
                    signature1: np.ndarray, 
                    signature2: np.ndarray,
                    is_training: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Augment a pair of signatures.
        
        Args:
            signature1: First signature image
            signature2: Second signature image
            is_training: Whether to apply training augmentations
            
        Returns:
            Tuple of augmented signature tensors
        """
        if is_training:
            # Apply individual augmentations
            aug1 = self.individual_transform(image=signature1)
            aug2 = self.individual_transform(image=signature2)
            
            # Apply shared transformations
            final1 = self.shared_transform(image=aug1['image'])
            final2 = self.shared_transform(image=aug2['image'])
        else:
            # Only apply shared transformations for validation
            final1 = self.shared_transform(image=signature1)
            final2 = self.shared_transform(image=signature2)
        
        return final1['image'], final2['image']


class OnlineAugmentation:
    """
    Online augmentation during training for dynamic augmentation.
    """
    
    def __init__(self, target_size: Tuple[int, int] = (224, 224)):
        """
        Initialize online augmentation.
        
        Args:
            target_size: Target size for signature images
        """
        self.target_size = target_size
        self.augmentation_pipeline = SignatureAugmentationPipeline(
            target_size=target_size, 
            augmentation_strength='medium'
        )
    
    def __call__(self, image: np.ndarray, is_training: bool = True) -> torch.Tensor:
        """
        Apply online augmentation.
        
        Args:
            image: Input signature image
            is_training: Whether to apply training augmentations
            
        Returns:
            Augmented image as torch tensor
        """
        return self.augmentation_pipeline.augment_image(image, is_training)
    
    def set_strength(self, strength: str):
        """
        Dynamically change augmentation strength.
        
        Args:
            strength: 'light', 'medium', or 'heavy'
        """
        self.augmentation_pipeline = SignatureAugmentationPipeline(
            target_size=self.target_size,
            augmentation_strength=strength
        )