import os import gc import json import math import torch import mlflow import logging import platform import numpy as np import pandas as pd from PIL import Image from tqdm import tqdm import torch.nn as nn import torch.optim as optim from torchvision import models import matplotlib.pyplot as plt import torch.nn.functional as F from sklearn.manifold import TSNE from torchvision import transforms from kymatio.torch import Scattering2D from torch.utils.data import Dataset, DataLoader from pytorch_metric_learning.miners import BatchHardMiner from pytorch_metric_learning.losses import MultiSimilarityLoss from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau from sklearn.metrics import roc_curve, auc, precision_recall_fscore_support from typing import Dict, List, Tuple, Optional, Union, Any from dataclasses import dataclass, asdict import warnings warnings.filterwarnings('ignore') # ---------------------------- # Configuration Management # ---------------------------- @dataclass @dataclass class TrainingConfig: # Model Architecture model_name: str = "resnet34" embedding_dim: int = 128 normalize_embeddings: bool = True pretrained_path: Optional[str] = "../../model/pretrained_model/ResNet34.pt" # Training Hyperparameters batch_size: int = 512 max_epochs: int = 20 grad_accum_steps: int = 10 device: str = "cuda" if torch.cuda.is_available() else "cpu" # Learning Rate Configuration head_lr: float = 1e-3 # Higher LR for embedding head (untrained) backbone_lr: float = 1e-4 # Lower LR for backbone (pretrained) lr_scheduler: str = "cosine" # "cosine" or "plateau" weight_decay: float = 1e-4 # Curriculum Learning Parameters - ADJUSTED FOR PRECISION curriculum_strategy: str = "progressive" # "progressive", "exponential", "linear" initial_hard_ratio: float = 0.6 # Increased from 0.1 for more hard negatives early final_hard_ratio: float = 0.9 # Increased from 0.8 for focus on hard cases curriculum_warmup_epochs: int = 1 # Reduced from 2 for faster hard sample exposure # Data Augmentation remove_bg: bool = False augmentation_strength: float = 0.5 # 0.0 = no aug, 1.0 = strong aug # Loss Configuration - ADJUSTED FOR PRECISION multisim_alpha: float = 2.5 # Increased from 2.0 (penalize false positives more) multisim_beta: float = 60.0 # Increased from 50.0 (larger margin) multisim_base: float = 0.4 # Decreased from 0.5 (stricter similarity) # Triplet Loss Parameters - NEW triplet_margin: float = 1.0 # Margin for triplet loss triplet_weight: float = 0.3 # Weight for triplet loss component false_positive_penalty_weight: float = 0.3 # Extra penalty for false positives # Mining Configuration use_hard_mining: bool = True # Precision Focus Parameters - NEW target_precision: float = 0.75 # Target precision for threshold selection negative_weight_multiplier: float = 2.5 # How much more to weight hard negatives # Checkpoint Configuration run_id: Optional[str] = None last_epoch_weights: Optional[str] = None save_frequency: int = 1 # Save every N epochs # Early Stopping patience: int = 15 min_delta: float = 0.001 # Logging log_frequency: int = 100 # Log every N steps visualize_frequency: int = 1 # Visualize every N epochs tracking_uri: str = "http://127.0.0.1:5555" def __post_init__(self): """Validate configuration parameters.""" assert 0.0 <= self.initial_hard_ratio <= 1.0, "Initial hard ratio must be in [0, 1]" assert 0.0 <= self.final_hard_ratio <= 1.0, "Final hard ratio must be in [0, 1]" assert self.curriculum_strategy in ["progressive", "exponential", "linear"] assert self.lr_scheduler in ["cosine", "plateau"] assert 0.0 <= self.target_precision <= 1.0, "Target precision must be in [0, 1]" # Global configuration CONFIG = TrainingConfig() # ---------------------------- # MLFlow Setup # ---------------------------- class MLFlowManager: """Centralized MLflow management for experiment tracking.""" def __init__(self, tracking_uri: str = "http://127.0.0.1:5555"): mlflow.set_tracking_uri(tracking_uri) self.experiment_name = "Signature Verification - Advanced Training" self._setup_experiment() def _setup_experiment(self): """Setup MLflow experiment.""" try: self.experiment_id = mlflow.create_experiment(self.experiment_name) except: self.experiment_id = mlflow.get_experiment_by_name(self.experiment_name).experiment_id def start_run(self, run_id: Optional[str] = None): """Start MLflow run with configuration logging.""" return mlflow.start_run(run_id=run_id, experiment_id=self.experiment_id) def log_config(self, config: TrainingConfig): """Log training configuration.""" config_dict = asdict(config) mlflow.log_params(config_dict) # ---------------------------- # Curriculum Learning Manager # ---------------------------- class CurriculumLearningManager: """Advanced curriculum learning for both hard positives and hard negatives.""" def __init__(self, config: TrainingConfig): self.config = config self.current_epoch = 0 def get_hard_ratio(self, epoch: int) -> float: """Get hard negative ratio (forgeries) for current epoch.""" if epoch < self.config.curriculum_warmup_epochs: return self.config.initial_hard_ratio # Target: reach final_hard_ratio by max_epochs // 2 target_epoch = max(self.config.max_epochs // 2, self.config.curriculum_warmup_epochs + 3) if epoch >= target_epoch: return self.config.final_hard_ratio # Aggressive progression to reach target by mid-training progress = (epoch - self.config.curriculum_warmup_epochs) / (target_epoch - self.config.curriculum_warmup_epochs) initial = self.config.initial_hard_ratio final = self.config.final_hard_ratio if self.config.curriculum_strategy == "progressive": # Very aggressive: exponential growth early, then plateau ratio = initial + (final - initial) * (progress ** 0.5) elif self.config.curriculum_strategy == "exponential": ratio = initial + (final - initial) * (progress ** 0.3) else: # linear ratio = initial + (final - initial) * progress return min(max(ratio, 0.0), 1.0) def get_hard_positive_ratio(self, epoch: int) -> float: """Get hard positive ratio for current epoch - increases more gradually.""" if epoch < self.config.curriculum_warmup_epochs: return 0.1 # Start with 10% hard positives # Hard positives should increase more gradually than hard negatives max_epochs = self.config.max_epochs progress = min(1.0, (epoch - self.config.curriculum_warmup_epochs) / (max_epochs - self.config.curriculum_warmup_epochs)) # Target 40% hard positives by end of training initial_ratio = 0.1 final_ratio = 0.4 if self.config.curriculum_strategy == "progressive": ratio = initial_ratio + (final_ratio - initial_ratio) * (progress ** 0.7) else: ratio = initial_ratio + (final_ratio - initial_ratio) * progress return min(max(ratio, 0.0), final_ratio) def get_mining_difficulty(self, epoch: int) -> Dict[str, float]: """Adaptive mining parameters for both hard positives and negatives.""" progress = min(1.0, epoch / self.config.max_epochs) # Separate ratios for hard positives and hard negatives hard_negative_ratio = self.get_hard_ratio(epoch) hard_positive_ratio = self.get_hard_positive_ratio(epoch) # Dynamic weights for different sample types hard_pos_weight = 1.0 + 2.0 * progress # 1.0 → 3.0 hard_neg_weight = 1.0 + 4.0 * progress # 1.0 → 5.0 (harder negatives more important) return { # Margin parameters "margin_multiplier": 1.0 + 0.5 * progress, # Hard sample ratios "hard_negative_ratio": hard_negative_ratio, "hard_positive_ratio": hard_positive_ratio, "current_hard_ratio": hard_negative_ratio, # For backward compatibility # Sample weights "hard_positive_weight": hard_pos_weight, "hard_negative_weight": hard_neg_weight, "semi_positive_weight": 1.0 + 1.0 * progress, "semi_negative_weight": 1.0 + 2.0 * progress, # Difficulty thresholds "difficulty_threshold": 0.05 + 0.15 * progress, "selectivity": 0.8 + 0.2 * progress, # Mining aggressiveness "mining_temperature": max(0.5, 1.0 - 0.5 * progress), # Decreases over time # Focus balance (0 = equal focus, 1 = focus on negatives) "negative_focus": 0.5 + 0.3 * progress } # ---------------------------- # Enhanced Dataset with Advanced Curriculum Learning # ---------------------------- class SignatureDataset(Dataset): """ Advanced signature dataset with curriculum learning and mining statistics. """ def __init__( self, folder_img: str, excel_data: pd.DataFrame, curriculum_manager: CurriculumLearningManager, transform: Optional[transforms.Compose] = None, is_train: bool = True, config: TrainingConfig = CONFIG ): self.folder_img = folder_img self.is_train = is_train self.config = config self.curriculum_manager = curriculum_manager self.transform = transform or self._default_transforms() self.excel_data = excel_data.reset_index(drop=True) self.current_epoch = 0 # Data preparation self._handle_excel_person_ids() self._categorize_difficulty() # Curriculum learning data self.epoch_data = [] self._prepare_epoch_data() def _handle_excel_person_ids(self): """Properly separate genuine vs forged signature IDs with compact offset.""" # Map genuine person IDs to 0, 1, 2, ... genuine_ids = pd.concat([ self.excel_data["anchor_id"], self.excel_data[self.excel_data["easy_or_hard"] == "easy"]["negative_id"] ]).unique() self.genuine_id_mapping = {val: idx for idx, val in enumerate(genuine_ids)} max_genuine_id = len(genuine_ids) # Create forgery ID space with SMALLER offset (just enough to avoid collisions) forged_data = self.excel_data[self.excel_data["easy_or_hard"] == "hard"] if len(forged_data) > 0: unique_forged_persons = forged_data["negative_id"].unique() self.forgery_id_mapping = { val: idx + max_genuine_id + 100 # Smaller offset: 100 instead of 1000 for idx, val in enumerate(unique_forged_persons) } else: self.forgery_id_mapping = {} # Apply mappings self.excel_data["anchor_id"] = self.excel_data["anchor_id"].map(self.genuine_id_mapping) # Handle negatives based on type new_negative_ids = [] for idx, row in self.excel_data.iterrows(): if row["easy_or_hard"] == "easy": # Genuine different person: use regular ID new_negative_ids.append(self.genuine_id_mapping[row["negative_id"]]) else: # Forged signature: use offset ID to prevent clustering with genuine new_negative_ids.append(self.forgery_id_mapping[row["negative_id"]]) self.excel_data["negative_id"] = new_negative_ids print(f"ID mapping: Genuine IDs 0-{max_genuine_id-1}, Forgery IDs {max_genuine_id+100}-{max_genuine_id+100+len(self.forgery_id_mapping)-1}") def _categorize_difficulty(self): """Categorize samples by difficulty if not already done.""" if self.is_train and "easy_or_hard" in self.excel_data.columns: self.easy_df = self.excel_data[self.excel_data["easy_or_hard"] == "easy"] self.hard_df = self.excel_data[self.excel_data["easy_or_hard"] == "hard"] else: # All samples treated as medium difficulty self.easy_df = self.excel_data self.hard_df = pd.DataFrame() # Empty hard samples def _prepare_epoch_data(self): """Prepare data for current epoch based on curriculum.""" if not self.is_train: # Validation data preparation with better error handling if "image_1_path" in self.excel_data.columns and "image_2_path" in self.excel_data.columns: # Standard pair format required_cols = ["image_1_path", "image_2_path", "label"] # Find ID columns id_cols = [col for col in self.excel_data.columns if "id" in col.lower()] if len(id_cols) >= 2: required_cols.extend(id_cols[-2:]) # Take last 2 ID columns else: # Create dummy IDs if none exist self.excel_data["dummy_id1"] = 0 self.excel_data["dummy_id2"] = 1 required_cols.extend(["dummy_id1", "dummy_id2"]) self.epoch_data = self.excel_data[required_cols].values.tolist() else: # Fallback: try to use all available columns print(f"Warning: Expected validation columns not found. Available: {list(self.excel_data.columns)}") self.epoch_data = self.excel_data.values.tolist() print(f"Validation data prepared: {len(self.epoch_data)} samples") return # Training data preparation (unchanged) hard_ratio = self.curriculum_manager.get_hard_ratio(self.current_epoch) if len(self.hard_df) > 0: n_total = len(self.excel_data) n_hard = int(n_total * hard_ratio) n_easy = n_total - n_hard hard_sample = self.hard_df.sample( n=min(n_hard, len(self.hard_df)), random_state=self.current_epoch, replace=(n_hard > len(self.hard_df)) ) easy_sample = self.easy_df.sample( n=min(n_easy, len(self.easy_df)), random_state=self.current_epoch, replace=(n_easy > len(self.easy_df)) ) epoch_df = pd.concat([hard_sample, easy_sample]).sample( frac=1, random_state=self.current_epoch ).reset_index(drop=True) print(f"Epoch {self.current_epoch}: {len(hard_sample)} hard + {len(easy_sample)} easy = {len(epoch_df)} total (target ratio: {hard_ratio:.2f})") else: epoch_df = self.excel_data.sample( frac=1, random_state=self.current_epoch ).reset_index(drop=True) required_cols = ["anchor_path", "positive_path", "negative_path", "anchor_id", "negative_id"] missing_cols = [col for col in required_cols if col not in epoch_df.columns] if missing_cols: raise ValueError(f"Missing required training columns: {missing_cols}") self.epoch_data = epoch_df[required_cols].values.tolist() def set_epoch(self, epoch: int): """Update epoch and regenerate data.""" self.current_epoch = epoch self._prepare_epoch_data() def get_curriculum_stats(self) -> Dict[str, Any]: """Get current curriculum learning statistics.""" hard_ratio = self.curriculum_manager.get_hard_ratio(self.current_epoch) mining_params = self.curriculum_manager.get_mining_difficulty(self.current_epoch) return { "epoch": self.current_epoch, "hard_ratio": hard_ratio, "easy_ratio": 1.0 - hard_ratio, "total_samples": len(self.epoch_data), **mining_params } def __len__(self) -> int: return len(self.epoch_data) def __getitem__(self, index: int) -> Tuple[torch.Tensor, ...]: if self.is_train: return self._get_train_item(index) else: return self._get_val_item(index) def _get_train_item(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: """Return triplet: anchor, positive, negative with their IDs.""" anchor_path, positive_path, negative_path, pid, nid = self.epoch_data[index] anchor = self._load_image(anchor_path) positive = self._load_image(positive_path) negative = self._load_image(negative_path) return anchor, positive, negative, int(pid), int(nid) def _get_val_item(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: """Return: img1, img2, label, id1, id2.""" data_row = self.epoch_data[index] # Handle different data formats robustly if len(data_row) >= 5: img1_path, img2_path, label, id1, id2 = data_row[:5] elif len(data_row) >= 3: img1_path, img2_path, label = data_row[:3] # Fallback IDs id1, id2 = 0, 1 else: raise ValueError(f"Invalid validation data format: expected at least 3 columns, got {len(data_row)}") try: img1 = self._load_image(img1_path) img2 = self._load_image(img2_path) return img1, img2, torch.tensor(float(label), dtype=torch.float32), int(id1), int(id2) except Exception as e: print(f"Error loading validation item {index}: {e}") print(f"Data row: {data_row}") raise def _load_image(self, path: str) -> torch.Tensor: """Load and transform image.""" image = replace_background_with_white( path, self.folder_img, remove_bg=self.config.remove_bg ) return self.transform(image) if self.transform else image def _default_transforms(self) -> transforms.Compose: """Get default transforms with configurable augmentation strength.""" normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ) if self.is_train: aug_strength = self.config.augmentation_strength return transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p=0.5 * aug_strength), transforms.RandomRotation(degrees=int(10 * aug_strength)), transforms.ColorJitter( brightness=0.2 * aug_strength, contrast=0.2 * aug_strength ), transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0 * aug_strength)), transforms.ToTensor(), normalize ]) return transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), normalize ]) # ---------------------------- # Enhanced Model Architecture # ---------------------------- class ResNetBackbone(nn.Module): """Enhanced ResNet backbone with better weight loading.""" def __init__(self, model_name: str = "resnet34", pretrained_path: Optional[str] = None): super().__init__() # Initialize the ResNet model if model_name == "resnet18": self.resnet = models.resnet18(weights=None) elif model_name == "resnet34": self.resnet = models.resnet34(weights=None) elif model_name == "resnet50": self.resnet = models.resnet50(weights=None) else: raise ValueError(f"Unsupported model_name: {model_name}") # Load pretrained weights if pretrained_path and os.path.exists(pretrained_path): self._load_pretrained_weights(pretrained_path) elif pretrained_path: print(f"Warning: Pretrained path {pretrained_path} not found, using random initialization") # Remove the fully connected layer self.resnet.fc = nn.Identity() # Get output dimension with torch.no_grad(): dummy = torch.randn(1, 3, 224, 224) self.output_dim = self.resnet(dummy).shape[1] def _load_pretrained_weights(self, pretrained_path: str): """Load pretrained weights with comprehensive error handling.""" try: checkpoint = torch.load(pretrained_path, map_location="cpu", weights_only=False) state_dict = checkpoint.get("state_dict", checkpoint) # Handle prefix issues if not any(key.startswith("resnet.") for key in state_dict.keys()): state_dict = {f"resnet.{k}": v for k, v in state_dict.items()} # Filter matching keys and sizes model_dict = self.state_dict() filtered_state_dict = { k: v for k, v in state_dict.items() if k in model_dict and v.size() == model_dict[k].size() } # Load filtered weights missing_keys = self.load_state_dict(filtered_state_dict, strict=False) print(f"[INFO] Loaded pretrained weights: {len(filtered_state_dict)}/{len(model_dict)} parameters") if missing_keys.missing_keys: print(f"[INFO] Missing keys: {len(missing_keys.missing_keys)}") except Exception as e: print(f"[ERROR] Failed to load pretrained weights: {e}") raise def forward(self, x: torch.Tensor) -> torch.Tensor: return self.resnet(x) class AdvancedEmbeddingHead(nn.Module): """Advanced embedding head with residual connections and normalization.""" def __init__(self, input_dim: int, embedding_dim: int, dropout: float = 0.5): super().__init__() self.input_dim = input_dim self.embedding_dim = embedding_dim # Multi-layer embedding head with residual connections if input_dim > embedding_dim * 4: hidden_dim = max(embedding_dim * 2, input_dim // 4) self.layers = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim, embedding_dim * 2), nn.LayerNorm(embedding_dim * 2), nn.GELU(), nn.Dropout(dropout / 2), nn.Linear(embedding_dim * 2, embedding_dim), nn.LayerNorm(embedding_dim) ) else: # Simple head for smaller dimensions self.layers = nn.Sequential( nn.Linear(input_dim, embedding_dim), nn.LayerNorm(embedding_dim) ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x.flatten(1) # Flatten spatial dimensions return self.layers(x) class SiameseSignatureNetwork(nn.Module): """Advanced Siamese network with precision-focused loss.""" def __init__(self, config: TrainingConfig = CONFIG): super().__init__() self.config = config # Initialize backbone if config.model_name.startswith("resnet"): self.backbone = ResNetBackbone( model_name=config.model_name, pretrained_path=config.pretrained_path if config.last_epoch_weights is None else None ) backbone_dim = self.backbone.output_dim else: raise ValueError(f"Unsupported model: {config.model_name}") # Initialize embedding head self.embedding_head = AdvancedEmbeddingHead( input_dim=backbone_dim, embedding_dim=config.embedding_dim, dropout=0.5 ) self.normalize_embeddings = config.normalize_embeddings self.distance_threshold = 0.5 # Will be updated during validation # Loss components self.criterion = MultiSimilarityLoss( alpha=config.multisim_alpha, beta=config.multisim_beta, base=config.multisim_base ) # Add triplet margin loss for better separation self.triplet_loss = nn.TripletMarginLoss( margin=config.triplet_margin, p=2, reduction='none' # We'll apply weights manually ) # Loss weights self.triplet_weight = config.triplet_weight self.fp_penalty_weight = config.false_positive_penalty_weight # Mining if config.use_hard_mining: self.miner = BatchHardMiner() else: self.miner = None def get_parameter_groups(self) -> List[Dict[str, Any]]: """Get parameter groups for differential learning rates.""" backbone_params = list(self.backbone.parameters()) head_params = list(self.embedding_head.parameters()) return [ { 'params': backbone_params, 'lr': self.config.backbone_lr, 'name': 'backbone', 'weight_decay': self.config.weight_decay }, { 'params': head_params, 'lr': self.config.head_lr, 'name': 'embedding_head', 'weight_decay': self.config.weight_decay } ] def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: Optional[torch.Tensor] = None) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Forward pass for training or inference.""" a_features = self.backbone(anchor) a_emb = self.embedding_head(a_features) p_features = self.backbone(positive) p_emb = self.embedding_head(p_features) if self.normalize_embeddings: a_emb = F.normalize(a_emb, p=2, dim=1) p_emb = F.normalize(p_emb, p=2, dim=1) if negative is not None: n_features = self.backbone(negative) n_emb = self.embedding_head(n_features) if self.normalize_embeddings: n_emb = F.normalize(n_emb, p=2, dim=1) return a_emb, p_emb, n_emb return a_emb, p_emb def compute_loss(self, embeddings: torch.Tensor, labels: torch.Tensor, anchors: Optional[torch.Tensor] = None, positives: Optional[torch.Tensor] = None, negatives: Optional[torch.Tensor] = None, distance_weights: Optional[Dict[str, torch.Tensor]] = None) -> torch.Tensor: """Enhanced loss computation with precision focus and distance weighting.""" # MultiSimilarity loss if self.miner is not None: hard_pairs = self.miner(embeddings, labels) ms_loss = self.criterion(embeddings, labels, hard_pairs) else: ms_loss = self.criterion(embeddings, labels) total_loss = ms_loss # Add triplet loss if embeddings provided if anchors is not None and positives is not None and negatives is not None: # Compute triplet losses for each sample triplet_losses = self.triplet_loss(anchors, positives, negatives) # Apply distance-based weights if provided if distance_weights is not None: neg_weights = distance_weights.get('negative_weights', torch.ones_like(triplet_losses)) weighted_triplet_loss = (triplet_losses * neg_weights).mean() else: weighted_triplet_loss = triplet_losses.mean() total_loss += self.triplet_weight * weighted_triplet_loss # Additional penalty for hard negatives (false positives) with torch.no_grad(): d_an = F.pairwise_distance(anchors, negatives) # Find negatives that are too close (potential false positives) hard_negative_mask = d_an < self.distance_threshold if hard_negative_mask.any(): # Apply distance-based weights for false positive penalty if distance_weights is not None: neg_weights = distance_weights.get('negative_weights', torch.ones_like(d_an)) # Extra penalty weighted by how bad the false positive is false_positive_distances = self.distance_threshold - d_an[hard_negative_mask] false_positive_weights = neg_weights[hard_negative_mask] fp_loss = (false_positive_distances * false_positive_weights).mean() else: fp_loss = (self.distance_threshold - d_an[hard_negative_mask]).mean() total_loss += self.fp_penalty_weight * fp_loss return total_loss def predict_pair(self, img1: torch.Tensor, img2: torch.Tensor, threshold: Optional[float] = None, return_dist: bool = False) -> torch.Tensor: """Predict similarity between image pairs.""" self.eval() with torch.no_grad(): emb1, emb2 = self(img1, img2) distances = F.pairwise_distance(emb1, emb2) if return_dist: return distances thresh = threshold if threshold is not None else self.distance_threshold return (distances < thresh).long() # ---------------------------- # Advanced Training Metrics and Statistics # ---------------------------- class TrainingMetrics: """Enhanced training metrics with adaptive mining for both hard positives and negatives.""" def __init__(self): self.reset() # Track distance statistics for adaptive thresholds self.distance_history = {"positive": [], "negative": []} self.adaptive_stats = {} def reset(self): """Reset all metrics.""" self.losses = [] self.genuine_distances = [] self.forged_distances = [] # Separate mining stats for positives and negatives self.positive_mining_stats = {"easy": 0, "semi": 0, "hard": 0} self.negative_mining_stats = {"easy": 0, "semi": 0, "hard": 0} # Hard sample counts self.hard_positive_count = 0 self.hard_negative_count = 0 self.total_positive_pairs = 0 self.total_negative_pairs = 0 # False positive/negative tracking self.false_positive_count = 0 self.false_negative_count = 0 self.learning_rates = {} def update_distance_statistics(self, d_positive: np.ndarray, d_negative: np.ndarray): """Update running statistics for adaptive thresholds.""" # Keep rolling window of recent distances self.distance_history["positive"].extend(d_positive.tolist()) self.distance_history["negative"].extend(d_negative.tolist()) # Keep only recent history (last 5000 samples) for key in self.distance_history: if len(self.distance_history[key]) > 5000: self.distance_history[key] = self.distance_history[key][-5000:] # Compute adaptive statistics if len(self.distance_history["positive"]) > 100 and len(self.distance_history["negative"]) > 100: pos_distances = np.array(self.distance_history["positive"]) neg_distances = np.array(self.distance_history["negative"]) self.adaptive_stats = { "pos_mean": np.mean(pos_distances), "pos_std": np.std(pos_distances), "pos_q25": np.percentile(pos_distances, 25), "pos_q50": np.percentile(pos_distances, 50), "pos_q75": np.percentile(pos_distances, 75), "pos_q90": np.percentile(pos_distances, 90), "neg_mean": np.mean(neg_distances), "neg_std": np.std(neg_distances), "neg_q10": np.percentile(neg_distances, 10), "neg_q25": np.percentile(neg_distances, 25), "neg_q50": np.percentile(neg_distances, 50), "neg_q75": np.percentile(neg_distances, 75), "separation": np.mean(neg_distances) - np.mean(pos_distances), "overlap_region": max(0, np.percentile(pos_distances, 95) - np.percentile(neg_distances, 5)) } def compute_precision_focused_weights(self, d_positive: np.ndarray, d_negative: np.ndarray, negative_weight_multiplier: float = 2.5) -> Tuple[torch.Tensor, torch.Tensor]: """Compute sample weights with focus on improving precision.""" pos_weights = np.ones_like(d_positive) neg_weights = np.ones_like(d_negative) if self.adaptive_stats: # Hard negatives (forged that look genuine) get MUCH higher weight neg_q10 = self.adaptive_stats["neg_q10"] neg_q25 = self.adaptive_stats["neg_q25"] # Very hard negatives (bottom 10%) - highest weight very_hard_neg_mask = d_negative < neg_q10 neg_weights[very_hard_neg_mask] = negative_weight_multiplier * 1.5 # Hard negatives (10-25%) - high weight hard_neg_mask = (d_negative >= neg_q10) & (d_negative < neg_q25) neg_weights[hard_neg_mask] = negative_weight_multiplier # Semi-hard negatives (25-50%) - moderate weight semi_neg_mask = (d_negative >= neg_q25) & (d_negative < self.adaptive_stats["neg_q50"]) neg_weights[semi_neg_mask] = negative_weight_multiplier * 0.6 # Hard positives get moderate weight (but less than hard negatives) pos_q75 = self.adaptive_stats["pos_q75"] pos_q90 = self.adaptive_stats["pos_q90"] # Very hard positives (top 10%) very_hard_pos_mask = d_positive > pos_q90 pos_weights[very_hard_pos_mask] = 1.8 # Hard positives (75-90%) hard_pos_mask = (d_positive > pos_q75) & (d_positive <= pos_q90) pos_weights[hard_pos_mask] = 1.5 return torch.tensor(pos_weights, dtype=torch.float32), torch.tensor(neg_weights, dtype=torch.float32) def update_mining_stats(self, d_positive: np.ndarray, d_negative: np.ndarray, margin: float, difficulty_params: Dict[str, float]): """Intelligent adaptive mining for both hard positives and hard negatives.""" # Update distance statistics first self.update_distance_statistics(d_positive, d_negative) # Update totals self.total_positive_pairs += len(d_positive) self.total_negative_pairs += len(d_negative) # Use adaptive thresholds if available, otherwise fallback to fixed if self.adaptive_stats: self._adaptive_mining(d_positive, d_negative, difficulty_params) else: self._fixed_mining(d_positive, d_negative, margin) def _adaptive_mining(self, d_positive: np.ndarray, d_negative: np.ndarray, difficulty_params: Dict[str, float]): """Adaptive mining based on current distance distributions.""" stats = self.adaptive_stats # Get difficulty parameters hard_positive_ratio = difficulty_params.get("hard_positive_ratio", 0.3) hard_negative_ratio = difficulty_params.get("hard_negative_ratio", 0.3) # Dynamic thresholds for hard positives (far apart genuine pairs) # Use percentile based on desired hard positive ratio hard_pos_percentile = 100 - (hard_positive_ratio * 100) hard_pos_threshold = np.percentile(self.distance_history["positive"][-1000:], hard_pos_percentile) semi_pos_threshold = stats["pos_q50"] # Dynamic thresholds for hard negatives (close together impostor pairs) # Use percentile based on desired hard negative ratio hard_neg_percentile = hard_negative_ratio * 100 hard_neg_threshold = np.percentile(self.distance_history["negative"][-1000:], hard_neg_percentile) semi_neg_threshold = stats["neg_q50"] # Mine hard positives for dp in d_positive: if dp >= hard_pos_threshold: self.positive_mining_stats["hard"] += 1 self.hard_positive_count += 1 elif dp >= semi_pos_threshold: self.positive_mining_stats["semi"] += 1 else: self.positive_mining_stats["easy"] += 1 # Mine hard negatives for dn in d_negative: if dn <= hard_neg_threshold: self.negative_mining_stats["hard"] += 1 self.hard_negative_count += 1 elif dn <= semi_neg_threshold: self.negative_mining_stats["semi"] += 1 else: self.negative_mining_stats["easy"] += 1 def _fixed_mining(self, d_positive: np.ndarray, d_negative: np.ndarray, margin: float): """Fallback fixed mining for early epochs.""" # Fixed thresholds hard_pos_threshold = 0.5 # Far genuine pairs hard_neg_threshold = 0.3 # Close impostor pairs for dp in d_positive: if dp >= hard_pos_threshold: self.positive_mining_stats["hard"] += 1 self.hard_positive_count += 1 elif dp >= hard_pos_threshold * 0.7: self.positive_mining_stats["semi"] += 1 else: self.positive_mining_stats["easy"] += 1 for dn in d_negative: if dn <= hard_neg_threshold: self.negative_mining_stats["hard"] += 1 self.hard_negative_count += 1 elif dn <= hard_neg_threshold * 1.5: self.negative_mining_stats["semi"] += 1 else: self.negative_mining_stats["easy"] += 1 def get_mining_percentages(self) -> Dict[str, float]: """Get mining statistics as percentages with debugging info.""" total_pos = sum(self.positive_mining_stats.values()) total_neg = sum(self.negative_mining_stats.values()) percentages = {} # Positive pair mining stats if total_pos > 0: percentages.update({ "pos_mining_easy_pct": 100.0 * self.positive_mining_stats["easy"] / total_pos, "pos_mining_semi_pct": 100.0 * self.positive_mining_stats["semi"] / total_pos, "pos_mining_hard_pct": 100.0 * self.positive_mining_stats["hard"] / total_pos, }) else: percentages.update({ "pos_mining_easy_pct": 0.0, "pos_mining_semi_pct": 0.0, "pos_mining_hard_pct": 0.0, }) # Negative pair mining stats if total_neg > 0: percentages.update({ "neg_mining_easy_pct": 100.0 * self.negative_mining_stats["easy"] / total_neg, "neg_mining_semi_pct": 100.0 * self.negative_mining_stats["semi"] / total_neg, "neg_mining_hard_pct": 100.0 * self.negative_mining_stats["hard"] / total_neg, }) else: percentages.update({ "neg_mining_easy_pct": 0.0, "neg_mining_semi_pct": 0.0, "neg_mining_hard_pct": 0.0, }) # Overall hard sample ratios if self.total_positive_pairs > 0: percentages["hard_positive_ratio"] = 100.0 * self.hard_positive_count / self.total_positive_pairs else: percentages["hard_positive_ratio"] = 0.0 if self.total_negative_pairs > 0: percentages["hard_negative_ratio"] = 100.0 * self.hard_negative_count / self.total_negative_pairs else: percentages["hard_negative_ratio"] = 0.0 # False positive/negative rates total_samples = self.total_positive_pairs + self.total_negative_pairs if total_samples > 0: percentages["false_positive_rate"] = 100.0 * self.false_positive_count / self.total_negative_pairs if self.total_negative_pairs > 0 else 0.0 percentages["false_negative_rate"] = 100.0 * self.false_negative_count / self.total_positive_pairs if self.total_positive_pairs > 0 else 0.0 # Add adaptive stats if available if self.adaptive_stats: percentages.update({ "adaptive_separation": self.adaptive_stats["separation"], "adaptive_overlap": self.adaptive_stats["overlap_region"], "adaptive_pos_spread": self.adaptive_stats["pos_std"], "adaptive_neg_spread": self.adaptive_stats["neg_std"], }) return percentages def compute_separation_metrics(self) -> Dict[str, float]: """Compute distance separation metrics.""" if not self.genuine_distances or not self.forged_distances: return { "genuine_dist_mean": 0.0, "forged_dist_mean": 0.0, "genuine_dist_std": 0.0, "forged_dist_std": 0.0, "separation": 0.0, "overlap": 0.0, "separation_ratio": 0.0, "cohesion_ratio": 0.0 } gen_mean = np.mean(self.genuine_distances) forg_mean = np.mean(self.forged_distances) gen_std = np.std(self.genuine_distances) forg_std = np.std(self.forged_distances) separation = forg_mean - gen_mean overlap = max(0, gen_mean + 2*gen_std - (forg_mean - 2*forg_std)) # Cohesion ratio: how tight are genuine pairs relative to separation cohesion_ratio = gen_std / (separation + 1e-8) return { "genuine_dist_mean": gen_mean, "forged_dist_mean": forg_mean, "genuine_dist_std": gen_std, "forged_dist_std": forg_std, "separation": separation, "overlap": overlap, "separation_ratio": separation / (gen_std + forg_std + 1e-8), "cohesion_ratio": cohesion_ratio } # ---------------------------- # Enhanced Training Loop # ---------------------------- class SignatureTrainer: """Research-grade signature verification trainer.""" def __init__(self, config: TrainingConfig = CONFIG): self.config = config self.device = torch.device(config.device) # Initialize managers self.mlflow_manager = MLFlowManager(tracking_uri=self.config.tracking_uri) self.curriculum_manager = CurriculumLearningManager(config) # Training state self.current_epoch = 0 self.best_eer = float('inf') self.patience_counter = 0 self.global_step = 0 # Setup logging self._setup_logging() def _setup_logging(self): """Setup comprehensive logging.""" logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler('training.log'), logging.StreamHandler() ] ) self.logger = logging.getLogger(__name__) def _prepare_datasets(self) -> Tuple[SignatureDataset, SignatureDataset]: """Prepare training and validation datasets.""" # Load datasets train_data = pd.read_excel("../../data/classify/preprared_data/labels/train_triplets_balanced_v12.xlsx") val_data = pd.read_excel("../../data/classify/preprared_data/labels/valid_pairs_balanced_v12.xlsx") train_dataset = SignatureDataset( folder_img="../../data/classify/preprared_data/images/", excel_data=train_data, curriculum_manager=self.curriculum_manager, is_train=True, config=self.config ) val_dataset = SignatureDataset( folder_img="../../data/classify/preprared_data/images/", excel_data=val_data, curriculum_manager=self.curriculum_manager, is_train=False, config=self.config ) self.logger.info(f"Training samples: {len(train_dataset)}") self.logger.info(f"Validation samples: {len(val_dataset)}") return train_dataset, val_dataset def _compute_precision_optimized_threshold(self, distances: np.ndarray, labels: np.ndarray, target_precision: float = None) -> float: """Find threshold that achieves target precision while maximizing F1.""" if target_precision is None: target_precision = self.config.target_precision thresholds = np.linspace(distances.min(), distances.max(), 1000) best_threshold = thresholds[0] best_f1 = 0 best_precision = 0 best_recall = 0 for thresh in thresholds: predictions = (distances < thresh).astype(int) # Calculate metrics tp = np.sum((predictions == 1) & (labels == 1)) fp = np.sum((predictions == 1) & (labels == 0)) fn = np.sum((predictions == 0) & (labels == 1)) precision = tp / (tp + fp + 1e-8) recall = tp / (tp + fn + 1e-8) f1 = 2 * precision * recall / (precision + recall + 1e-8) # Prioritize precision while maintaining reasonable recall if precision >= target_precision and f1 > best_f1: best_f1 = f1 best_threshold = thresh best_precision = precision best_recall = recall # If we can't achieve target precision, get best precision with recall > 0.5 elif precision > best_precision and recall > 0.5: best_f1 = f1 best_threshold = thresh best_precision = precision best_recall = recall print(f" Precision-optimized threshold: {best_threshold:.4f} " f"(P: {best_precision:.3f}, R: {best_recall:.3f}, F1: {best_f1:.3f})") return best_threshold def _setup_model_and_optimizer(self) -> Tuple[SiameseSignatureNetwork, torch.optim.Optimizer, Any]: """Setup model, optimizer, and scheduler.""" # Initialize model model = SiameseSignatureNetwork(self.config) # Compile model if available if hasattr(torch, "compile") and platform.system() != "Windows": self.logger.info("Compiling model with torch.compile") model = torch.compile(model) model = model.to(self.device) # Count parameters total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) self.logger.info(f"Total parameters: {total_params:,}") self.logger.info(f"Trainable parameters: {trainable_params:,}") # Setup optimizer with parameter groups param_groups = model.get_parameter_groups() optimizer = torch.optim.AdamW(param_groups) # Log learning rates for group in param_groups: self.logger.info(f"Parameter group '{group['name']}': LR = {group['lr']:.2e}") # Setup scheduler if self.config.lr_scheduler == "cosine": scheduler = CosineAnnealingLR(optimizer, T_max=self.config.max_epochs) else: scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5) return model, optimizer, scheduler def _setup_checkpoint_management(self, run_id: str) -> Tuple[str, str]: """Setup checkpoint directories.""" checkpoint_dir = os.path.join("../../model/models_checkpoints/", run_id) figures_dir = os.path.join(checkpoint_dir, "figures") os.makedirs(checkpoint_dir, exist_ok=True) os.makedirs(figures_dir, exist_ok=True) return checkpoint_dir, figures_dir def _load_checkpoint(self, model: nn.Module, optimizer: torch.optim.Optimizer, scheduler: Any, scaler: torch.amp.GradScaler) -> int: """Load checkpoint if specified.""" if not self.config.last_epoch_weights: return 1 checkpoint_path = self.config.last_epoch_weights self.logger.info(f"Loading checkpoint from {checkpoint_path}") try: checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False) model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) scaler.load_state_dict(checkpoint.get("scaler_state_dict", scaler.state_dict())) start_epoch = checkpoint["epoch"] + 1 self.best_eer = checkpoint.get("best_eer", self.best_eer) model.distance_threshold = checkpoint.get("prediction_threshold", 0.5) self.logger.info(f"Resumed from epoch {start_epoch}, best EER: {self.best_eer:.4f}") return start_epoch except Exception as e: self.logger.error(f"Failed to load checkpoint: {e}") return 1 def train_epoch(self, model: nn.Module, train_loader: DataLoader, optimizer: torch.optim.Optimizer, scaler: torch.amp.GradScaler, epoch: int) -> TrainingMetrics: """Enhanced training with intelligent adaptive mining for both hard positives and negatives.""" model.train() metrics = TrainingMetrics() curriculum_stats = train_loader.dataset.get_curriculum_stats() # INTELLIGENT MARGIN CALCULATION base_margin = 0.5 # Base margin for normalized embeddings margin_multiplier = curriculum_stats["margin_multiplier"] adaptive_margin = base_margin * margin_multiplier # Progressive margin adjustment based on epoch epoch_progress = epoch / self.config.max_epochs progressive_factor = 1.2 - 0.4 * epoch_progress # 1.2 → 0.8 final_margin = adaptive_margin * progressive_factor # Tracking counters forgery_batch_count = 0 genuine_batch_count = 0 batch_fp_count = 0 batch_fn_count = 0 # Debug info debug_printed = False pbar = tqdm(train_loader, desc=f"[Train] Epoch {epoch}") for step, (anchors, positives, negatives, anchor_ids, negative_ids) in enumerate(pbar): # Move to device anchors = anchors.to(self.device, non_blocking=True) positives = positives.to(self.device, non_blocking=True) negatives = negatives.to(self.device, non_blocking=True) anchor_ids = anchor_ids.to(self.device, non_blocking=True) negative_ids = negative_ids.to(self.device, non_blocking=True) # Count forgery vs genuine negatives max_genuine_id = len(train_loader.dataset.genuine_id_mapping) forgery_mask = negative_ids >= max_genuine_id + 100 forgery_batch_count += forgery_mask.sum().item() genuine_batch_count += (~forgery_mask).sum().item() if not debug_printed and step == 0: print(f"\n[DEBUG Epoch {epoch}]") print(f" Final margin: {final_margin:.3f}") print(f" Hard negative ratio target: {curriculum_stats['hard_negative_ratio']:.3f}") print(f" Hard positive ratio target: {curriculum_stats['hard_positive_ratio']:.3f}") print(f" Negative weight multiplier: {self.config.negative_weight_multiplier:.2f}") print(f" Triplet weight: {self.config.triplet_weight:.2f}") print(f" FP penalty weight: {self.config.false_positive_penalty_weight:.2f}") debug_printed = True # Forward pass to get embeddings first with torch.amp.autocast(device_type=self.device.type): a_emb, p_emb, n_emb = model(anchors, positives, negatives) # Compute distances and weights BEFORE loss computation with torch.no_grad(): d_ap = F.pairwise_distance(a_emb, p_emb).cpu().numpy() d_an = F.pairwise_distance(a_emb, n_emb).cpu().numpy() # Get precision-focused weights pos_weights, neg_weights = metrics.compute_precision_focused_weights( d_ap, d_an, negative_weight_multiplier=self.config.negative_weight_multiplier ) pos_weights = pos_weights.to(self.device) neg_weights = neg_weights.to(self.device) # Track false positives/negatives fp_mask = d_an < model.distance_threshold fn_mask = d_ap > model.distance_threshold batch_fp_count = fp_mask.sum() batch_fn_count = fn_mask.sum() metrics.false_positive_count += batch_fp_count metrics.false_negative_count += batch_fn_count # Prepare distance weights for loss distance_weights = { 'positive_weights': pos_weights, 'negative_weights': neg_weights } # Now compute loss with weights with torch.amp.autocast(device_type=self.device.type): all_embeddings = torch.cat([a_emb, p_emb, n_emb], dim=0) all_labels = torch.cat([anchor_ids, anchor_ids, negative_ids], dim=0) # Compute loss with triplet component and distance weights batch_loss = model.compute_loss( all_embeddings, all_labels, anchors=a_emb, positives=p_emb, negatives=n_emb, distance_weights=distance_weights ) # Gradient accumulation loss = batch_loss / self.config.grad_accum_steps scaler.scale(loss).backward() if (step + 1) % self.config.grad_accum_steps == 0 or (step + 1) == len(train_loader): scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) self.global_step += 1 # Update metrics metrics.losses.append(batch_loss.item()) metrics.genuine_distances.extend(d_ap.tolist()) metrics.forged_distances.extend(d_an.tolist()) # Use enhanced mining with difficulty parameters metrics.update_mining_stats(d_ap, d_an, final_margin, curriculum_stats) # Store learning rates for i, group in enumerate(optimizer.param_groups): metrics.learning_rates[f"lr_{group.get('name', i)}"] = group['lr'] # Enhanced progress bar with precision focus sep = np.mean(d_an) - np.mean(d_ap) actual_forgery_ratio = forgery_batch_count / (forgery_batch_count + genuine_batch_count) if (forgery_batch_count + genuine_batch_count) > 0 else 0 # Get current mining stats mining_pcts = metrics.get_mining_percentages() pbar.set_postfix({ "loss": f"{batch_loss.item():.3f}", "h_neg%": f"{mining_pcts.get('neg_mining_hard_pct', 0):.0f}", "h_pos%": f"{mining_pcts.get('pos_mining_hard_pct', 0):.0f}", "d_sep": f"{sep:.3f}", "FP": f"{batch_fp_count}", "FN": f"{batch_fn_count}", "margin": f"{final_margin:.3f}" }) # Periodic logging if self.global_step % self.config.log_frequency == 0: enhanced_stats = { **curriculum_stats, **mining_pcts, "actual_forgery_ratio": actual_forgery_ratio, "batch_false_positives": int(batch_fp_count), "batch_false_negatives": int(batch_fn_count), "final_margin": final_margin, "epoch_progress": epoch_progress } self._log_training_step(metrics, enhanced_stats, self.global_step) # Memory cleanup del anchors, positives, negatives, a_emb, p_emb, n_emb torch.cuda.empty_cache() # End-of-epoch mining summary mining_pcts = metrics.get_mining_percentages() print(f"\n[Epoch {epoch} Mining Summary]") print(f" Hard Negatives: {mining_pcts.get('neg_mining_hard_pct', 0):.1f}% | Semi: {mining_pcts.get('neg_mining_semi_pct', 0):.1f}% | Easy: {mining_pcts.get('neg_mining_easy_pct', 0):.1f}%") print(f" Hard Positives: {mining_pcts.get('pos_mining_hard_pct', 0):.1f}% | Semi: {mining_pcts.get('pos_mining_semi_pct', 0):.1f}% | Easy: {mining_pcts.get('pos_mining_easy_pct', 0):.1f}%") print(f" Overall Hard Ratios - Positives: {mining_pcts.get('hard_positive_ratio', 0):.1f}% | Negatives: {mining_pcts.get('hard_negative_ratio', 0):.1f}%") print(f" False Positive Rate: {mining_pcts.get('false_positive_rate', 0):.1f}% | False Negative Rate: {mining_pcts.get('false_negative_rate', 0):.1f}%") if "adaptive_separation" in mining_pcts: print(f" Adaptive separation: {mining_pcts['adaptive_separation']:.3f} | Overlap: {mining_pcts['adaptive_overlap']:.3f}") return metrics def validate_epoch(self, model: nn.Module, val_loader: DataLoader, epoch: int) -> Tuple[float, float, Dict[str, float]]: """Validate for one epoch.""" model.eval() val_distances = [] val_labels = [] val_embeddings = [] val_person_ids = [] val_loss_total = 0.0 with torch.no_grad(): pbar = tqdm(val_loader, desc=f"[Val] Epoch {epoch}") for img1, img2, labels, id1, id2 in pbar: # Move to device img1 = img1.to(self.device, non_blocking=True) img2 = img2.to(self.device, non_blocking=True) labels = labels.to(self.device, non_blocking=True) id1 = id1.to(self.device, non_blocking=True) id2 = id2.to(self.device, non_blocking=True) # Forward pass emb1, emb2 = model(img1, img2) distances = F.pairwise_distance(emb1, emb2) # Compute validation loss val_loss = self._compute_validation_loss(emb1, emb2, labels, id1, id2, model.criterion) val_loss_total += val_loss.item() # Collect results val_distances.extend(distances.cpu().numpy()) val_labels.extend(labels.cpu().numpy()) val_embeddings.append(emb1.cpu().numpy()) val_embeddings.append(emb2.cpu().numpy()) val_person_ids.extend(id1.cpu().numpy()) val_person_ids.extend(id2.cpu().numpy()) # Update progress pos_mask = labels == 1 neg_mask = labels == 0 pos_dist = distances[pos_mask].mean().item() if pos_mask.any() else 0.0 neg_dist = distances[neg_mask].mean().item() if neg_mask.any() else 0.0 pbar.set_postfix({ "loss": f"{val_loss.item():.4f}", "d_pos": f"{pos_dist:.3f}", "d_neg": f"{neg_dist:.3f}", "sep": f"{neg_dist - pos_dist:.3f}" }) # Memory cleanup del img1, img2, emb1, emb2 torch.cuda.empty_cache() # Process results val_distances = np.array(val_distances) val_labels = np.array(val_labels) val_embeddings = np.concatenate(val_embeddings) val_person_ids = np.array(val_person_ids) avg_val_loss = val_loss_total / len(val_loader) # Compute metrics threshold, eer, metrics_dict = self._compute_validation_metrics( val_distances, val_labels, avg_val_loss ) # Update model threshold model.distance_threshold = threshold return threshold, eer, { "metrics": metrics_dict, "embeddings": val_embeddings, "labels": np.repeat(val_labels, 2), "person_ids": val_person_ids, "distances": np.repeat(val_distances, 2) } def _compute_validation_loss(self, emb1: torch.Tensor, emb2: torch.Tensor, binary_labels: torch.Tensor, person_ids1: torch.Tensor, person_ids2: torch.Tensor, criterion) -> torch.Tensor: """Compute validation loss using MultiSimilarityLoss.""" labels1 = person_ids1.clone() labels2 = person_ids2.clone() # Handle forged pairs forged_mask = binary_labels == 0 if forged_mask.any(): max_person_id = torch.max(torch.cat([person_ids1, person_ids2])).item() labels2[forged_mask] = labels2[forged_mask] + max_person_id + 1 # Handle genuine pairs genuine_mask = binary_labels == 1 labels2[genuine_mask] = labels1[genuine_mask] # Combine embeddings and labels all_embeddings = torch.cat([emb1, emb2], dim=0) all_labels = torch.cat([labels1, labels2], dim=0) return criterion(all_embeddings, all_labels) def _compute_validation_metrics(self, distances: np.ndarray, labels: np.ndarray, val_loss: float) -> Tuple[float, float, Dict[str, float]]: """Compute comprehensive validation metrics with precision focus.""" # Compute EER and threshold similarity_scores = 1.0 / (distances + 1e-8) fpr, tpr, thresholds = roc_curve(labels, similarity_scores, pos_label=1) fnr = 1 - tpr eer_idx = np.nanargmin(np.abs(fpr - fnr)) eer = fpr[eer_idx] eer_threshold = 1.0 / thresholds[eer_idx] # Get precision-optimized threshold precision_threshold = self._compute_precision_optimized_threshold(distances, labels) # Use precision-optimized threshold instead of EER threshold threshold = precision_threshold # Compute metrics with precision-optimized threshold predictions = (distances < threshold).astype(int) precision, recall, f1, _ = precision_recall_fscore_support( labels, predictions, average='binary', zero_division=0 ) accuracy = (predictions == labels).mean() roc_auc = auc(fpr, tpr) # Distance statistics genuine_dist = np.mean([d for d, l in zip(distances, labels) if l == 1]) forged_dist = np.mean([d for d, l in zip(distances, labels) if l == 0]) separation = forged_dist - genuine_dist # Confidence scores confidences = 1.0 / (distances + 1e-8) conf_genuine = np.mean([c for c, l in zip(confidences, labels) if l == 1]) conf_forged = np.mean([c for c, l in zip(confidences, labels) if l == 0]) metrics_dict = { "val_loss": val_loss, "val_EER": eer, "val_f1": f1, "val_auc": roc_auc, "val_accuracy": accuracy, "val_precision": precision, "val_recall": recall, "val_separation": separation, "val_genuine_dist": genuine_dist, "val_forged_dist": forged_dist, "val_genuine_conf": conf_genuine, "val_forged_conf": conf_forged, "threshold": threshold, "eer_threshold": eer_threshold, "precision_threshold": precision_threshold } return threshold, eer, metrics_dict def _log_training_step(self, metrics: TrainingMetrics, curriculum_stats: Dict, step: int): """Log training step metrics.""" if not metrics.losses: return try: # Compute separation metrics sep_metrics = metrics.compute_separation_metrics() # Get mining percentages mining_percentages = metrics.get_mining_percentages() # Log to MLflow log_dict = { "train_loss": np.mean(metrics.losses[-10:]), # Last 10 batches **sep_metrics, **mining_percentages, **curriculum_stats, **metrics.learning_rates } mlflow.log_metrics(log_dict, step=step) except Exception as e: print(f"Warning: Failed to log training step metrics: {e}") def _log_epoch_metrics(self, train_metrics: TrainingMetrics, val_metrics: Dict, epoch: int): """Log comprehensive epoch metrics.""" try: # Training metrics train_sep = train_metrics.compute_separation_metrics() train_mining = train_metrics.get_mining_percentages() log_dict = { "epoch": epoch, "train_loss_epoch": np.mean(train_metrics.losses), **train_sep, **train_mining, **val_metrics["metrics"], **train_metrics.learning_rates } mlflow.log_metrics(log_dict, step=epoch) # Log key metrics to console self.logger.info(f"Epoch {epoch}/{self.config.max_epochs} Summary:") self.logger.info(f" Train Loss: {log_dict['train_loss_epoch']:.4f}") self.logger.info(f" Val EER: {log_dict['val_EER']:.4f}") self.logger.info(f" Val F1: {log_dict['val_f1']:.4f}") self.logger.info(f" Separation: {log_dict['separation']:.4f}") except Exception as e: self.logger.error(f"Failed to log epoch metrics: {e}") # Log minimal metrics as fallback mlflow.log_metrics({ "epoch": epoch, "train_loss_epoch": np.mean(train_metrics.losses) if train_metrics.losses else 0.0, **val_metrics["metrics"] }, step=epoch) def _save_checkpoint(self, model: nn.Module, optimizer: torch.optim.Optimizer, scheduler: Any, scaler: torch.amp.GradScaler, epoch: int, threshold: float, eer: float, checkpoint_dir: str, is_best: bool = False): """Save model checkpoint.""" checkpoint = { "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "scaler_state_dict": scaler.state_dict(), "prediction_threshold": threshold, "best_eer": self.best_eer, "eer": eer, "config": asdict(self.config) } # Save regular checkpoint if epoch % self.config.save_frequency == 0: torch.save(checkpoint, os.path.join(checkpoint_dir, f"epoch_{epoch}.pth")) # Save best checkpoint if is_best: torch.save(checkpoint, os.path.join(checkpoint_dir, "best_model.pth")) self.logger.info(f"New best model saved with EER: {eer:.4f}") def _create_visualizations(self, val_results: Dict, epoch: int, figures_dir: str): """Create comprehensive visualizations.""" if epoch % self.config.visualize_frequency != 0: return # Distance distribution plot self._plot_distance_distribution( val_results["distances"][:len(val_results["distances"])//2], val_results["labels"][:len(val_results["labels"])//2], epoch, figures_dir ) # t-SNE embedding visualization self._plot_tsne_embeddings( val_results["embeddings"], val_results["labels"], val_results["person_ids"], val_results["distances"], epoch, figures_dir ) def _plot_distance_distribution(self, distances: np.ndarray, labels: np.ndarray, epoch: int, figures_dir: str): """Plot distance distribution.""" genuine_dists = distances[labels == 1] forged_dists = distances[labels == 0] plt.figure(figsize=(12, 8)) plt.hist(genuine_dists, bins=50, alpha=0.6, color='blue', label=f'Genuine (μ={np.mean(genuine_dists):.4f}±{np.std(genuine_dists):.4f})') plt.hist(forged_dists, bins=50, alpha=0.6, color='red', label=f'Forged (μ={np.mean(forged_dists):.4f}±{np.std(forged_dists):.4f})') separation = np.mean(forged_dists) - np.mean(genuine_dists) plt.axvline(np.mean(genuine_dists), color='blue', linestyle='--', alpha=0.7) plt.axvline(np.mean(forged_dists), color='red', linestyle='--', alpha=0.7) plt.title(f'Distance Distribution - Epoch {epoch}\nSeparation: {separation:.4f}', fontsize=14) plt.xlabel('Embedding Distance', fontsize=12) plt.ylabel('Frequency', fontsize=12) plt.legend(fontsize=12) plt.grid(alpha=0.3) plt.savefig(os.path.join(figures_dir, f"distance_dist_epoch_{epoch}.png"), dpi=150, bbox_inches='tight') plt.close() def _plot_tsne_embeddings(self, embeddings: np.ndarray, labels: np.ndarray, person_ids: np.ndarray, distances: np.ndarray, epoch: int, figures_dir: str, n_samples: int = 3000): """Create comprehensive t-SNE visualization.""" # Sample for computational efficiency if len(embeddings) > n_samples: indices = np.random.choice(len(embeddings), n_samples, replace=False) embeddings = embeddings[indices] labels = labels[indices] person_ids = person_ids[indices] distances = distances[indices] # Compute t-SNE tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(embeddings)-1)) embeddings_2d = tsne.fit_transform(embeddings) fig, axes = plt.subplots(1, 3, figsize=(20, 6)) # 1. Genuine vs Forged for label_val, color, name in [(0, 'red', 'Forged'), (1, 'blue', 'Genuine')]: mask = labels == label_val if mask.any(): axes[0].scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1], c=color, label=name, alpha=0.6, s=20) axes[0].set_title(f'Genuine vs Forged - Epoch {epoch}') axes[0].legend() axes[0].grid(alpha=0.3) # 2. Person clusters unique_ids = np.unique(person_ids) colors = plt.cm.tab20(np.linspace(0, 1, min(20, len(unique_ids)))) # Show top 15 most frequent IDs id_counts = {pid: np.sum(person_ids == pid) for pid in unique_ids} top_ids = sorted(id_counts.items(), key=lambda x: x[1], reverse=True)[:15] for idx, (pid, count) in enumerate(top_ids): mask = person_ids == pid color = colors[idx % len(colors)] # Plot the cluster points axes[1].scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1], c=[color], label=f'ID {pid} (n={count})', alpha=0.7, s=25) # Compute the centroid (mean position) of the points in this cluster centroid = np.mean(embeddings_2d[mask], axis=0) # Add the person ID text at the centroid axes[1].text(centroid[0], centroid[1], f'ID {pid}', fontsize=10, color='black', alpha=0.8, ha='center') # Plot others in gray other_mask = ~np.isin(person_ids, [pid for pid, _ in top_ids]) if other_mask.any(): axes[1].scatter(embeddings_2d[other_mask, 0], embeddings_2d[other_mask, 1], c='gray', label='Others', alpha=0.3, s=15) axes[1].set_title(f'Person Clusters - Epoch {epoch}') axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=8) axes[1].grid(alpha=0.3) # 3. Distance-based coloring scatter = axes[2].scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=distances, cmap='viridis', alpha=0.7, s=20) plt.colorbar(scatter, ax=axes[2], label='Distance') axes[2].set_title(f'Distance Visualization - Epoch {epoch}') axes[2].grid(alpha=0.3) plt.tight_layout() plt.savefig(os.path.join(figures_dir, f"tsne_epoch_{epoch}.png"), dpi=150, bbox_inches='tight') plt.close() def train(self): """Main training loop.""" torch.backends.cudnn.benchmark = True self.logger.info(f"Starting training on device: {self.device}") # Prepare components train_dataset, val_dataset = self._prepare_datasets() model, optimizer, scheduler = self._setup_model_and_optimizer() scaler = torch.amp.GradScaler(self.device.type, enabled=(self.device.type == "cuda")) # MLflow setup with self.mlflow_manager.start_run(run_id=self.config.run_id): run_id = mlflow.active_run().info.run_id self.mlflow_manager.log_config(self.config) # Setup checkpoints checkpoint_dir, figures_dir = self._setup_checkpoint_management(run_id) # Load checkpoint if specified start_epoch = self._load_checkpoint(model, optimizer, scheduler, scaler) # Data loaders val_loader = DataLoader( val_dataset, batch_size=self.config.batch_size, shuffle=False, num_workers=4, pin_memory=True, prefetch_factor=2 ) # Training loop for epoch in range(start_epoch, self.config.max_epochs + 1): self.current_epoch = epoch # Update curriculum train_dataset.set_epoch(epoch) train_loader = DataLoader( train_dataset, batch_size=self.config.batch_size, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True, prefetch_factor=2 ) # Training phase train_metrics = self.train_epoch(model, train_loader, optimizer, scaler, epoch) # Validation phase threshold, eer, val_results = self.validate_epoch(model, val_loader, epoch) # Logging self._log_epoch_metrics(train_metrics, val_results, epoch) # Visualizations self._create_visualizations(val_results, epoch, figures_dir) # Model checkpoint management is_best = eer < self.best_eer if is_best: self.best_eer = eer self.patience_counter = 0 else: self.patience_counter += 1 self._save_checkpoint( model, optimizer, scheduler, scaler, epoch, threshold, eer, checkpoint_dir, is_best ) # Early stopping if self.patience_counter >= self.config.patience: self.logger.info(f"Early stopping after {self.config.patience} epochs without improvement") break # Learning rate scheduling if self.config.lr_scheduler == "cosine": scheduler.step() else: scheduler.step(eer) # Memory cleanup gc.collect() torch.cuda.empty_cache() # Final logging mlflow.log_metric("final_best_eer", self.best_eer) self.logger.info(f"Training completed. Best EER: {self.best_eer:.4f}") # ---------------------------- # Image Processing Utilities # ---------------------------- def estimate_background_color_pil(image: Image.Image, border_width: int = 10, method: str = "median") -> np.ndarray: """Estimate background color from image borders.""" if image.mode != 'RGB': image = image.convert('RGB') np_img = np.array(image) h, w, _ = np_img.shape # Extract border pixels top = np_img[:border_width, :, :].reshape(-1, 3) bottom = np_img[-border_width:, :, :].reshape(-1, 3) left = np_img[:, :border_width, :].reshape(-1, 3) right = np_img[:, -border_width:, :].reshape(-1, 3) all_border_pixels = np.concatenate([top, bottom, left, right], axis=0) if method == "mean": return np.mean(all_border_pixels, axis=0).astype(np.uint8) else: return np.median(all_border_pixels, axis=0).astype(np.uint8) def replace_background_with_white(image_name: str, folder_img: str, tolerance: int = 40, method: str = "median", remove_bg: bool = False) -> Image.Image: """Replace background with white based on border color estimation.""" image_path = os.path.join(folder_img, image_name) image = Image.open(image_path).convert("RGB") if not remove_bg: return image np_img = np.array(image) bg_color = estimate_background_color_pil(image, method=method) # Create mask for background pixels diff = np.abs(np_img.astype(np.int32) - bg_color.astype(np.int32)) mask = np.all(diff < tolerance, axis=2) # Replace background with white result = np_img.copy() result[mask] = [255, 255, 255] return Image.fromarray(result) # ---------------------------- # Main Execution # ---------------------------- def main(): """Main execution function with aggressive curriculum.""" # Test distance ranges first print("\n[INFO] Testing distance ranges for margin calibration...") dummy_emb1 = F.normalize(torch.randn(1000, 128), p=2, dim=1) dummy_emb2 = F.normalize(torch.randn(1000, 128), p=2, dim=1) dummy_distances = F.pairwise_distance(dummy_emb1, dummy_emb2).numpy() print(f"Random embeddings: mean={dummy_distances.mean():.3f}, std={dummy_distances.std():.3f}") print(f"Expected margin range: {dummy_distances.std() * 0.5:.3f} - {dummy_distances.std() * 1.5:.3f}") # Aggressive curriculum configuration CONFIG.model_name = "resnet34" CONFIG.embedding_dim = 128 CONFIG.max_epochs = 20 # Shorter with aggressive curriculum CONFIG.head_lr = 2e-3 # Higher for faster adaptation CONFIG.backbone_lr = 1e-4 CONFIG.curriculum_strategy = "progressive" # AGGRESSIVE SETTINGS CONFIG.initial_hard_ratio = 0.4 # Start much higher CONFIG.final_hard_ratio = 0.85 # Target very high CONFIG.curriculum_warmup_epochs = 1 # Very short warmup CONFIG.batch_size = 256 # Smaller batches for more frequent updates CONFIG.grad_accum_steps = 8 # Smaller accumulation CONFIG.tracking_uri = "http://127.0.0.1:5555" #CONFIG.run_id = "aa58e3a1f3314351bc1dd2b82ab156ad" #CONFIG.last_epoch_weights = "../../model/models_checkpoints/aa58e3a1f3314351bc1dd2b82ab156ad/best_model.pth" trainer = SignatureTrainer(CONFIG) trainer.train() if __name__ == "__main__": main()