SR-Forge

Datasets

Dataset

Bases: Dataset

Base dataset class with built-in disk caching and automatic transform application.

Subclasses implement __getitem__ and __len__. The framework wraps __getitem__ so that:

1. If caching is enabled and a cached file exists for the requested index, the fully-transformed entry is loaded from disk and returned immediately.
2. On cache miss, the original __getitem__ is called, transforms are applied, and the result is stored to disk for future hits.
3. If caching is disabled, the original __getitem__ is called and transforms are applied every time.

Caching can be enabled in two ways:

• At construction time via the cache_dir parameter (useful in YAML configs).
• After construction via the :meth:cache method (useful in scripts).

cache_path: Optional[Path] property

Return the cache directory path, or None if caching is disabled.

meta_item(index) -> Entry

Return the entry at index as it would leave __getitem__ — but with meta tensors (correct shapes, zero data).

Builds the raw meta entry via :meth:_meta_raw_item, then runs this dataset's actual transform pipeline on it via shape tracing (abstract interpretation on device='meta'). The result's tensor shapes are exactly what a real __getitem__ would produce, obtained without reading pixel data or executing tensor math.

Raises:

Type	Description
ShapeTraceUnsupported	When raw shapes aren't cheaply available or a transform can't be traced (see the trace ladder in :mod:srforge.utils.shape_trace).

cache(path: Union[str, Path]) -> Dataset

Enable disk caching of fully-transformed entries.

Cached entries are stored as pickle files at {path}/{index}.pkl. On cache hit, the entry is returned directly — no transforms run again.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Directory where cached pickle files are stored. Created automatically on first cache write.	required

Returns:

Type	Description
Dataset	self, for method chaining.