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