SR-Forge

Datasets¶

You have files on disk — images, targets, metadata — and you need to feed them to a neural network. A Dataset bridges that gap: you tell it how to load one sample, and the framework takes care of batching, transforms, and caching.

How It Works¶

A Dataset is an object with a length and an index. Given an integer index, it loads the corresponding files from disk and returns an Entry — SR-Forge's data container. You implement two methods: __getitem__ (load one sample) and __len__ (how many samples exist).

During training, PyTorch's DataLoader iterates over the dataset — calling dataset[i] for each sample in a batch — then collates them into a single batched Entry that the model receives:

dataset[0] → Entry(image=[C,H,W], target=[C,H,W], name="scene_001")
dataset[1] → Entry(image=[C,H,W], target=[C,H,W], name="scene_002")
dataset[2] → Entry(image=[C,H,W], target=[C,H,W], name="scene_003")
dataset[3] → Entry(image=[C,H,W], target=[C,H,W], name="scene_004")
         │
         └─ DataLoader collates ─→ Entry(image=[4,C,H,W], target=[4,C,H,W], ...)
                                          │
                                          └─→ model(batch)

Your __getitem__ returns tensors in their natural shapes — a single image is [C, H, W], not [1, C, H, W]. The batch dimension is added later by collation, so you never add it yourself.

What the Base Class Provides¶

SR-Forge's Dataset extends PyTorch's torch.utils.data.Dataset with three features, controlled by constructor parameters:

name — a string identifier for the dataset (e.g. "train", "site_A"). ConcatDataset uses it to prefix entry names so they stay unique when combining multiple sources. Required when concatenating datasets.
transforms — a list of transforms applied automatically after every __getitem__ call. You don't call transforms yourself — the framework runs them in order and returns the result.
cache_dir — a path for disk caching. When set, fully-transformed entries are stored as pickle files so expensive preprocessing only runs once. Subsequent reads return the cached result instantly.

All three can be set from Python or from YAML configs.

Creating a Custom Dataset¶

Subclass Dataset and implement four things:

__getitem__ — load files and return an Entry with natural tensor shapes ([C, H, W]).
__len__ — return the number of samples. DataLoader needs this to know how many batches to create.
Set the name field on every Entry — used for saving predictions and per-scene benchmark scores.
Accept **kwargs and pass them to super().__init__() — this forwards name, transforms, cache_dir, and recache to the base class. recache=True wipes the existing cache_dir at construction time before reusing it — set this when transforms or input data have changed and stale cached entries need to be invalidated.

from srforge.dataset import Dataset
from srforge.data import Entry
from srforge.utils.io import load_image

class MyDataset(Dataset):
    def __init__(self, image_paths, target_paths, **kwargs):
        super().__init__(**kwargs)                        # (4)
        self._images = image_paths
        self._targets = target_paths

    def __len__(self):                                   # (2)
        return len(self._images)

    def __getitem__(self, idx) -> Entry:                 # (1)
        return Entry(
            name=self._images[idx].stem,                 # (3)
            image=load_image(self._images[idx]),          # [C, H, W]
            target=load_image(self._targets[idx]),        # [C, H, W]
        )

That's it. You don't call collate() and you don't apply transforms — the framework handles both automatically.

What you don't need to worry about¶

Bare scalars are fine — name="scene", scale=4, flag=True are all passed through as-is. Collation handles converting them to batched form later.
Nested dict metadata is fine — meta={"a": 1, "b": "x"} stays as-is until collation.
Don't collate — PyTorch's DataLoader (the standard tool for iterating over datasets in batches) does this for you.
Don't apply transforms — they are applied automatically by the Dataset wrapper.

What happens behind the scenes¶

When you define __getitem__ on a Dataset subclass, the framework's __init_subclass__ hook automatically wraps it. The wrapper handles both caching and transforms:

Cache check — if caching is enabled and a cached file exists for this index, return it immediately (skip everything below).
Calls your original __getitem__ — loads files, builds the raw Entry.
Applies transforms from self._transforms in order.
Cache write — if caching is enabled, stores the fully-transformed Entry to disk for next time.

Transforms see data in its natural form — bare strings, bare ints, natural tensor shapes:

# You return this from __getitem__:
Entry(name="scene_042", image=tensor([3, 64, 64]), target=tensor([3, 64, 64]))

# Transforms see exactly this — natural values, no wrapping

Tensor Conventions¶

Tensors returned from a Dataset should use natural shapes — no batch dimension:

# Correct — natural shape
image = torch.randn(3, 64, 64)   # [C, H, W]

# Wrong — batch dimension shouldn't be present in __getitem__
image = torch.randn(1, 3, 64, 64)  # [1, C, H, W] — collation will add another dim

SR-Forge's load_image() returns [C, H, W]:

Reads with OpenCV → [H, W, C] numpy array
Permutes to [C, H, W]
Converts to float tensor

If you load data through other means, use the natural shape:

# Loading from a numpy array
array = np.load("data.npy")               # [C, H, W]
tensor = torch.from_numpy(array)          # [C, H, W]

# Creating synthetic data
noise = torch.randn(1, 32, 32)            # [1, 32, 32]

Transforms on Datasets¶

Transforms are passed at construction time and applied automatically after every __getitem__ call:

from srforge.transform.data import Downsample, Normalize

ds = MyDataset(
    image_paths=paths,
    target_paths=targets,
    transforms=[
        Downsample(scale_factor=2).set_io({"inputs": {"image": "image"}}),
        Normalize().set_io({"inputs": {"image": "target"}}),
    ],
)

entry = ds[0]  # image is downsampled, target is normalized

.set_io({"inputs": {"image": "image"}}) tells the transform which Entry field to operate on. Transforms and IO binding are covered in detail in the next sections — for now, just know that each transform needs to be pointed at a field.

Caching preprocessed entries¶

If your transforms are expensive (e.g., resizing, normalization), you can cache the fully-processed results so they're only computed once:

ds = MyDataset(
    image_paths=paths,
    target_paths=targets,
    transforms=[Downsample(scale_factor=2).set_io({"inputs": {"image": "image"}})],
).cache("/tmp/my_cache")

# First epoch: loads + transforms + caches each entry
# Subsequent epochs: returns cached entries instantly

The .cache(path) method stores fully-transformed entries as pickle files at {path}/{index}.pkl. On cache hit, the entry is returned directly — no transforms run again.

In YAML configs, pass cache_dir as a constructor parameter:

train_dataset:
  _target: LazyDataset
  params:
    root: data/train
    mappings:
      image: input
      target: ground_truth
    cache_dir: /tmp/cache/train

Built-in Dataset Classes¶

Which one do I need?

Your situation	Use
One folder per scene, one file per field	LazyDataset
Multispectral — one subfolder per band inside each scene	LazyMultiSpectralDataset
You have explicit lists of file paths (e.g. from a JSON manifest)	LazyConfigDataset
None of the above (custom format, database, synthetic data)	Write your own
Take a slice / shuffle / filter of an existing dataset	SubsetDataset
Combine several datasets into one	ConcatDataset
Images too big for GPU — split into patches	PatchedDataset / MultiSpectralPatchedDataset

LazyDataset¶

Loads data from a directory structure where each example is a subdirectory:

root/
├── scene_001/
│   ├── input.tif          # matched by mappings["image"] = "input"
│   └── ground_truth.tif   # matched by mappings["target"] = "ground_truth"
├── scene_002/
│   ├── input.tif
│   └── ground_truth.tif
└── ...

from srforge.dataset.lazy_datasets import LazyDataset

ds = LazyDataset(
    root="data/train",
    mappings={"image": "input", "target": "ground_truth"},
    name="my_train",
    depth=0,            # 0 = root/scene/*  ;  1 = root/group/scene/*  ; etc.
    allowed_dirs=None,  # optional list — restrict scan to these top-level subdirs
)
# ds[0] → Entry(name="scene_001", image=tensor([C,H,W]), target=tensor([C,H,W]))

Files are matched by glob pattern (mappings value + *). Supported file types:

Extension	Loader
`.jpg` / `.jpeg` / `.png` / `.bmp` / `.tif` / `.tiff` / `.dib` / `.jp2`	OpenCV → `[C, H, W]` float tensor
`.pkl` / `.pickle`	`pickle.load` (unchanged)
`.json`	`json.load` with `list`/`tuple` values auto-converted to float tensors

Extra constructor params:

depth: int = 0 — how many directory levels above the scene directory to descend into before scanning. depth=0 treats the immediate children of root as scenes; depth=1 treats root/<group>/<scene> as the per-scene layout.
allowed_dirs: list[str] | None = None — when set, only the named top-level subdirectories under root are scanned. Useful when a shared root has both train and val data side-by-side and you want one LazyDataset per split without cloning files.

Fields with no matching file in a scene directory load as None — collation handles this by emitting a list of values instead of a stacked tensor.

LazyMultiSpectralDataset¶

For multispectral data where each scene has one subdirectory per spectral band:

root/
├── scene_001/
│   ├── b2/
│   │   ├── image.tif
│   │   └── target.tif
│   ├── b8/
│   │   ├── image.tif
│   │   └── target.tif
│   └── meta.json       # optional scene-level metadata
└── ...

from srforge.dataset.lazy_datasets import LazyMultiSpectralDataset

ds = LazyMultiSpectralDataset(
    root="data/train",
    mappings={"image": "image", "target": "target"},
    name="multispectral_train",
)
# ds[0].image  → {"b2": tensor([C,H,W]), "b8": tensor([C,H,W])}
# ds[0].target → {"b2": tensor([C,H,W]), "b8": tensor([C,H,W])}
# ds[0].bands  → ["b2", "b8"]

Each field becomes a dict keyed by band name. Band-level and scene-level metadata (from meta.json files) are automatically merged into the Entry as additional fields.

LazyConfigDataset¶

Driven by a JSON config or dict mapping field names to lists of file paths:

from srforge.dataset.lazy_datasets import LazyConfigDataset

ds = LazyConfigDataset(
    config={
        "image": ["data/image/001.tif", "data/image/002.tif"],
        "target": ["data/target/001.tif", "data/target/002.tif"],
    },
    name="config_dataset",
)

All lists must have the same length, and files at the same index must share the same stem (used as the entry name).

Wrapper Datasets¶

Wrapper datasets wrap an existing dataset to modify its behavior — selecting subsets, concatenating sources, or splitting entries into patches. They delegate to the inner dataset's __getitem__, which means the original dataset is never modified.

This creates a powerful pattern: layered caching with per-layer transforms. Each layer in the chain can have its own transforms, and caching at any layer freezes everything below it:

# Layer 1: base dataset with expensive preprocessing — cached
base = LazyDataset(
    root="data/train",
    mappings={"image": "input", "target": "ground_truth"},
    transforms=[Normalize().set_io({"inputs": {"image": "image"}})],
)
base.cache("/tmp/cache/train")

# Layer 2: wrapper that extracts patches and applies random augmentations — NOT cached
patched = PatchedDataset(
    base,
    field_sizes={"image": 32, "target": 128},
    transforms=[RandomFlip().set_io({"inputs": [{"image": "image"}, {"image": "target"}]})],
)

On the first epoch, base.__getitem__ loads files, applies Normalize, and caches the result. On subsequent epochs, cached entries are returned instantly — no file I/O, no preprocessing. PatchedDataset then extracts patches from the cached full-size images, and its transforms (RandomFlip) run every time, producing different augmentations each epoch.

This separation is why dataset transforms and wrapper datasets exist as distinct concepts:

Cached layer (inner): expensive, deterministic preprocessing — computed once
Uncached layer (wrapper): cheap, random augmentations — different every epoch

You can stack multiple layers, and each .cache() call creates a caching boundary. Transforms on the cached dataset are frozen (computed once and stored); transforms on the wrapper run fresh every time.

SubsetDataset¶

Select a subset of entries by index:

# Take first 100 entries
subset = ds.take(100)

# Take 50 entries starting at offset 10
subset = ds.take(50, offset=10)

# Shuffle with a seed
shuffled = ds.shuffle(seed=42)

# Reverse order
reversed_ds = ds.reverse()

# Filter by predicate
filtered = ds.filter(lambda entry: entry.scale > 2)

All of these return a SubsetDataset that references the original dataset with reordered indices.

ConcatDataset¶

Combine multiple datasets into one. Each entry's name is prefixed with the source dataset's name to stay unique:

ds_a = MyDataset(paths_a, name="site_A")
ds_b = MyDataset(paths_b, name="site_B")

combined = ds_a + ds_b
# or: ConcatDataset([ds_a, ds_b])

# combined[0].name → "site_A_scene_001"
# combined[len(ds_a)].name → "site_B_scene_001"

All datasets must have a name set (not None).

PatchedDataset¶

Splits each entry's spatial fields into non-overlapping patches. Useful when full images are too large for GPU memory:

from srforge.dataset.patching import PatchedDataset

patched = PatchedDataset(
    dataset=ds,
    field_sizes={"image": 32, "target": 128},  # patch size per field
    same_shapes=True,   # True if all entries have the same spatial dims
)
# len(patched) == len(ds) * num_patches_per_entry
# patched[0].name → "scene_001_000"  (original name + patch index)

With same_shapes=True, the patch grid is computed once from the first entry (fast). With same_shapes=False, each entry's grid is computed independently (slower init, handles variable sizes).

MultiSpectralPatchedDataset¶

Advanced patching for multispectral data with per-band spatial resolutions. A reference band defines the patch grid, and coordinates are mapped to other bands by their spatial ratio:

from srforge.dataset.patching import MultiSpectralPatchedDataset

patched = MultiSpectralPatchedDataset(
    dataset=ms_dataset,
    field_sizes={"image": 64},
    ref_band="b2",        # reference band for grid computation
    rounding="round",     # coordinate mapping: "round", "floor", or "ceil"
)

Supports optional internal disk caching (use_internal_cache=True) with cross-platform file locking for multi-process data loading.

DataLoader Integration¶

DataLoaderFactory creates a PyTorch DataLoader with the right collation for your dataset. It takes the same batch_size, shuffle, num_workers, pin_memory, and other keyword arguments you'd normally pass to DataLoader, plus an optional device parameter for multi-GPU detection:

from srforge.data.loader import DataLoaderFactory

train_loader = DataLoaderFactory(
    train_dataset,
    batch_size=cfg.training.batch_size,
    shuffle=True,
    num_workers=cfg.training.num_workers,
    device=cfg.system.device,
    pin_memory_device=str(device),
    pin_memory=True,
).get_loader()

val_loader = DataLoaderFactory(
    val_dataset,
    batch_size=1,
    shuffle=False,
    num_workers=0,
    device=cfg.system.device,
).get_loader()

get_loader() inspects the first entry in the dataset (dataset[0]) to decide which loader and collation to use:

Dataset returns	Loader	Collation
`Entry`	`torch.utils.data.DataLoader`	`GeneralCollation` — stacks tensors along a new dim 0
`GraphEntry` (single-GPU)	`torch_geometric.loader.DataLoader`	PyG native — concatenates node features, offsets edge indices
`GraphEntry` (multi-GPU)	`torch_geometric.loader.DataListLoader`	PyG native — returns a list of individual graphs per device

The two collation strategies are fundamentally different:

Entry (stacking): Each sample has shape [C, H, W]. Stacking 4 of them gives [4, C, H, W] — a new batch dimension is added. Non-tensor fields are collected: bare strings "s1", "s2" become ["s1", "s2"], lists ["b2", "b8"] become [["b2", "b8"], ["b2", "b8"]]. See Collation Rules for all types.
GraphEntry (concatenation): Graphs have variable numbers of nodes and edges, so they can't be stacked. PyG concatenates node features along dim 0 and offsets edge indices so the combined graph is a single disconnected graph. This is standard PyTorch Geometric batching.

Multi-GPU is detected when device is a list with more than one element (e.g., device=[0, 1]). For Entry datasets, the same DataLoader + GeneralCollation is used regardless — multi-GPU distribution happens at the model level via DataParallel. For GraphEntry, PyG provides a dedicated DataListLoader for multi-GPU.

for batch in train_loader:
    batch.image.shape   # [B, C, H, W]  — stacked tensors
    batch.target.shape  # [B, C, H, W]
    batch.name          # ["scene_001", "scene_002", ...]  — collected strings
    batch.is_batched    # True

Inheritance and the Wrapper Mechanism¶

Understanding how the automatic wrapper works helps when building advanced datasets.

How `__init_subclass__` wrapping works¶

Every class that inherits from Dataset and defines __getitem__ gets its method automatically wrapped at class definition time (not at instantiation). The wrapper applies transforms after your code runs:

class MyDataset(Dataset):
    def __getitem__(self, idx):  # ← wrapped at class definition time
        return Entry(...)

class ChildDataset(MyDataset):
    pass  # No __getitem__ → inherits MyDataset's wrapped version (no double-wrap)

class OverridingChild(MyDataset):
    def __getitem__(self, idx):  # ← gets its OWN wrapper (independent)
        return Entry(...)

Rules: - If a subclass defines __getitem__, it gets its own wrapper. - If a subclass inherits __getitem__ without overriding, no new wrapper is added — the parent's wrapper handles everything. - This prevents double-wrapping in inheritance chains.

Wrapper datasets and delegation¶

SubsetDataset and ConcatDataset work by delegating to the inner dataset:

class SubsetDataset(Dataset):
    def __getitem__(self, item):
        return self._dataset[item]  # Calls inner dataset's wrapped __getitem__

The inner dataset's __getitem__ already applies its own transforms (and returns from cache if enabled). The wrapper's __getitem__ is also wrapped, so if the wrapper has its own transforms, those are applied on top — creating the layered pipeline described in Wrapper Datasets.

This means caching and transforms compose naturally through delegation:

wrapper.__getitem__(idx)
    │
    ├─ inner.__getitem__(mapped_idx)
    │      ├─ cache hit? → return cached (already preprocessed)
    │      └─ cache miss? → load → preprocess → cache → return
    │
    └─ apply wrapper's transforms (augmentations, etc.)

The same cached base dataset can be wrapped differently for different purposes — one wrapper for training (with augmentations), another for validation (without), both sharing the same cache.

Next: Transforms — Learn how to process and transform Entry data