Dataset

Load and manage image datasets for training

Overview

The dataset module provides functionality for loading and managing image datasets. It automatically discovers images, loads captions, and provides convenient iteration and batching.

dataset.load()

Load a dataset from a folder of images.

from hypergen import dataset

ds = dataset.load("./my_images")

Parameters

path · string | Path · required

Path to folder containing images. The folder should contain image files and optionally text files with captions.

extensions · list[string]

List of file extensions to include. Defaults to common image formats: [".jpg", ".jpeg", ".png", ".webp", ".bmp"]

**kwargs · dict

Additional arguments (reserved for future use)

Returns

dataset · Dataset

Dataset instance containing loaded images and captions

Directory Structure

The dataset loader expects a directory structure like this:

my_images/
�� image1.jpg
�� image1.txt      # Optional caption for image1.jpg
�� image2.png
�� image2.txt      # Optional caption for image2.png
�� image3.jpg      # No caption file = None

Captions are loaded from .txt files with the same name as the image. If a caption file doesn’t exist, the caption will be None.

Example

from hypergen import dataset

# Load dataset from folder
ds = dataset.load("./training_images")
print(f"Loaded {len(ds)} images")

# Load with custom extensions
ds = dataset.load(
    "./my_images",
    extensions=[".jpg", ".png"]  # Only load JPG and PNG
)

# Check dataset size
print(f"Dataset contains {len(ds)} images")

Errors

FileNotFoundError · error

Raised if the specified path does not exist

ValueError · error

Raised if:

• The path is not a directory
• No images found with specified extensions

---

Dataset Operations

Length

Get the number of items in the dataset using len().

ds = dataset.load("./images")
print(len(ds))  # e.g., 100

---

Indexing

Access individual items by index. Returns a tuple of (image, caption).

ds = dataset.load("./images")

# Get first item
image, caption = ds[0]

print(f"Image size: {image.size}")
print(f"Caption: {caption}")

Returns

image · PIL.Image.Image

PIL Image object loaded from disk

caption · string | None

Caption text from corresponding .txt file, or None if no caption exists

---

Iteration

Iterate over all items in the dataset.

ds = dataset.load("./images")

for image, caption in ds:
    print(f"Image: {image.size}, Caption: {caption}")
    # Process image and caption

Example

from hypergen import dataset

ds = dataset.load("./training_images")

# Process all images
for image, caption in ds:
    # Resize images
    resized = image.resize((512, 512))

    # Print caption
    if caption:
        print(f"Caption: {caption}")

---

Dataset.batch()

Iterate over the dataset in batches.

for batch in ds.batch(32):
    # Process batch of 32 images
    pass

Parameters

batch_size · integer · required

Number of items per batch

Yields

batch · list[tuple[PIL.Image.Image, string | None]]

List of (image, caption) tuples. The last batch may be smaller than batch_size if the dataset size is not evenly divisible.

Example

from hypergen import dataset

ds = dataset.load("./images")

# Process in batches of 32
for batch in ds.batch(32):
    print(f"Processing batch of {len(batch)} images")

    for image, caption in batch:
        # Process each item in batch
        pass

# Training example
from hypergen import model

m = model.load("stabilityai/sdxl-turbo").to("cuda")

# Process dataset in batches during training
for batch in ds.batch(4):
    images = [item[0] for item in batch]
    captions = [item[1] for item in batch]

    # Training step with batch
    # ...

Notes

The last batch may contain fewer items than batch_size if the dataset size is not evenly divisible.

---

Caption Files

HyperGen automatically loads captions from text files matching image names.

Format

For each image file (e.g., photo.jpg), create a text file with the same name (e.g., photo.txt) containing the caption:

photo.jpg       # Image file
photo.txt       # Caption: "A beautiful sunset over mountains"

The caption file should contain plain text. The entire contents of the file will be used as the caption, with leading/trailing whitespace removed.

Example Caption File

A photorealistic portrait of a person in natural lighting, professional photography, high detail, sharp focus

Behavior

• If a caption file exists, its contents are loaded as a string
• If no caption file exists, the caption will be None
• Caption files are optional - you can have some images with captions and others without
• Captions are automatically associated with their corresponding images by filename

---

Use Cases

Training

from hypergen import model, dataset

# Load model and dataset
m = model.load("stabilityai/stable-diffusion-xl-base-1.0").to("cuda")
ds = dataset.load("./my_style_images")

# Train LoRA
lora = m.train_lora(ds, steps=1000)

Data Analysis

from hypergen import dataset

ds = dataset.load("./images")

# Analyze dataset
print(f"Total images: {len(ds)}")

# Count images with captions
captioned = sum(1 for _, caption in ds if caption is not None)
print(f"Images with captions: {captioned}")

# Check image sizes
sizes = set()
for image, _ in ds:
    sizes.add(image.size)

print(f"Unique image sizes: {sizes}")

Preprocessing

from hypergen import dataset
from PIL import Image

ds = dataset.load("./raw_images")

# Preprocess and save
output_dir = "./processed_images"

for i, (image, caption) in enumerate(ds):
    # Resize to square
    size = min(image.size)
    image = image.crop((0, 0, size, size))
    image = image.resize((512, 512))

    # Save processed image
    image.save(f"{output_dir}/image_{i:04d}.jpg")

    # Save caption if exists
    if caption:
        with open(f"{output_dir}/image_{i:04d}.txt", "w") as f:
            f.write(caption)

---

Type Reference

Dataset

The main dataset class returned by dataset.load().

class Dataset:
    image_paths: list[Path]      # Paths to image files
    captions: list[str | None]   # Corresponding captions

    def load(path: str | Path, **kwargs) -> Dataset: ...
    def __len__() -> int: ...
    def __getitem__(idx: int) -> tuple[PIL.Image.Image, str | None]: ...
    def __iter__() -> Iterator[tuple[PIL.Image.Image, str | None]]: ...
    def batch(batch_size: int) -> Iterator[list[tuple[PIL.Image.Image, str | None]]]: ...