Skip to content

feat: shift the scan inside the patch and fill what lies outside it with EM-synthesised background #144

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

Draft
Hendrik-code wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: shift the scan inside the patch and fill what lies outside it with EM-synthesised background #144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 2 additions & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ MANIFEST
# configs beside it are tracked only because they predate the blanket rule above;
# a new one has to say so.
!smauglab/configs/transform_params_paper-synthseg.json
# The paper's Ours plus RandomFOVExtensionGPU (README "Field of view beyond the scan").
!smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json
*.yaml
*.nii
*.nii.gz
Expand Down
63 changes: 63 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,68 @@ transfer augmentations (`RandomScharrGPU`, `RandomUnsharpMaskGPU`, `RandomRandCo
`p=0` and `RandomSynthSegGPU` at `p=1.0` in their place. The general-enhancement half is
identical between the two.

### Field of view beyond the scan

A model trained on a narrow field of view has never been shown what lies outside it. A
sagittal spine slab is about 8 cm wide, and nnU-Net centres it in a 128 mm patch and pads
the rest with 0. So every training window contains spine, and beyond about ±4 cm there is
only a flat slab labelled background. On a whole-body axial scan the model then finds
vertebrae in muscle, gluteals, fat and bowel. For the Ours setup trained on SpineGAN in-phase,
60% of the false-positive voxels on NAKO VIBE lie more than 5 cm from the spine. The contrast
augmentations do not touch this, because the problem is context, not intensity.

`RandomFOVExtensionGPU` shows the network windows it can only get right by rejecting
context. For each selected sample it does six things:

1. **Shift.** Image and segmentation move by the same whole-voxel shift, up to `shift` of
the patch per axis. At `shift: 1.0` the spine can leave the window entirely.
2. **Find the fill region.** That is everything that is not scan: the padding (every
channel exactly `pad_value`) plus whatever the shift vacated.
3. **Split the background with EM.** The patch's own background is split into a random 3
to 10 intensity regions, the same way SynthSeg's label completion splits its
background. Only background voxels count: label 0, more than `margin` voxels from any
label, and off the padding's interpolated edge.
4. **Post-process the region map.** A 3×3×3 majority filter removes the EM speckle, and a
random diffeomorphic field (SynthSeg's) deforms the map.
5. **Composite into the fill.** `layers` copies of the map, each randomly rotated, scaled,
mirrored and translated, are composited into the fill region.
6. **Render.** Each fill voxel is drawn from its region's mean and standard deviation,
measured on the patch itself. The within-region variation is spatially correlated
(σ = 1 voxel), because iid draws read as grain next to the smooth scan. The result is
then blurred as SynthSeg blurs, and its label stays 0.

| parameter | default | meaning |
| --- | --- | --- |
| `shift` | `0.5` | maximum shift as a fraction of the patch, one value or one per axis; `0` only fills the existing padding |
| `regions` | `[3, 10]` | range for the number of EM regions, drawn per sample |
| `margin` | `4` | voxels between a label and the background used to build the fill |
| `layers` | `6` | warped copies of the region map composited into the fill |
| `pad_value` | `0.0` | the intensity that marks padding; nnU-Net pads with 0 |
| `p` | `0.5` | probability per sample |

Things to know before using it:

* **Selection is always per sample.** The transform keeps this even when the pipeline sets
`same_on_batch`. Without a `pipeline` section a config runs with `same_on_batch: true`,
which would make `p` one coin per batch and could take the spine out of both samples of
a step.
* **The logged training loss rises.** A window the shift has emptied of spine has no
foreground, so its Dice term is about 0 and cross-entropy does the work. The loss curve
is therefore not comparable with a run without the transform.
* **Pipeline position matters.** The transform runs right after `RandomAffineGPU` and
before every intensity transform, so the synthesised background is augmented like real
tissue. A label-to-image generator later in the pipeline (SynthSeg, PALETTE synthesis)
regenerates the fill from label 0, which overwrites it.
* **Padding is found by exact value.** Two cases follow from that:
* Background that the scanner itself zeroed (the masked air of a Dixon reconstruction)
normalises to a constant other than 0. It counts as scan. It stays where it is, and
it can become one of the EM regions.
* With `use_mask_for_norm`, everything outside the body is 0, so it gets filled.

`smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json` (`ae065824`)
is the paper's Ours config (`c69b9872`) plus this transform at the defaults above. Nothing
else changes.

## Available augmentations

Which augmentations exist, and which of the two backends implements each one. An empty
Expand All @@ -222,6 +284,7 @@ column -- see "Run Monai training with SmaugLab augmentations" above. Regenerate
| affine | GEO | `RandomAffineGPU` | — |
| crop | GEO | `RandomCropTransformGPU` | — |
| spatial | GEO | — | `SpatialTransform` |
| fov_extension | GEO | `RandomFOVExtensionGPU` | — |
| gaussian_noise | GE | `RandomGaussianNoiseGPU` | `GaussianNoiseTransform` |
| gaussian_blur | GE | `RandomGaussianBlurGPU` | `GaussianBlurTransform` |
| brightness | GE | `RandomBrightnessGPU` | `MultiplicativeBrightnessTransform` |
Expand Down
13 changes: 13 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 @@ -35,6 +35,19 @@
0.1
]
},
"RandomFOVExtensionGPU": {
"keepdim": true,
"layers": 6,
"margin": 4,
"p": 0.5,
"p_batch": 1.0,
"pad_value": 0.0,
"regions": [
3,
10
],
"shift": 0.5
},
"RandomSynthSegGPU": {
"apply_bias_field": true,
"apply_intensity_augmentation": true,
Expand Down
241 changes: 241 additions & 0 deletions smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,241 @@
{
"GPU": {
"RandomFlipTransformGPU": {
"flip_axis": [
0
],
"same_on_batch": false,
"keepdim": true,
"p": 0.5
},
"RandomAffineGPU": {
"degrees": 5,
"translate": [
0.1,
0.1,
0.1
],
"scale": [
0.7,
1.4
],
"shears": [
-5,
5,
-5,
5,
-5,
5
],
"resample": "bilinear",
"p": 0.0
},
"RandomFOVExtensionGPU": {
"shift": 0.5,
"regions": [
3,
10
],
"margin": 4,
"layers": 6,
"pad_value": 0.0,
"p": 0.5
},
"RandomInverseGPU": {
"retain_stats": true,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"mix_prob": 0.0,
"p": 0.25
},
"RandomHistogramEqualizationGPU": {
"retain_stats": true,
"in_seg": 0,
"out_seg": 0,
"mix_in_out": true,
"mix_prob": 0.6,
"p": 0.25
},
"RandomRedistributeSegGPU": {
"in_seg": 0.25,
"retain_stats": true,
"p": 0.4
},
"RandomScharrGPU": {
"absolute": true,
"retain_stats": true,
"mix_prob": 0.5,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.25
},
"RandomUnsharpMaskGPU": {
"sigma": 1.0,
"unsharp_amount": 1.5,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"mix_prob": 0.5,
"p": 0.25
},
"RandomRandConvGPU": {
"kernel_sizes": [
3,
5,
7
],
"retain_stats": true,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"mix_prob": 0.5,
"p": 0.2
},
"RandomClampGPU": {
"max_clamp_amount": 0.2,
"retain_stats": true,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.0
},
"RandomGaussianNoiseGPU": {
"mean": 0.0,
"std": 1.0,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.1
},
"RandomGaussianBlurGPU": {
"sigma": 1.0,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.2
},
"RandomBrightnessGPU": {
"brightness_range": [
0.75,
1.25
],
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.15
},
"RandomGammaGPU": {
"gamma_range": [
0.7,
1.5
],
"retain_stats": true,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.3
},
"RandomInvGammaGPU": {
"gamma_range": [
0.7,
1.5
],
"retain_stats": true,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.1
},
"RandomContrastGPU": {
"contrast_range": [
0.75,
1.25
],
"retain_stats": false,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.15
},
"RandomLog1pGPU": {
"retain_stats": true,
"in_seg": 0,
"out_seg": 0,
"mix_in_out": false,
"p": 0.1
},
"RandomSqrtGPU": {
"retain_stats": true,
"in_seg": 0,
"out_seg": 0,
"mix_in_out": false,
"p": 0.1
},
"RandomSinGPU": {
"retain_stats": true,
"in_seg": 0,
"out_seg": 0,
"mix_in_out": false,
"p": 0.1
},
"RandomExpGPU": {
"retain_stats": true,
"in_seg": 0,
"out_seg": 0,
"mix_in_out": false,
"p": 0.1
},
"RandomSigmoidGPU": {
"retain_stats": true,
"in_seg": 0,
"out_seg": 0,
"mix_in_out": false,
"p": 0.1
},
"RandomLowResTransformGPU": {
"scale": [
0.5,
1.0
],
"same_on_batch": false,
"p": 0.25
},
"RandomAcqTransformGPU": {
"scale": [
0.6,
1.0
],
"same_on_batch": false,
"p": 0.0
},
"RandomBiasFieldGPU": {
"retain_stats": true,
"coefficients": 0.2,
"in_seg": 0.0,
"out_seg": 0.0,
"mix_in_out": true,
"p": 0.15
},
"ZscoreNormalizationGPU": {
"p": 0.0
}
},
"CPU": {
"SpatialTransform": {
"patch_center_dist_from_border": 80,
"random_crop": false,
"p_elastic_deform": 0.0,
"p_rotation": 0.2,
"p_scaling": 0.2,
"scaling": [
0.7,
1.4
],
"p_synchronize_scaling_across_axes": 1,
"bg_style_seg_sampling": false,
"mode_seg": "nearest"
}
}
}
5 changes: 5 additions & 0 deletions smauglab/registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ class AugId(str, Enum):
AFFINE = "affine"
CROP = "crop"
SPATIAL = "spatial" # nnU-Net's own SpatialTransform
FOV_EXTENSION = "fov_extension"
# -- general enhancement
GAUSSIAN_NOISE = "gaussian_noise"
GAUSSIAN_BLUR = "gaussian_blur"
Expand Down Expand Up @@ -200,6 +201,10 @@ def __init__(self, source: str, problems: list[str]) -> None:
Backend.GPU: (
"RandomFlipTransformGPU", # FlipTransform
"RandomAffineGPU", # AffineTransform
# No ladder key. After the last transform that can open padding (the affine's zero
# fill) and before everything that rewrites intensities, so the synthesised
# background is augmented exactly like the anatomy around it.
"RandomFOVExtensionGPU",
"RandomSynthSegGPU", # SynthSeg
"RandomPaletteGPU", # RandomPALETTETransform
"PaletteSynthesisGPU", # same slot: it is the same augmentation, composed
Expand Down
1 change: 1 addition & 0 deletions smauglab/transforms/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ class carries an `@register(...)` decorator, so the registry is complete once th
"""

from smauglab.transforms.cpu import artifact, contrast, external, fromSeg, spatial # noqa: F401
from smauglab.transforms.gpu import context as gpu_context # noqa: F401
from smauglab.transforms.gpu import contrast as gpu_contrast # noqa: F401
from smauglab.transforms.gpu import domain_transfer # noqa: F401
from smauglab.transforms.gpu import fromSeg as gpu_fromSeg
Expand Down
Loading
Loading