dataset.py

"""
dataset.py
~~~~~~~~~~
Infinite on-the-fly TensorFlow dataset pipeline for CAPTCHA training.
No images are ever saved to disk — every sample is generated fresh each epoch.

Image size: 200 × 64 px  (width × height)
  Why 200 wide?
    CTC needs time_steps >= 2 * max_label_len - 1.
    After 2× stride-2 pooling on the width axis: 200 / 4 = 50 time steps.
    For a 7-char label that's 50 vs the minimum 13 — comfortable margin.
    128px would give only 32 steps (still valid but tight for longer labels).

Character set: 62 chars  (digits + lowercase + uppercase)
  Indices 0-61  →  '0'-'9', 'a'-'z', 'A'-'Z'
  Index  62     →  CTC blank token

Usage:
    from dataset import make_dataset, make_combined_dataset, CHARS, NUM_CLASSES

    # One dataset per type (17 separate, fully on-the-fly)
    ds_type1 = make_dataset(captcha_type=1, batch_size=32)
    for images, labels, label_lengths in ds_type1.take(10):
        ...   # images: (B, 64, 200, 3) float32  labels: (B, pad) int32

    # All 17 types mixed together
    ds_all = make_combined_dataset(batch_size=32)

    # Keras model.fit() helper (returns dict y)
    ds_keras = make_dataset(1, batch_size=32, keras_format=True)
    model.fit(ds_keras, steps_per_epoch=500, epochs=50)
"""

from __future__ import annotations

import numpy as np
import tensorflow as tf
from PIL import Image as PILImage

from captcha_generators import generate, generate_random, TYPES

# ── Character set ─────────────────────────────────────────────────────────────

CHARS       = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
CHAR_TO_IDX = {c: i for i, c in enumerate(CHARS)}
IDX_TO_CHAR = {i: c for i, c in enumerate(CHARS)}
BLANK_IDX   = len(CHARS)          # 62
NUM_CLASSES = len(CHARS) + 1      # 63  (62 chars + 1 blank)
MAX_LABEL_LEN = 8                 # max chars any generator produces

# ── Image dimensions ──────────────────────────────────────────────────────────

IMG_W = 200   # width  → determines CTC time steps after CNN
IMG_H = 64    # height
IMG_C = 3     # RGB channels

# ── Encoding helpers ──────────────────────────────────────────────────────────

def encode_label(text: str) -> np.ndarray:
    """
    Map each character in *text* to its integer index.
    Unknown characters are silently skipped.
    Returns a 1-D int32 numpy array.
    """
    return np.array([CHAR_TO_IDX[c] for c in text if c in CHAR_TO_IDX],
                    dtype=np.int32)


def decode_label(indices) -> str:
    """
    Map a sequence of integer indices back to a string.
    CTC blank tokens (index 62) and padding (-1) are stripped.
    """
    return "".join(
        IDX_TO_CHAR[int(i)]
        for i in indices
        if int(i) not in (-1, BLANK_IDX)
    )


def preprocess_image(pil_img) -> np.ndarray:
    """
    Resize a PIL image to (IMG_H, IMG_W) **preserving aspect ratio** and
    pad the remaining width with the estimated background colour.
    Returns a float32 array in [0, 1] of shape (IMG_H, IMG_W, IMG_C).

    Real CAPTCHAs vary widely in aspect ratio (100×40 … 280×54).  Stretching
    them all to a fixed box distorts the glyphs; instead we scale to the target
    height, fit the width, and pad so characters keep their natural shape.
    """
    pil_img = pil_img.convert("RGB")
    w, h = pil_img.size

    # Scale so height == IMG_H, keep aspect ratio
    new_w = max(1, int(round(w * (IMG_H / h))))
    pil_img = pil_img.resize((new_w, IMG_H), PILImage.LANCZOS)
    arr = np.array(pil_img, dtype=np.float32) / 255.0      # (IMG_H, new_w, 3)

    if new_w == IMG_W:
        return arr
    if new_w > IMG_W:
        # Too wide after height-scaling: squeeze width to fit (rare)
        pil_img = pil_img.resize((IMG_W, IMG_H), PILImage.LANCZOS)
        return np.array(pil_img, dtype=np.float32) / 255.0

    # Pad width to IMG_W, centered, using the median border colour as background
    border = np.concatenate([arr[0, :, :], arr[-1, :, :],
                             arr[:, 0, :], arr[:, -1, :]], axis=0)
    bg = np.median(border, axis=0)                          # (3,)
    canvas = np.ones((IMG_H, IMG_W, IMG_C), dtype=np.float32) * bg
    left = (IMG_W - new_w) // 2
    canvas[:, left:left + new_w, :] = arr
    return canvas



# ── Core Python generators ────────────────────────────────────────────────────

def _single_type_generator(captcha_type: int):
    """
    Infinite Python generator for one CAPTCHA type.
    Yields (image_array, label_indices, label_length) tuples.
    """
    while True:
        pil_img, label = generate(captcha_type)
        image   = preprocess_image(pil_img)
        indices = encode_label(label)
        yield image, indices, np.int32(len(indices))


def _combined_generator():
    """
    Infinite Python generator that cycles through all 17 types randomly.
    Yields (image_array, label_indices, label_length, type_id) tuples.
    """
    import random
    while True:
        pil_img, label, t = generate_random()
        image   = preprocess_image(pil_img)
        indices = encode_label(label)
        yield image, indices, np.int32(len(indices)), np.int32(t)


# ── tf.data.Dataset factories ─────────────────────────────────────────────────

def make_dataset(
    captcha_type: int,
    batch_size: int = 32,
    shuffle_buffer: int = 256,
    prefetch: int = tf.data.AUTOTUNE,
    keras_format: bool = False,
) -> tf.data.Dataset:
    """
    Build an infinite on-the-fly tf.data.Dataset for a single CAPTCHA type.

    Args:
        captcha_type:   Integer 1-17.
        batch_size:     Samples per batch.
        shuffle_buffer: Size of the shuffle buffer (0 = no shuffle).
        prefetch:       Number of batches to prefetch (AUTOTUNE recommended).
        keras_format:   If True, each batch is (images, y_dict) where
                        y_dict = {'labels': ..., 'label_lengths': ...}.
                        Use this with model.fit().

    Returns:
        Batched tf.data.Dataset.
        Each batch (keras_format=False):
            images        : (B, 64, 200, 3)  float32
            labels        : (B, pad_len)      int32  (padded with -1)
            label_lengths : (B,)              int32
    """
    if captcha_type not in TYPES:
        raise ValueError(f"captcha_type must be 1-17, got {captcha_type!r}")

    output_sig = (
        tf.TensorSpec(shape=(IMG_H, IMG_W, IMG_C), dtype=tf.float32),
        tf.TensorSpec(shape=(None,),               dtype=tf.int32),
        tf.TensorSpec(shape=(),                    dtype=tf.int32),
    )

    ds = tf.data.Dataset.from_generator(
        lambda: _single_type_generator(captcha_type),
        output_signature=output_sig,
    )

    if shuffle_buffer > 0:
        ds = ds.shuffle(buffer_size=shuffle_buffer, reshuffle_each_iteration=True)

    ds = ds.padded_batch(
        batch_size,
        padded_shapes=((IMG_H, IMG_W, IMG_C), (None,), ()),
        padding_values=(
            tf.constant(0.0, tf.float32),
            tf.constant(-1,  tf.int32),   # -1 = padding sentinel for labels
            tf.constant(0,   tf.int32),
        ),
        drop_remainder=False,
    )

    if keras_format:
        ds = ds.map(
            lambda imgs, lbl, llen: (
                imgs,
                {"labels": lbl, "label_lengths": llen},
            ),
            num_parallel_calls=tf.data.AUTOTUNE,
        )

    return ds.prefetch(prefetch)


def make_combined_dataset(
    batch_size: int = 32,
    shuffle_buffer: int = 256,
    prefetch: int = tf.data.AUTOTUNE,
    include_type_id: bool = False,
    keras_format: bool = False,
) -> tf.data.Dataset:
    """
    Build an infinite dataset that mixes all 17 CAPTCHA types randomly.

    Args:
        include_type_id: If True, each batch includes the captcha_type integer
                         (useful for multi-task learning or analysis).
        keras_format:    Same as make_dataset().

    Returns:
        Batched tf.data.Dataset.
        Each batch (include_type_id=False, keras_format=False):
            images        : (B, 64, 200, 3)  float32
            labels        : (B, pad_len)      int32
            label_lengths : (B,)              int32
    """
    output_sig = (
        tf.TensorSpec(shape=(IMG_H, IMG_W, IMG_C), dtype=tf.float32),
        tf.TensorSpec(shape=(None,),               dtype=tf.int32),
        tf.TensorSpec(shape=(),                    dtype=tf.int32),
        tf.TensorSpec(shape=(),                    dtype=tf.int32),  # type_id
    )

    ds = tf.data.Dataset.from_generator(
        _combined_generator,
        output_signature=output_sig,
    )

    if shuffle_buffer > 0:
        ds = ds.shuffle(buffer_size=shuffle_buffer, reshuffle_each_iteration=True)

    if include_type_id:
        ds = ds.padded_batch(
            batch_size,
            padded_shapes=((IMG_H, IMG_W, IMG_C), (None,), (), ()),
            padding_values=(
                tf.constant(0.0, tf.float32),
                tf.constant(-1,  tf.int32),
                tf.constant(0,   tf.int32),
                tf.constant(0,   tf.int32),
            ),
            drop_remainder=False,
        )
    else:
        # Drop the type_id column for simplicity
        ds = ds.map(lambda img, lbl, llen, _t: (img, lbl, llen),
                    num_parallel_calls=tf.data.AUTOTUNE)
        ds = ds.padded_batch(
            batch_size,
            padded_shapes=((IMG_H, IMG_W, IMG_C), (None,), ()),
            padding_values=(
                tf.constant(0.0, tf.float32),
                tf.constant(-1,  tf.int32),
                tf.constant(0,   tf.int32),
            ),
            drop_remainder=False,
        )

    if keras_format:
        ds = ds.map(
            lambda imgs, lbl, llen: (
                imgs,
                {"labels": lbl, "label_lengths": llen},
            ),
            num_parallel_calls=tf.data.AUTOTUNE,
        )

    return ds.prefetch(prefetch)


def make_all_datasets(
    batch_size: int = 32,
    shuffle_buffer: int = 256,
    keras_format: bool = False,
) -> dict[int, tf.data.Dataset]:
    """
    Convenience function — returns a dict of 17 separate tf.data.Datasets,
    one per CAPTCHA type.

    Returns:
        {1: ds_type1, 2: ds_type2, ..., 17: ds_type17}
    """
    return {
        t: make_dataset(t, batch_size=batch_size,
                        shuffle_buffer=shuffle_buffer,
                        keras_format=keras_format)
        for t in TYPES
    }


# ── Quick sanity check ────────────────────────────────────────────────────────

def verify_pipeline(captcha_type: int = 1, n_batches: int = 3, batch_size: int = 4) -> None:
    """
    Print shape and label info for a few batches.  Use to confirm the pipeline
    is working before kicking off a full training run.

    Example:
        from dataset import verify_pipeline
        verify_pipeline(captcha_type=3)
    """
    print(f"\n── Verifying pipeline for type {captcha_type} ──")
    print(f"   IMG_W={IMG_W}  IMG_H={IMG_H}  NUM_CLASSES={NUM_CLASSES}  BLANK={BLANK_IDX}\n")

    ds = make_dataset(captcha_type, batch_size=batch_size, shuffle_buffer=16)
    for batch_idx, (images, labels, lengths) in enumerate(ds.take(n_batches)):
        print(f"  Batch {batch_idx + 1}:")
        print(f"    images shape  : {images.shape}  dtype={images.dtype}"
              f"  min={images.numpy().min():.3f}  max={images.numpy().max():.3f}")
        print(f"    labels shape  : {labels.shape}  dtype={labels.dtype}")
        print(f"    label_lengths : {lengths.numpy().tolist()}")
        for i in range(len(lengths)):
            raw    = labels[i].numpy()
            length = int(lengths[i])
            text   = decode_label(raw[:length])
            print(f"      sample {i}: encoded={raw[:length].tolist()}  decoded='{text}'")
        print()


if __name__ == "__main__":
    for t in TYPES:
        verify_pipeline(captcha_type=t, n_batches=1, batch_size=2)