Skip to content

Give every GPU transfer augmentation a configurable mix range #137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Hendrik-code wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Give every GPU transfer augmentation a configurable mix range #137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
984f037
feat: give every GPU transfer augmentation a configurable mix range
Hendrik-code Oct 15, 2026
0d7e5ac
fix: make the mix attributes visible to mypy, and reject a bool range
Hendrik-code Oct 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 83 additions & 0 deletions smauglab/configs/all_augmentations.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,11 @@
"keepdim": true,
"max_res_aniso": 8.0,
"max_res_iso": 4.0,
"mix_alpha": [
0.0,
1.0
],
"mix_prob": 0.0,
"n_channels": 1,
"n_neutral_labels": null,
"nonlin_scale": 0.04,
Expand Down Expand Up @@ -109,6 +114,11 @@
"label_classes": null,
"label_remap_prob": 0.5,
"min_label_voxels": 4,
"mix_alpha": [
0.0,
1.0
],
"mix_prob": 0.0,
"n_kmeans_subsample": 10000,
"p": 1.0,
"p_batch": 1.0,
Expand Down Expand Up @@ -151,6 +161,11 @@
"dark_threshold": 0.01,
"initial_partitioner": null,
"keepdim": false,
"mix_alpha": [
0.0,
1.0
],
"mix_prob": 0.0,
"overlay": null,
"p": 1.0,
"p_batch": 1.0,
Expand All @@ -167,6 +182,11 @@
"blend_targets": 1,
"include_self": true,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_prob": 0.0,
"p": 0.2,
"p_batch": 1.0,
"p_class_mix": 0.0,
Expand All @@ -186,6 +206,10 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
Expand All @@ -200,6 +224,10 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
Expand All @@ -218,6 +246,11 @@
],
"in_seg": 0.2,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_prob": 0.0,
"p": 1.0,
"p_batch": 1.0,
"retain_stats": true,
Expand All @@ -234,6 +267,10 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
Expand All @@ -249,6 +286,10 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
Expand All @@ -263,6 +304,10 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
Expand All @@ -285,6 +330,10 @@
5,
7
],
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
Expand Down Expand Up @@ -327,6 +376,10 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
Expand Down Expand Up @@ -409,7 +462,12 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
"p": 1.0,
"p_batch": 1.0,
Expand All @@ -422,7 +480,12 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
"p": 1.0,
"p_batch": 1.0,
Expand All @@ -435,7 +498,12 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
"p": 1.0,
"p_batch": 1.0,
Expand All @@ -448,7 +516,12 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
"p": 1.0,
"p_batch": 1.0,
Expand All @@ -461,7 +534,12 @@
],
"in_seg": 0.0,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"out_seg": 0.0,
"p": 1.0,
"p_batch": 1.0,
Expand Down Expand Up @@ -514,7 +592,12 @@
"in_seg": 0.0,
"invert": false,
"keepdim": true,
"mix_alpha": [
0.0,
1.0
],
"mix_in_out": false,
"mix_prob": 0.0,
"order": 3,
"out_seg": 0.0,
"p": 1.0,
Expand Down
11 changes: 11 additions & 0 deletions smauglab/transforms/gpu/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
from torch.distributions import RelaxedBernoulli
from torch.nn import Module

from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA

try: # private to kornia, and the only part of the faster parameter draw that needs it
from kornia.augmentation.container.image import ImageSequentialBase, _get_new_batch_shape

Expand Down Expand Up @@ -84,6 +86,15 @@ class ImageOnlyTransform(RigidAffineAugmentationBase3D):

"""

#: Blending the result back towards the input; see `gpu/mixing.py`. Declared here,
#: with the values that mean "off", rather than only assigned by `configure_mix` in
#: the constructors that expose them: an attribute a type checker cannot see on the
#: class resolves through `nn.Module.__getattr__`, which is typed `Tensor | Module`,
#: and neither of those is a number. Only the transforms that take them as
#: constructor arguments can be configured, since the signature is the config schema.
mix_prob: float = 0.0
mix_alpha: tuple[float, float] = DEFAULT_MIX_ALPHA

def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor:
# These transforms never move anything, so the matrix is always the
# identity -- shared and cached rather than rebuilt per module. See
Expand Down
Loading
Loading