| """ |
| 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 |
| |
| |
| 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: |
| 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() |
| ]) |
| |
| |
| 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 |
| |
| |
| 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() |
| ]) |
| |
| |
| 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: |
| |
| aug1 = self.individual_transform(image=signature1) |
| aug2 = self.individual_transform(image=signature2) |
| |
| |
| final1 = self.shared_transform(image=aug1['image']) |
| final2 = self.shared_transform(image=aug2['image']) |
| else: |
| |
| 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 |
| ) |
|
|
