InklyAI / src /data /preprocessing.py

Upload folder using huggingface_hub

8eab354 verified about 1 year ago

8.57 kB

"""
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)

	"""
	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)