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