dew.data.images

Image datasets: TFDS, Hugging Face hub and arrayrecord shards, one transform.

Every image dataset resizes, augments and captions its records the same way. What differs is where the records come from and how one is read, which is the three hooks a subclass fills in. Records leave as {"image": uint8 [size, size, 3], "caption": str}, plus "label" when the source carries a class index. load(tokenize=) is where a run’s own condition reads the captions: the dataset carries the text, the encoder decides what tokens it becomes. cv2, tensorflow_datasets and HF datasets are imported on use, so import dew.data costs none of them.

Name	Summary
Augmentation
import_opencv	Import OpenCV in the thread that opens a loader, before its readers start.
unpack_dict_of_byte_arrays	The str -> bytes entries of one packed arrayrecord record.
pack_dict_of_byte_arrays	unpacked’s entries length-prefixed in dict order, the layout unpack_dict_of_byte_arrays reads.
decode_image	An encoded image as RGB uint8, in the orientation its pixels are stored.
resize_image	image at size square; area interpolation down, cubic up.
Augment	Says which augmentations a mode applies: flip for flip_only, both for flip_jitter.
image_augmentations	The augmentations mode names: flip_only (DiT style), flip_jitter, or none (deterministic evaluation and debugging).
augment_image	Flips and colour-jitters image, seeded by the record’s own rng.
PROMPT_TEMPLATES
class_names	The class names of a labels file, one per line, read once per process.
record_caption	The caption a record already carries, for datasets that ship their text.
ImageTransform	Resizes, augments and captions one record, seeded by the record’s own rng.
ImageDataset	Reads captioned images through grain, resized to image_size.
OxfordFlowers	Reads prepared Oxford Flowers ArrayRecords, captioned from their class names.
HFImages	Reads a Hugging Face hub dataset of images by index, captioned from its ‘caption’ or ‘text’ column.
ArrayRecordImages	Reads image and caption pairs from arrayrecord shards under path/<shard>/, each record a packed dict.
Laion12mCoco	laion-aesthetics-12M (score >= 6) plus MS-COCO 2017: 228 shards, 236 GiB, about 15M samples.
Laion2bAesthetic	laion-2B-en aesthetic >= 4.2 subset: 569 shards, 550 GiB, larger but noisier.
DiffusionDB	diffusiondb (SD synthetic images and prompts): 31 shards, 60 GiB, 1.97M samples.
CC3M	Conceptual Captions 3M: 50 shards, 37 GiB, about 3.3M samples (shard 00039 missing).
CombinedMsml612	The four msml612 datasets together, about 883 GiB and 20M samples.
CC12M	Conceptual Captions 12M, in the arrayrecord2 layout of the msml612 bucket.
Combined30M	Four arrayrecord2 shard sets of the msml612 bucket, about 30M samples.

Augmentation

attribute source

Augmentation = Literal['none', 'flip_only', 'flip_jitter']

import_opencv

function source

def import_opencv() -> None

Import OpenCV in the thread that opens a loader, before its readers start.

The reader threads reach their first decode together, so each would otherwise make the first import of cv2 at once. An import that fails in one of them leaves the others the half-built module, which surfaces as module 'cv2' has no attribute 'INTER_AREA' instead of the failure itself. Imported here, a broken install raises its own error when the loader opens. The import stays out of the module’s top so that reading text never loads OpenCV.

unpack_dict_of_byte_arrays

function source

def unpack_dict_of_byte_arrays(packed_data: bytes) -> dict[str, bytes]

The str -> bytes entries of one packed arrayrecord record.

Each entry is a uint32 key length, the utf-8 key, a uint32 value length and the value, in that order. pack_dict_of_byte_arrays writes it.

pack_dict_of_byte_arrays

function source

def pack_dict_of_byte_arrays(unpacked: dict) -> bytes

unpacked’s entries length-prefixed in dict order, the layout unpack_dict_of_byte_arrays reads.

decode_image

function source

def decode_image(encoded: bytes, *, at_least: int | None = None) -> np.ndarray

An encoded image as RGB uint8, in the orientation its pixels are stored.

Grey is replicated and a 16-bit sample kept to its high byte. An image with transparency is composited onto white, as img2dataset does for the url shards the online loader streams. EXIF orientation is ignored, as PIL’s Image.open ignores it, on the reduced decodes too, which would otherwise apply it.

With at_least, an opaque image is decoded at the largest 1/2, 1/4 or 1/8 reduction that keeps both sides >= at_least (the DCT scale of a JPEG), so the resize after it still only shrinks and most of the decode is skipped.

Every failure is a ValueError. PIL reads the header first, since it refuses a decompression bomb from the header alone where cv2 would decode up to 2**30 pixels, and cv2 hands back None for a half-written file.

resize_image

function source

def resize_image(image: np.ndarray, size: int) -> np.ndarray

image at size square; area interpolation down, cubic up.

Augment

dataclass source

class Augment(flip: bool, jitter: bool)

Says which augmentations a mode applies: flip for flip_only, both for flip_jitter. ‘none’ maps to no Augment at all.

image_augmentations

function source

def image_augmentations(mode: Augmentation) -> Augment | None

The augmentations mode names: flip_only (DiT style), flip_jitter, or none (deterministic evaluation and debugging).

augment_image

function source

def augment_image(
    augment: Augment | None,
    image: np.ndarray,
    rng: np.random.Generator,
) -> np.ndarray

Flips and colour-jitters image, seeded by the record’s own rng.

Every draw comes from grain’s per-record rng, a Philox keyed by the record index, so a record’s augmentation is the same however many workers, threads or processes produced its batch. uint8 pixels go through float32 and are rounded and clipped once at the end.

PROMPT_TEMPLATES

attribute source

PROMPT_TEMPLATES = ('a photo of a {}', 'a photo of a {} flower', 'This is a photo of a {}', 'This is a photo of a {} flower', 'A photo of a {} flower')

class_names

function source

def class_names(path: str) -> tuple[str, ...]

The class names of a labels file, one per line, read once per process.

record_caption

function source

def record_caption(element) -> str

The caption a record already carries, for datasets that ship their text.

Hub image datasets keep it in a ‘caption’ or a ‘text’ column.

ImageTransform

class source

class ImageTransform(spec: ImageDataset)

Resizes, augments and captions one record, seeded by the record’s own rng.

It is built where its loader opens and unpickled where a spawned worker starts, both before any reader thread runs, and both import OpenCV (import_opencv).

ImageTransform.random_map

def random_map(element: Batch | bytes, rng: np.random.Generator) -> Batch

ImageDataset

dataclass source

class ImageDataset(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads captioned images through grain, resized to image_size.

Validation comes from one of two places. val_split names a split of the dataset’s own, which is opened as a second source and scored in record order, val_batches batches of it or all of it when that is None. Without one, val_batches batches of records are held out of the head of the training source, in canonical order, so FID and CLIP are never measured on records the model trained on. None or 0 holds nothing out and validates nothing.

count takes that many records from the head of the source. A source that reports no length needs it set.

ImageDataset.source

def source(split: str | None = None) -> Records

Opens the records by index (__getitem__, and __len__ unless count says how many there are).

split names a split other than the one this spec reads, which is how val_split opens a second source. A dataset whose records are one pile refuses it.

ImageDataset.record

def record(
    element: Batch | bytes,
    rng: np.random.Generator,
) -> tuple[np.ndarray | bytes, str, int | None]

One record as (image, caption, class index or None).

The image is RGB uint8, or the encoded bytes for the transform to decode at the size it needs.

ImageDataset.records

def records(source: Records) -> int

The records the run uses, from the head of the source.

A source that cannot count itself is Counted’s other case, where the spec’s own count is the whole record of how many there are.

ImageDataset.load

def load(*, batch: int, tokenize: Tokenize | None = None) -> Dataset

OxfordFlowers

dataclass source

class OxfordFlowers(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    split: str = 'all',
    labels: str | None = None,
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads prepared Oxford Flowers ArrayRecords, captioned from their class names.

Preparation runs separately. Reading uses TFDS metadata and NumPy image decoding through its read-only builder, without TensorFlow or dataset generation code in the training process.

path: str | None

Prepared version directory containing dataset_info.json and ArrayRecords.

labels: str | None

Class-name file override; unset reads label.labels.txt in path.

OxfordFlowers.source

def source(split: str | None = None)

OxfordFlowers.record

def record(element: Batch | bytes, rng)

HFImages

dataclass source

class HFImages(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    name: str = '',
    split: str = 'train',
    options: HubOptions = HFOptions(),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads a Hugging Face hub dataset of images by index, captioned from its ‘caption’ or ‘text’ column.

name is the repo id and split the split to read. options is everything else datasets.load_dataset takes, the same value the hf provider holds, so a dataset behind a config name, a revision, its own data_files or a token is read here too.

Its images arrive decoded by datasets rather than by decode_image, so a JPEG’s EXIF orientation is applied, where decode_image keeps the stored one.

HFImages.source

def source(split: str | None = None)

HFImages.record

def record(element: Batch | bytes, rng)

ArrayRecordImages

dataclass source

class ArrayRecordImages(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = (),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Reads image and caption pairs from arrayrecord shards under path/<shard>/, each record a packed dict.

Two layouts are read. ‘jpg’/‘txt’ entries hold an encoded image, decoded on read. The prepare_images.py layout holds ‘image’/‘shape’/‘caption’, where the image is uint8 HxWx3 already at training size, the shape is two little-endian int32s, and ‘label’ is optional.

path is the bucket mount or directory the shards live under. An empty shards reads every arrayrecord file in path itself, which is the layout prepare_images.py writes.

ArrayRecordImages.source

def source(split: str | None = None)

ArrayRecordImages.record

def record(element: Batch | bytes, rng)

Laion12mCoco

dataclass source

class Laion12mCoco(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = ('arrayrecord2/laion12m_coco',),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

laion-aesthetics-12M (score >= 6) plus MS-COCO 2017: 228 shards, 236 GiB, about 15M samples.

Laion2bAesthetic

dataclass source

class Laion2bAesthetic(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = ('arrayrecord2/laion2B-en-aesthetic',),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

laion-2B-en aesthetic >= 4.2 subset: 569 shards, 550 GiB, larger but noisier.

DiffusionDB

dataclass source

class DiffusionDB(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = ('arrayrecord2/diffusiondb',),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

diffusiondb (SD synthetic images and prompts): 31 shards, 60 GiB, 1.97M samples.

CC3M

dataclass source

class CC3M(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = ('arrayrecord2/cc3m',),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Conceptual Captions 3M: 50 shards, 37 GiB, about 3.3M samples (shard 00039 missing).

CombinedMsml612

dataclass source

class CombinedMsml612(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = ('arrayrecord2/laion12m_coco', 'arrayrecord2/laion2B-en-aesthetic', 'arrayrecord2/diffusiondb', 'arrayrecord2/cc3m'),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

The four msml612 datasets together, about 883 GiB and 20M samples.

CC12M

dataclass source

class CC12M(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = ('arrayrecord2/cc12m',),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Conceptual Captions 12M, in the arrayrecord2 layout of the msml612 bucket.

Combined30M

dataclass source

class Combined30M(
    image_size: int = 128,
    augmentation: Augmentation = 'flip_jitter',
    val_batches: int | None = 4,
    val_split: str | None = None,
    count: int | None = None,
    path: str | None = None,
    shards: tuple[str, ...] = ('arrayrecord2/laion-aesthetics-12m+mscoco-2017', 'arrayrecord2/cc12m', 'arrayrecord2/aestheticCoyo_0.26_clip_5.5aesthetic_256plus', 'arrayrecord2/playground+leonardo_x4+cc3m.parquet'),
    *,
    seed: int = 0,
    loading: Loading = Loading(),
)

Four arrayrecord2 shard sets of the msml612 bucket, about 30M samples.