DataTransform¶
A DataTransform processes the content of Entry fields. You define a transform() method (batched) or transform_unbatched() method (per-sample) that takes values in and returns values out — SR-Forge handles everything else: extracting values from Entry fields, calling your method, and storing results back. It works with any data type: tensors, strings, integers, dicts, and more. Depending on IO binding, results can overwrite the original field or be written to a new field, so a DataTransform can also grow the Entry.
This is the most common type of transform. Use it for normalization, augmentation, resizing, color conversion, and any other per-field operation.
Your First DataTransform¶
A DataTransform is a class that overrides transform() (for batched processing) or transform_unbatched() (for per-sample processing). Let's start with transform():
from srforge.transform import DataTransform
import torch
class ZScoreNormalize(DataTransform):
def transform(self, image: torch.Tensor) -> torch.Tensor:
mean = image.mean(dim=(-2, -1), keepdim=True)
std = image.std(dim=(-2, -1), keepdim=True)
return (image - mean) / (std + 1e-8)
This computes the mean and standard deviation from the data itself across the last two (spatial) dimensions. Using negative indices (-2, -1) makes the code shape-agnostic — it works with [B, C, H, W], [B, C, N, H, W], [C, H, W], or any other layout.
You can call it directly with a value — no Entry, no config, no setup:
zscore = ZScoreNormalize()
result = zscore(torch.randn(4, 3, 64, 64)) # via __call__
result = zscore.transform(torch.randn(4, 3, 64, 64)) # same thing
ZScoreNormalize has no __init__() because it computes everything from the input. But DataTransform extends torch.nn.Module, so you can store configuration as self.* attributes and use them inside transform(). The transform() parameters are only for values that come from Entry fields — everything else (hyperparameters, thresholds, flags, learned weights) lives on self:
class Resize(DataTransform):
def __init__(self, size: tuple[int, int], mode: str = "bilinear"):
super().__init__()
self.size = size # Fixed config — not an Entry field
self.mode = mode
def transform(self, image: torch.Tensor) -> torch.Tensor:
return F.interpolate(image, size=self.size, mode=self.mode)
The transform() method's parameter names and type annotations define the transform's interface. The parameter name and the Entry field name don't have to match. For example, ZScoreNormalize has a parameter called image, but you can point it at any field: set_io({"inputs": {"image": "target"}}) binds the Entry field "target" to the image parameter. The method signature is the contract — set_io() connects it to actual data.
Using with Entry¶
In a pipeline, your transform reads from and writes to Entry fields. The mechanism that connects a transform to specific Entry fields is called IO binding — it tells the framework which fields to read from and where to write results. This section covers the essentials; for the complete syntax reference across all component types, see IO Binding.
The simplest way to bind a DataTransform is set_io():
from srforge.data import Entry
zscore = ZScoreNormalize()
zscore.set_io({"inputs": {"image": "image"}})
entry = Entry(image=torch.randn(1, 3, 64, 64))
entry = zscore(entry)
set_io() only configures the binding — it doesn't call anything yet. The transform runs when you call zscore(entry). That call is equivalent to:
You can apply the same transform to multiple fields at once:
The call zscore(entry) is now equivalent to:
entry["image"] = zscore.transform(image=entry["image"])
entry["target"] = zscore.transform(image=entry["target"])
To keep the originals and store results under new names:
zscore.set_io({
"inputs": [{"image": "image"}, {"image": "target"}],
"outputs": ["image_norm", "target_norm"],
})
entry = zscore(entry)
Now zscore(entry) is equivalent to:
entry["image_norm"] = zscore.transform(image=entry["image"])
entry["target_norm"] = zscore.transform(image=entry["target"])
How Type Annotations Control Container Handling¶
When an Entry field contains a container (dict, list, tuple) instead of a plain value, the framework decides whether to recurse into it or pass it whole based on the transform() method's type annotations.
Suppose entry["data"] is a dict of tensors: {"a": tensor_a, "b": tensor_b}. Here's what transform() receives depending on the annotation:
| Annotation | Entry field value | What transform() receives |
Calls |
|---|---|---|---|
x: torch.Tensor |
{"a": tensor_a, "b": tensor_b} |
tensor_a, then tensor_b (each leaf separately) |
One per leaf |
x: Union[torch.Tensor, np.ndarray] |
{"a": tensor_a, "b": tensor_b} |
tensor_a, then tensor_b |
One per leaf |
x: dict |
{"a": tensor_a, "b": tensor_b} |
{"a": tensor_a, "b": tensor_b} (the whole dict) |
One total |
x: list |
[tensor_a, tensor_b] |
[tensor_a, tensor_b] (the whole list) |
One total |
| No annotation | {"a": tensor_a, "b": tensor_b} |
tensor_a, then tensor_b |
One per leaf |
Why annotation-driven recursion?
This design means you write your transform once for a single tensor, and it automatically works when the Entry field is a dict of tensors (e.g., {"b2": tensor, "b8": tensor}). Without this, every transform would need manual dict iteration. The annotation is the opt-in signal: torch.Tensor says "I expect leaf values, recurse for me", while dict says "I need the whole container."
The annotation controls whether the framework recurses into containers, not what the leaf values must be. Leaf values can be anything — tensors, strings, integers, etc. The same recursion behavior applies regardless of what's inside the container. For example, x: torch.Tensor with {"a": "hello", "b": "world"} would still recurse and pass each string to transform() individually.
Annotations control recursion, not type checking
The annotation does not currently validate types — it only controls whether the framework recurses into containers. If the actual value doesn't match the annotation, transform() still gets called. For example, x: torch.Tensor with a string field still works — the string is passed directly (it's not a container). A future release will add type validation.
Recurses into Containers¶
class Downsample(DataTransform):
def transform(self, image: torch.Tensor) -> torch.Tensor:
return F.interpolate(image, scale_factor=0.5)
entry = Entry(lrs={"b1": tensor_b1, "b2": tensor_b2})
ds = Downsample()
ds.set_io({"inputs": {"image": "lrs"}})
entry = ds(entry)
Because image is annotated as torch.Tensor, the framework recurses into the dict and calls transform() on each leaf value. The output preserves the same container structure:
Passes Container as a Whole¶
When you annotate as dict (or list), the framework passes the entire container — useful when you need to see all values at once:
class SumAllBands(DataTransform):
def transform(self, bands: dict) -> torch.Tensor:
# Receives the ENTIRE dict — can combine across bands
return sum(bands.values())
entry = Entry(bands={"b2": tensor_b2, "b8": tensor_b8})
entry = SumAllBands().set_io({"inputs": {"bands": "bands"}, "outputs": "total"})(entry)
A more complex real-world example — resizing bands with different spatial resolutions to a common size:
class ResizeBandsToCommonResolution(DataTransform):
def transform(self, field: dict) -> dict:
max_h = max(v.shape[-2] for v in field.values() if isinstance(v, torch.Tensor))
max_w = max(v.shape[-1] for v in field.values() if isinstance(v, torch.Tensor))
return {k: resize(v, (max_h, max_w)) for k, v in field.items()}
The annotation controls behavior, not the config format — the same set_io() syntax works regardless.
Multi-Input DataTransform¶
Some operations need multiple inputs from different Entry fields — like histogram matching that needs both an image and a reference:
class MatchReference(DataTransform):
def transform(self, image: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
return histogram_match(image, reference)
Multi-input transforms use the same set_io() format — a dict with inputs and outputs keys that maps each transform() parameter to an Entry field:
matcher = MatchReference()
matcher.set_io({
"inputs": {"image": "lrs", "reference": "hr"},
"outputs": "matched",
})
entry = matcher(entry)
Equivalent to:
inputs and outputs are reserved keys
Avoid naming your transform() parameters inputs or outputs. These names are used by the config syntax — when the framework sees a dict with an inputs or outputs key, it interprets it as the IO config structure, not as a parameter-to-field mapping.
Parallel Recursion¶
When both parameters are Tensor-annotated and the fields contain dicts, the framework performs parallel recursion — iterating over matching dict keys:
entry = Entry(
lrs={"b1": tensor_b1, "b2": tensor_b2},
hr={"b1": ref_b1, "b2": ref_b2},
)
entry = MatchReference().set_io({
"inputs": {"image": "lrs", "reference": "hr"},
"outputs": "matched",
})(entry)
# Result: entry["matched"] = {
# "b1": histogram_match(tensor_b1, ref_b1),
# "b2": histogram_match(tensor_b2, ref_b2),
# }
For more on multi-input config syntax, see the IO Binding Reference.
Optional Parameters¶
Parameters with default values in transform() are optional — you can omit them from the IO config, and the default value applies. But if you do map them, the Entry field must exist:
class AddNoise(DataTransform):
def transform(self, image: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
noise = torch.randn_like(image)
if mask is not None:
noise = noise * mask
return image + noise
noise = AddNoise()
# OK — mask not in IO config, default (None) applies
noise.set_io({"inputs": {"image": "x"}, "outputs": "noisy"})
noise(Entry(x=tensor))
# OK — mask mapped and field exists
noise.set_io({"inputs": {"image": "x", "mask": "m"}, "outputs": "noisy"})
noise(Entry(x=tensor, m=mask_tensor))
# FAILS — mask mapped but field missing → KeyError
noise.set_io({"inputs": {"image": "x", "mask": "m"}, "outputs": "noisy"})
noise(Entry(x=tensor)) # KeyError: field 'm' not found
This applies to both transform() and transform_unbatched(). The rule is simple: if you map it, it must exist.
IO Binding Options¶
The sections below show all the ways to configure IO binding — in Python, YAML, and the flow DSL. Some reference concepts covered later (YAML configuration, SequentialModel). Feel free to skim and return later.
There are three ways to configure which Entry fields a DataTransform reads from and writes to. set_io() only configures the binding — the transform runs when you call zscore(entry).
In Python (programmatic):
ZScoreNormalize().set_io({"inputs": {"image": "image"}}) # In-place
ZScoreNormalize().set_io({"inputs": [{"image": "image"}, {"image": "target"}]}) # Multiple fields
ZScoreNormalize().set_io({"inputs": {"image": "image"}, "outputs": "image_norm"}) # New output name
In YAML config:
zscore:
_target: ZScoreNormalize
io:
inputs:
image: image # Read from entry["image"]
outputs: image_norm # Write to entry["image_norm"]
In SequentialModel flow DSL:
All three configure the same binding. When the transform is called, it executes: entry["image_norm"] = zscore.transform(image=entry["image"]).
All three are interchangeable. For the complete syntax reference, see IO Binding Reference: DataTransform.
Works with unbatched entries too
Even though transform() receives batched data, it works transparently with unbatched entries (from a Dataset). The framework auto-wraps via collate([entry]), calls transform() with [1, C, H, W], and unwraps via result[0]. The wrapping overhead is negligible. Alternatively, you can implement transform_unbatched() instead — see Per-Sample Processing.
In SequentialModel¶
Inside a SequentialModel, DataTransforms get their IO from the flow DSL — no set_io() needed:
DataTransforms are fully reusable — the same instance can appear in multiple flow lines with different field mappings:
mul = Multiply(2.0)
flow = [
"x -> mul -> x_scaled",
"y -> mul -> y_scaled", # Same instance, different fields
]
For multi-parameter transforms, use positional or named mapping — the same syntax as Model:
# Positional — maps by transform() signature order
flow = ["x, ref -> match -> matched"]
# Named — explicit parameter mapping
flow = ["(image=x, reference=ref) -> match -> matched"]
See SequentialModel: DataTransforms in Flow for details.
Examples¶
Image Resizing¶
class Resize(DataTransform):
def __init__(self, size: tuple[int, int]):
super().__init__()
self.size = size
def transform(self, image: torch.Tensor) -> torch.Tensor:
return F.interpolate(image, size=self.size, mode='bilinear')
resize:
_target: Resize
params:
size: [256, 256]
io:
inputs: [{image: image}, {image: target}]
outputs: [image_resized, target_resized]
Adding Gaussian Noise¶
class AddGaussianNoise(DataTransform):
def __init__(self, sigma: float = 0.1):
super().__init__()
self.sigma = sigma
def transform(self, image: torch.Tensor) -> torch.Tensor:
noise = torch.randn_like(image) * self.sigma
return image + noise
Color Space Conversion¶
class RGBtoYCbCr(DataTransform):
def transform(self, image: torch.Tensor) -> torch.Tensor:
return self._convert_rgb_to_ycbcr(image)
Clipping Values (Multi-Field)¶
class ClipValues(DataTransform):
def __init__(self, min_val: float = 0.0, max_val: float = 1.0):
super().__init__()
self.min_val = min_val
self.max_val = max_val
def transform(self, image: torch.Tensor) -> torch.Tensor:
return torch.clamp(image, self.min_val, self.max_val)
clip:
_target: ClipValues
io:
inputs: [{image: image_norm}, {image: target_norm}]
outputs: [image_clipped, target_clipped]
Stacking Multi-Temporal Data¶
For multi-temporal data where each sample has a list of images (e.g. lrs=[t1, t2]), StackSequence converts the collected list into a single tensor after collation:
from srforge.transform.data import StackSequence
stack = StackSequence(dim=0)
stack.set_io({"inputs": {"image": "lrs"}})
# Input: entry.lrs == [[img1, img2], [img3, img4]] — 2 samples, 2 images each (C,H,W)
# Output: entry.lrs == tensor [2, 2, C, H, W] — stacked along dim=0, batch stacked
entry = stack(entry)
For simpler cases where the field is a flat list of tensors, use StackTensors:
from srforge.transform.data import StackTensors
stack = StackTensors(dim=1)
result = stack.transform([torch.randn(3, 8, 8) for _ in range(4)])
# result.shape == [4, 3, 8, 8]
Common Mistakes¶
Wrong Annotation for Container-Level Operations¶
# BAD: Tensor annotation causes recursion into dicts
class ResizeBands(DataTransform):
def transform(self, field: torch.Tensor) -> torch.Tensor:
# This receives INDIVIDUAL tensors, not the whole dict!
...
# GOOD: dict annotation passes the container as-is
class ResizeBands(DataTransform):
def transform(self, field: dict) -> dict:
# This receives the WHOLE dict
...
The transform() parameter annotation controls recursion behavior. Use torch.Tensor for per-element operations, dict (or other non-Tensor types) when you need the whole container.
Mapped Field Must Exist¶
If you map a parameter (even an optional one) to an Entry field, that field must be present at runtime. This catches misconfigured pipelines where a field you expected is missing:
noise.set_io({"inputs": {"image": "x", "mask": "m"}, "outputs": "noisy"})
noise(Entry(x=tensor)) # KeyError: field 'm' not found in entry
If you don't need mask, omit it from the IO config entirely — the Python default applies:
noise.set_io({"inputs": {"image": "x"}, "outputs": "noisy"})
noise(Entry(x=tensor)) # OK — mask defaults to None
Per-Sample Processing: transform_unbatched¶
Everything above uses transform(), which receives batched data. If you prefer to work with single samples, implement transform_unbatched() instead — write your logic for one sample, and the framework handles batching automatically:
class PadToSquare(DataTransform):
def transform_unbatched(self, image: torch.Tensor) -> torch.Tensor:
# image shape: [C, H, W] — no batch dim
_, h, w = image.shape
size = max(h, w)
padded = torch.zeros(image.shape[0], size, size)
padded[:, :h, :w] = image
return padded
Per-sample code is often simpler to write — no batch dimension to manage, shapes are intuitive, and you can use plain indexing. The framework handles looping over the batch and re-stacking the results.
This is a complete DataTransform. All IO binding, annotation-driven recursion, and SequentialModel integration work unchanged — the only difference is that the framework loops over each sample when the Entry is batched.
What transform_unbatched Receives¶
Values arrive without a batch dimension — each sample is processed individually:
Batched (what transform sees) |
Per-sample (what transform_unbatched sees) |
|---|---|
tensor([B, C, H, W]) |
tensor([C, H, W]) |
["s1", "s2", ...] |
"s1" |
[["b2","b8"], ...] |
["b2", "b8"] |
tensor([4, 8]) |
tensor(4) (0-dim scalar) |
dict {"a": tensor[B,...]} |
dict {"a": tensor[...]} (recursed) |
The unbatching uses integer indexing (tensor[i], list[i]) — so the batch dim is removed. Return values are re-batched via torch.stack (tensors), list collection (everything else), or per-key/per-position recursion (dicts/tuples).
IO Binding¶
IO binding works identically regardless of which method you implemented. The framework inspects the overridden method's signature for parameter names and type annotations:
class PerSampleAdd(DataTransform):
def transform_unbatched(self, image: torch.Tensor, offset: torch.Tensor) -> torch.Tensor:
return image + offset
t = PerSampleAdd()
t.set_io({"inputs": {"image": "lr", "offset": "bias"}, "outputs": "result"})
result = t(entry) # Works with Entry, in SequentialModel, and standalone
All IO binding syntax (flat, structured, YAML, flow DSL) works the same way.
Multiple Outputs¶
Both transform() and transform_unbatched() can return tuples for multiple outputs. Use the structured set_io() syntax with a list of output field names:
class SplitChannels(DataTransform):
def transform(self, image: torch.Tensor):
return image[:, :3], image[:, 3:] # first 3 channels, rest
class SplitPerSample(DataTransform):
def transform_unbatched(self, image: torch.Tensor):
return image[:3], image[3:] # first 3 channels, rest
# Both use the same structured IO config:
t.set_io({"inputs": {"image": "image"}, "outputs": ["rgb", "extra"]})
For transform_unbatched, each tuple element is re-batched independently.
Note
Multiple outputs require a list in the outputs key. The flow DSL (x -> t -> y) pairs inputs to outputs 1:1 for single-input transforms. For multi-param transforms, use named mapping: (image=x, ref=r) -> t -> out.
In SequentialModel¶
transform_unbatched-only DataTransforms work transparently in SequentialModel flows — no special syntax or configuration needed:
seq = SequentialModel(
modules={"double": PerSampleDouble()},
flow="image -> double -> image",
)
e1 = Entry(name="s1", image=torch.ones(3, 4, 4))
e2 = Entry(name="s2", image=torch.ones(3, 4, 4))
batch = Entry.collate([e1, e2])
result = seq(batch) # Each sample doubled independently, then re-batched
The SequentialModel doesn't need to know which method the transform implements. The framework handles dispatch internally.
Next: EntryTransform — When you need full control over Entry fields