model.py

"""
model.py
~~~~~~~~
CRNN (Convolutional Recurrent Neural Network) for CAPTCHA text recognition.

Architecture — Shi et al. 2016, "An End-to-End Trainable Neural Network
for Image-based Sequence Recognition" (modernised with BatchNorm + Dropout):

  Input     (B, H=64, W=200, C=3)
  ──────────────────────────────────────────────────────────────
  CNN       6 conv blocks           → (B, 1, 50, 512)
  Reshape   squeeze height          → (B, T=50, 512)
  BiLSTM×2  bidirectional context   → (B, T=50, 512)
  Dense     per-step projection     → (B, T=50, 63)   ← logits
  ──────────────────────────────────────────────────────────────
  CTC loss  (training)   tf.nn.ctc_loss
  CTC decode(inference)  greedy or beam search

Why CRNN + CTC?
  • CNN  — learns local visual features (stroke curves, serifs, noise patterns)
  • Width axis maps naturally to the time axis CTC needs
  • BiLSTM — captures left ↔ right context across characters
  • CTC — no need to pre-segment characters; handles variable-length text

CNN width reduction detail:
  Blocks 1-2: MaxPool(2×2)  → width halved twice: 200 → 100 → 50
  Blocks 3-6: MaxPool(2×1)  → width unchanged at 50, height halved to 1
  After CNN:  T = 50 time steps  (≥ 2×max_label_len − 1 = 15, safe margin)
"""

from __future__ import annotations

import keras
import tensorflow as tf
import numpy as np

from dataset import NUM_CLASSES, BLANK_IDX, IMG_H, IMG_W, IMG_C, decode_label

# ── CNN building block ────────────────────────────────────────────────────────

def _conv_block(
    x: keras.KerasTensor,
    filters: int,
    pool: tuple[int, int],
    double_conv: bool = False,
    name: str = "",
) -> keras.KerasTensor:
    """Conv → BN → ReLU (optionally repeated) → MaxPool."""
    x = keras.layers.Conv2D(
        filters, (3, 3), padding="same", use_bias=False, name=f"{name}_conv1"
    )(x)
    x = keras.layers.BatchNormalization(name=f"{name}_bn1")(x)
    x = keras.layers.Activation("relu", name=f"{name}_relu1")(x)

    if double_conv:
        x = keras.layers.Conv2D(
            filters, (3, 3), padding="same", use_bias=False, name=f"{name}_conv2"
        )(x)
        x = keras.layers.BatchNormalization(name=f"{name}_bn2")(x)
        x = keras.layers.Activation("relu", name=f"{name}_relu2")(x)

    x = keras.layers.MaxPooling2D(pool_size=pool, strides=pool, name=f"{name}_pool")(x)
    return x


# ── Model builder ─────────────────────────────────────────────────────────────

def build_crnn_model(
    img_h: int = IMG_H,
    img_w: int = IMG_W,
    img_c: int = IMG_C,
    num_classes: int = NUM_CLASSES,
    rnn_units: int = 256,
    num_rnn_layers: int = 2,
    rnn_type: str = "lstm",       # 'lstm' or 'gru'
    dropout: float = 0.25,
) -> keras.Model:
    """
    Build and return the CRNN model.

    Args:
        img_h, img_w, img_c : Input image dimensions (height, width, channels).
        num_classes          : Total classes including CTC blank (63).
        rnn_units            : Hidden units per direction in each BiRNN layer.
        num_rnn_layers       : Number of stacked BiRNN layers (1 or 2).
        rnn_type             : 'lstm' (default, slightly better) or 'gru' (faster).
        dropout              : Dropout rate applied after each RNN layer.

    Returns:
        keras.Model  inputs=(B, H, W, C)  outputs=(B, T, num_classes)
        The output is RAW LOGITS — no softmax, no activation.
        Pass directly to ctc_loss() during training.
        Pass through tf.nn.softmax then ctc_decode() during inference.
    """
    inputs = keras.Input(shape=(img_h, img_w, img_c), name="image")

    # ── CNN backbone ───────────────────────────────────────────────────────────
    # Block 1: H 64→32, W 200→100, channels 3→64
    x = _conv_block(inputs, 64,  pool=(2, 2), name="block1")
    # Block 2: H 32→16, W 100→50, channels 64→128
    x = _conv_block(x,     128, pool=(2, 2), name="block2")
    # Block 3: H 16→8,  W 50  (unchanged), double conv, channels 128→256
    x = _conv_block(x,     256, pool=(2, 1), double_conv=True, name="block3")
    # Block 4: H 8→4,   W 50  (unchanged), channels 256→512
    x = _conv_block(x,     512, pool=(2, 1), name="block4")
    # Block 5: H 4→2,   W 50  (unchanged), channels 512→512
    x = _conv_block(x,     512, pool=(2, 1), name="block5")
    # Block 6: H 2→1,   W 50  (unchanged), channels 512→512
    x = _conv_block(x,     512, pool=(2, 1), name="block6")
    # x shape: (B, 1, 50, 512)

    # ── Height squeeze → sequence ──────────────────────────────────────────────
    # Remove the height=1 axis → (B, 50, 512)
    x = keras.layers.Lambda(lambda t: tf.squeeze(t, axis=1), name="squeeze")(x)

    # ── Bidirectional RNN ──────────────────────────────────────────────────────
    RNNCell = keras.layers.LSTM if rnn_type.lower() == "lstm" else keras.layers.GRU

    for i in range(num_rnn_layers):
        return_seq = True   # always True — we need per-timestep output
        x = keras.layers.Bidirectional(
            RNNCell(rnn_units, return_sequences=return_seq,
                    dropout=dropout, name=f"rnn{i + 1}"),
            merge_mode="concat",
            name=f"birnn{i + 1}",
        )(x)
        # x shape: (B, 50, rnn_units*2)

    # ── Output projection ──────────────────────────────────────────────────────
    # Raw logits — shape (B, T=50, num_classes=63)
    # DO NOT apply softmax here; tf.nn.ctc_loss expects unnormalised log-probs.
    logits = keras.layers.Dense(num_classes, name="logits")(x)

    model = keras.Model(inputs=inputs, outputs=logits, name="CRNN_CTC")
    return model


# ── CTC loss ──────────────────────────────────────────────────────────────────

def ctc_loss(
    logits: tf.Tensor,
    labels: tf.Tensor,
    label_lengths: tf.Tensor,
    logit_lengths: tf.Tensor | None = None,
) -> tf.Tensor:
    """
    Compute mean CTC loss over a batch.

    Args:
        logits        : (B, T, C) raw logits from the model.
        labels        : (B, pad_len) int32 padded label indices (-1 = padding).
        label_lengths : (B,) int32 true length of each label sequence.
        logit_lengths : (B,) int32 length of each logit sequence.
                        Defaults to T (full width) for every sample.

    Returns:
        Scalar tensor — mean CTC loss over the batch.
    """
    batch_size = tf.shape(logits)[0]
    time_steps = tf.shape(logits)[1]

    if logit_lengths is None:
        logit_lengths = tf.fill([batch_size], time_steps)

    # tf.nn.ctc_loss expects labels without padding tokens.
    # We pass the dense padded tensor + label_lengths; TF handles the masking.
    loss = tf.nn.ctc_loss(
        labels=tf.cast(labels, tf.int32),
        logits=logits,
        label_length=tf.cast(label_lengths, tf.int32),
        logit_length=tf.cast(logit_lengths, tf.int32),
        logits_time_major=False,   # logits is (B, T, C)
        blank_index=BLANK_IDX,     # 62 = last class
    )
    return tf.reduce_mean(loss)


# ── CTC decode ────────────────────────────────────────────────────────────────

def ctc_decode(
    logits: tf.Tensor,
    logit_lengths: tf.Tensor | None = None,
    method: str = "greedy",
    beam_width: int = 5,
) -> list[str]:
    """
    Decode model logits into text strings.

    Args:
        logits        : (B, T, C) raw logits (no softmax needed — applied here).
        logit_lengths : (B,) int32. Defaults to T for all samples.
        method        : 'greedy' (fast) or 'beam' (slightly more accurate).
        beam_width    : Beam width when method='beam'.

    Returns:
        List of decoded text strings, length = B.
    """
    batch_size = tf.shape(logits)[0]
    time_steps = tf.shape(logits)[1]

    if logit_lengths is None:
        logit_lengths = tf.fill([batch_size], time_steps)

    # Apply softmax and convert to log-probs for the CTC decoder
    log_probs = tf.nn.log_softmax(logits, axis=-1)

    # TF CTC decoders expect (T, B, C) — time-major
    log_probs_tm = tf.transpose(log_probs, perm=[1, 0, 2])

    seq_len = tf.cast(logit_lengths, tf.int32)

    if method == "beam":
        decoded, _ = tf.nn.ctc_beam_search_decoder(
            log_probs_tm,
            sequence_length=seq_len,
            beam_width=beam_width,
            top_paths=1,
        )
        sparse = decoded[0]
    else:
        decoded, _ = tf.nn.ctc_greedy_decoder(
            log_probs_tm,
            sequence_length=seq_len,
            merge_repeated=True,
        )
        sparse = decoded[0]

    # Convert SparseTensor to dense padded with -1
    dense = tf.sparse.to_dense(sparse, default_value=-1).numpy()

    return [decode_label(row) for row in dense]


# ── Accuracy metrics ──────────────────────────────────────────────────────────

def character_accuracy(preds: list[str], targets: list[str]) -> float:
    """
    Character-level accuracy: fraction of correctly predicted characters
    across all samples (aligned by position, length-padded).
    """
    total = correct = 0
    for pred, tgt in zip(preds, targets):
        for p, t in zip(pred.ljust(len(tgt)), tgt):
            total += 1
            if p == t:
                correct += 1
    return correct / total if total else 0.0


def sequence_accuracy(preds: list[str], targets: list[str]) -> float:
    """
    Sequence-level accuracy: fraction of samples where the entire
    predicted string exactly matches the ground truth.
    """
    if not preds:
        return 0.0
    return sum(p == t for p, t in zip(preds, targets)) / len(preds)


# ── Summary ───────────────────────────────────────────────────────────────────

def model_summary(model: keras.Model) -> None:
    """Print a clean model summary with parameter count."""
    model.summary(line_length=80)
    total   = model.count_params()
    trainable = sum(
        int(tf.reduce_prod(v.shape)) for v in model.trainable_variables
    )
    print(f"\n  Total params     : {total:,}")
    print(f"  Trainable params : {trainable:,}")
    print(f"  Input  shape     : {model.input_shape}")
    print(f"  Output shape     : {model.output_shape}  (B, T, num_classes)")
    print(f"  T (time steps)   : {model.output_shape[1]}")
    print(f"  num_classes      : {model.output_shape[2]}  "
          f"(chars 0-{NUM_CLASSES-2} + blank={BLANK_IDX})\n")


# ── Quick build test ──────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("Building CRNN model...")
    m = build_crnn_model()
    model_summary(m)

    # Forward pass smoke test
    dummy = tf.zeros((2, IMG_H, IMG_W, IMG_C))
    out   = m(dummy, training=False)
    print(f"Smoke test — input {dummy.shape} → output {out.shape}")
    assert out.shape == (2, 50, NUM_CLASSES), f"Unexpected output shape: {out.shape}"
    print("All checks passed.")