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 to folder containing images. The folder should contain image files and optionally text files with captions.
List of file extensions to include. Defaults to common image formats: [".jpg", ".jpeg", ".png", ".webp", ".bmp"]
Additional arguments (reserved for future use)
Returns
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
.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
Raised if the specified path does not exist
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
PIL Image object loaded from disk
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
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.
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"
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]]]: ...