advanced_model.py

#!/usr/bin/env python3
"""
Enhanced OCR Model Architecture
Advanced CNN-Transformer model with multi-scale features, spatial attention, and robust decoding
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import List, Tuple, Optional

# ============================================================
# Residual Block Implementation
# ============================================================

class ResidualBlock(nn.Module):
    """Residual block with batch normalization and ReLU activation"""
    
    def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride, 1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
        # Shortcut connection
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        return F.relu(out)

# ============================================================
# Spatial Attention Mechanism
# ============================================================

class SpatialAttention(nn.Module):
    """Spatial attention mechanism for focusing on important regions"""
    
    def __init__(self, in_channels: int, reduction: int = 16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        
        self.fc = nn.Sequential(
            nn.Conv2d(in_channels, in_channels // reduction, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels // reduction, in_channels, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        attention = self.sigmoid(avg_out + max_out)
        return x * attention

# ============================================================
# Multi-Scale Feature Extraction
# ============================================================

class MultiScaleBlock(nn.Module):
    """Multi-scale feature extraction block"""
    
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.scale1 = nn.Conv2d(in_channels, out_channels // 4, 1, 1, 0)
        self.scale2 = nn.Conv2d(in_channels, out_channels // 4, 3, 1, 1)
        self.scale3 = nn.Conv2d(in_channels, out_channels // 4, 5, 1, 2)
        self.scale4 = nn.Conv2d(in_channels, out_channels // 4, 7, 1, 3)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        s1 = self.scale1(x)
        s2 = self.scale2(x)
        s3 = self.scale3(x)
        s4 = self.scale4(x)
        out = torch.cat([s1, s2, s3, s4], dim=1)
        return self.relu(self.bn(out))

# ============================================================
# Positional Encoding
# ============================================================

class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding for transformer"""
    
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Add positional encoding to input tensor"""
        return x + self.pe[:, :x.size(1), :]

# ============================================================
# Advanced CNN-Transformer Model
# ============================================================

class AdvancedFastPlateOCR(nn.Module):
    """
    Enhanced OCR model with:
    - Multi-scale CNN backbone with spatial attention
    - Transformer decoder with improved attention
    - Positional encoding
    - Robust decoding strategies
    - Confidence-based filtering
    """
    
    def __init__(self, vocab_size: int, hidden: int = 256, num_layers: int = 4, 
                 nhead: int = 8, use_pe: bool = True, dropout: float = 0.1,
                 use_attention: bool = True, use_multiscale: bool = True):
        super().__init__()
        
        self.vocab_size = vocab_size
        self.hidden = hidden
        self.use_pe = use_pe
        self.use_attention = use_attention
        self.use_multiscale = use_multiscale
        
        # Enhanced CNN backbone with multi-scale features and attention
        self.cnn = nn.Sequential(
            # Initial conv block
            nn.Conv2d(3, 64, 7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, stride=2, padding=1),
            
            # Multi-scale blocks with attention
            self._make_enhanced_layer(64, 128, 2, stride=2),
            self._make_enhanced_layer(128, 256, 2, stride=2),
            self._make_enhanced_layer(256, hidden, 2, stride=2),
            
            # Global average pooling
            nn.AdaptiveAvgPool2d((1, None))
        )
        
        # Transformer decoder
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden, 
            nhead=nhead, 
            dim_feedforward=hidden * 4,
            dropout=dropout, 
            activation="relu", 
            batch_first=True
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        
        # Embedding and output layers
        self.embedding = nn.Embedding(vocab_size, hidden)
        self.fc = nn.Linear(hidden, vocab_size)
        
        # Positional encoding
        if use_pe:
            self.pos_enc = PositionalEncoding(hidden)
        
        # Initialize weights
        self._init_weights()
    
    def _make_layer(self, in_channels: int, out_channels: int, blocks: int, stride: int = 1):
        """Create a layer with multiple residual blocks"""
        layers = []
        layers.append(ResidualBlock(in_channels, out_channels, stride))
        for _ in range(1, blocks):
            layers.append(ResidualBlock(out_channels, out_channels))
        return nn.Sequential(*layers)
    
    def _make_enhanced_layer(self, in_channels: int, out_channels: int, blocks: int, stride: int = 1):
        """Create enhanced layer with multi-scale features and attention"""
        layers = []
        
        # First block with multi-scale features
        if self.use_multiscale:
            layers.append(MultiScaleBlock(in_channels, out_channels))
        else:
            layers.append(ResidualBlock(in_channels, out_channels, stride))
        
        # Add spatial attention
        if self.use_attention:
            layers.append(SpatialAttention(out_channels))
        
        # Additional residual blocks
        for _ in range(1, blocks):
            layers.append(ResidualBlock(out_channels, out_channels))
            if self.use_attention:
                layers.append(SpatialAttention(out_channels))
        
        return nn.Sequential(*layers)
    
    def _init_weights(self):
        """Initialize model weights"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, imgs: torch.Tensor, tgt_inp: torch.Tensor) -> torch.Tensor:
        """
        Forward pass
        
        Args:
            imgs: Input images [B, 3, H, W]
            tgt_inp: Target input sequences [B, T]
            
        Returns:
            Logits [B, T, vocab_size]
        """
        # Extract features with CNN
        feats = self.cnn(imgs).flatten(2).permute(0, 2, 1)  # [B, Seq, C]
        
        # Add positional encoding to features
        if self.use_pe:
            feats = self.pos_enc(feats)
        
        # Decoder forward pass
        tgt_emb = self.embedding(tgt_inp)  # [B, T, C]
        out = self.decoder(tgt_emb, feats)  # [B, T, C]
        
        return self.fc(out)  # [B, T, vocab_size]
    
    def greedy_decode(self, imgs: torch.Tensor, sos_id: int = 1, eos_id: int = 2, 
                     max_len: int = 32, device: str = "cpu") -> List[List[int]]:
        """
        Greedy decoding for inference
        
        Args:
            imgs: Input images [B, 3, H, W]
            sos_id: Start of sequence token ID
            eos_id: End of sequence token ID
            max_len: Maximum sequence length
            device: Device to run on
            
        Returns:
            List of predicted sequences
        """
        self.eval()
        with torch.no_grad():
            # Extract features
            feats = self.cnn(imgs).flatten(2).permute(0, 2, 1)
            if self.use_pe:
                feats = self.pos_enc(feats)
            
            B = feats.size(0)
            ys = torch.full((B, 1), sos_id, dtype=torch.long, device=device)
            finished = [False] * B
            
            for _ in range(max_len - 1):
                tgt_emb = self.embedding(ys)
                out = self.decoder(tgt_emb, feats)
                logits = self.fc(out[:, -1, :])
                next_word = logits.argmax(dim=-1, keepdim=True)
                ys = torch.cat([ys, next_word], dim=1)
                
                # Check for EOS tokens
                for i in range(B):
                    if next_word[i].item() == eos_id:
                        finished[i] = True
                
                if all(finished):
                    break
            
            return [ys[i].tolist() for i in range(B)]
    
    def beam_decode(self, imgs: torch.Tensor, beam_width: int = 5, max_len: int = 32,
                    sos_id: int = 1, eos_id: int = 2, device: str = "cpu") -> List[int]:
        """
        Beam search decoding for better accuracy
        
        Args:
            imgs: Input images [B, 3, H, W] (batch size must be 1)
            beam_width: Beam search width
            max_len: Maximum sequence length
            sos_id: Start of sequence token ID
            eos_id: End of sequence token ID
            device: Device to run on
            
        Returns:
            Best predicted sequence
        """
        self.eval()
        with torch.no_grad():
            # Extract features
            feats = self.cnn(imgs).flatten(2).permute(0, 2, 1)
            if self.use_pe:
                feats = self.pos_enc(feats)
            
            B = feats.size(0)
            if B != 1:
                # Fallback to greedy for batch > 1
                return self.greedy_decode(imgs, sos_id, eos_id, max_len, device)[0]
            
            memory = feats[0:1]
            sequences = [(torch.tensor([sos_id], device=device), 0.0)]
            min_len = 4  # Minimum sequence length

            for _ in range(max_len):
                all_candidates = []
                
                for seq, score in sequences:
                    if seq[-1].item() == eos_id:
                        all_candidates.append((seq, score))
                        continue
                    
                    # Get next token probabilities
                    tgt_emb = self.embedding(seq.unsqueeze(0))
                    tgt_mask = nn.Transformer.generate_square_subsequent_mask(seq.size(0)).to(device)
                    out = self.decoder(tgt_emb, memory, tgt_mask=tgt_mask)
                    logits = self.fc(out[:, -1, :])
                    logp = F.log_softmax(logits, dim=-1)[0]
                    
                    # Get top-k candidates
                    topv, topi = torch.topk(logp, k=min(beam_width, logp.size(0)))
                    
                    for k in range(topv.size(0)):
                        token = int(topi[k].item())
                        
                        # Skip EOS if sequence is too short
                        if len(seq) < min_len and token == eos_id:
                            continue
                        
                        new_seq = torch.cat([seq, torch.tensor([token], device=device)])
                        length_penalty = ((5 + len(new_seq)) / 6) ** 0.7
                        new_score = (score + float(topv[k].item())) / length_penalty
                        all_candidates.append((new_seq, new_score))
                
                if not all_candidates:
                    break
                
                # Keep top beam_width sequences
                sequences = sorted(all_candidates, key=lambda t: t[1], reverse=True)[:beam_width]

            return sequences[0][0].tolist()
    
    def robust_decode(self, imgs: torch.Tensor, sos_id: int = 1, eos_id: int = 2, 
                     max_len: int = 32, device: str = "cpu", min_confidence: float = 0.7,
                     method: str = "beam") -> Tuple[List[List[int]], List[float]]:
        """
        Robust decoding with confidence scoring
        
        Args:
            imgs: Input images [B, 3, H, W]
            sos_id: Start of sequence token ID
            eos_id: End of sequence token ID
            max_len: Maximum sequence length
            device: Device to run on
            min_confidence: Minimum confidence threshold
            method: Decoding method ("greedy" or "beam")
            
        Returns:
            Tuple of (predicted_sequences, confidence_scores)
        """
        self.eval()
        with torch.no_grad():
            # Extract features
            feats = self.cnn(imgs).flatten(2).permute(0, 2, 1)
            if self.use_pe:
                feats = self.pos_enc(feats)
            
            B = feats.size(0)
            all_sequences = []
            all_confidences = []
            
            for b in range(B):
                # Get predictions
                if method == "beam":
                    pred_ids = self.beam_decode(imgs[b:b+1], device=device)[0]
                else:
                    pred_ids = self.greedy_decode(imgs[b:b+1], device=device)[0]
                
                # Calculate confidence
                confidence = self._calculate_sequence_confidence(
                    imgs[b:b+1], pred_ids, feats[b:b+1], device
                )
                
                # Filter by confidence
                if confidence >= min_confidence:
                    all_sequences.append(pred_ids)
                    all_confidences.append(confidence)
                else:
                    # Return uncertain prediction
                    all_sequences.append([sos_id, 0, eos_id])  # 0 = uncertain token
                    all_confidences.append(confidence)
            
            return all_sequences, all_confidences
    
    def _calculate_sequence_confidence(self, img: torch.Tensor, pred_ids: List[int], 
                                     feats: torch.Tensor, device: str) -> float:
        """Calculate confidence score for a predicted sequence"""
        try:
            # Convert prediction to tensor
            pred_tensor = torch.tensor(pred_ids, device=device).unsqueeze(0)
            
            # Get model predictions
            tgt_emb = self.embedding(pred_tensor)
            tgt_mask = nn.Transformer.generate_square_subsequent_mask(pred_tensor.size(1)).to(device)
            out = self.decoder(tgt_emb, feats, tgt_mask=tgt_mask)
            logits = self.fc(out)
            
            # Calculate average confidence
            probs = F.softmax(logits, dim=-1)
            max_probs = torch.max(probs, dim=-1)[0]
            
            # Exclude SOS and EOS tokens
            valid_probs = max_probs[0, 1:-1]  # Skip first and last
            if len(valid_probs) > 0:
                confidence = float(torch.mean(valid_probs))
            else:
                confidence = 0.0
            
            return min(confidence, 1.0)  # Cap at 1.0
            
        except Exception:
            return 0.0
    
    def get_model_info(self) -> dict:
        """Get model information"""
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        
        return {
            "vocab_size": self.vocab_size,
            "hidden_dim": self.hidden,
            "use_pe": self.use_pe,
            "total_parameters": total_params,
            "trainable_parameters": trainable_params,
            "model_size_mb": total_params * 4 / (1024 * 1024)  # Assuming float32
        }

# ============================================================
# Model Factory Functions
# ============================================================

def create_model(vocab_size: int, config: dict = None) -> AdvancedFastPlateOCR:
    """
    Create model with given configuration
    
    Args:
        vocab_size: Size of vocabulary
        config: Model configuration dictionary
        
    Returns:
        Initialized model
    """
    default_config = {
        "hidden": 256,
        "num_layers": 4,
        "nhead": 8,
        "use_pe": True,
        "dropout": 0.1
    }
    
    if config:
        default_config.update(config)
    
    return AdvancedFastPlateOCR(vocab_size=vocab_size, **default_config)

def load_model_from_checkpoint(checkpoint_path: str, vocab_size: int, 
                              config: dict = None, device: str = "cpu") -> AdvancedFastPlateOCR:
    """
    Load model from checkpoint
    
    Args:
        checkpoint_path: Path to model checkpoint
        vocab_size: Size of vocabulary
        config: Model configuration
        device: Device to load model on
        
    Returns:
        Loaded model
    """
    model = create_model(vocab_size, config)
    
    checkpoint = torch.load(checkpoint_path, map_location=device)
    
    if "model" in checkpoint:
        model.load_state_dict(checkpoint["model"])
    else:
        model.load_state_dict(checkpoint)
    
    model.to(device)
    model.eval()
    
    return model

# ============================================================
# Testing and Validation
# ============================================================

def test_model():
    """Test model functionality"""
    print("🧪 Testing Advanced OCR Model...")
    
    # Create model
    vocab_size = 39  # 26 letters + 10 digits + 3 special tokens
    model = create_model(vocab_size)
    
    # Test forward pass
    batch_size = 2
    imgs = torch.randn(batch_size, 3, 96, 512)
    tgt = torch.randint(0, vocab_size, (batch_size, 10))
    
    with torch.no_grad():
        output = model(imgs, tgt)
        print(f"✅ Forward pass: {imgs.shape} -> {output.shape}")
    
    # Test greedy decoding
    pred_ids = model.greedy_decode(imgs)
    print(f"✅ Greedy decode: {len(pred_ids)} sequences")
    
    # Test beam search (single image)
    single_img = imgs[0:1]
    beam_pred = model.beam_decode(single_img)
    print(f"✅ Beam search: {len(beam_pred)} tokens")
    
    # Model info
    info = model.get_model_info()
    print(f"✅ Model info: {info['total_parameters']:,} parameters ({info['model_size_mb']:.1f} MB)")
    
    print("🎉 All tests passed!")

if __name__ == "__main__":
    test_model()