Skip to content
SR-Forge

Transforms

Transformation

Bases: ABC

Base for all transforms.

Shape-tracing declarations (see :mod:srforge.utils.shape_trace):

  • SHAPE_PRESERVING — set to True on subclasses whose output tensor shapes always equal their input shapes (most numpy/scipy/ cv2-roundtrip transforms: filters, denoising, histogram ops). Lets shape tracing skip executing the transform. Verified against reality on the first real pass — a wrong declaration raises.
  • infer_shape(shape_in) -> shape_out — optionally implement on subclasses that change shapes but can't run on meta tensors (e.g. a scipy-based resize). Same runtime verification applies.

Torch-native transforms need neither: their real code executes on meta tensors directly.

DataTransform

Bases: IOModule, Transformation, ABC

Process the content of Entry fields.

Subclasses override one of two methods:

  • transform(**kwargs) — batched processing. Receives values with a leading batch dimension (tensors are [B, C, H, W], names are lists). This is the default and preferred path for performance.
  • transform_unbatched(**kwargs) — unbatched processing. Receives values without the batch dimension (tensors are [C, H, W], names are bare strings).

Dispatch: transform() is always preferred when it exists — for both batched and unbatched entries. For unbatched entries, the framework auto-wraps via collate([entry]) → transform() → result[0]. The wrapping overhead for a single entry is negligible compared to having a single optimized code path. transform_unbatched() is only used when transform() is not implemented:

  • transform() exists → always used (wrap/unwrap for unbatched).
  • Only transform_unbatched() → called directly for unbatched entries, looped per sample for batched entries.

IO binding works identically for both paths. The framework inspects the overridden method's signature (transform or transform_unbatched) for parameter names and type annotations (which control annotation-driven recursion into containers).

Annotation-driven recursion: when a field contains a container (dict, list, tuple) and the corresponding transform/ transform_unbatched parameter is annotated as torch.Tensor (or Union[..., torch.Tensor]), the framework recurses into the container and calls the method on each leaf value. Annotating as dict, list, or omitting the annotation for non-Tensor types passes the container as a whole.

transform(*args, **kwargs) -> Any

Batched transform — override for batch-parallel processing.

Subclasses should override either this method (for batched processing) or transform_unbatched (for unbatched processing). If only transform_unbatched is overridden, the framework automatically loops per sample at runtime.

transform_unbatched(**kwargs) -> Any

Unbatched transform — override this for single-sample logic.

Receives kwargs without batch dim: tensors are [C,H,W], strings are bare, scalars are Python types. Return values in the same convention.

When transform_unbatched is overridden and transform is not, the framework automatically loops over batch samples, calling this method once per sample, and re-batches the results.

set_io(io_cfg, *, strict=True, require_all=True)

Bind IO mapping using the unified config format.

Accepts: {"inputs": {param: field}, "outputs": <string|list|dict>}

EntryTransform

Bases: IOModule, Transformation, ABC

Operate on Entry structure — add, remove, rename fields, change type.

Subclasses override one of two methods:

  • transform(entry) — batched processing. Receives an Entry with is_batched=True (tensors are [B, C, H, W], names are lists).
  • transform_unbatched(entry) — unbatched processing. Receives an Entry with is_batched=False (tensors are [C, H, W], names are bare strings).

Dispatch depends on the Entry's is_batched flag:

  • Batched entry + transform() → call directly (fast path).
  • Batched entry + only transform_unbatched() → unbatch, loop per sample, re-collate.
  • Unbatched entry + transform_unbatched() → call directly (natural fit).
  • Unbatched entry + only transform() → collate([entry]) → batched transform → result[0] (auto-wrap/unwrap).

If both are overridden, transform() wins for batched entries, transform_unbatched() wins for unbatched entries. Unlike DataTransform (where the framework mediates field extraction), EntryTransform methods receive the Entry directly — so each method must match the Entry's batched state.

Field keys: Subclasses accept Entry field names as constructor parameters with a _key suffix (e.g., field_key="lrs"). This makes field binding explicit and avoids the need for IOSpec or set_io() — each transform knows its own field mappings.

transform(entry: Entry) -> Entry

Batched transform — override for batch-parallel processing.

Receives an Entry with is_batched=True. Tensors have shape [B, C, H, W], names are lists of strings.

transform_unbatched(entry: Entry) -> Entry

Unbatched transform — override this for single-sample logic.

Receives an Entry with is_batched=False. Tensors have shape [C, H, W], names are bare strings.