diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000..a693997 --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,9 @@ +# Revisions to skip in `git blame`. +# +# Enable locally with: +# git config blame.ignoreRevsFile .git-blame-ignore-revs +# +# GitHub picks this file up automatically. + +# style: apply ruff autofix and formatter (#34) +4c78918e1be261ca2373f1fbfb302654c7803a03 diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..a627ce4 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,12 @@ +# Owners are requested for review automatically on every pull request. +# Combined with a branch protection rule requiring one approval, this is the +# "requiring code-review" half of issue #34. +# +# See https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners + +* @Hendrik-code @NathanMolinier + +# Changing the CI, packaging or linting setup affects every contributor. +/.github/ @Hendrik-code +/pyproject.toml @Hendrik-code +/.pre-commit-config.yaml @Hendrik-code diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..9c36a88 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,45 @@ +--- +name: Bug report +about: Something does not work as expected +title: '' +labels: bug +assignees: '' +--- + +**Describe the bug** +A clear and concise description of what goes wrong. + +**To reproduce** +Steps to reproduce the behaviour, ideally with the exact command: + +```bash +# e.g. AUGLAB_PARAMS_GPU_JSON=/abs/path/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU +``` + +**Config JSON** +If the problem involves augmentation parameters, paste the relevant part of your +transform params JSON (or attach the file). + +```json + +``` + +**Expected behaviour** +What you expected to happen instead. + +**Error output** +The full traceback, not just the last line. + +``` + +``` + +**Environment** +- AugLab version or commit: +- Python version: +- PyTorch version and CUDA build: +- nnU-Net version (if applicable): +- OS: + +**Additional context** +Anything else that might matter — dataset, image orientation, patch size. diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..b8c41cd --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,25 @@ +--- +name: Feature request +about: Suggest a new augmentation, option, or improvement +title: '' +labels: enhancement +assignees: '' +--- + +**What problem does this solve?** +A clear description of the limitation you are hitting. + +**Proposed solution** +What you would like AugLab to do. + +**If this is a new augmentation** +- What does it simulate (acquisition artefact, contrast change, anatomy change)? +- Reference or paper, if there is one: +- Should it run on GPU, CPU, or both? +- What parameters should the config JSON expose? + +**Alternatives considered** +Other approaches you thought about. + +**Additional context** +Screenshots, example images, or links. diff --git a/.github/ISSUE_TEMPLATE/question.md b/.github/ISSUE_TEMPLATE/question.md new file mode 100644 index 0000000..095917b --- /dev/null +++ b/.github/ISSUE_TEMPLATE/question.md @@ -0,0 +1,17 @@ +--- +name: Question +about: Ask about usage, configuration, or results +title: '' +labels: question +assignees: '' +--- + +**Your question** +What you are trying to do and where you are stuck. + +**What you have tried** +Commands, config JSONs, or documentation you already looked at. + +**Environment (if relevant)** +- AugLab version or commit: +- Python / PyTorch version: diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..7db9b2b --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,26 @@ +## What does this change? + + + +## Why? + + + +## How was it tested? + + + +- [ ] `pytest` passes locally +- [ ] `pre-commit run --all-files` passes +- [ ] Added or updated tests covering the change +- [ ] Ran a training / augmentation job end to end + +## Anything reviewers should look at closely? + + + +## Checklist + +- [ ] Augmentation behaviour is unchanged, or the change is intentional and described above +- [ ] New transforms are reachable from a config JSON +- [ ] `project.version` in `pyproject.toml` bumped, if this should be released diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..adc0ea4 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,33 @@ +name: lint + +on: + push: + branches: [main] + pull_request: + branches: [main] + +permissions: + contents: read + +concurrency: + group: lint-${{ github.ref }} + cancel-in-progress: true + +jobs: + # Runs the pre-commit hooks rather than a bare `ruff check`, so the hooks + # contributors run locally and the ones CI enforces cannot drift apart. + # + # Deliberately does NOT auto-commit fixes back to the branch: that breaks on + # pull requests from forks and rewrites contributors' branches under them. + # A red check plus `pre-commit run --all-files` locally is the fix. + pre-commit: + name: pre-commit + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - uses: pre-commit/action@v3.0.1 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..04cba4c --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,125 @@ +name: publish to PyPI + +on: + release: + types: [created] + workflow_dispatch: + inputs: + repository: + description: Which index to upload to + required: true + default: testpypi + type: choice + options: [testpypi, pypi] + +permissions: + contents: read + +jobs: + publish: + name: build and upload + runs-on: ubuntu-latest + # Requires the PYPI_API_TOKEN (and, for dry runs, TEST_PYPI_API_TOKEN) + # repository secret. Until that is set by an admin this job cannot upload. + environment: pypi + + steps: + - uses: actions/checkout@v4 + with: + # poetry-dynamic-versioning derives the version from the git tag; + # a shallow clone has no tags and would build as 0.0.0. + fetch-depth: 0 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install build tooling + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build sdist and wheel + run: python -m build + + - name: Check the built version came from the release tag + # poetry-dynamic-versioning derives the version from the tag, so the two cannot + # disagree the way a hand-maintained project.version could. What can + # still go wrong is the tag not matching [tool.poetry-dynamic-versioning].pattern, + # or the checkout arriving without tags -- both of which silently yield + # the 0.0.0 fallback or a .dev version. Catch that before uploading. + if: github.event_name == 'release' + run: | + python - <<'PY' + import glob, os, pathlib, re, sys + + from packaging.version import InvalidVersion, Version + + tag = os.environ["RELEASE_TAG"] + # Same prefixes as [tool.poetry-dynamic-versioning].pattern. + expected = re.sub(r"^(?:[rvV]|release[-_])", "", tag) + built = {pathlib.Path(p).name.split("-")[1] for p in glob.glob("dist/*.whl")} + + print(f"tag={tag!r} expected={expected!r} built={sorted(built)}") + + # Compare parsed versions, not strings: a tag like v2.0.0-beta1 is + # legitimately normalised to 2.0.0b1 in the artifact name. + try: + if {Version(b) for b in built} != {Version(expected)}: + raise InvalidVersion + except InvalidVersion: + sys.exit( + f"built version {sorted(built)} does not match release tag '{tag}'. " + "Either the tag does not match [tool.poetry-dynamic-versioning].pattern in " + "pyproject.toml, or the checkout has no tags (needs fetch-depth: 0)." + ) + PY + env: + RELEASE_TAG: ${{ github.event.release.tag_name }} + + - name: Refuse to publish a fallback version + run: | + python - <<'PY' + import glob, pathlib, sys + + built = {pathlib.Path(p).name.split("-")[1] for p in glob.glob("dist/*.whl")} + if "0.0.0" in built: + sys.exit( + "refusing to publish 0.0.0 -- poetry-dynamic-versioning found no git tag, so the " + "version is the fallback rather than a real release" + ) + print(f"version looks real: {sorted(built)}") + PY + + - name: Check distribution metadata + run: twine check dist/* + + - name: Verify the wheel ships the package data + run: | + python - <<'PY' + import glob, sys, zipfile + + names = zipfile.ZipFile(glob.glob("dist/*.whl")[0]).namelist() + if "auglab/configs/transform_params_gpu.json" not in names: + sys.exit("refusing to publish: wheel has no config JSONs") + if len([n for n in names if n.endswith(".py")]) < 20: + sys.exit("refusing to publish: wheel is missing modules") + print("wheel contents look complete") + PY + + - name: Upload to TestPyPI + if: github.event_name == 'workflow_dispatch' && inputs.repository == 'testpypi' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }} + TWINE_REPOSITORY: testpypi + run: twine upload dist/* + + - name: Upload to PyPI + if: github.event_name == 'release' || (github.event_name == 'workflow_dispatch' && inputs.repository == 'pypi') + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + # Uploads the sdist as well as the wheel; auglab currently has no sdist + # on PyPI, which blocks anyone who needs to build from source. + run: twine upload dist/* diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..b193866 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,156 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + branches: [main] + +permissions: + contents: read + +concurrency: + group: tests-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + name: test (python ${{ matrix.python-version }}) + # Linux only. The suite is pure CPU torch/kornia, and paying the torch + # install cost on a Windows runner buys nothing for a pipeline that only + # ever runs on Linux clusters. + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # pyproject sets requires-python = ">=3.10". + python-version: ["3.10", "3.11", "3.12"] + + steps: + - uses: actions/checkout@v4 + with: + # poetry-dynamic-versioning derives the version from the git tag; + # a shallow clone has no tags and would build as 0.0.0. + fetch-depth: 0 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + cache-dependency-path: pyproject.toml + + - name: Install CPU-only PyTorch + # Must come first. Installing from the default index pulls the CUDA + # build (several GB of nvidia-* wheels), which is slow and can fill the + # runner disk. Nothing here needs a GPU. + run: | + python -m pip install --upgrade pip + pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu + + - name: Install AugLab + run: pip install -e ".[dev]" + + - name: Run tests + run: pytest -v --cov=auglab --cov-report=xml --cov-report=term-missing + + - name: Upload coverage to Codecov + # Only after a merge to main, and only once per matrix. + if: matrix.python-version == '3.12' && github.event_name == 'push' + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: false + + kornia-compat: + # auglab subclasses kornia's private augmentation internals, which move + # between minor releases -- kornia 0.8.3 dropped kornia.core.Module and the + # whole kornia.utils.helpers module. The `test` job only ever installs the + # newest kornia, so it cannot catch a break at the other end of the + # supported range. This job pins both ends. + name: kornia ${{ matrix.kornia-version }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # Oldest and newest supported, matching the pin in pyproject.toml. + kornia-version: ["0.7.3", "0.8.3"] + + steps: + - uses: actions/checkout@v4 + with: + # poetry-dynamic-versioning derives the version from the git tag; + # a shallow clone has no tags and would build as 0.0.0. + fetch-depth: 0 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + + - name: Install CPU-only PyTorch + run: | + python -m pip install --upgrade pip + pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu + + - name: Install AugLab with a pinned kornia + run: | + pip install -e ".[dev]" + pip install "kornia==${{ matrix.kornia-version }}" + + - name: Run tests + run: pytest -q -m "not slow" + + build: + name: build distribution + # Catches packaging breakage on every PR instead of on release day. The + # published 20260109 wheel shipped without any config JSONs; this job plus + # unit_tests/test_packaging.py is what stops that recurring. + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + # poetry-dynamic-versioning derives the version from the git tag; + # a shallow clone has no tags and would build as 0.0.0. + fetch-depth: 0 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install build tooling + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build sdist and wheel + run: python -m build + + - name: Check distribution metadata + run: twine check dist/* + + - name: Verify the wheel ships the package data + run: | + python - <<'PY' + import glob, sys, zipfile + + wheel = glob.glob("dist/*.whl")[0] + names = zipfile.ZipFile(wheel).namelist() + modules = [n for n in names if n.endswith(".py")] + configs = [n for n in names if n.startswith("auglab/configs/") and n.endswith(".json")] + + print(f"{wheel}: {len(modules)} modules, {len(configs)} configs") + problems = [] + if len(modules) < 20: + problems.append(f"only {len(modules)} modules in the wheel") + if "auglab/configs/transform_params_gpu.json" not in configs: + problems.append("default transform_params_gpu.json is missing") + if problems: + sys.exit("wheel is incomplete: " + "; ".join(problems)) + PY + + - uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/ diff --git a/.gitignore b/.gitignore index 62f7dfb..5707ab7 100644 --- a/.gitignore +++ b/.gitignore @@ -40,6 +40,9 @@ share/python-wheels/ .installed.cfg *.egg MANIFEST +*.json +*.yaml +auglab/configs_paul/* # PyInstaller # Usually these files are written by a python script from a template @@ -164,6 +167,9 @@ dmypy.json # pytype static type analyzer .pytype/ +# Ruff +.ruff_cache/ + # Cython debug symbols cython_debug/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..0667716 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,30 @@ +# See https://pre-commit.com for more information +# See https://pre-commit.com/hooks.html for more hooks +# +# Install once per clone with: +# pip install -e ".[dev]" +# pre-commit install +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v6.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-json + - id: check-merge-conflict + - id: check-added-large-files + # Augmentation work produces large NIfTI/weight files; .gitignore covers + # the usual suspects, this catches the rest before they reach a PR. + args: [--maxkb=1000] + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.16.1 + hooks: + # Run the linter. + - id: ruff + types_or: [python, pyi] + args: [--fix] + # Run the formatter. + - id: ruff-format + types_or: [python, pyi] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..c7c16ab --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,128 @@ +# Contributing to AugLab + +Thanks for contributing. This page covers the development setup, the checks CI +runs, and how versioning and releases work. + +## Development setup + +```bash +git clone git@github.com:neuropoly/AugLab.git +cd AugLab + +python3 -m venv venv +source venv/bin/activate + +# PyTorch first, matching your CUDA version (see https://pytorch.org). +# For development and running the tests, the CPU build is enough: +pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu + +pip install -e ".[dev]" +pre-commit install +``` + +`pre-commit install` is the important step: it wires the same Ruff lint and +format hooks CI enforces into your local `git commit`, so you find problems +before pushing. + +## Running the checks + +```bash +pytest # the full suite, ~10 seconds +pytest -m "not slow" # skip the wheel-building packaging tests +pre-commit run --all-files # everything CI's lint job runs +ruff check . # lint only +ruff format . # format in place +``` + +## The test suite + +`unit_tests/` runs entirely on CPU with 24×24×24 volumes and needs no image +data on disk, so it is fast enough to gate every pull request. + +| File | What it covers | +| --- | --- | +| `test_imports.py` | Every module under `auglab/` imports cleanly | +| `test_configs.py` | Every shipped config parses, builds a pipeline, runs a forward pass, and is reproducible under a fixed seed | +| `test_transforms_gpu.py` | Each GPU transform in isolation | +| `test_packaging.py` | Builds the real wheel and checks its contents | + +Transforms in `test_transforms_gpu.py` are discovered by introspection, so a +new transform class is covered as soon as it lands — as long as it can be built +with default arguments. If yours needs configuration, cover it by adding a +config JSON under `auglab/configs/`, which `test_configs.py` picks up +automatically. + +Note that these are smoke and contract tests: they check that transforms run, +preserve shape, stay finite, and do not corrupt the segmentation labels. They +do not verify that an augmentation is *visually* or *statistically* correct. + +## Style + +Ruff handles both linting and formatting; the configuration lives in +`pyproject.toml`. Line length is 140. + +If a rule genuinely fights a deliberate choice, add a narrow `# noqa: RULE` +with a short reason on the line rather than widening the global ignore list. + +`git blame` is configured to skip the bulk reformatting commit: + +```bash +git config blame.ignoreRevsFile .git-blame-ignore-revs +``` + +## Pull requests + +1. Branch off `main` (`yourinitials/short-description`). +2. Make the change, with tests. +3. Make sure `pytest` and `pre-commit run --all-files` pass. +4. Open a PR. CODEOWNERS requests reviewers automatically. +5. One approval and green checks are required before merge. + +## Versioning + +The version comes from the git tag via +[poetry-dynamic-versioning](https://github.com/mtkennerly/poetry-dynamic-versioning), +the same arrangement as [TPTBox](https://github.com/Hendrik-code/TPTBox). The +`version = "0.0.0"` in `pyproject.toml` is a placeholder — **never bump it by +hand**; it is substituted at build time. + +To release, tag a commit and publish a GitHub release; `publish.yml` does the +rest. + +| Tag | Version built | +| --- | --- | +| `r20260801` | `20260801` | +| `v20260801` / `20260801` | `20260801` | +| `v1.2.3` | `1.2.3` | +| `v1.0.0rc1` | `1.0.0rc1` | +| `v2.0.0-beta1` | `2.0.0b1` (PEP 440 normalised) | +| *(23 commits past `r20260615`)* | `20260616.dev23` | + +An optional `r`/`v`/`release-` prefix is stripped; see +`[tool.poetry-dynamic-versioning].pattern`. A `.post` tag is not supported and +fails the build loudly rather than silently dropping the suffix. + +Two things to know: + +- Anything that builds or installs the package needs the **tags** present, so + every workflow checkout uses `fetch-depth: 0`. Building outside a git + checkout fails with *"Unable to detect version control system"* rather than + producing a wrong version. Published sdists are unaffected — the concrete + version is baked into their `pyproject.toml` at build time. +- The build backend is poetry's, but **you do not need the `poetry` CLI or a + `poetry.lock`**. Optional dependencies are declared as extras rather than + poetry groups precisely so that `pip install -e ".[dev]"` keeps working. + +## Dependency pins + +`kornia` is capped at `>=0.7.3,<0.9`. AugLab subclasses kornia's *private* +augmentation internals (`_AugmentationBase`, `RigidAffineAugmentationBase3D`, +`augmentation.container.ops`, `_adapted_rsampling`, `_tuple_range_reader`), +which move between minor releases — 0.8.3 removed `kornia.core.Module` and the +whole `kornia.utils.helpers` module. The `kornia-compat` CI job runs the suite +against both ends of the supported range, so a break shows up here rather than +in a user's training run. + +`auglab/transforms/gpu/contrast.py` imports the private +`torchvision.transforms._functional_tensor`. It still exists as of torchvision +0.28, but carries the same risk. diff --git a/README.md b/README.md index 5020bc8..c48edc4 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,9 @@ [![arXiv](https://img.shields.io/badge/Preprint-arXiv:2605.03098-orange)](https://arxiv.org/abs/2605.03098) -[![Python Versions](https://img.shields.io/pypi/pyversions/spineps)](https://pypi.org/project/spineps/) +[![PyPI](https://img.shields.io/pypi/v/auglab)](https://pypi.org/project/auglab/) +[![Python Versions](https://img.shields.io/pypi/pyversions/auglab)](https://pypi.org/project/auglab/) +[![tests](https://github.com/neuropoly/AugLab/actions/workflows/tests.yml/badge.svg)](https://github.com/neuropoly/AugLab/actions/workflows/tests.yml) +[![lint](https://github.com/neuropoly/AugLab/actions/workflows/lint.yml/badge.svg)](https://github.com/neuropoly/AugLab/actions/workflows/lint.yml) +[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) [![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0) # AugLab @@ -129,6 +133,20 @@ python scripts/train_monai.py --config /config.json --transforms BasicTransform: transforms = [] @@ -72,62 +66,63 @@ def get_training_transforms( # Load transform parameters from json file configs_path = importlib.resources.files(configs) json_path = os.environ.get("AUGLAB_PARAMS_CPU_JSON", str(configs_path / "transform_params.json")) - transforms.append(AugTransforms(json_path=json_path, do_dummy_2d_data_aug=do_dummy_2d_data_aug, patch_size=patch_size, rotation_for_DA=rotation_for_DA, mirror_axes=mirror_axes)) + transforms.append( + AugTransforms( + json_path=json_path, + do_dummy_2d_data_aug=do_dummy_2d_data_aug, + patch_size=patch_size, + rotation_for_DA=rotation_for_DA, + mirror_axes=mirror_axes, + ) + ) if do_dummy_2d_data_aug: - ignore_axes = (0,) transforms.append(Convert3DTo2DTransform()) patch_size_spatial = patch_size[1:] else: patch_size_spatial = patch_size - ignore_axes = None transforms.append( SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=0, - random_crop=False, + patch_size_spatial, + patch_center_dist_from_border=0, + random_crop=False, p_elastic_deform=0, p_rotation=0, - rotation=rotation_for_DA, - p_scaling=0, - scaling=(0.7, 1.4), + rotation=rotation_for_DA, + p_scaling=0, + scaling=(0.7, 1.4), p_synchronize_scaling_across_axes=1, - bg_style_seg_sampling=False, - mode_seg='nearest' + bg_style_seg_sampling=False, + mode_seg="nearest", ) ) if do_dummy_2d_data_aug: transforms.append(Convert2DTo3DTransform()) - + if use_mask_for_norm is not None and any(use_mask_for_norm): - transforms.append(MaskImageTransform( - apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], - channel_idx_in_seg=0, - set_outside_to=0, - )) + transforms.append( + MaskImageTransform( + apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], + channel_idx_in_seg=0, + set_outside_to=0, + ) + ) - transforms.append( - RemoveLabelTansform(-1, 0) - ) + transforms.append(RemoveLabelTansform(-1, 0)) # The following augmentations are related to special nnunet executions if is_cascaded: - assert foreground_labels is not None, 'We need foreground_labels for cascade augmentations' + assert foreground_labels is not None, "We need foreground_labels for cascade augmentations" transforms.append( - MoveSegAsOneHotToDataTransform( - source_channel_idx=1, - all_labels=foreground_labels, - remove_channel_from_source=True - ) + MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True) ) transforms.append( RandomTransform( ApplyRandomBinaryOperatorTransform( - channel_idx=list(range(-len(foreground_labels), 0)), - strel_size=(1, 8), - p_per_label=1 - ), apply_probability=0.4 + channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1 + ), + apply_probability=0.4, ) ) transforms.append( @@ -136,8 +131,9 @@ def get_training_transforms( channel_idx=list(range(-len(foreground_labels), 0)), fill_with_other_class_p=0, dont_do_if_covers_more_than_x_percent=0.15, - p_per_label=1 - ), apply_probability=0.2 + p_per_label=1, + ), + apply_probability=0.2, ) ) @@ -145,8 +141,7 @@ def get_training_transforms( # the ignore label must also be converted transforms.append( ConvertSegmentationToRegionsTransform( - regions=list(regions) + [ignore_label] if ignore_label is not None else regions, - channel_in_seg=0 + regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0 ) ) @@ -157,8 +152,7 @@ def get_training_transforms( class nnUNetTrainerDAExtGPU(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, - device: torch.device = torch.device('cuda')): + def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): super().__init__(plans, configuration, fold, dataset_json, device) self.num_epochs = 1000 @@ -167,17 +161,22 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic configs_path = importlib.resources.files(configs) json_path = os.environ.get("AUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) self.transforms = AugTransformsGPU(json_path=json_path).to(self.device) - print(f'Using AugLab GPU transforms with parameters from: {json_path}') + print(f"Using AugLab GPU transforms with parameters from: {json_path}") + + # Load JSON parameters for validation augmentation checkpoints + with open(json_path) as f: + config = json.load(f) + self.validation_augmentation_checkpoints = config.get("ValidationAugmentationCheckpoints", {}).get("checkpoints", [0]) + self.ema_dice_validation = [None] * len(self.validation_augmentation_checkpoints) + self.best_ema_dice_validation = [None] * len(self.validation_augmentation_checkpoints) # Copy json transfrom parameters to output folder - shutil.copy( - json_path, - os.path.join(self.output_folder, 'transform_params_gpu_used_for_training.json') - ) + shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_gpu_used_for_training.json")) def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): - rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ + rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() + ) # Remove mirroring mirror_axes = None self.inference_allowed_mirroring_axes = None @@ -185,52 +184,47 @@ def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): @staticmethod def get_training_transforms( - patch_size: Union[np.ndarray, Tuple[int]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[List, Tuple, None], - mirror_axes: Tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: List[bool] = None, - is_cascaded: bool = False, - foreground_labels: Union[Tuple[int, ...], List[int]] = None, - regions: List[Union[List[int], Tuple[int, ...], int]] = None, - ignore_label: int = None, - retain_stats: bool = False + patch_size: Union[np.ndarray, tuple[int]], + rotation_for_DA: RandomScalar, + deep_supervision_scales: Union[list, tuple, None], + mirror_axes: tuple[int, ...], + do_dummy_2d_data_aug: bool, + use_mask_for_norm: list[bool] | None = None, + is_cascaded: bool = False, + foreground_labels: Union[tuple[int, ...], list[int]] | None = None, + regions: list[Union[list[int], tuple[int, ...], int]] | None = None, + ignore_label: int | None = None, + retain_stats: bool = False, ) -> BasicTransform: transforms = [] configs_path = importlib.resources.files(configs) json_path = os.environ.get("AUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) - with open(json_path, 'r') as f: + with open(json_path) as f: config = json.load(f) ### Keep some nnunet transforms if do_dummy_2d_data_aug: - ignore_axes = (0,) transforms.append(Convert3DTo2DTransform()) patch_size_spatial = patch_size[1:] else: patch_size_spatial = patch_size - ignore_axes = None - if 'nnUNetSpatialTransform' in config: - spatial_params = config['nnUNetSpatialTransform'] - else: - spatial_params = {} + spatial_params = config.get("nnUNetSpatialTransform", {}) transforms.append( SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=spatial_params.get('patch_center_dist_from_border', 0), - random_crop=spatial_params.get('random_crop', False), - p_elastic_deform=spatial_params.get('p_elastic_deform', 0), - p_rotation=spatial_params.get('p_rotation', 0), - rotation=rotation_for_DA, - p_scaling=spatial_params.get('p_scaling', 0), - scaling=spatial_params.get('scaling', (0.7, 1.4)), - p_synchronize_scaling_across_axes=spatial_params.get('p_synchronize_scaling_across_axes', 1), - bg_style_seg_sampling=False, - mode_seg='nearest' + patch_size_spatial, + patch_center_dist_from_border=spatial_params.get("patch_center_dist_from_border", 0), + random_crop=spatial_params.get("random_crop", False), + p_elastic_deform=spatial_params.get("p_elastic_deform", 0), + p_rotation=spatial_params.get("p_rotation", 0), + rotation=rotation_for_DA, + p_scaling=spatial_params.get("p_scaling", 0), + scaling=spatial_params.get("scaling", (0.7, 1.4)), + p_synchronize_scaling_across_axes=spatial_params.get("p_synchronize_scaling_across_axes", 1), + bg_style_seg_sampling=False, + mode_seg="nearest", ) ) @@ -238,33 +232,28 @@ def get_training_transforms( transforms.append(Convert2DTo3DTransform()) if use_mask_for_norm is not None and any(use_mask_for_norm): - transforms.append(MaskImageTransform( - apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], - channel_idx_in_seg=0, - set_outside_to=0, - )) + transforms.append( + MaskImageTransform( + apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], + channel_idx_in_seg=0, + set_outside_to=0, + ) + ) - transforms.append( - RemoveLabelTansform(-1, 0) - ) + transforms.append(RemoveLabelTansform(-1, 0)) # The following augmentations are related to special nnunet executions if is_cascaded: - assert foreground_labels is not None, 'We need foreground_labels for cascade augmentations' + assert foreground_labels is not None, "We need foreground_labels for cascade augmentations" transforms.append( - MoveSegAsOneHotToDataTransform( - source_channel_idx=1, - all_labels=foreground_labels, - remove_channel_from_source=True - ) + MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True) ) transforms.append( RandomTransform( ApplyRandomBinaryOperatorTransform( - channel_idx=list(range(-len(foreground_labels), 0)), - strel_size=(1, 8), - p_per_label=1 - ), apply_probability=0.4 + channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1 + ), + apply_probability=0.4, ) ) transforms.append( @@ -273,8 +262,9 @@ def get_training_transforms( channel_idx=list(range(-len(foreground_labels), 0)), fill_with_other_class_p=0, dont_do_if_covers_more_than_x_percent=0.15, - p_per_label=1 - ), apply_probability=0.2 + p_per_label=1, + ), + apply_probability=0.2, ) ) @@ -282,8 +272,7 @@ def get_training_transforms( # the ignore label must also be converted transforms.append( ConvertSegmentationToRegionsTransform( - regions=list(regions) + [ignore_label] if ignore_label is not None else regions, - channel_in_seg=0 + regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0 ) ) @@ -297,44 +286,38 @@ def get_training_transforms( @staticmethod def get_validation_transforms( - deep_supervision_scales: Union[List, Tuple, None], - is_cascaded: bool = False, - foreground_labels: Union[Tuple[int, ...], List[int]] = None, - regions: List[Union[List[int], Tuple[int, ...], int]] = None, - ignore_label: int = None, + deep_supervision_scales: Union[list, tuple, None], + is_cascaded: bool = False, + foreground_labels: Union[tuple[int, ...], list[int]] | None = None, + regions: list[Union[list[int], tuple[int, ...], int]] | None = None, + ignore_label: int | None = None, ) -> BasicTransform: transforms = [] - transforms.append( - RemoveLabelTansform(-1, 0) - ) + transforms.append(RemoveLabelTansform(-1, 0)) if is_cascaded: transforms.append( - MoveSegAsOneHotToDataTransform( - source_channel_idx=1, - all_labels=foreground_labels, - remove_channel_from_source=True - ) + MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True) ) if regions is not None: # the ignore label must also be converted transforms.append( ConvertSegmentationToRegionsTransform( - regions=list(regions) + [ignore_label] if ignore_label is not None else regions, - channel_in_seg=0 + regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0 ) ) # transforms.append(ZscoreNormalization()) - if deep_supervision_scales is not None: - transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) + # NOTE: DownsampleSegForDSTransform is now handled in train_step for GPU augmentations + # if deep_supervision_scales is not None: + # transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) return ComposeTransforms(transforms) def train_step(self, batch: dict) -> dict: - data = batch['data'] - target = batch['target'] + data = batch["data"] + target = batch["target"] data = data.to(self.device, non_blocking=True) # Now target should be a single tensor, not a list @@ -349,7 +332,7 @@ def train_step(self, batch: dict) -> dict: # If the device_type is 'cpu' then it's slow as heck and needs to be disabled. # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) # So autocast will only be active if we have a cuda device. - with autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context(): + with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): # Apply GPU augmentations to full-resolution data/target data, target = self.transforms(data, target) @@ -373,29 +356,205 @@ def train_step(self, batch: dict) -> dict: l.backward() torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12) self.optimizer.step() - return {'loss': l.detach().cpu().numpy()} + return {"loss": l.detach().cpu().numpy()} + + def validation_step(self, batch: dict, augmentation_prob: float) -> dict: + data = batch["data"] + target = batch["target"] + + data = data.to(self.device, non_blocking=True) + # Now target should be a single tensor, not a list + target = target.to(self.device, non_blocking=True) + # if isinstance(target, list): + # target = [i.to(self.device, non_blocking=True) for i in target] + # else: + # target = target.to(self.device, non_blocking=True) + + # Autocast can be annoying + # If the device_type is 'cpu' then it's slow as heck and needs to be disabled. + # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) + # So autocast will only be active if we have a cuda device. + with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): + # Apply GPU augmentations to full-resolution data/target + if torch.rand(1).item() < augmentation_prob: + data, target = self.transforms(data, target) + + # Create multi-scale targets for deep supervision after augmentation + deep_supervision_scales = self._get_deep_supervision_scales() + if deep_supervision_scales is not None: + ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) + target = ds_transform(target) + + output = self.network(data) + del data + l = self.loss(output, target) + + # we only need the output with the highest output resolution (if DS enabled) + if self.enable_deep_supervision: + output = output[0] + target = target[0] + + # the following is needed for online evaluation. Fake dice (green line) + axes = [0, *list(range(2, output.ndim))] + + if self.label_manager.has_regions: + predicted_segmentation_onehot = (torch.sigmoid(output) > 0.5).long() + else: + # no need for softmax + output_seg = output.argmax(1)[:, None] + predicted_segmentation_onehot = torch.zeros(output.shape, device=output.device, dtype=torch.float32) + predicted_segmentation_onehot.scatter_(1, output_seg, 1) + del output_seg + + if self.label_manager.has_ignore_label: + if not self.label_manager.has_regions: + mask = (target != self.label_manager.ignore_label).float() + # CAREFUL that you don't rely on target after this line! + target[target == self.label_manager.ignore_label] = 0 + else: + mask = ~target[:, -1:] if target.dtype == torch.bool else 1 - target[:, -1:] + # CAREFUL that you don't rely on target after this line! + target = target[:, :-1] + else: + mask = None + + tp, fp, fn, _ = get_tp_fp_fn_tn(predicted_segmentation_onehot, target, axes=axes, mask=mask) + + tp_hard = tp.detach().cpu().numpy() + fp_hard = fp.detach().cpu().numpy() + fn_hard = fn.detach().cpu().numpy() + if not self.label_manager.has_regions: + # if we train with regions all segmentation heads predict some kind of foreground. In conventional + # (softmax training) there needs tobe one output for the background. We are not interested in the + # background Dice + # [1:] in order to remove background + tp_hard = tp_hard[1:] + fp_hard = fp_hard[1:] + fn_hard = fn_hard[1:] + + return {"loss": l.detach().cpu().numpy(), "tp_hard": tp_hard, "fp_hard": fp_hard, "fn_hard": fn_hard} + + @staticmethod + def _valaug_sidecar_path(checkpoint_path: str) -> str: + root, ext = os.path.splitext(checkpoint_path) + return f"{root}_valaug{ext}" + + def save_checkpoint(self, filename: str) -> None: + super().save_checkpoint(filename) + if self.local_rank == 0 and not self.disable_checkpointing: + torch.save( + { + "ema_dice_validation": self.ema_dice_validation, + "best_ema_dice_validation": self.best_ema_dice_validation, + "validation_augmentation_checkpoints": self.validation_augmentation_checkpoints, + }, + self._valaug_sidecar_path(filename), + ) + + def load_checkpoint(self, filename_or_checkpoint: Union[dict, str]) -> None: + super().load_checkpoint(filename_or_checkpoint) + if isinstance(filename_or_checkpoint, str): + sidecar = self._valaug_sidecar_path(filename_or_checkpoint) + if os.path.isfile(sidecar): + state = torch.load(sidecar, map_location="cpu") + # Only adopt persisted EMA state if the configured checkpoints list is unchanged; + # otherwise the per-index slots no longer correspond and we restart tracking. + if state.get("validation_augmentation_checkpoints") == self.validation_augmentation_checkpoints: + self.ema_dice_validation = state["ema_dice_validation"] + self.best_ema_dice_validation = state["best_ema_dice_validation"] + + def compute_validation_metrics(self, val_outputs: list[dict]): + """ + Based on on_validation_epoch_end nnUNetTrainer + """ + outputs_collated = collate_outputs(val_outputs) + tp = np.sum(outputs_collated["tp_hard"], 0) + fp = np.sum(outputs_collated["fp_hard"], 0) + fn = np.sum(outputs_collated["fn_hard"], 0) + + if self.is_ddp: + world_size = dist.get_world_size() + + tps = [None for _ in range(world_size)] + dist.all_gather_object(tps, tp) + tp = np.vstack([i[None] for i in tps]).sum(0) + + fps = [None for _ in range(world_size)] + dist.all_gather_object(fps, fp) + fp = np.vstack([i[None] for i in fps]).sum(0) + + fns = [None for _ in range(world_size)] + dist.all_gather_object(fns, fn) + fn = np.vstack([i[None] for i in fns]).sum(0) + + losses_val = [None for _ in range(world_size)] + dist.all_gather_object(losses_val, outputs_collated["loss"]) + loss_here = np.vstack(losses_val).mean() + else: + loss_here = np.mean(outputs_collated["loss"]) + + global_dc_per_class = [2 * i / (2 * i + j + k) for i, j, k in zip(tp, fp, fn)] + mean_fg_dice = np.nanmean(global_dc_per_class) + val_losses = loss_here + return mean_fg_dice, global_dc_per_class, val_losses + + def run_training(self): + self.on_train_start() + + for _epoch in range(self.current_epoch, self.num_epochs): + self.on_epoch_start() + + self.on_train_epoch_start() + train_outputs = [self.train_step(next(self.dataloader_train)) for _batch_id in range(self.num_iterations_per_epoch)] + self.on_train_epoch_end(train_outputs) + + with torch.no_grad(): + self.on_validation_epoch_start() + val_outputs = [[] for _ in range(len(self.validation_augmentation_checkpoints))] + for _batch_id in range(self.num_val_iterations_per_epoch): + batch = next(self.dataloader_val) + for prob_id, augmentation_prob in enumerate(self.validation_augmentation_checkpoints): + val_outputs[prob_id].append(self.validation_step(batch, augmentation_prob)) + for val_idx, val_output in enumerate(val_outputs): + mean_fg_dice, global_dc_per_class, val_losses = self.compute_validation_metrics(val_output) + self.ema_dice_validation[val_idx] = ( + self.ema_dice_validation[val_idx] * 0.9 + 0.1 * mean_fg_dice + if self.ema_dice_validation[val_idx] is not None + else mean_fg_dice + ) + if val_idx == 0: + self.logger.log("mean_fg_dice", mean_fg_dice, self.current_epoch) + self.logger.log("dice_per_class_or_region", global_dc_per_class, self.current_epoch) + self.logger.log("val_losses", val_losses, self.current_epoch) + + self.on_epoch_end() + for val_idx, augmentation_prob in enumerate(self.validation_augmentation_checkpoints): + ema_dice = self.ema_dice_validation[val_idx] + best = self.best_ema_dice_validation[val_idx] + if ema_dice is not None and (best is None or ema_dice > best): + self.best_ema_dice_validation[val_idx] = ema_dice + self.save_checkpoint(join(self.output_folder, f"checkpoint_best_validation_aug_{augmentation_prob}.pth")) + + self.on_train_end() class nnUNetTrainerDAExtHybrid(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, - device: torch.device = torch.device('cuda')): + def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): super().__init__(plans, configuration, fold, dataset_json, device) # Load transform parameters from json file configs_path = importlib.resources.files(configs) json_path = os.environ.get("AUGLAB_PARAMS_HYBRID_JSON", str(configs_path / "transform_params_hybrid.json")) self.transforms = AugTransformsGPU(json_path=json_path).to(self.device) - print(f'Using AugLab hybrid transforms with parameters from: {json_path}') + print(f"Using AugLab hybrid transforms with parameters from: {json_path}") # Copy json transfrom parameters to output folder - shutil.copy( - json_path, - os.path.join(self.output_folder, 'transform_params_hybrid_used_for_training.json') - ) + shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_hybrid_used_for_training.json")) def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): - rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = \ + rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes = ( super().configure_rotation_dummyDA_mirroring_and_inital_patch_size() + ) # Remove mirroring mirror_axes = None self.inference_allowed_mirroring_axes = None @@ -403,17 +562,17 @@ def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self): @staticmethod def get_training_transforms( - patch_size: Union[np.ndarray, Tuple[int]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[List, Tuple, None], - mirror_axes: Tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: List[bool] = None, - is_cascaded: bool = False, - foreground_labels: Union[Tuple[int, ...], List[int]] = None, - regions: List[Union[List[int], Tuple[int, ...], int]] = None, - ignore_label: int = None, - retain_stats: bool = False + patch_size: Union[np.ndarray, tuple[int]], + rotation_for_DA: RandomScalar, + deep_supervision_scales: Union[list, tuple, None], + mirror_axes: tuple[int, ...], + do_dummy_2d_data_aug: bool, + use_mask_for_norm: list[bool] | None = None, + is_cascaded: bool = False, + foreground_labels: Union[tuple[int, ...], list[int]] | None = None, + regions: list[Union[list[int], tuple[int, ...], int]] | None = None, + ignore_label: int | None = None, + retain_stats: bool = False, ) -> BasicTransform: transforms = [] @@ -421,36 +580,39 @@ def get_training_transforms( # Load transform parameters from json file configs_path = importlib.resources.files(configs) json_path = os.environ.get("AUGLAB_PARAMS_HYBRID_JSON", str(configs_path / "transform_params_hybrid.json")) - transforms.append(AugTransforms(json_path=json_path, do_dummy_2d_data_aug=do_dummy_2d_data_aug, patch_size=patch_size, rotation_for_DA=rotation_for_DA, mirror_axes=mirror_axes)) - - if use_mask_for_norm is not None and any(use_mask_for_norm): - transforms.append(MaskImageTransform( - apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], - channel_idx_in_seg=0, - set_outside_to=0, - )) - transforms.append( - RemoveLabelTansform(-1, 0) + AugTransforms( + json_path=json_path, + do_dummy_2d_data_aug=do_dummy_2d_data_aug, + patch_size=patch_size, + rotation_for_DA=rotation_for_DA, + mirror_axes=mirror_axes, + ) ) + if use_mask_for_norm is not None and any(use_mask_for_norm): + transforms.append( + MaskImageTransform( + apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], + channel_idx_in_seg=0, + set_outside_to=0, + ) + ) + + transforms.append(RemoveLabelTansform(-1, 0)) + # The following augmentations are related to special nnunet executions if is_cascaded: - assert foreground_labels is not None, 'We need foreground_labels for cascade augmentations' + assert foreground_labels is not None, "We need foreground_labels for cascade augmentations" transforms.append( - MoveSegAsOneHotToDataTransform( - source_channel_idx=1, - all_labels=foreground_labels, - remove_channel_from_source=True - ) + MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True) ) transforms.append( RandomTransform( ApplyRandomBinaryOperatorTransform( - channel_idx=list(range(-len(foreground_labels), 0)), - strel_size=(1, 8), - p_per_label=1 - ), apply_probability=0.4 + channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1 + ), + apply_probability=0.4, ) ) transforms.append( @@ -459,8 +621,9 @@ def get_training_transforms( channel_idx=list(range(-len(foreground_labels), 0)), fill_with_other_class_p=0, dont_do_if_covers_more_than_x_percent=0.15, - p_per_label=1 - ), apply_probability=0.2 + p_per_label=1, + ), + apply_probability=0.2, ) ) @@ -468,20 +631,19 @@ def get_training_transforms( # the ignore label must also be converted transforms.append( ConvertSegmentationToRegionsTransform( - regions=list(regions) + [ignore_label] if ignore_label is not None else regions, - channel_in_seg=0 + regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0 ) ) - + # NOTE: DownsampleSegForDSTransform is now handled in train_step for GPU augmentations # if deep_supervision_scales is not None: # transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) return ComposeTransforms(transforms) - + def train_step(self, batch: dict) -> dict: - data = batch['data'] - target = batch['target'] + data = batch["data"] + target = batch["target"] data = data.to(self.device, non_blocking=True) # Now target should be a single tensor, not a list @@ -496,16 +658,16 @@ def train_step(self, batch: dict) -> dict: # If the device_type is 'cpu' then it's slow as heck and needs to be disabled. # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) # So autocast will only be active if we have a cuda device. - with autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context(): + with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): # Apply GPU augmentations to full-resolution data/target data, target = self.transforms(data, target) - + # Create multi-scale targets for deep supervision after augmentation deep_supervision_scales = self._get_deep_supervision_scales() if deep_supervision_scales is not None: ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) target = ds_transform(target) - + output = self.network(data) # del data l = self.loss(output, target) @@ -520,4 +682,4 @@ def train_step(self, batch: dict) -> dict: l.backward() torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12) self.optimizer.step() - return {'loss': l.detach().cpu().numpy()} + return {"loss": l.detach().cpu().numpy()} diff --git a/auglab/trainers/nnUNetTrainerTest.py b/auglab/trainers/nnUNetTrainerTest.py index ed732a5..12b18bb 100644 --- a/auglab/trainers/nnUNetTrainerTest.py +++ b/auglab/trainers/nnUNetTrainerTest.py @@ -1,49 +1,48 @@ -from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer -from nnunetv2.utilities.helpers import dummy_context +import importlib +from typing import Union +import numpy as np +import torch from batchgeneratorsv2.helpers.scalar_type import RandomScalar from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform from batchgeneratorsv2.transforms.nnunet.random_binary_operator import ApplyRandomBinaryOperatorTransform from batchgeneratorsv2.transforms.nnunet.remove_connected_components import RemoveRandomConnectedComponentFromOneHotEncodingTransform from batchgeneratorsv2.transforms.nnunet.seg_to_onehot import MoveSegAsOneHotToDataTransform +from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms from batchgeneratorsv2.transforms.utils.deep_supervision_downsampling import DownsampleSegForDSTransform from batchgeneratorsv2.transforms.utils.nnunet_masking import MaskImageTransform +from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform from batchgeneratorsv2.transforms.utils.random import RandomTransform from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform -from batchgeneratorsv2.transforms.utils.pseudo2d import Convert3DTo2DTransform, Convert2DTo3DTransform -from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform - -import torch +from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer +from nnunetv2.utilities.helpers import dummy_context from torch import autocast -import importlib -from typing import Tuple, Union, List -import numpy as np -import auglab.configs as configs +from auglab import configs +from auglab.trainers.utils import DownsampleSegForDSTransformCustom from auglab.transforms.cpu.transforms import AugTransformsTest from auglab.transforms.gpu.transforms import AugTransformsGPU -from auglab.trainers.utils import DownsampleSegForDSTransformCustom + class nnUNetTrainerTest(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, - device: torch.device = torch.device('cuda')): + def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): super().__init__(plans, configuration, fold, dataset_json, device) - + @staticmethod def get_training_transforms( - patch_size: Union[np.ndarray, Tuple[int]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[List, Tuple, None], - mirror_axes: Tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: List[bool] = None, - is_cascaded: bool = False, - foreground_labels: Union[Tuple[int, ...], List[int]] = None, - regions: List[Union[List[int], Tuple[int, ...], int]] = None, - ignore_label: int = None, - retain_stats: bool = False + patch_size: Union[np.ndarray, tuple[int]], + rotation_for_DA: RandomScalar, + deep_supervision_scales: Union[list, tuple, None], + mirror_axes: tuple[int, ...], + do_dummy_2d_data_aug: bool, + use_mask_for_norm: list[bool] | None = None, + is_cascaded: bool = False, + foreground_labels: Union[tuple[int, ...], list[int]] | None = None, + regions: list[Union[list[int], tuple[int, ...], int]] | None = None, + ignore_label: int | None = None, + retain_stats: bool = False, ) -> BasicTransform: transforms = [] @@ -52,59 +51,52 @@ def get_training_transforms( ### Keep some nnunet transforms if do_dummy_2d_data_aug: - ignore_axes = (0,) transforms.append(Convert3DTo2DTransform()) patch_size_spatial = patch_size[1:] else: patch_size_spatial = patch_size - ignore_axes = None transforms.append( SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=0, - random_crop=False, + patch_size_spatial, + patch_center_dist_from_border=0, + random_crop=False, p_elastic_deform=0, p_rotation=0, - rotation=rotation_for_DA, - p_scaling=0, - scaling=(0.7, 1.4), + rotation=rotation_for_DA, + p_scaling=0, + scaling=(0.7, 1.4), p_synchronize_scaling_across_axes=1, - bg_style_seg_sampling=False, - mode_seg='nearest' + bg_style_seg_sampling=False, + mode_seg="nearest", ) ) if do_dummy_2d_data_aug: transforms.append(Convert2DTo3DTransform()) - + if use_mask_for_norm is not None and any(use_mask_for_norm): - transforms.append(MaskImageTransform( - apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], - channel_idx_in_seg=0, - set_outside_to=0, - )) + transforms.append( + MaskImageTransform( + apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], + channel_idx_in_seg=0, + set_outside_to=0, + ) + ) - transforms.append( - RemoveLabelTansform(-1, 0) - ) + transforms.append(RemoveLabelTansform(-1, 0)) # The following augmentations are related to special nnunet executions if is_cascaded: - assert foreground_labels is not None, 'We need foreground_labels for cascade augmentations' + assert foreground_labels is not None, "We need foreground_labels for cascade augmentations" transforms.append( - MoveSegAsOneHotToDataTransform( - source_channel_idx=1, - all_labels=foreground_labels, - remove_channel_from_source=True - ) + MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True) ) transforms.append( RandomTransform( ApplyRandomBinaryOperatorTransform( - channel_idx=list(range(-len(foreground_labels), 0)), - strel_size=(1, 8), - p_per_label=1 - ), apply_probability=0.4 + channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1 + ), + apply_probability=0.4, ) ) transforms.append( @@ -113,8 +105,9 @@ def get_training_transforms( channel_idx=list(range(-len(foreground_labels), 0)), fill_with_other_class_p=0, dont_do_if_covers_more_than_x_percent=0.15, - p_per_label=1 - ), apply_probability=0.2 + p_per_label=1, + ), + apply_probability=0.2, ) ) @@ -122,8 +115,7 @@ def get_training_transforms( # the ignore label must also be converted transforms.append( ConvertSegmentationToRegionsTransform( - regions=list(regions) + [ignore_label] if ignore_label is not None else regions, - channel_in_seg=0 + regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0 ) ) @@ -132,9 +124,9 @@ def get_training_transforms( return ComposeTransforms(transforms) + class nnUNetTrainerTestGPU(nnUNetTrainer): - def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, - device: torch.device = torch.device('cuda')): + def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): super().__init__(plans, configuration, fold, dataset_json, device) # Load transform parameters from json file @@ -144,75 +136,68 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic @staticmethod def get_training_transforms( - patch_size: Union[np.ndarray, Tuple[int]], - rotation_for_DA: RandomScalar, - deep_supervision_scales: Union[List, Tuple, None], - mirror_axes: Tuple[int, ...], - do_dummy_2d_data_aug: bool, - use_mask_for_norm: List[bool] = None, - is_cascaded: bool = False, - foreground_labels: Union[Tuple[int, ...], List[int]] = None, - regions: List[Union[List[int], Tuple[int, ...], int]] = None, - ignore_label: int = None, - retain_stats: bool = False + patch_size: Union[np.ndarray, tuple[int]], + rotation_for_DA: RandomScalar, + deep_supervision_scales: Union[list, tuple, None], + mirror_axes: tuple[int, ...], + do_dummy_2d_data_aug: bool, + use_mask_for_norm: list[bool] | None = None, + is_cascaded: bool = False, + foreground_labels: Union[tuple[int, ...], list[int]] | None = None, + regions: list[Union[list[int], tuple[int, ...], int]] | None = None, + ignore_label: int | None = None, + retain_stats: bool = False, ) -> BasicTransform: transforms = [] ### Keep some nnunet transforms if do_dummy_2d_data_aug: - ignore_axes = (0,) transforms.append(Convert3DTo2DTransform()) patch_size_spatial = patch_size[1:] else: patch_size_spatial = patch_size - ignore_axes = None transforms.append( SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=0, - random_crop=False, + patch_size_spatial, + patch_center_dist_from_border=0, + random_crop=False, p_elastic_deform=0, p_rotation=0, - rotation=rotation_for_DA, - p_scaling=0, - scaling=(0.7, 1.4), + rotation=rotation_for_DA, + p_scaling=0, + scaling=(0.7, 1.4), p_synchronize_scaling_across_axes=1, - bg_style_seg_sampling=False, - mode_seg='nearest' + bg_style_seg_sampling=False, + mode_seg="nearest", ) ) if do_dummy_2d_data_aug: transforms.append(Convert2DTo3DTransform()) - + if use_mask_for_norm is not None and any(use_mask_for_norm): - transforms.append(MaskImageTransform( - apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], - channel_idx_in_seg=0, - set_outside_to=0, - )) + transforms.append( + MaskImageTransform( + apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]], + channel_idx_in_seg=0, + set_outside_to=0, + ) + ) - transforms.append( - RemoveLabelTansform(-1, 0) - ) + transforms.append(RemoveLabelTansform(-1, 0)) # The following augmentations are related to special nnunet executions if is_cascaded: - assert foreground_labels is not None, 'We need foreground_labels for cascade augmentations' + assert foreground_labels is not None, "We need foreground_labels for cascade augmentations" transforms.append( - MoveSegAsOneHotToDataTransform( - source_channel_idx=1, - all_labels=foreground_labels, - remove_channel_from_source=True - ) + MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True) ) transforms.append( RandomTransform( ApplyRandomBinaryOperatorTransform( - channel_idx=list(range(-len(foreground_labels), 0)), - strel_size=(1, 8), - p_per_label=1 - ), apply_probability=0.4 + channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1 + ), + apply_probability=0.4, ) ) transforms.append( @@ -221,8 +206,9 @@ def get_training_transforms( channel_idx=list(range(-len(foreground_labels), 0)), fill_with_other_class_p=0, dont_do_if_covers_more_than_x_percent=0.15, - p_per_label=1 - ), apply_probability=0.2 + p_per_label=1, + ), + apply_probability=0.2, ) ) @@ -230,8 +216,7 @@ def get_training_transforms( # the ignore label must also be converted transforms.append( ConvertSegmentationToRegionsTransform( - regions=list(regions) + [ignore_label] if ignore_label is not None else regions, - channel_in_seg=0 + regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0 ) ) @@ -240,10 +225,10 @@ def get_training_transforms( # transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) return ComposeTransforms(transforms) - + def train_step(self, batch: dict) -> dict: - data = batch['data'] - target = batch['target'] + data = batch["data"] + target = batch["target"] data = data.to(self.device, non_blocking=True) # Now target should be a single tensor, not a list @@ -254,16 +239,16 @@ def train_step(self, batch: dict) -> dict: # If the device_type is 'cpu' then it's slow as heck and needs to be disabled. # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) # So autocast will only be active if we have a cuda device. - with autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context(): + with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): # Apply GPU augmentations to full-resolution data/target data, target = self.transforms(data, target) - + # Create multi-scale targets for deep supervision after augmentation if self.enable_deep_supervision: deep_supervision_scales = self._get_deep_supervision_scales() ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) target = ds_transform(target) - + output = self.network(data) # del data l = self.loss(output, target) @@ -278,4 +263,4 @@ def train_step(self, batch: dict) -> dict: l.backward() torch.nn.utils.clip_grad_norm_(self.network.parameters(), 12) self.optimizer.step() - return {'loss': l.detach().cpu().numpy()} + return {"loss": l.detach().cpu().numpy()} diff --git a/auglab/trainers/utils.py b/auglab/trainers/utils.py index 99da7e4..e950b23 100644 --- a/auglab/trainers/utils.py +++ b/auglab/trainers/utils.py @@ -1,55 +1,56 @@ +from typing import Union + import torch from torch.nn.functional import interpolate -from typing import Tuple, Union, List class DownsampleSegForDSTransformCustom: """ Custom deep supervision downsampling transform that handles batched tensors properly. Unlike the original DownsampleSegForDSTransform, this handles tensors with batch dimension. - - Input: [batch, channels, spatial_dims...] + + Input: [batch, channels, spatial_dims...] Output: List of [batch, channels, spatial_dims...] at different scales """ - def __init__(self, ds_scales: Union[List, Tuple]): + + def __init__(self, ds_scales: Union[list, tuple]): self.ds_scales = ds_scales - def __call__(self, segmentation: torch.Tensor) -> List[torch.Tensor]: + def __call__(self, segmentation: torch.Tensor) -> list[torch.Tensor]: """ Apply downsampling to segmentation tensor with batch dimension. - + Args: segmentation: [batch, channels, spatial_dims...] tensor - + Returns: List of downsampled tensors, each with shape [batch, channels, spatial_dims...] """ results = [] - for s in self.ds_scales: - if not isinstance(s, (tuple, list)): + for ds_scale in self.ds_scales: + if not isinstance(ds_scale, (tuple, list)): # If single scale value, apply to all spatial dimensions - s = [s] * (segmentation.ndim - 2) # -2 for batch and channel dims + s = [ds_scale] * (segmentation.ndim - 2) # -2 for batch and channel dims else: - assert len(s) == segmentation.ndim - 2, f"Scale length {len(s)} doesn't match spatial dims {segmentation.ndim - 2}" + assert len(ds_scale) == segmentation.ndim - 2, ( + f"Scale length {len(ds_scale)} doesn't match spatial dims {segmentation.ndim - 2}" + ) + s = ds_scale - if all([i == 1 for i in s]): + if all(i == 1 for i in s): # No downsampling needed results.append(segmentation) else: # Calculate new spatial shape spatial_shape = segmentation.shape[2:] # Skip batch and channel dims new_shape = [round(i * j) for i, j in zip(spatial_shape, s)] - + # Store original dtype dtype = segmentation.dtype - + # Interpolate (convert to float for interpolation, then back to original dtype) - downsampled = interpolate( - segmentation.float(), - size=new_shape, - mode='nearest-exact' - ).to(dtype) - + downsampled = interpolate(segmentation.float(), size=new_shape, mode="nearest-exact").to(dtype) + results.append(downsampled) - - return results \ No newline at end of file + + return results diff --git a/auglab/transforms/cpu/artifact.py b/auglab/transforms/cpu/artifact.py index 83cf246..24bc978 100644 --- a/auglab/transforms/cpu/artifact.py +++ b/auglab/transforms/cpu/artifact.py @@ -1,20 +1,19 @@ -import torch - -from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform +import gc +import random +import torch import torchio as tio -import gc +from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform -import random class ArtifactTransform(BasicTransform): def __init__(self, motion=False, ghosting=False, spike=False, bias_field=False, blur=False, noise=False, swap=False, random_pick=False): - ''' - Apply all selected artifacts (motion, ghosting, spike, bias field, blur, noise, and swap) to the image if they are enabled (set to True). + """ + Apply all selected artifacts (motion, ghosting, spike, bias field, blur, noise, and swap) to the image if they are enabled (set to True). If `random_pick` is True, randomly select and apply ONE of the enabled artifacts. Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py - ''' + """ super().__init__() self.motion = motion self.ghosting = ghosting @@ -34,169 +33,168 @@ def get_parameters(self, **data_dict) -> dict: "bias_field": self.bias_field, "blur": self.blur, "noise": self.noise, - "swap": self.swap + "swap": self.swap, } - enabled_artifacts = {k:v for k,v in artifacts.items() if v} + enabled_artifacts = {k: v for k, v in artifacts.items() if v} if self.random_pick and enabled_artifacts: selected_artifact = random.choice(list(enabled_artifacts.keys())) - artifacts = {k: (k == selected_artifact) for k,v in artifacts.items()} + artifacts = {k: (k == selected_artifact) for k, v in artifacts.items()} return artifacts - + def apply(self, data_dict: dict, **params) -> dict: - if data_dict.get('image') is not None and data_dict.get('segmentation') is not None: - data_dict['image'], data_dict['segmentation'] = self._apply_to_image(data_dict['image'], data_dict['segmentation'], **params) + if data_dict.get("image") is not None and data_dict.get("segmentation") is not None: + data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params) return data_dict def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor: - if params['motion']: + if params["motion"]: img, seg = aug_motion(img, seg) - if params['ghosting']: + if params["ghosting"]: img, seg = aug_ghosting(img, seg) - if params['spike']: + if params["spike"]: img, seg = aug_spike(img, seg) - if params['bias_field']: + if params["bias_field"]: img, seg = aug_bias_field(img, seg) - if params['blur']: + if params["blur"]: img, seg = aug_blur(img, seg) - if params['noise']: + if params["noise"]: img, seg = aug_noise(img, seg) - if params['swap']: + if params["swap"]: img, seg = aug_swap(img, seg) return img, seg + def aug_motion(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomMotion()(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomMotion()( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomMotion()(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomMotion()(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_ghosting(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomGhosting()(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomGhosting()( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomGhosting()(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomGhosting()(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_spike(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomSpike(intensity=(1, 2))(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomSpike(intensity=(1, 2))( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomSpike(intensity=(1, 2))(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomSpike(intensity=(1, 2))(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_bias_field(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomBiasField()(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomBiasField()( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomBiasField()(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomBiasField()(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_blur(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomBlur()(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomBlur()( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomBlur()(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomBlur()(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_noise(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomNoise()(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomNoise()( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomNoise()(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomNoise()(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_swap(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomSwap()(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomSwap()( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomSwap()(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomSwap()(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out - \ No newline at end of file diff --git a/auglab/transforms/cpu/contrast.py b/auglab/transforms/cpu/contrast.py index e7cfe4f..5918ade 100644 --- a/auglab/transforms/cpu/contrast.py +++ b/auglab/transforms/cpu/contrast.py @@ -1,17 +1,18 @@ import torch import torch.nn.functional as F +from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform -from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform, BasicTransform class ConvTransform(ImageOnlyTransform): - ''' + """ Applies a Laplace/Scharr filter to the image to highlight edges. Based on https://github.com/spinalcordtoolbox/disc-labeling-playground/blob/main/src/ply/models/transform.py - ''' - def __init__(self, kernel_type: str = 'Laplace', absolute: bool = False, retain_stats: bool = False): + """ + + def __init__(self, kernel_type: str = "Laplace", absolute: bool = False, retain_stats: bool = False): super().__init__() - if kernel_type not in ["Laplace","Scharr"]: + if kernel_type not in ["Laplace", "Scharr"]: raise NotImplementedError('Currently only "Laplace" and "Scharr" are supported.') else: self.kernel_type = kernel_type @@ -19,7 +20,7 @@ def __init__(self, kernel_type: str = 'Laplace', absolute: bool = False, retain_ self.retain_stats = retain_stats def get_parameters(self, **data_dict) -> dict: - spatial_dims = len(data_dict['image'].shape) - 1 + spatial_dims = len(data_dict["image"].shape) - 1 if spatial_dims == 2: if self.kernel_type == "Laplace": kernel = torch.tensor([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype=torch.float32) @@ -32,97 +33,82 @@ def get_parameters(self, **data_dict) -> dict: kernel = -1.0 * torch.ones(3, 3, 3, dtype=torch.float32) kernel[1, 1, 1] = 26.0 elif self.kernel_type == "Scharr": - kernel_x = torch.tensor([[[ 9, 0, -9], - [ 30, 0, -30], - [ 9, 0, -9]], - - [[ 30, 0, -30], - [100, 0, -100], - [ 30, 0, -30]], - - [[ 9, 0, -9], - [ 30, 0, -30], - [ 9, 0, -9]]], dtype=torch.float32) - - kernel_y = torch.tensor([[[ 9, 30, 9], - [ 0, 0, 0], - [ -9, -30, -9]], - - [[ 30, 100, 30], - [ 0, 0, 0], - [ -30, -100, -30]], - - [[ 9, 30, 9], - [ 0, 0, 0], - [ -9, -30, -9]]], dtype=torch.float32) - - kernel_z = torch.tensor([[[ 9, 30, 9], - [ 30, 100, 30], - [ 9, 30, 9]], - - [[ 0, 0, 0], - [ 0, 0, 0], - [ 0, 0, 0]], - - [[ -9, -30, -9], - [ -30, -100, -30], - [ -9, -30, -9]]], dtype=torch.float32) + kernel_x = torch.tensor( + [ + [[9, 0, -9], [30, 0, -30], [9, 0, -9]], + [[30, 0, -30], [100, 0, -100], [30, 0, -30]], + [[9, 0, -9], [30, 0, -30], [9, 0, -9]], + ], + dtype=torch.float32, + ) + + kernel_y = torch.tensor( + [ + [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], + [[30, 100, 30], [0, 0, 0], [-30, -100, -30]], + [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], + ], + dtype=torch.float32, + ) + + kernel_z = torch.tensor( + [ + [[9, 30, 9], [30, 100, 30], [9, 30, 9]], + [[0, 0, 0], [0, 0, 0], [0, 0, 0]], + [[-9, -30, -9], [-30, -100, -30], [-9, -30, -9]], + ], + dtype=torch.float32, + ) kernel = [kernel_x, kernel_y, kernel_z] else: raise ValueError(f"{self.__class__} can only handle 2D or 3D images.") - return { - 'kernel_type': self.kernel_type, - 'kernel': kernel, - 'absolute': self.absolute, - 'retain_stats': self.retain_stats - } - + return {"kernel_type": self.kernel_type, "kernel": kernel, "absolute": self.absolute, "retain_stats": self.retain_stats} + def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: - ''' + """ We expect (C, X, Y) or (C, X, Y, Z) shaped inputs for image and seg - ''' - for c in range(1): # Works on the first channel only - if params['retain_stats']: + """ + for c in range(1): # Works on the first channel only + if params["retain_stats"]: orig_mean = torch.mean(img[c]) orig_std = torch.std(img[c]) img_ = img[c].unsqueeze(0).unsqueeze(0) # adds temp batch and channel dim - if params['kernel_type'] == 'Laplace': - tot_ = apply_filter(img_, params['kernel']) - elif params['kernel_type'] == 'Scharr': + if params["kernel_type"] == "Laplace": + tot_ = apply_filter(img_, params["kernel"]) + elif params["kernel_type"] == "Scharr": tot_ = torch.zeros_like(img_) - for kernel in params['kernel']: - if params['absolute']: + for kernel in params["kernel"]: + if params["absolute"]: tot_ += torch.abs(apply_filter(img_, kernel)) else: tot_ += apply_filter(img_, kernel) - img[c] = tot_[0,0] - if params['retain_stats']: + img[c] = tot_[0, 0] + if params["retain_stats"]: mean = torch.mean(img[c]) std = torch.std(img[c]) - img[c] = (img[c] - mean)/torch.clamp(std, min=1e-7) - img[c] = img[c]*orig_std + orig_mean # return to original distribution + img[c] = (img[c] - mean) / torch.clamp(std, min=1e-7) + img[c] = img[c] * orig_std + orig_mean # return to original distribution return img class HistogramEqualTransform(ImageOnlyTransform): - ''' + """ Update image intensity using histogram manipulations Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py - ''' + """ + def __init__(self, retain_stats: bool = False): super().__init__() self.retain_stats = retain_stats - + def get_parameters(self, **data_dict) -> dict: - return { - 'retain_stats': self.retain_stats - } - + return {"retain_stats": self.retain_stats} + def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: for c in range(1): # Works on the first channel only - if params['retain_stats']: + if params["retain_stats"]: orig_mean = torch.mean(img[c]) orig_std = torch.std(img[c]) img_min, img_max = img[c].min(), img[c].max() @@ -143,36 +129,34 @@ def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: indices = torch.searchsorted(bin_edges[:-1], img_flattened) img_eq = torch.index_select(cdf, dim=0, index=torch.clamp(indices, 0, 255)) img[c] = img_eq.reshape(img[c].shape) - - if params['retain_stats']: + + if params["retain_stats"]: # Return to original distribution mean = torch.mean(img[c]) std = torch.std(img[c]) - img[c] = (img[c] - mean)/torch.clamp(std, min=1e-7) - img[c] = img[c]*orig_std + orig_mean + img[c] = (img[c] - mean) / torch.clamp(std, min=1e-7) + img[c] = img[c] * orig_std + orig_mean return img class FunctionTransform(ImageOnlyTransform): - ''' + """ Apply different functions to image pixels Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py - ''' - def __init__(self, function, retain_stats : bool = False): + """ + + def __init__(self, function, retain_stats: bool = False): super().__init__() self.function = function self.retain_stats = retain_stats def get_parameters(self, **data_dict) -> dict: - return { - 'function': self.function, - 'retain_stats': self.retain_stats - } - + return {"function": self.function, "retain_stats": self.retain_stats} + def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: for c in range(1): # Works on the first channel only - if params['retain_stats']: + if params["retain_stats"]: orig_mean = torch.mean(img[c]) orig_std = torch.std(img[c]) @@ -180,16 +164,17 @@ def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: img[c] = (img[c] - img.min()) / (img.max() - img.min() + 0.00001) # Apply function - img[c] = params['function'](img[c]) + img[c] = params["function"](img[c]) - if params['retain_stats']: + if params["retain_stats"]: # Return to original distribution mean = torch.mean(img[c]) std = torch.std(img[c]) - img[c] = (img[c] - mean)/torch.clamp(std, min=1e-7) - img[c] = img[c]*orig_std + orig_mean + img[c] = (img[c] - mean) / torch.clamp(std, min=1e-7) + img[c] = img[c] * orig_std + orig_mean return img + def apply_filter(x: torch.Tensor, kernel: torch.Tensor, **kwargs) -> torch.Tensor: """ Copied from https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/layers/simplelayers.py @@ -225,9 +210,7 @@ def apply_filter(x: torch.Tensor, kernel: torch.Tensor, **kwargs) -> torch.Tenso raise NotImplementedError(f"Only spatial dimensions up to 3 are supported but got {n_spatial}.") k_size = len(kernel.shape) if k_size < n_spatial or k_size > n_spatial + 2: - raise ValueError( - f"kernel must have {n_spatial} ~ {n_spatial + 2} dimensions to match the input shape {x.shape}." - ) + raise ValueError(f"kernel must have {n_spatial} ~ {n_spatial + 2} dimensions to match the input shape {x.shape}.") kernel = kernel.to(x) # broadcast kernel size to (batch chns, spatial_kernel_size) kernel = kernel.expand(batch, chns, *kernel.shape[(k_size - n_spatial) :]) @@ -242,10 +225,12 @@ def apply_filter(x: torch.Tensor, kernel: torch.Tensor, **kwargs) -> torch.Tenso output = conv(x, kernel, groups=kernel.shape[0], bias=None, **kwargs) return output.view(batch, chns, *output.shape[2:]) + class ZscoreNormalization(ImageOnlyTransform): - ''' + """ Z-score normalization of image - ''' + """ + def __init__(self) -> None: super().__init__() @@ -253,5 +238,5 @@ def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: for c in range(1): mean = torch.mean(img[c]) std = torch.std(img[c]) - img[c] = (img[c] - mean)/torch.clamp(std, min=1e-8) - return img \ No newline at end of file + img[c] = (img[c] - mean) / torch.clamp(std, min=1e-8) + return img diff --git a/auglab/transforms/cpu/fromSeg.py b/auglab/transforms/cpu/fromSeg.py index f0bb11d..cf7c528 100644 --- a/auglab/transforms/cpu/fromSeg.py +++ b/auglab/transforms/cpu/fromSeg.py @@ -1,18 +1,18 @@ -import torch -import torch.nn.functional as F - -from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform +from functools import partial import scipy.ndimage as ndi +import torch +from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform from scipy.stats import norm -from functools import partial + class RedistributeTransform(BasicTransform): - ''' + """ Redistribute image values using segmentation regions. Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py - ''' + """ + def __init__(self, classes=None, in_seg=0.2, retain_stats=False): super().__init__() self.classes = classes @@ -20,22 +20,21 @@ def __init__(self, classes=None, in_seg=0.2, retain_stats=False): self.retain_stats = retain_stats def get_parameters(self, **data_dict) -> dict: - return { - 'classes': self.classes, - 'in_seg': self.in_seg, - 'retain_stats': self.retain_stats - } - + return {"classes": self.classes, "in_seg": self.in_seg, "retain_stats": self.retain_stats} + def apply(self, data_dict: dict, **params) -> dict: - if data_dict.get('image') is not None and data_dict.get('segmentation') is not None: - data_dict['image'], data_dict['segmentation'] = self._apply_to_image(data_dict['image'], data_dict['segmentation'], **params) + if data_dict.get("image") is not None and data_dict.get("segmentation") is not None: + data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params) return data_dict def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor: for c in range(1): # Works on the first channel only - img[c], seg[c] = aug_redistribute_seg(img[c], seg[c], classes=params['classes'], in_seg=params['in_seg'], retain_stats=params['retain_stats']) + img[c], seg[c] = aug_redistribute_seg( + img[c], seg[c], classes=params["classes"], in_seg=params["in_seg"], retain_stats=params["retain_stats"] + ) return img, seg + def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False): """ Augment the image by redistributing the values of the image within the @@ -49,7 +48,7 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) if classes: _seg = combine_classes(_seg, classes) - + if retain_stats: # Compute original mean, std and min/max values original_mean, original_std = img.mean(), img.std() @@ -67,7 +66,7 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) # Loop over each label value for l in labels: # Get the mask for the current label - l_mask = (_seg == l) + l_mask = _seg == l # Get mean and std of the current label l_mean, l_std = img[l_mask].mean(), img[l_mask].std() @@ -89,8 +88,11 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) l_std_dilate = img[l_mask_dilate_excl].std() else: l_mean_dilate, l_std_dilate = l_mean, l_std # Fallback to original values - - redist_std = max(torch.rand(1, device=device) * 0.2 + 0.4 * abs((l_mean - l_mean_dilate) * l_std / (l_std_dilate + 1e-6)), torch.tensor([0.01], device=device)) + + redist_std = max( + torch.rand(1, device=device) * 0.2 + 0.4 * abs((l_mean - l_mean_dilate) * l_std / (l_std_dilate + 1e-6)), + torch.tensor([0.01], device=device), + ) redist = partial(norm.pdf, loc=l_mean.cpu().numpy(), scale=redist_std.cpu().numpy()) @@ -107,13 +109,14 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) # Return to original range mean = torch.mean(img) std = torch.std(img) - img = (img - mean)/torch.clamp(std, min=1e-7) - img = img*original_std + original_mean + img = (img - mean) / torch.clamp(std, min=1e-7) + img = img * original_std + original_mean return img, seg + def combine_classes(seg, classes): _seg = torch.zeros_like(seg) for i, c in enumerate(classes): _seg[torch.isin(seg, c)] = i + 1 - return _seg \ No newline at end of file + return _seg diff --git a/auglab/transforms/cpu/spatial.py b/auglab/transforms/cpu/spatial.py index e6c58c2..cbcd5a7 100644 --- a/auglab/transforms/cpu/spatial.py +++ b/auglab/transforms/cpu/spatial.py @@ -1,20 +1,19 @@ -import torch - -from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform, BasicTransform +import gc +import random +import torch import torchio as tio -import gc +from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform, ImageOnlyTransform -import random class SpatialCustomTransform(BasicTransform): def __init__(self, flip=False, affine=False, elastic=False, anisotropy=False, random_pick=False): - ''' - Apply all selected spatial transformation (flip, affine, elastic and anisotropy) to the image if they are enabled (set to True). + """ + Apply all selected spatial transformation (flip, affine, elastic and anisotropy) to the image if they are enabled (set to True). If `random_pick` is True, randomly select and apply ONE of the enabled transformation. Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py - ''' + """ super().__init__() self.flip = flip self.affine = affine @@ -23,144 +22,144 @@ def __init__(self, flip=False, affine=False, elastic=False, anisotropy=False, ra self.random_pick = random_pick def get_parameters(self, **data_dict) -> dict: - transfo = { - "flip" : self.flip, - "affine" : self.affine, - "elastic" : self.elastic, - "anisotropy" : self.anisotropy - } + transfo = {"flip": self.flip, "affine": self.affine, "elastic": self.elastic, "anisotropy": self.anisotropy} - enabled_transfo = {k:v for k,v in transfo.items() if v} + enabled_transfo = {k: v for k, v in transfo.items() if v} if self.random_pick and enabled_transfo: selected_transfo = random.choice(list(enabled_transfo.keys())) - transfo = {k: (k == selected_transfo) for k,v in transfo.items()} - + transfo = {k: (k == selected_transfo) for k, v in transfo.items()} + return transfo - + def apply(self, data_dict: dict, **params) -> dict: - if data_dict.get('image') is not None and data_dict.get('segmentation') is not None: - data_dict['image'], data_dict['segmentation'] = self._apply_to_image(data_dict['image'], data_dict['segmentation'], **params) + if data_dict.get("image") is not None and data_dict.get("segmentation") is not None: + data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params) return data_dict def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor: - if params['flip']: + if params["flip"]: img, seg = aug_flip(img, seg) - if params['affine']: + if params["affine"]: img, seg = aug_affine(img, seg) - if params['elastic']: + if params["elastic"]: img, seg = aug_elastic(img, seg) - if params['anisotropy']: + if params["anisotropy"]: img, seg = aug_anisotropy(img, seg) return img, seg + def aug_flip(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomFlip(axes=('LR',))(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomFlip(axes=("LR",))( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomFlip(axes=('LR',))(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomFlip(axes=("LR",))(tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg))) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_affine(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomAffine(degrees=10, translation=(0.1, 0.1, 0.1), scales=(0.9, 1.1))(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomAffine(degrees=10, translation=(0.1, 0.1, 0.1), scales=(0.9, 1.1))( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomAffine(degrees=10, translation=(0.1, 0.1, 0.1), scales=(0.9, 1.1))(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomAffine(degrees=10, translation=(0.1, 0.1, 0.1), scales=(0.9, 1.1))( + tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg)) + ) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_elastic(img, seg): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomElasticDeformation(max_displacement=40)(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomElasticDeformation(max_displacement=40)( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomElasticDeformation(max_displacement=40)(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg) - )) + subject = tio.RandomElasticDeformation(max_displacement=40)( + tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg)) + ) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + def aug_anisotropy(img, seg, downsampling=7): - if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation - subject = tio.RandomAnisotropy(downsampling=downsampling)(tio.Subject( - image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), - discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), - seg=tio.LabelMap(tensor=seg) - )) + if img.shape[0] == 2: # Step2: channel 1 --> image / channel 2 --> odd discs segmentation + subject = tio.RandomAnisotropy(downsampling=downsampling)( + tio.Subject( + image=tio.ScalarImage(tensor=torch.unsqueeze(img[0], dim=0)), + discs=tio.LabelMap(tensor=torch.unsqueeze(img[1], dim=0)), + seg=tio.LabelMap(tensor=seg), + ) + ) img_out = torch.cat((subject.image.data, subject.discs.data), axis=0) seg_out = subject.seg.data else: - subject = tio.RandomAnisotropy(downsampling=downsampling)(tio.Subject( - image=tio.ScalarImage(tensor=img), - seg=tio.LabelMap(tensor=seg, axis=0) - )) + subject = tio.RandomAnisotropy(downsampling=downsampling)( + tio.Subject(image=tio.ScalarImage(tensor=img), seg=tio.LabelMap(tensor=seg, axis=0)) + ) img_out, seg_out = subject.image.data, subject.seg.data del subject gc.collect() # Force garbage collection return img_out, seg_out + ### Shape transform + class ShapeTransform(ImageOnlyTransform): def __init__(self, shape_min=1, ignore_axes=()): - ''' + """ shape_min: minimal shape size along allowed axis Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py - ''' + """ super().__init__() self.shape_min = shape_min self.ignore_axes = ignore_axes def get_parameters(self, **data_dict) -> dict: - return { - 'shape_min': self.shape_min, - 'ignore_axes': self.ignore_axes - } - + return {"shape_min": self.shape_min, "ignore_axes": self.ignore_axes} + def apply(self, data_dict: dict, **params) -> dict: - if data_dict.get('image') is not None and data_dict.get('segmentation') is not None: - data_dict['image'], data_dict['segmentation'] = self._apply_to_image(data_dict['image'], data_dict['segmentation'], **params) + if data_dict.get("image") is not None and data_dict.get("segmentation") is not None: + data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params) return data_dict def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor: # Compute random shape img_shape = img.shape[1:] - new_shape = [random.randint(params["shape_min"], s) if i not in params["ignore_axes"] else s for i,s in enumerate(img_shape)] + new_shape = [random.randint(params["shape_min"], s) if i not in params["ignore_axes"] else s for i, s in enumerate(img_shape)] # Find image center - img_center = [s//2 for s in img_shape] + img_center = [s // 2 for s in img_shape] # Compute start and end crop indices per axis starts = [max(0, c - ns // 2) for c, ns in zip(img_center, new_shape)] @@ -170,4 +169,4 @@ def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tor slices = tuple(slice(start, end) for start, end in zip(starts, ends)) img_cropped = img[(slice(None), *slices)] # Keep channel dim intact seg_cropped = seg[(slice(None), *slices)] - return img_cropped, seg_cropped \ No newline at end of file + return img_cropped, seg_cropped diff --git a/auglab/transforms/cpu/transforms.py b/auglab/transforms/cpu/transforms.py index 5508073..32bf3e4 100644 --- a/auglab/transforms/cpu/transforms.py +++ b/auglab/transforms/cpu/transforms.py @@ -1,63 +1,71 @@ -import os import json -import torch -import numpy as np -from typing import Union, Tuple +import os +from typing import Union +import numpy as np +import torch from batchgeneratorsv2.helpers.scalar_type import RandomScalar from batchgeneratorsv2.transforms.intensity.brightness import MultiplicativeBrightnessTransform -from batchgeneratorsv2.transforms.intensity.contrast import ContrastTransform, BGContrast +from batchgeneratorsv2.transforms.intensity.contrast import BGContrast, ContrastTransform from batchgeneratorsv2.transforms.intensity.gamma import GammaTransform from batchgeneratorsv2.transforms.intensity.gaussian_noise import GaussianNoiseTransform from batchgeneratorsv2.transforms.noise.gaussian_blur import GaussianBlurTransform from batchgeneratorsv2.transforms.spatial.low_resolution import SimulateLowResolutionTransform from batchgeneratorsv2.transforms.spatial.mirroring import MirrorTransform from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform -from batchgeneratorsv2.transforms.utils.random import RandomTransform from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms -from batchgeneratorsv2.transforms.utils.pseudo2d import Convert3DTo2DTransform, Convert2DTo3DTransform +from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform +from batchgeneratorsv2.transforms.utils.random import RandomTransform from auglab.transforms.cpu.artifact import ArtifactTransform -from auglab.transforms.cpu.contrast import ConvTransform, HistogramEqualTransform, FunctionTransform +from auglab.transforms.cpu.contrast import ConvTransform, FunctionTransform, HistogramEqualTransform from auglab.transforms.cpu.fromSeg import RedistributeTransform -from auglab.transforms.cpu.spatial import SpatialCustomTransform, ShapeTransform +from auglab.transforms.cpu.spatial import ShapeTransform, SpatialCustomTransform + class AugTransforms(ComposeTransforms): - def __init__(self, json_path: str, do_dummy_2d_data_aug: bool, patch_size: Union[np.ndarray, Tuple[int]], - rotation_for_DA: RandomScalar, mirror_axes: Tuple[int]): + def __init__( + self, + json_path: str, + do_dummy_2d_data_aug: bool, + patch_size: Union[np.ndarray, tuple[int]], + rotation_for_DA: RandomScalar, + mirror_axes: tuple[int], + ): # Load transform parameters from JSON config_path = os.path.join(json_path) - with open(config_path, 'r') as f: + with open(config_path) as f: config = json.load(f) - - if 'CPU' in config.keys(): - self.transform_params = config['CPU'] + + if "CPU" in config.keys(): + self.transform_params = config["CPU"] else: self.transform_params = config - + self.transforms = self._build_transforms( - do_dummy_2d_data_aug=do_dummy_2d_data_aug, - patch_size=patch_size, - rotation_for_DA=rotation_for_DA, - mirror_axes=mirror_axes + do_dummy_2d_data_aug=do_dummy_2d_data_aug, patch_size=patch_size, rotation_for_DA=rotation_for_DA, mirror_axes=mirror_axes ) super().__init__(transforms=self.transforms) - def _build_transforms(self, do_dummy_2d_data_aug: bool, patch_size: Union[np.ndarray, Tuple[int]], - rotation_for_DA: RandomScalar, mirror_axes: Tuple[int]): + def _build_transforms( + self, do_dummy_2d_data_aug: bool, patch_size: Union[np.ndarray, tuple[int]], rotation_for_DA: RandomScalar, mirror_axes: tuple[int] + ): transform_params = self.transform_params transforms = [] # Scharr filter - conv_params = transform_params.get('ConvTransform') + conv_params = transform_params.get("ConvTransform") if conv_params is not None: - transforms.append(RandomTransform( - ConvTransform( - kernel_type=conv_params.get('kernel_type', 'Scharr'), - absolute=conv_params.get('absolute', True), - retain_stats=transform_params.get('retain_stats', False) - ), apply_probability=conv_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + ConvTransform( + kernel_type=conv_params.get("kernel_type", "Scharr"), + absolute=conv_params.get("absolute", True), + retain_stats=transform_params.get("retain_stats", False), + ), + apply_probability=conv_params.get("probability", 0), + ) + ) # Apply functions func_list = [ @@ -65,202 +73,230 @@ def _build_transforms(self, do_dummy_2d_data_aug: bool, patch_size: Union[np.nda torch.sqrt, torch.sin, torch.exp, - lambda x: 1/(1 + torch.exp(-x)), + lambda x: 1 / (1 + torch.exp(-x)), ] - func_params = transform_params.get('FunctionTransform') + func_params = transform_params.get("FunctionTransform") if func_params is not None: - for func in func_list: - transforms.append(RandomTransform( - FunctionTransform( - function=func, - retain_stats=transform_params.get('retain_stats', False) - ), apply_probability=func_params.get('probability', 0) - )) - + transforms.extend( + RandomTransform( + FunctionTransform(function=func, retain_stats=transform_params.get("retain_stats", False)), + apply_probability=func_params.get("probability", 0), + ) + for func in func_list + ) + # Histogram manipulations - hist_params = transform_params.get('HistogramEqualTransform') + hist_params = transform_params.get("HistogramEqualTransform") if hist_params is not None: - transforms.append(RandomTransform( - HistogramEqualTransform( - retain_stats=transform_params.get('retain_stats', False) - ), apply_probability=hist_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + HistogramEqualTransform(retain_stats=transform_params.get("retain_stats", False)), + apply_probability=hist_params.get("probability", 0), + ) + ) # Redistribute segmentation values - redist_params = transform_params.get('RedistributeTransform') + redist_params = transform_params.get("RedistributeTransform") if redist_params is not None: - transforms.append(RandomTransform( - RedistributeTransform( - in_seg=redist_params.get('in_seg', 0), - retain_stats=transform_params.get('retain_stats', False) - ), apply_probability=redist_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + RedistributeTransform(in_seg=redist_params.get("in_seg", 0), retain_stats=transform_params.get("retain_stats", False)), + apply_probability=redist_params.get("probability", 0), + ) + ) # Resolution transforms - shape_params = transform_params.get('ShapeTransform') + shape_params = transform_params.get("ShapeTransform") if shape_params is not None: - transforms.append(RandomTransform( - ShapeTransform( - shape_min=shape_params.get('shape_min'), - ignore_axes=tuple(shape_params.get('ignore_axes', None)) if shape_params.get('ignore_axes', None) is not None else None, - ), apply_probability=shape_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + ShapeTransform( + shape_min=shape_params.get("shape_min"), + ignore_axes=tuple(shape_params.get("ignore_axes", None)) + if shape_params.get("ignore_axes", None) is not None + else None, + ), + apply_probability=shape_params.get("probability", 0), + ) + ) # Artifacts generation - artifact_params = transform_params.get('ArtifactTransform') + artifact_params = transform_params.get("ArtifactTransform") if artifact_params is not None: - transforms.append(RandomTransform( - ArtifactTransform( - motion=artifact_params.get('motion', False), - ghosting=artifact_params.get('ghosting', False), - spike=artifact_params.get('spike', False), - bias_field=artifact_params.get('bias_field', False), - blur=artifact_params.get('blur', False), - noise=artifact_params.get('noise', False), - swap=artifact_params.get('swap', False), - random_pick=artifact_params.get('random_pick', False) - ), apply_probability=artifact_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + ArtifactTransform( + motion=artifact_params.get("motion", False), + ghosting=artifact_params.get("ghosting", False), + spike=artifact_params.get("spike", False), + bias_field=artifact_params.get("bias_field", False), + blur=artifact_params.get("blur", False), + noise=artifact_params.get("noise", False), + swap=artifact_params.get("swap", False), + random_pick=artifact_params.get("random_pick", False), + ), + apply_probability=artifact_params.get("probability", 0), + ) + ) # Spatial transforms - spatial_custom_params = transform_params.get('SpatialCustomTransform') + spatial_custom_params = transform_params.get("SpatialCustomTransform") if spatial_custom_params is not None: - transforms.append(RandomTransform( - SpatialCustomTransform( - flip=spatial_custom_params.get('flip', False), - affine=spatial_custom_params.get('affine', False), - elastic=spatial_custom_params.get('elastic', False), - anisotropy=spatial_custom_params.get('anisotropy', False), - random_pick=spatial_custom_params.get('random_pick', False) - ), apply_probability=spatial_custom_params.get('probability', 0) - )) - + transforms.append( + RandomTransform( + SpatialCustomTransform( + flip=spatial_custom_params.get("flip", False), + affine=spatial_custom_params.get("affine", False), + elastic=spatial_custom_params.get("elastic", False), + anisotropy=spatial_custom_params.get("anisotropy", False), + random_pick=spatial_custom_params.get("random_pick", False), + ), + apply_probability=spatial_custom_params.get("probability", 0), + ) + ) + # Spatial nnunet transform if do_dummy_2d_data_aug: - ignore_axes = (0,) transforms.append(Convert3DTo2DTransform()) patch_size_spatial = patch_size[1:] else: patch_size_spatial = patch_size - ignore_axes = None - - spatial_params = transform_params.get('SpatialTransform') + + spatial_params = transform_params.get("SpatialTransform") if spatial_params is not None: transforms.append( SpatialTransform( - patch_size_spatial, - patch_center_dist_from_border=spatial_params.get('patch_center_dist_from_border', 0), - random_crop=spatial_params.get('random_crop', False), - p_elastic_deform=spatial_params.get('p_elastic_deform', 0), - p_rotation=spatial_params.get('p_rotation', 0), - rotation=rotation_for_DA, - p_scaling=spatial_params.get('p_scaling', 0), - scaling=spatial_params.get('scaling', (0.7, 1.4)), - p_synchronize_scaling_across_axes=spatial_params.get('p_synchronize_scaling_across_axes', 1), - bg_style_seg_sampling=spatial_params.get('bg_style_seg_sampling', False), - mode_seg='nearest' + patch_size_spatial, + patch_center_dist_from_border=spatial_params.get("patch_center_dist_from_border", 0), + random_crop=spatial_params.get("random_crop", False), + p_elastic_deform=spatial_params.get("p_elastic_deform", 0), + p_rotation=spatial_params.get("p_rotation", 0), + rotation=rotation_for_DA, + p_scaling=spatial_params.get("p_scaling", 0), + scaling=spatial_params.get("scaling", (0.7, 1.4)), + p_synchronize_scaling_across_axes=spatial_params.get("p_synchronize_scaling_across_axes", 1), + bg_style_seg_sampling=spatial_params.get("bg_style_seg_sampling", False), + mode_seg="nearest", ) ) if do_dummy_2d_data_aug: transforms.append(Convert2DTo3DTransform()) - + # Noise transforms - noise_params = transform_params.get('GaussianNoiseTransform') + noise_params = transform_params.get("GaussianNoiseTransform") if noise_params is not None: - transforms.append(RandomTransform( - GaussianNoiseTransform( - noise_variance=tuple(noise_params.get('noise_variance', (0, 0.1))), - p_per_channel=noise_params.get('p_per_channel', 1), - synchronize_channels=noise_params.get('synchronize_channels', True) - ), apply_probability=noise_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + GaussianNoiseTransform( + noise_variance=tuple(noise_params.get("noise_variance", (0, 0.1))), + p_per_channel=noise_params.get("p_per_channel", 1), + synchronize_channels=noise_params.get("synchronize_channels", True), + ), + apply_probability=noise_params.get("probability", 0), + ) + ) # Gaussian blur - blur_params = transform_params.get('GaussianBlurTransform') + blur_params = transform_params.get("GaussianBlurTransform") if blur_params is not None: - transforms.append(RandomTransform( - GaussianBlurTransform( - blur_sigma=tuple(blur_params.get('blur_sigma', (0.5, 1.))), - synchronize_channels=blur_params.get('synchronize_channels', False), - synchronize_axes=blur_params.get('synchronize_axes', False), - p_per_channel=blur_params.get('p_per_channel', 0.5), - benchmark=blur_params.get('benchmark', True) - ), apply_probability=blur_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + GaussianBlurTransform( + blur_sigma=tuple(blur_params.get("blur_sigma", (0.5, 1.0))), + synchronize_channels=blur_params.get("synchronize_channels", False), + synchronize_axes=blur_params.get("synchronize_axes", False), + p_per_channel=blur_params.get("p_per_channel", 0.5), + benchmark=blur_params.get("benchmark", True), + ), + apply_probability=blur_params.get("probability", 0), + ) + ) # Brightness transforms - bright_params = transform_params.get('MultiplicativeBrightnessTransform') + bright_params = transform_params.get("MultiplicativeBrightnessTransform") if bright_params is not None: - transforms.append(RandomTransform( - MultiplicativeBrightnessTransform( - multiplier_range=BGContrast(tuple(bright_params.get('multiplier_range', (0.75, 1.25)))), - synchronize_channels=bright_params.get('synchronize_channels', False), - p_per_channel=bright_params.get('p_per_channel', 1) - ), apply_probability=bright_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + MultiplicativeBrightnessTransform( + multiplier_range=BGContrast(tuple(bright_params.get("multiplier_range", (0.75, 1.25)))), + synchronize_channels=bright_params.get("synchronize_channels", False), + p_per_channel=bright_params.get("p_per_channel", 1), + ), + apply_probability=bright_params.get("probability", 0), + ) + ) # Contrast transforms - contrast_params = transform_params.get('ContrastTransform') + contrast_params = transform_params.get("ContrastTransform") if contrast_params is not None: - transforms.append(RandomTransform( - ContrastTransform( - contrast_range=BGContrast(tuple(contrast_params.get('contrast_range', (0.75, 1.25)))), - preserve_range=contrast_params.get('preserve_range', True), - synchronize_channels=contrast_params.get('synchronize_channels', False), - p_per_channel=contrast_params.get('p_per_channel', 1) - ), apply_probability=contrast_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + ContrastTransform( + contrast_range=BGContrast(tuple(contrast_params.get("contrast_range", (0.75, 1.25)))), + preserve_range=contrast_params.get("preserve_range", True), + synchronize_channels=contrast_params.get("synchronize_channels", False), + p_per_channel=contrast_params.get("p_per_channel", 1), + ), + apply_probability=contrast_params.get("probability", 0), + ) + ) # Simulate low resolution - lowres_params = transform_params.get('SimulateLowResolutionTransform') + lowres_params = transform_params.get("SimulateLowResolutionTransform") if lowres_params is not None: - transforms.append(RandomTransform( - SimulateLowResolutionTransform( - scale=tuple(lowres_params.get('scale', (0.3, 1))), - synchronize_channels=lowres_params.get('synchronize_channels', True), - synchronize_axes=lowres_params.get('synchronize_axes', False), - ignore_axes=tuple(lowres_params.get('ignore_axes', ())), - allowed_channels=lowres_params.get('allowed_channels', None), - p_per_channel=lowres_params.get('p_per_channel', 0.5) - ), apply_probability=lowres_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + SimulateLowResolutionTransform( + scale=tuple(lowres_params.get("scale", (0.3, 1))), + synchronize_channels=lowres_params.get("synchronize_channels", True), + synchronize_axes=lowres_params.get("synchronize_axes", False), + ignore_axes=tuple(lowres_params.get("ignore_axes", ())), + allowed_channels=lowres_params.get("allowed_channels", None), + p_per_channel=lowres_params.get("p_per_channel", 0.5), + ), + apply_probability=lowres_params.get("probability", 0), + ) + ) # Gamma transforms - gamma_inv_params = transform_params.get('GammaTransform_invert') + gamma_inv_params = transform_params.get("GammaTransform_invert") if gamma_inv_params is not None: - transforms.append(RandomTransform( - GammaTransform( - gamma=BGContrast(tuple(gamma_inv_params.get('gamma', (0.7, 1.5)))), - p_invert_image=gamma_inv_params.get('p_invert_image', 1), - synchronize_channels=gamma_inv_params.get('synchronize_channels', False), - p_per_channel=gamma_inv_params.get('p_per_channel', 1), - p_retain_stats=gamma_inv_params.get('p_retain_stats', 1) - ), apply_probability=gamma_inv_params.get('probability', 0) - )) - - gamma_params = transform_params.get('GammaTransform') - if gamma_params is not None: - transforms.append(RandomTransform( - GammaTransform( - gamma=BGContrast(tuple(gamma_params.get('gamma', (0.7, 1.5)))), - p_invert_image=gamma_params.get('p_invert_image', 0), - synchronize_channels=gamma_params.get('synchronize_channels', False), - p_per_channel=gamma_params.get('p_per_channel', 1), - p_retain_stats=gamma_params.get('p_retain_stats', 1) - ), apply_probability=gamma_params.get('probability', 0) - )) + transforms.append( + RandomTransform( + GammaTransform( + gamma=BGContrast(tuple(gamma_inv_params.get("gamma", (0.7, 1.5)))), + p_invert_image=gamma_inv_params.get("p_invert_image", 1), + synchronize_channels=gamma_inv_params.get("synchronize_channels", False), + p_per_channel=gamma_inv_params.get("p_per_channel", 1), + p_retain_stats=gamma_inv_params.get("p_retain_stats", 1), + ), + apply_probability=gamma_inv_params.get("probability", 0), + ) + ) - # Mirroring transforms - if transform_params.get('mirror_axes') is not None and len(transform_params['mirror_axes']) > 0: + gamma_params = transform_params.get("GammaTransform") + if gamma_params is not None: transforms.append( - MirrorTransform( - allowed_axes=transform_params.get('mirror_axes') + RandomTransform( + GammaTransform( + gamma=BGContrast(tuple(gamma_params.get("gamma", (0.7, 1.5)))), + p_invert_image=gamma_params.get("p_invert_image", 0), + synchronize_channels=gamma_params.get("synchronize_channels", False), + p_per_channel=gamma_params.get("p_per_channel", 1), + p_retain_stats=gamma_params.get("p_retain_stats", 1), + ), + apply_probability=gamma_params.get("probability", 0), ) ) + # Mirroring transforms + if transform_params.get("mirror_axes") is not None and len(transform_params["mirror_axes"]) > 0: + transforms.append(MirrorTransform(allowed_axes=transform_params.get("mirror_axes"))) + return transforms + class AugTransformsTest(ComposeTransforms): def __init__(self): self.transforms = self._build_transforms() @@ -270,57 +306,62 @@ def _build_transforms(self): transforms = [] # Scharr filter - transforms.append(RandomTransform( - ConvTransform( - kernel_type="Scharr", - absolute=True, - ), apply_probability=0.9 - )) + transforms.append( + RandomTransform( + ConvTransform( + kernel_type="Scharr", + absolute=True, + ), + apply_probability=0.9, + ) + ) # Affine transforms - transforms.append(RandomTransform( - SpatialCustomTransform( - affine=True, - ), apply_probability=0.9 - )) + transforms.append( + RandomTransform( + SpatialCustomTransform( + affine=True, + ), + apply_probability=0.9, + ) + ) return transforms + if __name__ == "__main__": # Example usage import importlib - import auglab.configs as configs - from auglab.utils.image import Image, resample_nib + import cv2 - from auglab.utils.utils import normalize + + from auglab import configs from auglab.transforms.gpu.transforms import AugTransformsGPU - + from auglab.utils.image import Image, resample_nib + from auglab.utils.utils import normalize + configs_path = importlib.resources.files(configs) json_path = configs_path / "transform_params_hybrid_TAGE.json" # Load images and masks tensors - img_path = '/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz' - img = Image(img_path).change_orientation('RSP') - img = resample_nib(img, new_size=[1,1,1], new_size_type='mm', interpolation='linear') + img_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz" + img = Image(img_path).change_orientation("RSP") + img = resample_nib(img, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") img_tensor = torch.from_numpy(img.data.copy()).to(torch.float32).unsqueeze(0) - seg_path = '/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/derivatives/labels/sub-amu02/anat/sub-amu02_T1w_label-spine_dseg.nii.gz' - seg = Image(seg_path).change_orientation('RSP') - seg = resample_nib(seg, new_size=[1,1,1], new_size_type='mm', interpolation='nn') + seg_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/derivatives/labels/sub-amu02/anat/sub-amu02_T1w_label-spine_dseg.nii.gz" + seg = Image(seg_path).change_orientation("RSP") + seg = resample_nib(seg, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") seg_tensor_all = torch.from_numpy(seg.data.copy()) # Add segmentation values to different channels seg_tensor = torch.zeros((5, *seg_tensor_all.shape)) for i, value in enumerate([12, 13, 14, 15, 16]): - seg_tensor[i] = (seg_tensor_all == value) + seg_tensor[i] = seg_tensor_all == value # Example usage aug_transforms = AugTransforms( - json_path=json_path, - do_dummy_2d_data_aug=False, - patch_size=(128, 128, 128), - rotation_for_DA=(-10, 10), - mirror_axes=None + json_path=json_path, do_dummy_2d_data_aug=False, patch_size=(128, 128, 128), rotation_for_DA=(-10, 10), mirror_axes=None ) augmentor_gpu = AugTransformsGPU(json_path) @@ -329,31 +370,34 @@ def _build_transforms(self): tensor_dict = {} gpu = False for i in range(24): - tensor_dict[f'transfo_{str(i+1)}'] = aug_transforms(**{'image': img_tensor.detach().clone(), 'segmentation': seg_tensor.detach().clone()}) + tensor_dict[f"transfo_{i + 1!s}"] = aug_transforms(image=img_tensor.detach().clone(), segmentation=seg_tensor.detach().clone()) if gpu: - augmented_img, augmented_seg = augmentor_gpu(tensor_dict[f'transfo_{str(i+1)}']['image'].cuda().unsqueeze(0).clone(), tensor_dict[f'transfo_{str(i+1)}']['segmentation'].cuda().unsqueeze(0).clone()) - tensor_dict[f'transfo_{str(i+1)}']['image'] = augmented_img.cpu().squeeze(0) - tensor_dict[f'transfo_{str(i+1)}']['segmentation'] = augmented_seg.cpu().squeeze(0) - + augmented_img, augmented_seg = augmentor_gpu( + tensor_dict[f"transfo_{i + 1!s}"]["image"].cuda().unsqueeze(0).clone(), + tensor_dict[f"transfo_{i + 1!s}"]["segmentation"].cuda().unsqueeze(0).clone(), + ) + tensor_dict[f"transfo_{i + 1!s}"]["image"] = augmented_img.cpu().squeeze(0) + tensor_dict[f"transfo_{i + 1!s}"]["segmentation"] = augmented_seg.cpu().squeeze(0) + nb_img = len(tensor_dict.keys()) nb_col = 6 - for key in ['image', 'segmentation']: + for key in ["image", "segmentation"]: output = [] line = [] aug = [[]] - for idx, (augment, dic) in enumerate(tensor_dict.items()): + for _idx, (augment, _dic) in enumerate(tensor_dict.items()): if len(line) < nb_col: - img = 255*normalize(np.sum(tensor_dict[augment][key].detach().numpy(), axis=0, keepdims=True)[0,64]) + img = 255 * normalize(np.sum(tensor_dict[augment][key].detach().numpy(), axis=0, keepdims=True)[0, 64]) line.append(img) aug[-1].append(augment) else: output.append(np.concatenate(line, axis=1)) - img = 255*normalize(np.sum(tensor_dict[augment][key].detach().numpy(), axis=0, keepdims=True)[0,64]) + img = 255 * normalize(np.sum(tensor_dict[augment][key].detach().numpy(), axis=0, keepdims=True)[0, 64]) line = [img] aug.append([augment]) output.append(np.concatenate(line, axis=1)) out_img = np.concatenate(output, axis=0) - cv2.imwrite(f'img/transforms_default+plus_{key}.png', out_img) - print(aug_transforms) \ No newline at end of file + cv2.imwrite(f"img/transforms_default+plus_{key}.png", out_img) + print(aug_transforms) diff --git a/auglab/transforms/gpu/base.py b/auglab/transforms/gpu/base.py index b37d1d3..dc97c03 100644 --- a/auglab/transforms/gpu/base.py +++ b/auglab/transforms/gpu/base.py @@ -1,34 +1,35 @@ +import copy import warnings -from kornia.augmentation import RandomGamma +from collections.abc import Sequence +from typing import Any, Union +import kornia.augmentation as K +from kornia.augmentation import AugmentationSequential from kornia.augmentation._2d.base import RigidAffineAugmentationBase2D -from kornia.augmentation._3d.base import RigidAffineAugmentationBase3D from kornia.augmentation._3d.base import AugmentationBase3D, RigidAffineAugmentationBase3D from kornia.augmentation.base import _AugmentationBase -from kornia.constants import DataKey, Resample -from kornia.core import Tensor -from kornia.geometry.boxes import Boxes -from kornia.geometry.keypoints import Keypoints -from kornia.augmentation.container.patch import PatchSequential -from kornia.augmentation.container.video import VideoSequential from kornia.augmentation.container.image import ImageSequential -from kornia.augmentation.container.ops import AugmentationSequentialOps, SequentialOpsInterface, InputSequentialOps, BoxSequentialOps, KeypointSequentialOps, ClassSequentialOps - -from kornia.augmentation import AugmentationSequential -from kornia.augmentation.container.ops import MaskSequentialOps +from kornia.augmentation.container.ops import ( + AugmentationSequentialOps, + BoxSequentialOps, + ClassSequentialOps, + InputSequentialOps, + KeypointSequentialOps, + MaskSequentialOps, + SequentialOpsInterface, +) from kornia.augmentation.container.params import ParamItem -import kornia.augmentation as K -from kornia.augmentation.base import _AugmentationBase -from kornia.constants import DataKey -from kornia.core import Module, Tensor +from kornia.augmentation.container.patch import PatchSequential +from kornia.augmentation.container.video import VideoSequential +from kornia.constants import DataKey, Resample from kornia.geometry.boxes import Boxes from kornia.geometry.keypoints import Keypoints +from torch import Tensor +from torch.nn import Module -from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, Type -import copy +DataType = Union[Tensor, list[Tensor], Boxes, Keypoints] +SequenceDataType = Union[list[Tensor], list[list[Tensor]], list[Boxes], list[Keypoints]] -DataType = Union[Tensor, List[Tensor], Boxes, Keypoints] -SequenceDataType = Union[List[Tensor], List[List[Tensor]], List[Boxes], List[Keypoints]] class ImageOnlyTransform(RigidAffineAugmentationBase3D): r"""ImageOnlyTransform base class for customized image-only transformations. @@ -44,70 +45,72 @@ class ImageOnlyTransform(RigidAffineAugmentationBase3D): """ - def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: return self.identity_matrix(input) def apply_non_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: # For the images where batch_prob == False. return input def apply_non_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: return input def apply_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: return input def apply_non_transform_boxes( - self, input: Boxes, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Boxes, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Boxes: return input def apply_transform_boxes( - self, input: Boxes, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Boxes, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Boxes: return input def apply_non_transform_keypoint( - self, input: Keypoints, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Keypoints, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Keypoints: return input def apply_transform_keypoint( - self, input: Keypoints, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Keypoints, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Keypoints: return input def apply_non_transform_class( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: return input def apply_transform_class( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: return input + class AugmentationSequentialCustom(AugmentationSequential): """Custom AugmentationSequential to handle masks augmentations.""" + def __init__( self, *args: Union[_AugmentationBase, ImageSequential], - data_keys: Optional[Union[Sequence[str], Sequence[int], Sequence[DataKey]]] = (DataKey.INPUT,), - same_on_batch: Optional[bool] = None, - keepdim: Optional[bool] = None, - random_apply: Union[int, bool, Tuple[int, int]] = False, - random_apply_weights: Optional[List[float]] = None, + data_keys: Union[Sequence[str], Sequence[int], Sequence[DataKey]] | None = (DataKey.INPUT,), + same_on_batch: bool | None = None, + keepdim: bool | None = None, + random_apply: Union[int, bool, tuple[int, int]] = False, + random_apply_weights: list[float] | None = None, transformation_matrix_mode: str = "silent", - extra_args: Optional[Dict[DataKey, Dict[str, Any]]] = None, + extra_args: dict[DataKey, dict[str, Any]] | None = None, ) -> None: - self._transform_matrix: Optional[Tensor] - self._transform_matrices: List[Optional[Tensor]] = [] + self._transform_matrix: Tensor | None + self._transform_matrices: list[Tensor | None] = [] super().__init__( *args, @@ -119,13 +122,13 @@ def __init__( self._parse_transformation_matrix_mode(transformation_matrix_mode) - self._valid_ops_for_transform_computation: Tuple[Any, ...] = ( + self._valid_ops_for_transform_computation: tuple[Any, ...] = ( RigidAffineAugmentationBase2D, RigidAffineAugmentationBase3D, AugmentationSequential, ) - self.data_keys: Optional[List[DataKey]] + self.data_keys: list[DataKey] | None if data_keys is not None: self.data_keys = [DataKey.get(inp) for inp in data_keys] else: @@ -144,9 +147,7 @@ def __init__( self.contains_3d_augmentation: bool = False for arg in args: if isinstance(arg, PatchSequential) and not arg.is_intensity_only(): - warnings.warn( - "Geometric transformation detected in PatchSeqeuntial, which would break bbox, mask.", stacklevel=1 - ) + warnings.warn("Geometric transformation detected in PatchSeqeuntial, which would break bbox, mask.", stacklevel=1) if isinstance(arg, VideoSequential): self.contains_video_sequential = True # NOTE: only for images are supported for 3D. @@ -154,20 +155,17 @@ def __init__( self.contains_3d_augmentation = True self._transform_matrix = None self.extra_args = extra_args or {DataKey.MASK: {"resample": Resample.NEAREST, "align_corners": None}} - - def transform_masks( - self, input: Tensor, params: List[ParamItem], extra_args: Optional[Dict[str, Any]] = None - ) -> Tensor: + + def transform_masks(self, input: Tensor, params: list[ParamItem], extra_args: dict[str, Any] | None = None) -> Tensor: for param in params: module = self.get_submodule(param.name) input = MaskSequentialOpsCustom.transform(input, module=module, param=param, extra_args=extra_args) return input + class MaskSequentialOpsCustom(MaskSequentialOps): @classmethod - def transform( - cls, input: Tensor, module: Module, param: ParamItem, extra_args: Optional[Dict[str, Any]] = None - ) -> Tensor: + def transform(cls, input: Tensor, module: Module, param: ParamItem, extra_args: dict[str, Any] | None = None) -> Tensor: """Apply a transformation with respect to the parameters. Args: @@ -203,14 +201,11 @@ def transform( input = module(input, params=cls.get_instance_module_param(param), data_keys=[DataKey.MASK], **extra_args) elif isinstance(module, (_AugmentationBase)): - input = module.transform_masks( - input, params=cls.get_instance_module_param(param), flags=module.flags, **extra_args - ) - - elif isinstance(module, K.ImageSequential) and not module.is_intensity_only(): - input = module.transform_masks(input, params=cls.get_sequential_module_param(param), extra_args=extra_args) + input = module.transform_masks(input, params=cls.get_instance_module_param(param), flags=module.flags, **extra_args) - elif isinstance(module, K.container.ImageSequentialBase): + elif (isinstance(module, K.ImageSequential) and not module.is_intensity_only()) or isinstance( + module, K.container.ImageSequentialBase + ): input = module.transform_masks(input, params=cls.get_sequential_module_param(param), extra_args=extra_args) elif isinstance(module, (K.auto.operations.OperationBase,)): @@ -220,8 +215,8 @@ def transform( @classmethod def transform_list( - cls, input: List[Tensor], module: Module, param: ParamItem, extra_args: Optional[Dict[str, Any]] = None - ) -> List[Tensor]: + cls, input: list[Tensor], module: Module, param: ParamItem, extra_args: dict[str, Any] | None = None + ) -> list[Tensor]: """Apply a transformation with respect to the parameters. Args: @@ -233,27 +228,13 @@ def transform_list( """ if extra_args is None: extra_args = {} - if isinstance(module, (K.GeometricAugmentationBase2D,)): + if isinstance(module, (K.GeometricAugmentationBase2D, K.RigidAffineAugmentationBase3D)): tfm_input = [] params = cls.get_instance_module_param(param) params_i = copy.deepcopy(params) for i, inp in enumerate(input): params_i["batch_prob"] = params["batch_prob"][i] - tfm_inp = module.transform_masks( - inp, params=params_i, flags=module.flags, transform=module.transform_matrix, **extra_args - ) - tfm_input.append(tfm_inp) - input = tfm_input - - elif isinstance(module, (K.RigidAffineAugmentationBase3D,)): - tfm_input = [] - params = cls.get_instance_module_param(param) - params_i = copy.deepcopy(params) - for i, inp in enumerate(input): - params_i["batch_prob"] = params["batch_prob"][i] - tfm_inp = module.transform_masks( - inp, params=params_i, flags=module.flags, transform=module.transform_matrix, **extra_args - ) + tfm_inp = module.transform_masks(inp, params=params_i, flags=module.flags, transform=module.transform_matrix, **extra_args) tfm_input.append(tfm_inp) input = tfm_input @@ -267,15 +248,9 @@ def transform_list( tfm_input.append(tfm_inp) input = tfm_input - elif isinstance(module, K.ImageSequential) and not module.is_intensity_only(): - tfm_input = [] - seq_params = cls.get_sequential_module_param(param) - for inp in input: - tfm_inp = module.transform_masks(inp, params=seq_params, extra_args=extra_args) - tfm_input.append(tfm_inp) - input = tfm_input - - elif isinstance(module, K.container.ImageSequentialBase): + elif (isinstance(module, K.ImageSequential) and not module.is_intensity_only()) or isinstance( + module, K.container.ImageSequentialBase + ): tfm_input = [] seq_params = cls.get_sequential_module_param(param) for inp in input: @@ -285,13 +260,13 @@ def transform_list( elif isinstance(module, (K.auto.operations.OperationBase,)): raise NotImplementedError( - "The support for list of masks under auto operations are not yet supported. You are welcome to file a" - " PR in our repo." + "The support for list of masks under auto operations are not yet supported. You are welcome to file a PR in our repo." ) return input + class AugmentationSequentialOpsCustom(AugmentationSequentialOps): - def _get_op(self, data_key: DataKey) -> Type[SequentialOpsInterface[Any]]: + def _get_op(self, data_key: DataKey) -> type[SequentialOpsInterface[Any]]: """Return the corresponding operation given a data key.""" if data_key == DataKey.INPUT: return InputSequentialOps @@ -304,14 +279,14 @@ def _get_op(self, data_key: DataKey) -> Type[SequentialOpsInterface[Any]]: if data_key == DataKey.CLASS: return ClassSequentialOps raise RuntimeError(f"Operation for `{data_key.name}` is not found.") - + def transform( self, *arg: DataType, module: Module, param: ParamItem, - extra_args: Dict[DataKey, Dict[str, Any]], - data_keys: Optional[Union[List[str], List[int], List[DataKey]]] = None, + extra_args: dict[DataKey, dict[str, Any]], + data_keys: Union[list[str], list[int], list[DataKey]] | None = None, ) -> Union[DataType, SequenceDataType]: _data_keys = self.preproc_datakeys(data_keys) @@ -326,7 +301,7 @@ def transform( extra_args=extra_args, ), ) - + keys = [dk.name for dk in _data_keys] if "MASK" in keys: mask_index = keys.index("MASK") @@ -342,4 +317,4 @@ def transform( outputs.append(op.transform(inp, module, param=param, extra_args=extra_arg)) if len(outputs) == 1 and isinstance(outputs, (list, tuple)): return outputs[0] - return outputs \ No newline at end of file + return outputs diff --git a/auglab/transforms/gpu/contrast.py b/auglab/transforms/gpu/contrast.py index 570162b..1167440 100644 --- a/auglab/transforms/gpu/contrast.py +++ b/auglab/transforms/gpu/contrast.py @@ -1,39 +1,38 @@ +import math +import random +from typing import Any, Union + import torch -import torch.nn as nn -from torch.nn import functional as F import torchvision.transforms._functional_tensor as F_t - -from typing import Any, Dict, Optional -from kornia.core import Tensor -import random -import math +from torch import Tensor +from torch.nn import functional as F from auglab.transforms.gpu.base import ImageOnlyTransform -from typing import Any, Dict, Optional, Tuple, Union, List -def _choose_region_mode(p_in: float, p_out: float, seg_mask: Optional[torch.Tensor]) -> str: +def _choose_region_mode(p_in: float, p_out: float, seg_mask: torch.Tensor | None) -> str: # noqa: ARG001 -- seg_mask kept for signature symmetry with _apply_region_mode """Sample where to apply the transform: 'in', 'out', or 'all'. - p_in, p_out are probabilities in [0,1]. - - If seg_mask is None, or both probs are 0, return 'all'. - - If p_in + p_out > 1, renormalize so p_all=0. + - If both probs are 0, or both fire at once, return 'all'. + - seg_mask is accepted but unused here; _apply_region_mode treats a None + mask as 'all' regardless of the mode chosen. """ p_in = float(max(0.0, min(1.0, p_in))) p_out = float(max(0.0, min(1.0, p_out))) in_bool = torch.rand(()) < p_in out_bool = torch.rand(()) < p_out if in_bool and not out_bool: - return 'in' + return "in" if out_bool and not in_bool: - return 'out' - return 'all' + return "out" + return "all" def _apply_region_mode( orig: torch.Tensor, transformed: torch.Tensor, - seg_mask: Optional[torch.Tensor], + seg_mask: torch.Tensor | None, mode: str, normalize: bool = False, mix_in_out: bool = False, @@ -46,7 +45,7 @@ def _apply_region_mode( mix_in_out: if True, randomly apply transform to some of the segmentation, not all. """ - if seg_mask is None or mode == 'all': + if seg_mask is None or mode == "all": return transformed # Rescale transformed based on min max orig @@ -116,7 +115,7 @@ class RandomConvTransformGPU(ImageOnlyTransform): def __init__( self, kernel_type: str = "Laplace", - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default same_on_batch: bool = False, retain_stats: bool = False, in_seg: float = 0.0, @@ -126,6 +125,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) if kernel_type not in ["Laplace", "Scharr", "GaussianBlur", "UnsharpMask", "RandConv"]: raise NotImplementedError('Currently only "Laplace", "Scharr", "GaussianBlur", "UnsharpMask" and "RandConv" are supported.') @@ -199,14 +200,12 @@ def get_kernel(self, device: torch.device) -> torch.Tensor: return kernel @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Initialize kernel kernel = self.get_kernel(device=input.device) # Load segmentation - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") # Apply convolution for c in self.apply_to_channel: @@ -267,7 +266,7 @@ def apply_transform( x = (x - nm) / (ns + eps) * os + om # Apply region selection - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) @@ -292,7 +291,11 @@ def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torc padding = [kernel.shape[2] // 2, kernel.shape[2] // 2, kernel.shape[3] // 2, kernel.shape[3] // 2] elif dim == 3: kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1], kernel.shape[2]) - padding = [kernel.shape[2] // 2, kernel.shape[2] // 2, kernel.shape[3] // 2, kernel.shape[3] // 2] + [ + padding = [ + kernel.shape[2] // 2, + kernel.shape[2] // 2, + kernel.shape[3] // 2, + kernel.shape[3] // 2, kernel.shape[4] // 2, kernel.shape[4] // 2, ] @@ -303,7 +306,7 @@ def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torc # padding = (left, right, top, bottom) img = F.pad(img, padding, mode="reflect") - if dim == 2: + if dim == 2: # noqa: SIM108 -- the 2d/3d split reads better spelled out than as a ternary img = F.conv2d(img, kernel, groups=img.shape[-(1 + dim)]) else: # dim == 3 img = F.conv3d(img, kernel, groups=img.shape[-(1 + dim)]) @@ -369,7 +372,7 @@ def __init__( self, mean: float = 0.0, std: float = 0.1, - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, @@ -378,6 +381,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.apply_to_channel = apply_to_channel self.mean = mean @@ -387,11 +392,9 @@ def __init__( self.mix_in_out = mix_in_out @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Generate Gaussian noise with the same shape as input - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: if self.same_on_batch: std = torch.rand(1, device=input.device, dtype=input.dtype) * self.std @@ -402,10 +405,10 @@ def apply_transform( noise = torch.randn_like(input[:, c], device=input.device, dtype=input.dtype) for i in range(input.shape[0]): noise[i] = noise[i] * std[i] + self.mean - + orig = input[:, c] x = orig + noise - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) @@ -414,7 +417,7 @@ def apply_transform( print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[:, c] = x - + return input @@ -437,7 +440,7 @@ class RandomBrightnessGPU(ImageOnlyTransform): def __init__( self, brightness_range: list[float, float] = (0.9, 1.1), - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, @@ -446,6 +449,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.brightness_range = brightness_range self.apply_to_channel = apply_to_channel @@ -454,24 +459,29 @@ def __init__( self.mix_in_out = mix_in_out @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Apply brightness adjustment - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: channel_data = input[:, c] # [N, ...spatial...] orig = channel_data.clone() if self.same_on_batch: - factor = torch.rand(1, device=input.device, dtype=input.dtype) * (self.brightness_range[1] - self.brightness_range[0]) + self.brightness_range[0] + factor = ( + torch.rand(1, device=input.device, dtype=input.dtype) * (self.brightness_range[1] - self.brightness_range[0]) + + self.brightness_range[0] + ) x = channel_data * factor else: - factor = torch.rand(input.shape[0], device=input.device, dtype=input.dtype) * (self.brightness_range[1] - self.brightness_range[0]) + self.brightness_range[0] + factor = ( + torch.rand(input.shape[0], device=input.device, dtype=input.dtype) + * (self.brightness_range[1] - self.brightness_range[0]) + + self.brightness_range[0] + ) x = channel_data.clone() for i in range(input.shape[0]): x[i] = x[i] * factor[i] - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -479,7 +489,7 @@ def apply_transform( print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[:, c] = x - + return input @@ -505,7 +515,7 @@ def __init__( self, gamma_range: list[float, float] = (0.9, 1.1), invert_image: bool = False, - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default retain_stats: bool = False, same_on_batch: bool = False, in_seg: float = 0.0, @@ -515,6 +525,8 @@ def __init__( keepdim: bool = False, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.gamma_range = gamma_range self.invert_image = invert_image @@ -525,17 +537,13 @@ def __init__( self.mix_in_out = mix_in_out @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Apply gamma transform - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: - if self.invert_image: - channel_data = -input[:, c] # [N, ...spatial...] - else: - channel_data = input[:, c] # [N, ...spatial...] + # [N, ...spatial...] + channel_data = -input[:, c] if self.invert_image else input[:, c] orig_full = input[:, c].clone() if self.retain_stats: @@ -591,7 +599,7 @@ def apply_transform( if self.invert_image: channel_data = -channel_data - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) channel_data = _apply_region_mode(orig_full, channel_data, seg_mask, region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -599,7 +607,7 @@ def apply_transform( print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[:, c] = channel_data - + return input @@ -623,7 +631,7 @@ class RandomContrastGPU(ImageOnlyTransform): def __init__( self, contrast_range: list[float, float] = (0.9, 1.1), - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default retain_stats: bool = False, same_on_batch: bool = False, in_seg: float = 0.0, @@ -633,6 +641,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.contrast_range = contrast_range self.apply_to_channel = apply_to_channel @@ -642,12 +652,10 @@ def __init__( self.mix_in_out = mix_in_out @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Apply brightness adjustment - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: channel_data = input[:, c] # [N, ...spatial...] orig = channel_data.clone() @@ -658,24 +666,30 @@ def apply_transform( orig_stds = channel_data.std(dim=reduce_dims) if self.same_on_batch: - factor = torch.rand(1, device=input.device, dtype=input.dtype) * (self.contrast_range[1] - self.contrast_range[0]) + self.contrast_range[0] + factor = ( + torch.rand(1, device=input.device, dtype=input.dtype) * (self.contrast_range[1] - self.contrast_range[0]) + + self.contrast_range[0] + ) x = channel_data.clone() for i in range(input.shape[0]): mean = x[i].mean() x[i] = (x[i] - mean) * factor + mean else: - factor = torch.rand(input.shape[0], device=input.device, dtype=input.dtype) * (self.contrast_range[1] - self.contrast_range[0]) + self.contrast_range[0] + factor = ( + torch.rand(input.shape[0], device=input.device, dtype=input.dtype) * (self.contrast_range[1] - self.contrast_range[0]) + + self.contrast_range[0] + ) x = channel_data.clone() for i in range(input.shape[0]): mean = x[i].mean() x[i] = (x[i] - mean) * factor[i] + mean - + if self.retain_stats: # Adjust mean and std to match original eps = 1e-8 reduce_dims = tuple(range(1, x.dim())) new_mean = x.mean(dim=reduce_dims) # [N] - new_std = x.std(dim=reduce_dims) # [N] + new_std = x.std(dim=reduce_dims) # [N] # reshape stats to broadcast over spatial dims: [N,1,1,...] shape = [x.shape[0]] + [1] * (x.dim() - 1) nm = new_mean.view(shape) @@ -683,7 +697,7 @@ def apply_transform( om = orig_means.view(shape) os = orig_stds.view(shape) x = (x - nm) / (ns + eps) * os + om - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -691,7 +705,7 @@ def apply_transform( print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[:, c] = x - + return input @@ -715,7 +729,7 @@ class RandomFunctionGPU(ImageOnlyTransform): def __init__( self, func: callable = lambda x: x**2, - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default retain_stats: bool = False, same_on_batch: bool = False, in_seg: float = 0.0, @@ -725,6 +739,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.func = func self.retain_stats = retain_stats @@ -734,12 +750,10 @@ def __init__( self.mix_in_out = mix_in_out @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Apply function transform - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: x = input[:, c] # shape [N, ...spatial...] orig = x.clone() @@ -768,7 +782,7 @@ def apply_transform( om = orig_means.view(shape) os = orig_stds.view(shape) x = (x - nm) / (ns + eps) * os + om - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -797,7 +811,7 @@ class RandomInverseGPU(ImageOnlyTransform): def __init__( self, - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default retain_stats: bool = False, same_on_batch: bool = False, in_seg: float = 0.0, @@ -808,6 +822,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats @@ -817,15 +833,13 @@ def __init__( self.mix_in_out = mix_in_out @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Inverse image - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: for i in range(input.shape[0]): - x= input[i, c] # shape [...spatial...] + x = input[i, c] # shape [...spatial...] orig = x.clone() if self.retain_stats: orig_means = x.mean() @@ -845,7 +859,7 @@ def apply_transform( alpha = torch.rand(1, device=input.device) x = alpha * orig + (1 - alpha) * x - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask[i]) x = _apply_region_mode(orig, x, seg_mask[i], region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -875,7 +889,7 @@ class RandomHistogramEqualizationGPU(ImageOnlyTransform): def __init__( self, - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default retain_stats: bool = False, same_on_batch: bool = False, in_seg: float = 0.0, @@ -886,6 +900,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.retain_stats = retain_stats self.apply_to_channel = apply_to_channel @@ -895,12 +911,10 @@ def __init__( self.mix_prob = mix_prob @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Apply histogram equalization transform - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: channel_data = input[:, c] # shape [N, ...spatial...] orig = channel_data.clone() @@ -956,7 +970,7 @@ def apply_transform( os = orig_stds.view(shape) channel_data = (channel_data - nm) / (ns + eps) * os + om - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) channel_data = _apply_region_mode(orig, channel_data, seg_mask, region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -991,9 +1005,9 @@ class RandomBiasFieldGPU(ImageOnlyTransform): def __init__( self, - coefficients: Union[float, Tuple[float, float]] = 0.5, + coefficients: Union[float, tuple[float, float]] = 0.5, order: int = 3, - apply_to_channel: list[int] = [0], + apply_to_channel: list[int] | None = None, invert: bool = False, retain_stats: bool = False, in_seg: float = 0.0, @@ -1004,6 +1018,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) if isinstance(coefficients, (int, float)): self.coeff_range = (-float(coefficients), float(coefficients)) @@ -1027,11 +1043,11 @@ def _num_coeffs(self, dim: int) -> int: if dim == 3: for xo in range(self.order + 1): for yo in range(self.order + 1 - xo): - for zo in range(self.order + 1 - (xo + yo)): + for _zo in range(self.order + 1 - (xo + yo)): count += 1 elif dim == 2: for xo in range(self.order + 1): - for yo in range(self.order + 1 - xo): + for _yo in range(self.order + 1 - xo): count += 1 else: raise ValueError("Only 2D or 3D spatial dims supported for bias field") @@ -1047,7 +1063,7 @@ def _sample_coeffs(self, batch_size: int, device: torch.device, dtype: torch.dty coeff = torch.empty(n, batch_size, device=device, dtype=dtype).uniform_(low, high) return coeff # shape (n_coeffs, B) - def _make_grids(self, spatial_shape: Tuple[int, ...], device: torch.device, dtype: torch.dtype) -> List[torch.Tensor]: + def _make_grids(self, spatial_shape: tuple[int, ...], device: torch.device, dtype: torch.dtype) -> list[torch.Tensor]: # Create coordinate grids normalized to [-1, 1] if len(spatial_shape) == 2: h, w = spatial_shape @@ -1069,9 +1085,9 @@ def _make_grids(self, spatial_shape: Tuple[int, ...], device: torch.device, dtyp def apply_transform( self, input: Tensor, - params: Dict[str, Tensor], - flags: Dict[str, Any], - transform: Optional[Tensor] = None, + params: dict[str, Tensor], + flags: dict[str, Any], + transform: Tensor | None = None, ) -> Tensor: # input: (N, C, [D,] H, W) if input.dim() not in (4, 5): @@ -1084,7 +1100,7 @@ def apply_transform( coeffs = self._sample_coeffs(batch_size, device, dtype, dim) # (n_coeffs, B) grids = self._make_grids(spatial, device, dtype) - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") # Initialize bias map per batch element bias_map = torch.zeros((batch_size, *spatial), device=device, dtype=dtype) @@ -1140,7 +1156,7 @@ def apply_transform( nm = new_mean.view(shape) ns = new_std.view(shape) channel = (channel - nm) / (ns + eps) * os + om - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) channel = _apply_region_mode(orig, channel, seg_mask, region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -1148,9 +1164,10 @@ def apply_transform( print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[:, c] = channel - + return input + # Random clamping transform class RandomClampGPU(ImageOnlyTransform): """Apply random gamma adjustment to image. @@ -1167,11 +1184,11 @@ class RandomClampGPU(ImageOnlyTransform): Returns: Tensor: Image with adjusted brightness. """ - + def __init__( self, max_clamp_amount: float = 0.2, - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default retain_stats: bool = False, same_on_batch: bool = False, in_seg: float = 0.0, @@ -1181,6 +1198,8 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.max_clamp_amount = max_clamp_amount self.apply_to_channel = apply_to_channel @@ -1190,12 +1209,10 @@ def __init__( self.mix_in_out = mix_in_out @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Apply clamping - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: channel_data = input[:, c] # [N, ...spatial...] orig = channel_data.clone() @@ -1204,7 +1221,7 @@ def apply_transform( # store per-sample mean/std (shape [N]) orig_means = channel_data.mean(dim=reduce_dims) orig_stds = channel_data.std(dim=reduce_dims) - + if self.same_on_batch: min_percentile = torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount max_percentile = 1.0 - (torch.rand(1, device=input.device, dtype=input.dtype) * self.max_clamp_amount) @@ -1221,13 +1238,13 @@ def apply_transform( min_val = torch.quantile(x[i].flatten(), min_percentile) max_val = torch.quantile(x[i].flatten(), max_percentile) x[i] = torch.clamp(x[i], min_val, max_val) - + if self.retain_stats: # Adjust mean and std to match original eps = 1e-8 reduce_dims = tuple(range(1, x.dim())) new_mean = x.mean(dim=reduce_dims) # [N] - new_std = x.std(dim=reduce_dims) # [N] + new_std = x.std(dim=reduce_dims) # [N] # reshape stats to broadcast over spatial dims: [N,1,1,...] shape = [x.shape[0]] + [1] * (x.dim() - 1) nm = new_mean.view(shape) @@ -1235,7 +1252,7 @@ def apply_transform( om = orig_means.view(shape) os = orig_stds.view(shape) x = (x - nm) / (ns + eps) * os + om - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=self.mix_in_out) # Final safety: check if nan/inf appeared @@ -1243,7 +1260,7 @@ def apply_transform( print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[:, c] = x - + return input @@ -1258,13 +1275,15 @@ class ZscoreNormalizationGPU(ImageOnlyTransform): def __init__( self, - apply_to_channel: list[int] = [0], + apply_to_channel: list[int] | None = None, keepdim: bool = True, in_seg: float = 0.0, out_seg: float = 0.0, p: float = 1.0, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=False, keepdim=keepdim) self.apply_to_channel = apply_to_channel self.in_seg = in_seg @@ -1274,12 +1293,12 @@ def __init__( def apply_transform( self, input: Tensor, - params: Dict[str, Tensor], - flags: Dict[str, Any], - transform: Optional[Tensor] = None, + params: dict[str, Tensor], + flags: dict[str, Any], + transform: Tensor | None = None, ) -> Tensor: # input: (N, C, [D,] H, W) - seg_mask = params.get('seg', None) + seg_mask = params.get("seg") for c in self.apply_to_channel: if c < 0 or c >= input.shape[1]: continue # skip invalid channel index @@ -1290,7 +1309,7 @@ def apply_transform( # use unbiased=False for stability, and clamp std to avoid division by ~0 std = channel.std(dim=reduce_dims, keepdim=True, unbiased=False).clamp_min(1e-8) channel = (channel - mean) / std - if not seg_mask is None: + if seg_mask is not None: region_mode = _choose_region_mode(self.in_seg, self.out_seg, seg_mask) channel = _apply_region_mode(orig, channel, seg_mask, region_mode) # Final safety: check if nan/inf appeared @@ -1298,5 +1317,5 @@ def apply_transform( print(f"Warning nan: {self.__class__.__name__}", flush=True) continue input[:, c] = channel - + return input diff --git a/auglab/transforms/gpu/domain_transfer.py b/auglab/transforms/gpu/domain_transfer.py index 68a02f3..bf5f7e2 100644 --- a/auglab/transforms/gpu/domain_transfer.py +++ b/auglab/transforms/gpu/domain_transfer.py @@ -35,26 +35,22 @@ """ import math +from typing import Any import numpy as np import torch +from torch import Tensor from torch.distributions import Dirichlet from torch.nn import functional as F -from typing import Any, Dict, List, Optional, Tuple - -from kornia.core import Tensor from auglab.transforms.gpu.base import ImageOnlyTransform # Default transfer LUT bank (built by embeddaug/analysis/playground/build_transfer_bank.py). -DEFAULT_BANK_PATH = ( - "/DATA/NAS/ongoing_projects/hendrik/nathan-transferaug/" - "embeddaug/analysis/playground/results/domain_transfer_bank.npz" -) +DEFAULT_BANK_PATH = "/DATA/NAS/ongoing_projects/hendrik/nathan-transferaug/embeddaug/analysis/playground/results/domain_transfer_bank.npz" def _gaussian_kernel1d(sigma: float, device, dtype) -> torch.Tensor: - radius = max(1, int(round(3.0 * sigma))) + radius = max(1, round(3.0 * sigma)) x = torch.arange(-radius, radius + 1, device=device, dtype=dtype) k = torch.exp(-0.5 * (x / sigma) ** 2) return k / k.sum() @@ -68,9 +64,12 @@ def _gaussian_blur3d(x: torch.Tensor, sigma: float) -> torch.Tensor: k = _gaussian_kernel1d(sigma, x.device, x.dtype) r = (k.numel() - 1) // 2 for dim in (2, 3, 4): - shape = [1, 1, 1, 1, 1]; shape[dim] = k.numel() - ker = k.view(shape).repeat(c, 1, 1, 1, 1) # [C,1,kD,kH,kW] for separable conv - pad = [0, 0, 0, 0, 0, 0]; pad[(4 - dim) * 2] = r; pad[(4 - dim) * 2 + 1] = r + shape = [1, 1, 1, 1, 1] + shape[dim] = k.numel() + ker = k.view(shape).repeat(c, 1, 1, 1, 1) # [C,1,kD,kH,kW] for separable conv + pad = [0, 0, 0, 0, 0, 0] + pad[(4 - dim) * 2] = r + pad[(4 - dim) * 2 + 1] = r x = F.conv3d(F.pad(x, pad, mode="replicate"), ker, groups=c) return x @@ -84,7 +83,7 @@ def _random_bias_field3d(shape, std: float, scale: float, device, dtype) -> torc ``contrast.py::RandomBiasFieldGPU``; kept local so this module stays self-contained. """ d, h, w = shape - small = [max(2, int(math.ceil(s * scale))) for s in (d, h, w)] + small = [max(2, math.ceil(s * scale)) for s in (d, h, w)] s = torch.rand((), device=device) * std field = torch.randn(1, 1, *small, device=device, dtype=dtype) * s field = F.interpolate(field, size=(d, h, w), mode="trilinear", align_corners=True) @@ -101,10 +100,10 @@ def _random_smooth_field01(shape, scale: float, gain: float, device, dtype) -> t balance between the two domains per draw. """ d, h, w = shape - small = [max(2, int(math.ceil(s * scale))) for s in (d, h, w)] + small = [max(2, math.ceil(s * scale)) for s in (d, h, w)] field = torch.randn(1, 1, *small, device=device, dtype=dtype) field = F.interpolate(field, size=(d, h, w), mode="trilinear", align_corners=True)[0, 0] - offset = (torch.rand((), device=device, dtype=dtype) * 4.0 - 2.0) + offset = torch.rand((), device=device, dtype=dtype) * 4.0 - 2.0 return torch.sigmoid(gain * field + offset) @@ -113,13 +112,13 @@ class RandomDomainTransferGPU(ImageOnlyTransform): def __init__( self, - bank_path: Optional[str] = None, - source_label: Optional[str] = None, - targets: Optional[List[str]] = None, + bank_path: str | None = None, + source_label: str | None = None, + targets: list[str] | None = None, include_self: bool = True, any_source: bool = False, sigma: float = 2.0, - apply_to_channel: List[int] = [0], + apply_to_channel: list[int] | None = None, zscore_io: str = "auto", pct: float = 1.0, blend_targets: int = 1, @@ -135,10 +134,12 @@ def __init__( keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) bank_path = bank_path or DEFAULT_BANK_PATH data = np.load(bank_path, allow_pickle=True) - self.labels: List[str] = [str(x) for x in data["labels"].tolist()] + self.labels: list[str] = [str(x) for x in data["labels"].tolist()] self.L = int(data["L"]) self.num_classes = int(data["num_classes"]) self.any_source = bool(any_source) @@ -156,9 +157,9 @@ def __init__( x, y = k.split("__", 1) if x not in self.labels or y not in self.labels: continue - if not include_self and x == y: # identity transfers + if not include_self and x == y: # identity transfers continue - if targets is not None and y not in targets: # optional: restrict the target domain + if targets is not None and y not in targets: # optional: restrict the target domain continue keys.append(k) self.targets = keys @@ -202,7 +203,7 @@ def __init__( self.spatial_mix_scale = float(spatial_mix_scale) self.spatial_mix_gain = float(spatial_mix_gain) - def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> Tuple[Tensor, bool]: + def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> tuple[Tensor, bool]: """Build the per-class LUTs to use for one sample. Returns ``(lut_used, per_class)`` where ``lut_used`` is ``[n_draws, NC, L]`` with @@ -215,9 +216,9 @@ def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> Tuple[ K = T if (self.blend_targets <= 0 or self.blend_targets > T) else self.blend_targets per_class = (self.p_class_mix > 0.0) and (float(torch.rand((), device=device)) < self.p_class_mix) - if K == 1 and not per_class: # fast path == original behaviour + if K == 1 and not per_class: # fast path == original behaviour ti = int(torch.randint(T, (1,), device=device).item()) - return lut_bank[ti:ti + 1], False # [1, NC, L] + return lut_bank[ti : ti + 1], False # [1, NC, L] n_draws = n_seg_c if per_class else 1 beta = torch.zeros(n_draws, T, device=device, dtype=lut_bank.dtype) @@ -229,16 +230,15 @@ def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> Tuple[ conc = torch.full((K,), self.blend_concentration, device=device, dtype=torch.float32) wd = Dirichlet(conc).sample().to(lut_bank.dtype) beta[d, idx] = wd - lut_used = torch.einsum("dt,tcl->dcl", beta, lut_bank) # [n_draws, NC, L] + lut_used = torch.einsum("dt,tcl->dcl", beta, lut_bank) # [n_draws, NC, L] return lut_used, per_class @staticmethod - def _accumulate(lut_used: Tensor, per_class: bool, w_b: Tensor, - il: Tensor, ih: Tensor, xf: Tensor, n_seg_c: int) -> Tensor: + def _accumulate(lut_used: Tensor, per_class: bool, w_b: Tensor, il: Tensor, ih: Tensor, xf: Tensor, n_seg_c: int) -> Tensor: """Class-weighted LUT interpolation ``Σ_c w_c · interp(LUT_c, x)`` → ``[D, H, W]``.""" acc = torch.zeros_like(xf) for c in range(n_seg_c): - lut_c = lut_used[c if per_class else 0, c] # [L] + lut_c = lut_used[c if per_class else 0, c] # [L] acc += w_b[c] * (lut_c[il] * (1 - xf) + lut_c[ih] * xf) return acc @@ -253,25 +253,23 @@ def _to_unit(self, x: Tensor) -> Tensor: """Map a (z-scored) image into the LUT's [0,1] domain via percentile scaling (clip), matching how the bank's source histograms were normalised.""" flat = x.reshape(-1).float() - if flat.numel() > 1_000_000: # cap for torch.quantile + if flat.numel() > 1_000_000: # cap for torch.quantile flat = flat[torch.linspace(0, flat.numel() - 1, 1_000_000, device=x.device).long()] lo = torch.quantile(flat, self.pct / 100.0) hi = torch.quantile(flat, 1.0 - self.pct / 100.0) return ((x - lo) / (hi - lo).clamp_min(1e-6)).clamp(0.0, 1.0) @torch.no_grad() - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: if "seg" not in params: return input seg = params["seg"] - if seg.dim() != input.dim(): # accept [N, ...] integer seg → one-hot-ish + if seg.dim() != input.dim(): # accept [N, ...] integer seg → one-hot-ish if seg.dim() == input.dim() - 1: seg = seg.unsqueeze(1) else: return input - if input.dim() != 5: # this GPU pipeline is 3D: [N, C, D, H, W] + if input.dim() != 5: # this GPU pipeline is 3D: [N, C, D, H, W] return input out = input.clone() @@ -286,9 +284,9 @@ def apply_transform( N = input.shape[0] for ch in self.apply_to_channel: for b in range(N): - x = input[b, ch] # [D, H, W] — may be z-scored + x = input[b, ch] # [D, H, W] — may be z-scored zmode = self._is_zscore(x) - if zmode: # remember scale, map into LUT's [0,1] domain + if zmode: # remember scale, map into LUT's [0,1] domain mu, sd = x.mean(), x.std().clamp_min(1e-6) x01 = self._to_unit(x) else: @@ -303,24 +301,21 @@ def apply_transform( # hybridised per class (see _sample_blended_luts). With p_spatial_mix, two # independent domain transfers are blended across space by a smooth field, so # different regions look like different target sequences. - spatial = (self.p_spatial_mix > 0.0) and \ - (float(torch.rand((), device=input.device)) < self.p_spatial_mix) + spatial = (self.p_spatial_mix > 0.0) and (float(torch.rand((), device=input.device)) < self.p_spatial_mix) if spatial: lutA, pcA = self._sample_blended_luts(lut_bank, n_seg_c, input.device) lutB, pcB = self._sample_blended_luts(lut_bank, n_seg_c, input.device) accA = self._accumulate(lutA, pcA, w[b], il, ih, xf, n_seg_c) accB = self._accumulate(lutB, pcB, w[b], il, ih, xf, n_seg_c) - a = _random_smooth_field01(x01.shape, self.spatial_mix_scale, - self.spatial_mix_gain, input.device, x01.dtype) + a = _random_smooth_field01(x01.shape, self.spatial_mix_scale, self.spatial_mix_gain, input.device, x01.dtype) acc = (1.0 - a) * accA + a * accB else: lut_used, per_class = self._sample_blended_luts(lut_bank, n_seg_c, input.device) acc = self._accumulate(lut_used, per_class, w[b], il, ih, xf, n_seg_c) acc = acc.clamp(0.0, 1.0) - if self.bias_field_std > 0.0: # smooth multiplicative spatial inhomogeneity - field = _random_bias_field3d(acc.shape, self.bias_field_std, self.bias_scale, - acc.device, acc.dtype) + if self.bias_field_std > 0.0: # smooth multiplicative spatial inhomogeneity + field = _random_bias_field3d(acc.shape, self.bias_field_std, self.bias_scale, acc.device, acc.dtype) acc = (acc * field).clamp(0.0, 1.0) if zmode: diff --git a/auglab/transforms/gpu/fromSeg.py b/auglab/transforms/gpu/fromSeg.py index 8b0e149..8fc34d6 100644 --- a/auglab/transforms/gpu/fromSeg.py +++ b/auglab/transforms/gpu/fromSeg.py @@ -1,18 +1,16 @@ import random +from typing import Any import torch -from torch import nn -from torch.nn import functional as F - -from typing import Any, Dict, Optional, Tuple, Union, List, Protocol -from kornia.core import Tensor import torch.distributed as dist +from torch import Tensor, nn +from torch.nn import functional as F from auglab.transforms.gpu.base import ImageOnlyTransform - # ── PALETTE AUG helpers ────────────────────────────────────────────────── + def _kmeans_1d(values: torch.Tensor, C: int, n_iter: int = 10) -> torch.Tensor: """1-D K-means on foreground values. Returns (C,) centroids.""" centroids = torch.linspace(values.min().item(), values.max().item(), C, device=values.device) @@ -47,7 +45,7 @@ def _voronoi_region_ids( fg: torch.Tensor, C: int, device: torch.device, - s_choices: List[int], + s_choices: list[int], skip_sub_parc_prob: float, ) -> tuple[torch.Tensor, int]: """Spatially subdivide each K-means cluster into Voronoi sub-regions. @@ -83,12 +81,14 @@ def _voronoi_region_ids( # ───────────────────────────────────────────────────────────────────────────── + def _normal_pdf(x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: inv = 1.0 / (std + 1e-6) return (inv / (torch.sqrt(torch.tensor(2.0 * 3.141592653589793, device=x.device, dtype=x.dtype)))) * torch.exp( -0.5 * ((x - mean) * inv) ** 2 ) + ## Redistribute segmentation values transform (GPU) class RandomRedistributeSegGPU(ImageOnlyTransform): """Redistribute image values using segmentation regions (GPU version). @@ -100,15 +100,21 @@ class RandomRedistributeSegGPU(ImageOnlyTransform): def __init__( self, in_seg: float = 0.2, - apply_to_channel: list[int] = [0], + apply_to_channel: list[int] | None = None, retain_stats: bool = False, same_on_batch: bool = False, p: float = 1.0, keepdim: bool = True, - std_noise_range: list[float] = [0.1, 0.3], - dilation_iterations_range: list[int] = [1, 3], + std_noise_range: list[float] | None = None, + dilation_iterations_range: list[int] | None = None, **kwargs, ) -> None: + if dilation_iterations_range is None: + dilation_iterations_range = [1, 3] + if std_noise_range is None: + std_noise_range = [0.1, 0.3] + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.in_seg = in_seg self.apply_to_channel = apply_to_channel @@ -117,13 +123,11 @@ def __init__( self.dilation_iterations_range = dilation_iterations_range @torch.no_grad() - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # Expect segmentation provided in params: shape [N, 1, ...] or [N, C_seg, ...] - if 'seg' not in params: + if "seg" not in params: return input - seg = params['seg'] + seg = params["seg"] if seg.dim() != input.dim(): # Allow seg [N, ...] by adding channel dim if seg.dim() == input.dim() - 1: @@ -151,8 +155,8 @@ def apply_transform( orig_std = flat.std(dim=1, unbiased=False) # Normalize entire batch to [0,1] per sample - img_min = img_batch.view(N, -1).min(dim=1)[0].view(N, *([1] * (img_batch.dim()-1))) - img_max = img_batch.view(N, -1).max(dim=1)[0].view(N, *([1] * (img_batch.dim()-1))) + img_min = img_batch.view(N, -1).min(dim=1)[0].view(N, *([1] * (img_batch.dim() - 1))) + img_max = img_batch.view(N, -1).max(dim=1)[0].view(N, *([1] * (img_batch.dim() - 1))) denom = (img_max - img_min).clamp_min(1e-6) x_batch = (img_batch - img_min) / denom @@ -176,7 +180,9 @@ def apply_transform( # Vectorized dilation for all regions (3 iterations) dilated = masks.float() - dilation_iterations = torch.randint(self.dilation_iterations_range[0], self.dilation_iterations_range[1]+1, (1,), device=input.device)[0].item() + dilation_iterations = torch.randint( + self.dilation_iterations_range[0], self.dilation_iterations_range[1] + 1, (1,), device=input.device + )[0].item() for _ in range(dilation_iterations): if spatial_dims == 3: dilated = F.max_pool3d(dilated.unsqueeze(0), 3, 1, 1).squeeze(0) @@ -194,27 +200,29 @@ def apply_transform( # Means means = (mask_flat * x_flat).sum(dim=1) / counts # Std (compute variance then sqrt) avoid indexing overhead - diffs = (x_flat - means.view(R,1)) * mask_flat + diffs = (x_flat - means.view(R, 1)) * mask_flat vars = (diffs * diffs).sum(dim=1) / counts.clamp_min(1) stds = vars.sqrt() # Dilated stats dil_counts = dil_flat.sum(dim=1).clamp_min(1) dil_means = (dil_flat * x_flat).sum(dim=1) / dil_counts - dil_diffs = (x_flat - dil_means.view(R,1)) * dil_flat + dil_diffs = (x_flat - dil_means.view(R, 1)) * dil_flat dil_vars = (dil_diffs * dil_diffs).sum(dim=1) / dil_counts dil_stds = dil_vars.sqrt() # redist_std per region - std_noise_range = torch.rand(1, device=input.device)[0] * (self.std_noise_range[1] - self.std_noise_range[0]) + self.std_noise_range[0] + std_noise_range = ( + torch.rand(1, device=input.device)[0] * (self.std_noise_range[1] - self.std_noise_range[0]) + self.std_noise_range[0] + ) redist_std = torch.maximum( torch.rand(R, device=input.device) * std_noise_range + 0.4 * torch.abs((means - dil_means) * stds / (dil_stds + 1e-6)), - torch.full((R,), 0.01, device=input.device, dtype=input.dtype) + torch.full((R,), 0.01, device=input.device, dtype=input.dtype), ) # Build additive term to_add = torch.zeros_like(x) - rand_sign = (2 * torch.rand(R, device=input.device) - 1) # random sign factor per region + rand_sign = 2 * torch.rand(R, device=input.device) - 1 # random sign factor per region if in_seg_bool.item(): # Only inside region for r in range(R): @@ -291,20 +299,28 @@ class RandomPALETTEGPU(ImageOnlyTransform): def __init__( self, - c_choices: List[int] = [2, 3, 4, 5, 6], - s_choices: List[int] = [2, 3, 4, 5, 6, 7, 8, 9, 10], - blur_sigmas: List[float] = [0.0, 0.0, 0.0, 0.3, 0.5, 0.8], + c_choices: list[int] | None = None, + s_choices: list[int] | None = None, + blur_sigmas: list[float] | None = None, dark_threshold: float = 0.01, n_kmeans_subsample: int = 10_000, skip_parcellation_prob: float = 0.10, skip_sub_parc_prob: float = 0.40, - alpha_magnitude_range: List[float] = [0.5, 2.0], + alpha_magnitude_range: list[float] | None = None, label_remap_prob: float = 0.5, min_label_voxels: int = 4, - label_classes: Optional[List[int]] = None, + label_classes: list[int] | None = None, p: float = 1.0, **kwargs: Any, ) -> None: + if alpha_magnitude_range is None: + alpha_magnitude_range = [0.5, 2.0] + if blur_sigmas is None: + blur_sigmas = [0.0, 0.0, 0.0, 0.3, 0.5, 0.8] + if s_choices is None: + s_choices = [2, 3, 4, 5, 6, 7, 8, 9, 10] + if c_choices is None: + c_choices = [2, 3, 4, 5, 6] super().__init__(p=p, **kwargs) self.c_choices = c_choices self.s_choices = s_choices @@ -322,13 +338,13 @@ def __init__( def apply_transform( self, input: Tensor, - params: Dict[str, Any], - flags: Dict[str, Any], - transform: Optional[Tensor] = None, + params: dict[str, Any], + flags: dict[str, Any], + transform: Tensor | None = None, ) -> Tensor: - seg_raw: Optional[torch.Tensor] = params.get("seg", None) + seg_raw: torch.Tensor | None = params.get("seg") - labels: Optional[torch.Tensor] = None + labels: torch.Tensor | None = None if seg_raw is not None and seg_raw.ndim == 5 and seg_raw.shape[1] > 1: labels = collapse_onehot_to_index(seg_raw) elif seg_raw is not None and seg_raw.ndim == 5 and seg_raw.shape[1] == 1: @@ -349,11 +365,15 @@ def apply_transform( flat_m_all = (images_01 > self.dark_threshold).float() # foreground mask # Voxel coordinates (shared — same spatial dims for every sample) - coords = torch.stack(torch.meshgrid( - torch.arange(D, device=device, dtype=torch.float32), - torch.arange(H, device=device, dtype=torch.float32), - torch.arange(W, device=device, dtype=torch.float32), - indexing="ij"), dim=-1).reshape(N, 3) + coords = torch.stack( + torch.meshgrid( + torch.arange(D, device=device, dtype=torch.float32), + torch.arange(H, device=device, dtype=torch.float32), + torch.arange(W, device=device, dtype=torch.float32), + indexing="ij", + ), + dim=-1, + ).reshape(N, 3) # ── Step 1: PALETTE K-means + Voronoi per-region affine remap ────────── synth_list = [] @@ -374,9 +394,9 @@ def apply_transform( C_k = self.c_choices[int(torch.rand(1, device=device).item() * len(self.c_choices))] idx = torch.randint(0, N, (min(N, 40_000),), device=device) samp = flat[idx] - sub_fg = samp[samp > self.dark_threshold][:self.n_kmeans_subsample] + sub_fg = samp[samp > self.dark_threshold][: self.n_kmeans_subsample] if sub_fg.numel() < 4: - sub_fg = samp[:self.n_kmeans_subsample] + sub_fg = samp[: self.n_kmeans_subsample] centroids = _kmeans_1d(sub_fg, C_k) sorted_c, sort_idx = torch.sort(centroids) @@ -385,8 +405,13 @@ def apply_transform( lbl_l = sort_idx[lbl_s].long() rid, R = _voronoi_region_ids( - coords, lbl_l, flat_m, C_k, device, - self.s_choices, self.skip_sub_parc_prob, + coords, + lbl_l, + flat_m, + C_k, + device, + self.s_choices, + self.skip_sub_parc_prob, ) s_c = torch.zeros(R, device=device).scatter_add_(0, rid, flat * flat_m) @@ -402,7 +427,7 @@ def apply_transform( synth_list.append(synth_i) - synth = torch.stack(synth_list) # (B, N) + synth = torch.stack(synth_list) # (B, N) synth_01 = synth.reshape(B, 1, D, H, W) sigma = random.choice(self.blur_sigmas) @@ -424,13 +449,10 @@ def apply_transform( for c in unique_classes: c_val = int(c.item()) - c_mask = (lbl == c_val).float() # (B, N) - c_cnt = c_mask.sum(dim=1, keepdim=True) # (B, 1) + c_mask = (lbl == c_val).float() # (B, N) + c_cnt = c_mask.sum(dim=1, keepdim=True) # (B, 1) - apply = ( - (torch.rand(B, 1, device=device) < self.label_remap_prob) - & (c_cnt >= self.min_label_voxels) - ).float() + apply = ((torch.rand(B, 1, device=device) < self.label_remap_prob) & (c_cnt >= self.min_label_voxels)).float() if apply.sum() == 0: continue @@ -469,7 +491,7 @@ def apply_transform( def _next_shared_seed() -> int: - global _SHARED_RNG_COUNTER + global _SHARED_RNG_COUNTER # noqa: PLW0603 -- module-level counter is the point: it makes successive seeds distinct _SHARED_RNG_COUNTER += 1 seed = (int(torch.initial_seed()) + _SHARED_RNG_COUNTER) % (2**63 - 1) if dist.is_available() and dist.is_initialized(): @@ -480,8 +502,7 @@ def _next_shared_seed() -> int: @staticmethod -def _minmax_norm(x: torch.Tensor, eps: float = 1e-8 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +def _minmax_norm(x: torch.Tensor, eps: float = 1e-8) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Per-sample min-max normalise to [0, 1]. Returns (normed, min, max).""" B = x.shape[0] x_flat = x.view(B, -1) @@ -489,11 +510,12 @@ def _minmax_norm(x: torch.Tensor, eps: float = 1e-8 vmax = x_flat.max(dim=1).values.view(B, 1, 1, 1, 1) return (x - vmin) / (vmax - vmin + eps), vmin, vmax + @staticmethod -def _minmax_denorm(x_norm: torch.Tensor, vmin: torch.Tensor, - vmax: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: +def _minmax_denorm(x_norm: torch.Tensor, vmin: torch.Tensor, vmax: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: return x_norm * (vmax - vmin + eps) + vmin + @staticmethod def _zscore_renorm(x: torch.Tensor, bg_threshold: float = 1e-6) -> torch.Tensor: """Per-sample foreground-masked z-score. Mirrors nnUNet's use_mask_for_norm=True. @@ -502,19 +524,21 @@ def _zscore_renorm(x: torch.Tensor, bg_threshold: float = 1e-6) -> torch.Tensor: Eliminates the train/inference distribution mismatch that would occur because nnUNet always z-scores at inference time. """ - fg = x.abs() > bg_threshold + fg = x.abs() > bg_threshold fg_f = fg.float() - n = fg_f.sum(dim=(2, 3, 4), keepdim=True).clamp(min=1) + n = fg_f.sum(dim=(2, 3, 4), keepdim=True).clamp(min=1) mean = (x * fg_f).sum(dim=(2, 3, 4), keepdim=True) / n - var = ((x - mean).pow(2) * fg_f).sum(dim=(2, 3, 4), keepdim=True) / n - std = var.sqrt().clamp(min=1e-8) + var = ((x - mean).pow(2) * fg_f).sum(dim=(2, 3, 4), keepdim=True) / n + std = var.sqrt().clamp(min=1e-8) return torch.where(fg, (x - mean) / std, torch.zeros_like(x)) + def _shared_cpu_generator() -> torch.Generator: generator = torch.Generator(device="cpu") generator.manual_seed(_next_shared_seed()) return generator + def _shared_rand(shape: tuple[int, ...], device: torch.device, dtype: torch.dtype) -> torch.Tensor: if not (dist.is_available() and dist.is_initialized()): return torch.rand(shape, device=device, dtype=dtype) @@ -536,8 +560,8 @@ def collapse_onehot_to_index(seg_raw: torch.Tensor) -> torch.Tensor: Background voxels (all-zero across channels) map to 0. Foreground voxels map to argmax(seg_raw, dim=1) + 1. """ - foreground_mask = seg_raw.any(dim=1, keepdim=True) # [B,1,D,H,W] bool - labels = torch.argmax(seg_raw, dim=1, keepdim=True).long() + 1 # 0-based → 1-based + foreground_mask = seg_raw.any(dim=1, keepdim=True) # [B,1,D,H,W] bool + labels = torch.argmax(seg_raw, dim=1, keepdim=True).long() + 1 # 0-based → 1-based labels = torch.where(foreground_mask, labels, torch.zeros_like(labels)) return labels diff --git a/auglab/transforms/gpu/spatial.py b/auglab/transforms/gpu/spatial.py index 25708be..86462ed 100644 --- a/auglab/transforms/gpu/spatial.py +++ b/auglab/transforms/gpu/spatial.py @@ -1,16 +1,20 @@ -from kornia.constants import Resample -from kornia.core import Tensor -from kornia.augmentation._3d.base import RigidAffineAugmentationBase3D +from typing import Any, Union + +import torch +import torch.nn.functional as F from kornia.augmentation import random_generator as rg -from kornia.geometry import deg2rad, get_affine_matrix3d, warp_affine3d +from kornia.augmentation._3d.base import RigidAffineAugmentationBase3D from kornia.augmentation.random_generator.base import RandomGeneratorBase, UniformDistribution from kornia.augmentation.utils import _adapted_rsampling, _tuple_range_reader -from kornia.utils.helpers import _extract_device_dtype -from kornia.constants import DataKey -import torch -import torch.nn.functional as F +from kornia.constants import DataKey, Resample +from kornia.geometry import deg2rad, get_affine_matrix3d, warp_affine3d +from torch import Tensor + +try: # kornia < 0.8.3 + from kornia.utils.helpers import _extract_device_dtype +except ImportError: # kornia >= 0.8.3 moved it and dropped kornia.utils.helpers + from kornia.core.utils import _extract_device_dtype -from typing import Any, Dict, Optional, Tuple, Union from auglab.transforms.gpu.base import ImageOnlyTransform @@ -67,7 +71,7 @@ class RandomAffine3DCustom(RigidAffineAugmentationBase3D): >>> import torch >>> rng = torch.manual_seed(0) >>> input = torch.rand(1, 1, 3, 3, 3) - >>> aug = RandomAffine3D((15., 20., 20.), p=1.) + >>> aug = RandomAffine3D((15.0, 20.0, 20.0), p=1.0) >>> aug(input), aug.transform_matrix (tensor([[[[[0.4503, 0.4763, 0.1680], [0.2029, 0.4267, 0.3515], @@ -86,7 +90,7 @@ class RandomAffine3DCustom(RigidAffineAugmentationBase3D): To apply the exact augmenation again, you may take the advantage of the previous parameter state: >>> input = torch.rand(1, 3, 32, 32, 32) - >>> aug = RandomAffine3D((15., 20., 20.), p=1.) + >>> aug = RandomAffine3D((15.0, 20.0, 20.0), p=1.0) >>> (aug(input) == aug(input, params=aug._params)).all() tensor(True) @@ -97,26 +101,26 @@ def __init__( degrees: Union[ Tensor, float, - Tuple[float, float], - Tuple[float, float, float], - Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]], + tuple[float, float], + tuple[float, float, float], + tuple[tuple[float, float], tuple[float, float], tuple[float, float]], ], - translate: Optional[Union[Tensor, Tuple[float, float, float]]] = None, - scale: Optional[Union[Tensor, Tuple[float, float], Tuple[Tuple[float, float], Tuple[float, float], Tuple[float, float]]]] = None, + translate: Union[Tensor, tuple[float, float, float]] | None = None, + scale: Union[Tensor, tuple[float, float], tuple[tuple[float, float], tuple[float, float], tuple[float, float]]] | None = None, shears: Union[ - None, Tensor, float, - Tuple[float, float], - Tuple[float, float, float, float, float, float], - Tuple[ - Tuple[float, float], - Tuple[float, float], - Tuple[float, float], - Tuple[float, float], - Tuple[float, float], - Tuple[float, float], + tuple[float, float], + tuple[float, float, float, float, float, float], + tuple[ + tuple[float, float], + tuple[float, float], + tuple[float, float], + tuple[float, float], + tuple[float, float], + tuple[float, float], ], + None, ] = None, resample: Union[str, int, Resample] = Resample.BILINEAR.name, same_on_batch: bool = False, @@ -133,7 +137,7 @@ def __init__( self.flags = {"resample": Resample.get(resample), "align_corners": align_corners} self._param_generator = rg.AffineGenerator3D(degrees, translate, scale, shears) - def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: transform: Tensor = get_affine_matrix3d( params["translations"], params["center"], @@ -148,9 +152,7 @@ def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags ).to(input) return transform - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: if not isinstance(transform, Tensor): raise TypeError(f"Expected the transform to be a Tensor. Gotcha {type(transform)}") @@ -168,13 +170,13 @@ def apply_transform( ) def apply_non_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are no transformation applied.""" return input def apply_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are transformed. @@ -182,7 +184,7 @@ def apply_transform_mask( Convert "resample" arguments to "nearest" by default. """ - resample_method: Optional[Resample] + resample_method: Resample | None if "resample" in flags: resample_method = flags["resample"] flags["resample"] = Resample.get("nearest") @@ -200,7 +202,7 @@ class RandomLowResTransformGPU(RigidAffineAugmentationBase3D): def __init__( self, - scale: Tuple[float, float] = (0.3, 1.0), + scale: tuple[float, float] = (0.3, 1.0), same_on_batch: bool = False, p: float = 1.0, keepdim: bool = True, @@ -209,13 +211,11 @@ def __init__( super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self._param_generator = ScaleGenerator3D(scale=scale) - def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: return self.identity_matrix(input) @torch.no_grad() - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # input shape: (B, C, D, H, W) if not isinstance(input, torch.Tensor): raise TypeError(f"Expected input to be a Tensor. Got {type(input)}") @@ -228,9 +228,9 @@ def apply_transform( scales = params["scale"] # shape [B, 3] - if flags['data_keys'][0] is DataKey.IMAGE: + if flags["data_keys"][0] is DataKey.IMAGE: resample = "trilinear" - elif flags['data_keys'][0] is DataKey.MASK: + elif flags["data_keys"][0] is DataKey.MASK: resample = "nearest" else: raise ValueError(f"Unsupported data key {flags['data_keys'][0]} for RandomLowResTransformGPU. Expected IMAGE or MASK.") @@ -247,9 +247,9 @@ def apply_transform( sx, sy, sz = scales[b] # compute downsampled size - down_D = max(1, int(round(float(sz) * D))) - down_H = max(1, int(round(float(sy) * H))) - down_W = max(1, int(round(float(sx) * W))) + down_D = max(1, round(float(sz) * D)) + down_H = max(1, round(float(sy) * H)) + down_W = max(1, round(float(sx) * W)) # downsample x_down = F.interpolate( @@ -272,13 +272,13 @@ def apply_transform( return out def apply_non_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are no transformation applied.""" return input def apply_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are transformed. @@ -291,7 +291,7 @@ def apply_transform_mask( class ScaleGenerator3D(RandomGeneratorBase): - def __init__(self, scale: Tuple[float, float], one_dim: bool = False) -> None: + def __init__(self, scale: tuple[float, float], one_dim: bool = False) -> None: super().__init__() self.scale = scale self.one_dim = one_dim @@ -309,12 +309,10 @@ def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: self.scaley_sampler = UniformDistribution(scale[1, 0], scale[1, 1], validate_args=False) self.scalez_sampler = UniformDistribution(scale[2, 0], scale[2, 1], validate_args=False) - def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) -> Dict[str, torch.Tensor]: + def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> dict[str, torch.Tensor]: batch_size = batch_shape[0] - _device, _dtype = _extract_device_dtype( - [self.scalex_sampler, self.scaley_sampler, self.scalez_sampler] - ) + _device, _dtype = _extract_device_dtype([self.scalex_sampler, self.scaley_sampler, self.scalez_sampler]) scalex = _adapted_rsampling((batch_size,), self.scalex_sampler, same_on_batch) scaley = _adapted_rsampling((batch_size,), self.scaley_sampler, same_on_batch) @@ -332,23 +330,23 @@ class RandomAcqTransformGPU(ImageOnlyTransform): def __init__( self, - scale: Tuple[float, float] = (0.3, 1.0), + scale: tuple[float, float] = (0.3, 1.0), one_dim: bool = False, same_on_batch: bool = False, - apply_to_channel: list[int] = [0], # Apply to first channel by default + apply_to_channel: list[int] | None = None, # Apply to first channel by default p: float = 1.0, keepdim: bool = True, **kwargs, ) -> None: + if apply_to_channel is None: + apply_to_channel = [0] super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self.flags = {"resample": "trilinear"} self.apply_to_channel = apply_to_channel self._param_generator = ScaleGenerator3D(scale=scale, one_dim=one_dim) @torch.no_grad() - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # input shape: (B, C, D, H, W) if not isinstance(input, torch.Tensor): raise TypeError(f"Expected input to be a Tensor. Got {type(input)}") @@ -378,9 +376,9 @@ def apply_transform( sx, sy, sz = scales[b] # compute downsampled size - down_D = max(1, int(round(float(sz) * D))) - down_H = max(1, int(round(float(sy) * H))) - down_W = max(1, int(round(float(sx) * W))) + down_D = max(1, round(float(sz) * D)) + down_H = max(1, round(float(sy) * H)) + down_W = max(1, round(float(sx) * W)) # downsample x_down = F.interpolate( @@ -391,12 +389,16 @@ def apply_transform( ) # upsample back to original resolution - x_up = F.interpolate( - x_down, - size=(D, H, W), - mode=interp_up, - align_corners=False if "linear" in interp_up else None, - ).squeeze(0).squeeze(0) # [D, H, W] + x_up = ( + F.interpolate( + x_down, + size=(D, H, W), + mode=interp_up, + align_corners=False if "linear" in interp_up else None, + ) + .squeeze(0) + .squeeze(0) + ) # [D, H, W] # place patch back into the canvas for the correct channel only canvas[c] = x_up @@ -405,6 +407,7 @@ def apply_transform( return out + # Flip transforms class RandomFlipTransformGPU(RigidAffineAugmentationBase3D): """ @@ -429,13 +432,11 @@ def __init__( # generator creates per-batch flip flags for axes (z, y, x) self._param_generator = FlipGenerator3D(flip_axis=self.flip_axis) - def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: return self.identity_matrix(input) @torch.no_grad() - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # input shape: (B, C, D, H, W) if not isinstance(input, torch.Tensor): @@ -451,11 +452,8 @@ def apply_transform( out = input.clone() # For each batch element, build list of spatial dims to flip (D,H,W -> dims 2,3,4) for b in range(batch_size): - flip_dims = [] # fb expected as length-3 tensor for (z,y,x) - for axis in range(3): - if axis in self.flip_axis: - flip_dims.append(1 + axis) + flip_dims = [1 + axis for axis in range(3) if axis in self.flip_axis] if len(flip_dims) > 0: out[b] = torch.flip(input[b], dims=tuple(flip_dims)) @@ -463,13 +461,13 @@ def apply_transform( return out def apply_non_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are no transformation applied.""" return input def apply_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are transformed. @@ -501,7 +499,7 @@ def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: # use uniform samplers per axis and threshold at 0.5 self._samplers = [UniformDistribution(0.0, 1.0, validate_args=False) for _ in range(3)] - def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) -> Dict[str, torch.Tensor]: + def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> dict[str, torch.Tensor]: batch_size = batch_shape[0] _device, _dtype = _extract_device_dtype(self._samplers) @@ -527,6 +525,7 @@ def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) -> return {"flip": flips} + # Crop transform class RandomCropTransformGPU(RigidAffineAugmentationBase3D): """ @@ -535,8 +534,8 @@ class RandomCropTransformGPU(RigidAffineAugmentationBase3D): def __init__( self, - crop: Tuple[float, float] = (1.0, 1.0), - pos: Tuple[float, float, float] = (0.5, 1), # Fraction of the pos + crop: tuple[float, float] = (1.0, 1.0), + pos: tuple[float, float, float] = (0.5, 1), # Fraction of the pos same_on_batch: bool = False, p: float = 1.0, keepdim: bool = True, @@ -545,13 +544,11 @@ def __init__( super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) self._param_generator = CropGenerator3D(crop=crop, pos=pos) - def compute_transformation(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor: + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: return self.identity_matrix(input) @torch.no_grad() - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: # input shape: (B, C, D, H, W) if not isinstance(input, torch.Tensor): raise TypeError(f"Expected input to be a Tensor. Got {type(input)}") @@ -574,22 +571,22 @@ def apply_transform( # determine crop fraction and crop size on the image cx, cy, cz = crops[b] # interpret crop as fraction of upsampled size to keep - crop_D = max(1, int(round(float(cz) * D))) - crop_H = max(1, int(round(float(cy) * H))) - crop_W = max(1, int(round(float(cx) * W))) + crop_D = max(1, round(float(cz) * D)) + crop_H = max(1, round(float(cy) * H)) + crop_W = max(1, round(float(cx) * W)) # determine pos fraction of the image px, py, pz = pos[b] - + # center position center_z = float(pz) * D center_y = float(py) * H center_x = float(px) * W # choose top-left-front corner - start_z = int(round(center_z - crop_D / 2.0)) - start_y = int(round(center_y - crop_H / 2.0)) - start_x = int(round(center_x - crop_W / 2.0)) + start_z = round(center_z - crop_D / 2.0) + start_y = round(center_y - crop_H / 2.0) + start_x = round(center_x - crop_W / 2.0) # clamp to valid limits max_z = max(0, D - crop_D) @@ -615,13 +612,13 @@ def apply_transform( return out def apply_non_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are no transformation applied.""" return input def apply_transform_mask( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None ) -> Tensor: """Process masks corresponding to the inputs that are transformed. @@ -632,11 +629,12 @@ def apply_transform_mask( output = self.apply_transform(input, params, flags, transform) return output + class CropGenerator3D(RandomGeneratorBase): - def __init__(self, crop: Tuple[float, float], pos: Tuple[float, float], one_dim: bool = False) -> None: + def __init__(self, crop: tuple[float, float], pos: tuple[float, float], one_dim: bool = False) -> None: super().__init__() self.crop = crop - self.pos = pos # Position of the crop box center, as a fraction of the image dimensions (e.g. 0.5 for centered) + self.pos = pos # Position of the crop box center, as a fraction of the image dimensions (e.g. 0.5 for centered) self.one_dim = one_dim def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: @@ -651,7 +649,7 @@ def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: self.cropx_sampler = UniformDistribution(crop[0, 0], crop[0, 1], validate_args=False) self.cropy_sampler = UniformDistribution(crop[1, 0], crop[1, 1], validate_args=False) self.cropz_sampler = UniformDistribution(crop[2, 0], crop[2, 1], validate_args=False) - + pos = _tuple_range_reader(self.pos, 3, device, dtype) if self.one_dim: # Pick a random dimension to apply cropping @@ -664,7 +662,7 @@ def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: self.posy_sampler = UniformDistribution(pos[1, 0], pos[1, 1], validate_args=False) self.posz_sampler = UniformDistribution(pos[2, 0], pos[2, 1], validate_args=False) - def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) -> Dict[str, torch.Tensor]: + def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> dict[str, torch.Tensor]: batch_size = batch_shape[0] _device, _dtype = _extract_device_dtype( @@ -681,4 +679,4 @@ def forward(self, batch_shape: Tuple[int, ...], same_on_batch: bool = False) -> posz = _adapted_rsampling((batch_size,), self.posz_sampler, same_on_batch) pos = torch.stack([posx, posy, posz], dim=1) - return {"crop": torch.as_tensor(crop, device=_device, dtype=_dtype), "pos": torch.as_tensor(pos, device=_device, dtype=_dtype)} \ No newline at end of file + return {"crop": torch.as_tensor(crop, device=_device, dtype=_dtype), "pos": torch.as_tensor(pos, device=_device, dtype=_dtype)} diff --git a/auglab/transforms/gpu/transforms.py b/auglab/transforms/gpu/transforms.py index 956f087..ba5ca0a 100644 --- a/auglab/transforms/gpu/transforms.py +++ b/auglab/transforms/gpu/transforms.py @@ -1,75 +1,100 @@ -import os, json +import json +import os +from typing import Any -import torch.nn as nn -import torch import numpy as np - -from auglab.transforms.gpu.base import ImageOnlyTransform -from typing import Any, Dict, Optional, Tuple, Union, List -from kornia.core import Tensor - -from auglab.transforms.gpu.contrast import RandomConvTransformGPU, RandomGaussianNoiseGPU, RandomBrightnessGPU, RandomGammaGPU, RandomFunctionGPU, \ -RandomHistogramEqualizationGPU, RandomInverseGPU, RandomBiasFieldGPU, RandomContrastGPU, ZscoreNormalizationGPU, RandomClampGPU -from auglab.transforms.gpu.spatial import RandomAffine3DCustom, RandomLowResTransformGPU, RandomFlipTransformGPU, RandomAcqTransformGPU, RandomCropTransformGPU -from auglab.transforms.gpu.fromSeg import RandomRedistributeSegGPU, RandomPALETTEGPU +import torch +from torch import Tensor, nn + +from auglab.transforms.gpu.base import AugmentationSequentialCustom, ImageOnlyTransform +from auglab.transforms.gpu.contrast import ( + RandomBiasFieldGPU, + RandomBrightnessGPU, + RandomClampGPU, + RandomContrastGPU, + RandomConvTransformGPU, + RandomFunctionGPU, + RandomGammaGPU, + RandomGaussianNoiseGPU, + RandomHistogramEqualizationGPU, + RandomInverseGPU, + ZscoreNormalizationGPU, +) from auglab.transforms.gpu.domain_transfer import RandomDomainTransferGPU +from auglab.transforms.gpu.fromSeg import RandomPALETTEGPU, RandomRedistributeSegGPU +from auglab.transforms.gpu.spatial import ( + RandomAcqTransformGPU, + RandomAffine3DCustom, + RandomCropTransformGPU, + RandomFlipTransformGPU, + RandomLowResTransformGPU, +) from auglab.transforms.synthseg.transforms import RandomSynthSegGPU -from auglab.transforms.gpu.base import AugmentationSequentialCustom + class AugTransformsGPU(AugmentationSequentialCustom): """ Module to perform data augmentation on GPU. """ + def __init__(self, json_path: str): # Load transform parameters from JSON config_path = os.path.join(json_path) - with open(config_path, 'r') as f: + with open(config_path) as f: config = json.load(f) - if 'GPU' in config.keys(): - self.transform_params = config['GPU'] + if "GPU" in config.keys(): + self.transform_params = config["GPU"] else: self.transform_params = config transforms = self._build_transforms() - super().__init__(*transforms, data_keys=["input", "mask"], same_on_batch=True) # Same_on_batch to ensure mask are aligned with images correctly (custom) see AugmentationSequentialOpsCustom in base.py + super().__init__( + *transforms, data_keys=["input", "mask"], same_on_batch=True + ) # Same_on_batch to ensure mask are aligned with images correctly (custom) see AugmentationSequentialOpsCustom in base.py def _build_transforms(self) -> list[nn.Module]: transforms = [] # Flipping transforms - flip_params = self.transform_params.get('FlipTransform') + flip_params = self.transform_params.get("FlipTransform") if flip_params is not None: - transforms.append(RandomFlipTransformGPU( - flip_axis=flip_params.get('flip_axis', [0]), - p=flip_params.get('probability', 0), - same_on_batch=flip_params.get('same_on_batch', False), - keepdim=flip_params.get('keepdim', True) - )) + transforms.append( + RandomFlipTransformGPU( + flip_axis=flip_params.get("flip_axis", [0]), + p=flip_params.get("probability", 0), + same_on_batch=flip_params.get("same_on_batch", False), + keepdim=flip_params.get("keepdim", True), + ) + ) # Spatial transforms - affine_params = self.transform_params.get('AffineTransform') + affine_params = self.transform_params.get("AffineTransform") if affine_params is not None: - transforms.append(RandomAffine3DCustom( - degrees=affine_params.get('degrees', 10), - translate=affine_params.get('translate', [0.1, 0.1, 0.1]), - scale=affine_params.get('scale', [0.9, 1.1]), - shears=affine_params.get('shear', [-10, 10, -10, 10, -10, 10]), - resample=affine_params.get('resample', "bilinear"), - p=affine_params.get('probability', 0) - )) + transforms.append( + RandomAffine3DCustom( + degrees=affine_params.get("degrees", 10), + translate=affine_params.get("translate", [0.1, 0.1, 0.1]), + scale=affine_params.get("scale", [0.9, 1.1]), + shears=affine_params.get("shear", [-10, 10, -10, 10, -10, 10]), + resample=affine_params.get("resample", "bilinear"), + p=affine_params.get("probability", 0), + ) + ) # SynthSeg generative augmentation: replace the image with a GMM synthesis # of the segmentation (intensity-only here, so the mask stays consistent; # geometric transforms above deform the labels first). All SynthSeg # generator parameters are read straight from the config block. - synthseg_params = self.transform_params.get('SynthSeg') + synthseg_params = self.transform_params.get("SynthSeg") if synthseg_params is not None: - synthseg_kwargs = {k: v for k, v in synthseg_params.items() if k != 'probability'} - transforms.append(RandomSynthSegGPU( - p=synthseg_params.get('probability', 1.0), - **synthseg_kwargs, - )) + synthseg_kwargs = {k: v for k, v in synthseg_params.items() if k != "probability"} + transforms.append( + RandomSynthSegGPU( + p=synthseg_params.get("probability", 1.0), + **synthseg_kwargs, + ) + ) ## Transfer augmentations (TA) ######################### @@ -95,8 +120,7 @@ def _build_transforms(self) -> list[nn.Module]: # Domain transfer: randomly re-render the image as another sequence/cluster (TA) # Accept either the class-name key or the descriptive key. - domain_params = self.transform_params.get('RandomDomainTransferGPU') \ - or self.transform_params.get('DomainTransferTransform') + domain_params = self.transform_params.get("RandomDomainTransferGPU") or self.transform_params.get("DomainTransferTransform") if domain_params is not None: transforms.append( RandomDomainTransferGPU( @@ -123,7 +147,7 @@ def _build_transforms(self) -> list[nn.Module]: ) # Inverse transform (max - pixel_value) - inverse_params = self.transform_params.get('InverseTransform') + inverse_params = self.transform_params.get("InverseTransform") if inverse_params is not None: transforms.append( RandomInverseGPU( @@ -137,7 +161,7 @@ def _build_transforms(self) -> list[nn.Module]: ) # Histogram manipulations - histo_params = self.transform_params.get('HistogramEqualizationTransform') + histo_params = self.transform_params.get("HistogramEqualizationTransform") if histo_params is not None: transforms.append( RandomHistogramEqualizationGPU( @@ -151,142 +175,164 @@ def _build_transforms(self) -> list[nn.Module]: ) # Redistribute segmentation values transform - redistribute_params = self.transform_params.get('RedistributeSegTransform') + redistribute_params = self.transform_params.get("RedistributeSegTransform") if redistribute_params is not None: - transforms.append(RandomRedistributeSegGPU( - in_seg=redistribute_params.get('in_seg', 0.2), - retain_stats=redistribute_params.get('retain_stats', False), - p=redistribute_params.get('probability', 0), - std_noise_range=redistribute_params.get('std_noise_range', [0.1, 0.3]), - dilation_iterations_range=redistribute_params.get('dilation_iterations_range', [1, 3]), - )) + transforms.append( + RandomRedistributeSegGPU( + in_seg=redistribute_params.get("in_seg", 0.2), + retain_stats=redistribute_params.get("retain_stats", False), + p=redistribute_params.get("probability", 0), + std_noise_range=redistribute_params.get("std_noise_range", [0.1, 0.3]), + dilation_iterations_range=redistribute_params.get("dilation_iterations_range", [1, 3]), + ) + ) # Scharr filter - scharr_params = self.transform_params.get('ScharrTransform') + scharr_params = self.transform_params.get("ScharrTransform") if scharr_params is not None: - transforms.append(RandomConvTransformGPU( - kernel_type=scharr_params.get('kernel_type', 'Scharr'), - p=scharr_params.get('probability', 0), - in_seg=scharr_params.get('in_seg', 0.0), - out_seg=scharr_params.get('out_seg', 0.0), - mix_in_out=scharr_params.get('mix_in_out', False), - retain_stats=scharr_params.get('retain_stats', True), - absolute=scharr_params.get('absolute', True), - mix_prob=scharr_params.get('mix_prob', 0.0), - )) + transforms.append( + RandomConvTransformGPU( + kernel_type=scharr_params.get("kernel_type", "Scharr"), + p=scharr_params.get("probability", 0), + in_seg=scharr_params.get("in_seg", 0.0), + out_seg=scharr_params.get("out_seg", 0.0), + mix_in_out=scharr_params.get("mix_in_out", False), + retain_stats=scharr_params.get("retain_stats", True), + absolute=scharr_params.get("absolute", True), + mix_prob=scharr_params.get("mix_prob", 0.0), + ) + ) # Unsharp masking - unsharp_params = self.transform_params.get('UnsharpMaskTransform') + unsharp_params = self.transform_params.get("UnsharpMaskTransform") if unsharp_params is not None: - transforms.append(RandomConvTransformGPU( - kernel_type=unsharp_params.get('kernel_type', 'UnsharpMask'), - p=unsharp_params.get('probability', 0), - in_seg=unsharp_params.get('in_seg', 0.0), - out_seg=unsharp_params.get('out_seg', 0.0), - mix_in_out=unsharp_params.get('mix_in_out', False), - sigma=unsharp_params.get('sigma', 1.0), - unsharp_amount=unsharp_params.get('unsharp_amount', 1.5), - mix_prob=unsharp_params.get('mix_prob', 0.0), - )) + transforms.append( + RandomConvTransformGPU( + kernel_type=unsharp_params.get("kernel_type", "UnsharpMask"), + p=unsharp_params.get("probability", 0), + in_seg=unsharp_params.get("in_seg", 0.0), + out_seg=unsharp_params.get("out_seg", 0.0), + mix_in_out=unsharp_params.get("mix_in_out", False), + sigma=unsharp_params.get("sigma", 1.0), + unsharp_amount=unsharp_params.get("unsharp_amount", 1.5), + mix_prob=unsharp_params.get("mix_prob", 0.0), + ) + ) # RandomConv transform - randconv_params = self.transform_params.get('RandomConvTransform') + randconv_params = self.transform_params.get("RandomConvTransform") if randconv_params is not None: - transforms.append(RandomConvTransformGPU( - kernel_type=randconv_params.get('kernel_type', 'RandConv'), - p=randconv_params.get('probability', 0), - in_seg=randconv_params.get('in_seg', 0.0), - out_seg=randconv_params.get('out_seg', 0.0), - mix_in_out=randconv_params.get('mix_in_out', False), - retain_stats=randconv_params.get('retain_stats', False), - kernel_sizes=randconv_params.get('kernel_sizes', [1,3,5,7]), - mix_prob=randconv_params.get('mix_prob', 0.0), - )) + transforms.append( + RandomConvTransformGPU( + kernel_type=randconv_params.get("kernel_type", "RandConv"), + p=randconv_params.get("probability", 0), + in_seg=randconv_params.get("in_seg", 0.0), + out_seg=randconv_params.get("out_seg", 0.0), + mix_in_out=randconv_params.get("mix_in_out", False), + retain_stats=randconv_params.get("retain_stats", False), + kernel_sizes=randconv_params.get("kernel_sizes", [1, 3, 5, 7]), + mix_prob=randconv_params.get("mix_prob", 0.0), + ) + ) ## General enhancement (GE) # Clamping transform - clamp_params = self.transform_params.get('ClampTransform') + clamp_params = self.transform_params.get("ClampTransform") if clamp_params is not None: - transforms.append(RandomClampGPU( - max_clamp_amount=clamp_params.get('max_clamp_amount', 0.0), - in_seg=clamp_params.get('in_seg', 0.0), - out_seg=clamp_params.get('out_seg', 0.0), - mix_in_out=clamp_params.get('mix_in_out', False), - retain_stats=clamp_params.get('retain_stats', False), - p=clamp_params.get('probability', 0), - )) + transforms.append( + RandomClampGPU( + max_clamp_amount=clamp_params.get("max_clamp_amount", 0.0), + in_seg=clamp_params.get("in_seg", 0.0), + out_seg=clamp_params.get("out_seg", 0.0), + mix_in_out=clamp_params.get("mix_in_out", False), + retain_stats=clamp_params.get("retain_stats", False), + p=clamp_params.get("probability", 0), + ) + ) # Noise transforms - noise_params = self.transform_params.get('GaussianNoiseTransform') + noise_params = self.transform_params.get("GaussianNoiseTransform") if noise_params is not None: - transforms.append(RandomGaussianNoiseGPU( - mean=noise_params.get('mean', 0.0), - std=noise_params.get('std', 1.0), - in_seg=noise_params.get('in_seg', 0.0), - out_seg=noise_params.get('out_seg', 0.0), - mix_in_out=noise_params.get('mix_in_out', False), - p=noise_params.get('probability', 0), - )) + transforms.append( + RandomGaussianNoiseGPU( + mean=noise_params.get("mean", 0.0), + std=noise_params.get("std", 1.0), + in_seg=noise_params.get("in_seg", 0.0), + out_seg=noise_params.get("out_seg", 0.0), + mix_in_out=noise_params.get("mix_in_out", False), + p=noise_params.get("probability", 0), + ) + ) # Gaussian blur - gaussianblur_params = self.transform_params.get('GaussianBlurTransform') + gaussianblur_params = self.transform_params.get("GaussianBlurTransform") if gaussianblur_params is not None: - transforms.append(RandomConvTransformGPU( - kernel_type=gaussianblur_params.get('kernel_type', 'GaussianBlur'), - in_seg=gaussianblur_params.get('in_seg', 0.0), - out_seg=gaussianblur_params.get('out_seg', 0.0), - mix_in_out=gaussianblur_params.get('mix_in_out', False), - p=gaussianblur_params.get('probability', 0), - sigma=gaussianblur_params.get('sigma', 1.0), - )) + transforms.append( + RandomConvTransformGPU( + kernel_type=gaussianblur_params.get("kernel_type", "GaussianBlur"), + in_seg=gaussianblur_params.get("in_seg", 0.0), + out_seg=gaussianblur_params.get("out_seg", 0.0), + mix_in_out=gaussianblur_params.get("mix_in_out", False), + p=gaussianblur_params.get("probability", 0), + sigma=gaussianblur_params.get("sigma", 1.0), + ) + ) # Brightness transforms - brightness_params = self.transform_params.get('BrightnessTransform') + brightness_params = self.transform_params.get("BrightnessTransform") if brightness_params is not None: - transforms.append(RandomBrightnessGPU( - brightness_range=brightness_params.get('brightness_range', [0.5, 1.5]), - in_seg=brightness_params.get('in_seg', 0.0), - out_seg=brightness_params.get('out_seg', 0.0), - mix_in_out=brightness_params.get('mix_in_out', False), - p=brightness_params.get('probability', 0), - )) + transforms.append( + RandomBrightnessGPU( + brightness_range=brightness_params.get("brightness_range", [0.5, 1.5]), + in_seg=brightness_params.get("in_seg", 0.0), + out_seg=brightness_params.get("out_seg", 0.0), + mix_in_out=brightness_params.get("mix_in_out", False), + p=brightness_params.get("probability", 0), + ) + ) # Gamma transforms - gamma_params = self.transform_params.get('GammaTransform') + gamma_params = self.transform_params.get("GammaTransform") if gamma_params is not None: - transforms.append(RandomGammaGPU( - gamma_range=gamma_params.get('gamma_range', [0.7, 1.5]), - p=gamma_params.get('probability', 0), - invert_image=False, - in_seg=gamma_params.get('in_seg', 0.0), - out_seg=gamma_params.get('out_seg', 0.0), - mix_in_out=gamma_params.get('mix_in_out', False), - retain_stats=gamma_params.get('retain_stats', False), - )) - - inv_gamma_params = self.transform_params.get('InvGammaTransform') + transforms.append( + RandomGammaGPU( + gamma_range=gamma_params.get("gamma_range", [0.7, 1.5]), + p=gamma_params.get("probability", 0), + invert_image=False, + in_seg=gamma_params.get("in_seg", 0.0), + out_seg=gamma_params.get("out_seg", 0.0), + mix_in_out=gamma_params.get("mix_in_out", False), + retain_stats=gamma_params.get("retain_stats", False), + ) + ) + + inv_gamma_params = self.transform_params.get("InvGammaTransform") if inv_gamma_params is not None: - transforms.append(RandomGammaGPU( - gamma_range=inv_gamma_params.get('gamma_range', [0.7, 1.5]), - p=inv_gamma_params.get('probability', 0), - in_seg=inv_gamma_params.get('in_seg', 0.0), - out_seg=inv_gamma_params.get('out_seg', 0.0), - mix_in_out=inv_gamma_params.get('mix_in_out', False), - invert_image=True, - retain_stats=inv_gamma_params.get('retain_stats', False), - )) + transforms.append( + RandomGammaGPU( + gamma_range=inv_gamma_params.get("gamma_range", [0.7, 1.5]), + p=inv_gamma_params.get("probability", 0), + in_seg=inv_gamma_params.get("in_seg", 0.0), + out_seg=inv_gamma_params.get("out_seg", 0.0), + mix_in_out=inv_gamma_params.get("mix_in_out", False), + invert_image=True, + retain_stats=inv_gamma_params.get("retain_stats", False), + ) + ) # nnUNetV2 Contrast transforms - contrast_params = self.transform_params.get('ContrastTransform') + contrast_params = self.transform_params.get("ContrastTransform") if contrast_params is not None: - transforms.append(RandomContrastGPU( - contrast_range=contrast_params.get('contrast_range', [0.75, 1.25]), - p=contrast_params.get('probability', 0), - in_seg=contrast_params.get('in_seg', 0.0), - out_seg=contrast_params.get('out_seg', 0.0), - mix_in_out=contrast_params.get('mix_in_out', False), - retain_stats=contrast_params.get('retain_stats', False) - )) + transforms.append( + RandomContrastGPU( + contrast_range=contrast_params.get("contrast_range", [0.75, 1.25]), + p=contrast_params.get("probability", 0), + in_seg=contrast_params.get("in_seg", 0.0), + out_seg=contrast_params.get("out_seg", 0.0), + mix_in_out=contrast_params.get("mix_in_out", False), + retain_stats=contrast_params.get("retain_stats", False), + ) + ) # Apply functions func_list = [ @@ -294,68 +340,77 @@ def _build_transforms(self) -> list[nn.Module]: torch.sqrt, torch.sin, torch.exp, - lambda x: 1/(1 + torch.exp(-x)), + lambda x: 1 / (1 + torch.exp(-x)), ] - function_params = self.transform_params.get('FunctionTransform') + function_params = self.transform_params.get("FunctionTransform") if function_params is not None: - for func in func_list: - transforms.append(RandomFunctionGPU( + transforms.extend( + RandomFunctionGPU( func=func, - p=function_params.get('probability', 0), - in_seg=function_params.get('in_seg', 0.0), - out_seg=function_params.get('out_seg', 0.0), - mix_in_out=function_params.get('mix_in_out', False), - retain_stats=function_params.get('retain_stats', False), - )) + p=function_params.get("probability", 0), + in_seg=function_params.get("in_seg", 0.0), + out_seg=function_params.get("out_seg", 0.0), + mix_in_out=function_params.get("mix_in_out", False), + retain_stats=function_params.get("retain_stats", False), + ) + for func in func_list + ) # Shape transforms (Cropping and Simulating low resolution) - lowres_params = self.transform_params.get('SimulateLowResTransform') + lowres_params = self.transform_params.get("SimulateLowResTransform") if lowres_params is not None: - transforms.append(RandomLowResTransformGPU( - p=lowres_params.get('probability', 0), - scale=lowres_params.get('scale', [0.3, 1.0]), - same_on_batch=lowres_params.get('same_on_batch', False) - )) + transforms.append( + RandomLowResTransformGPU( + p=lowres_params.get("probability", 0), + scale=lowres_params.get("scale", [0.3, 1.0]), + same_on_batch=lowres_params.get("same_on_batch", False), + ) + ) - acq_params = self.transform_params.get('AcqTransform') + acq_params = self.transform_params.get("AcqTransform") if acq_params is not None: - transforms.append(RandomAcqTransformGPU( - p=acq_params.get('probability', 0), - scale=acq_params.get('scale', [0.3, 1.0]), - one_dim=True, - same_on_batch=acq_params.get('same_on_batch', False) - )) - - crop_params = self.transform_params.get('CropTransform') + transforms.append( + RandomAcqTransformGPU( + p=acq_params.get("probability", 0), + scale=acq_params.get("scale", [0.3, 1.0]), + one_dim=True, + same_on_batch=acq_params.get("same_on_batch", False), + ) + ) + + crop_params = self.transform_params.get("CropTransform") if crop_params is not None: - transforms.append(RandomCropTransformGPU( - p=crop_params.get('probability', 0), - crop=crop_params.get('crop', [1.0, 1.0]), - pos=crop_params.get('pos', [0.0, 1.0]), - same_on_batch=acq_params.get('same_on_batch', False) - )) + transforms.append( + RandomCropTransformGPU( + p=crop_params.get("probability", 0), + crop=crop_params.get("crop", [1.0, 1.0]), + pos=crop_params.get("pos", [0.0, 1.0]), + same_on_batch=acq_params.get("same_on_batch", False), + ) + ) # Bias field artifact - bias_field_params = self.transform_params.get('BiasFieldTransform') + bias_field_params = self.transform_params.get("BiasFieldTransform") if bias_field_params is not None: - transforms.append(RandomBiasFieldGPU( - p=bias_field_params.get('probability', 0), - in_seg=bias_field_params.get('in_seg', 0.0), - out_seg=bias_field_params.get('out_seg', 0.0), - mix_in_out=bias_field_params.get('mix_in_out', False), - retain_stats=bias_field_params.get('retain_stats', False), - coefficients=bias_field_params.get('coefficients', 0.5), - )) + transforms.append( + RandomBiasFieldGPU( + p=bias_field_params.get("probability", 0), + in_seg=bias_field_params.get("in_seg", 0.0), + out_seg=bias_field_params.get("out_seg", 0.0), + mix_in_out=bias_field_params.get("mix_in_out", False), + retain_stats=bias_field_params.get("retain_stats", False), + coefficients=bias_field_params.get("coefficients", 0.5), + ) + ) ## Random Z-score normalization - zscore_params = self.transform_params.get('ZscoreNormalizationTransform') + zscore_params = self.transform_params.get("ZscoreNormalizationTransform") if zscore_params is not None: - transforms.append(ZscoreNormalizationGPU( - p=zscore_params.get('probability', 0) - )) + transforms.append(ZscoreNormalizationGPU(p=zscore_params.get("probability", 0))) return transforms + class RandomChooseXTransformsGPU(ImageOnlyTransform): """Randomly choose X transforms to apply from a given list of ImageOnlyTransform transforms (GPU version). @@ -372,7 +427,7 @@ class RandomChooseXTransformsGPU(ImageOnlyTransform): def __init__( self, - transforms_list: List[ImageOnlyTransform], + transforms_list: list[ImageOnlyTransform], num_transforms: int = 1, same_on_batch: bool = False, p: float = 1.0, @@ -385,7 +440,7 @@ def __init__( self.transforms_list = nn.ModuleList(transforms_list) self.num_transforms = num_transforms - def _apply_mix(self, x: Tensor, seg: Optional[Tensor]) -> Tensor: + def _apply_mix(self, x: Tensor, seg: Tensor | None) -> Tensor: if self.num_transforms == 0 or len(self.transforms_list) == 0: return x @@ -393,7 +448,7 @@ def _apply_mix(self, x: Tensor, seg: Optional[Tensor]) -> Tensor: # sample without replacement idx = torch.randperm(len(self.transforms_list), device=x.device)[:k] - child_params: Dict[str, Tensor] = {} + child_params: dict[str, Tensor] = {} if seg is not None: child_params["seg"] = seg @@ -402,19 +457,15 @@ def _apply_mix(self, x: Tensor, seg: Optional[Tensor]) -> Tensor: if torch.rand(1, device=x.device, dtype=x.dtype) > t.p: continue if not hasattr(t, "apply_transform"): - raise TypeError( - f"All transforms must implement apply_transform like ImageOnlyTransform. Got {type(t)}" - ) + raise TypeError(f"All transforms must implement apply_transform like ImageOnlyTransform. Got {type(t)}") # Most contrast transforms perform their random sampling inside apply_transform. t_flags = getattr(t, "flags", {}) x = t.apply_transform(x, child_params, t_flags, transform=None) return x @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: - seg = params.get("seg", None) + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: + seg = params.get("seg") if self.same_on_batch: return self._apply_mix(input, seg) @@ -424,14 +475,12 @@ def apply_transform( for i in range(batch_size): xi = out[i : i + 1] seg_i = None - if seg is not None and isinstance(seg, torch.Tensor) and seg.shape[0] == batch_size: - seg_i = seg[i : i + 1] - else: - seg_i = seg + seg_i = seg[i : i + 1] if seg is not None and isinstance(seg, torch.Tensor) and seg.shape[0] == batch_size else seg xi = self._apply_mix(xi, seg_i) out[i : i + 1] = xi return out + def normalize(arr: np.ndarray) -> np.ndarray: """ Normalize a tensor to the range [0, 1]. @@ -441,19 +490,25 @@ def normalize(arr: np.ndarray) -> np.ndarray: normalized_arr = (arr - min_val) / (max_val - min_val + 1e-8) return normalized_arr + def pad_numpy_array(arr, shape): """ Pad a numpy array to the desired shape with zeros. """ # Calculate padding needed for each dimension - pad_width = [(max(0, shape[i] - arr.shape[i]) // 2, max(0, shape[i] - arr.shape[i]) - max(0, shape[i] - arr.shape[i]) // 2) for i in range(len(shape))] - padded_arr = np.pad(arr, pad_width, mode='constant', constant_values=0) + pad_width = [ + (max(0, shape[i] - arr.shape[i]) // 2, max(0, shape[i] - arr.shape[i]) - max(0, shape[i] - arr.shape[i]) // 2) + for i in range(len(shape)) + ] + padded_arr = np.pad(arr, pad_width, mode="constant", constant_values=0) return padded_arr + if __name__ == "__main__": # Example usage import importlib - import auglab.configs as configs + + from auglab import configs from auglab.utils.image import Image, resample_nib configs_path = importlib.resources.files(configs) @@ -461,24 +516,24 @@ def pad_numpy_array(arr, shape): augmentor = AugTransformsGPU(json_path) # Load images and masks tensors - img_path = '/home/ge.polymtl.ca/p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz' - img = Image(img_path).change_orientation('RSP') - img = resample_nib(img, new_size=[1,1,1], new_size_type='mm', interpolation='linear') + img_path = "/home/ge.polymtl.ca/p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz" + img = Image(img_path).change_orientation("RSP") + img = resample_nib(img, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") img_tensor = torch.from_numpy(img.data.copy()).to(torch.float32) - seg_path = '/home/ge.polymtl.ca/p118739/data/datasets/data-multi-subject/derivatives/labels/sub-amu02/anat/sub-amu02_T1w_label-spine_dseg.nii.gz' - seg = Image(seg_path).change_orientation('RSP') - seg = resample_nib(seg, new_size=[1,1,1], new_size_type='mm', interpolation='nn') + seg_path = "/home/ge.polymtl.ca/p118739/data/datasets/data-multi-subject/derivatives/labels/sub-amu02/anat/sub-amu02_T1w_label-spine_dseg.nii.gz" + seg = Image(seg_path).change_orientation("RSP") + seg = resample_nib(seg, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") seg_tensor_all = torch.from_numpy(seg.data.copy()) - - img2_path = '/home/ge.polymtl.ca/p118739/data/datasets/spider-challenge-2023/sub-002/anat/sub-002_acq-lowresSag_T2w.nii.gz' - img2 = Image(img2_path).change_orientation('RSP') - img2 = resample_nib(img2, new_size=[1,1,1], new_size_type='mm', interpolation='linear') + + img2_path = "/home/ge.polymtl.ca/p118739/data/datasets/spider-challenge-2023/sub-002/anat/sub-002_acq-lowresSag_T2w.nii.gz" + img2 = Image(img2_path).change_orientation("RSP") + img2 = resample_nib(img2, new_size=[1, 1, 1], new_size_type="mm", interpolation="linear") img2_tensor = torch.from_numpy(img2.data.copy()).to(torch.float32) - seg2_path = '/home/ge.polymtl.ca/p118739/data/datasets/spider-challenge-2023/derivatives/labels/sub-002/anat/sub-002_acq-lowresSag_T2w_label-spine_dseg.nii.gz' - seg2 = Image(seg2_path).change_orientation('RSP') - seg2 = resample_nib(seg2, new_size=[1,1,1], new_size_type='mm', interpolation='nn') + seg2_path = "/home/ge.polymtl.ca/p118739/data/datasets/spider-challenge-2023/derivatives/labels/sub-002/anat/sub-002_acq-lowresSag_T2w_label-spine_dseg.nii.gz" + seg2 = Image(seg2_path).change_orientation("RSP") + seg2 = resample_nib(seg2, new_size=[1, 1, 1], new_size_type="mm", interpolation="nn") seg2_tensor_all = torch.from_numpy(seg2.data.copy()) # Combine two images to same size @@ -488,7 +543,7 @@ def pad_numpy_array(arr, shape): size2 = img2_tensor.shape[dim] min_size = min(size1, size2) new_shape.append(min_size) - + new_img_tensor = torch.zeros(new_shape) new_img2_tensor = torch.zeros(new_shape) new_seg_tensor_all = torch.zeros(new_shape) @@ -496,25 +551,31 @@ def pad_numpy_array(arr, shape): gap = (torch.tensor(img_tensor.shape) - torch.tensor(new_shape)) // 2 gap2 = (torch.tensor(img2_tensor.shape) - torch.tensor(new_shape)) // 2 - new_img_tensor = img_tensor[gap[0]:gap[0]+new_shape[0], gap[1]:gap[1]+new_shape[1], gap[2]:gap[2]+new_shape[2]] - new_img2_tensor = img2_tensor[gap2[0]:gap2[0]+new_shape[0], gap2[1]:gap2[1]+new_shape[1], gap2[2]:gap2[2]+new_shape[2]] - new_seg_tensor_all = seg_tensor_all[gap[0]:gap[0]+new_shape[0], gap[1]:gap[1]+new_shape[1], gap[2]:gap[2]+new_shape[2]] - new_seg2_tensor_all = seg2_tensor_all[gap2[0]:gap2[0]+new_shape[0], gap2[1]:gap2[1]+new_shape[1], gap2[2]:gap2[2]+new_shape[2]] + new_img_tensor = img_tensor[gap[0] : gap[0] + new_shape[0], gap[1] : gap[1] + new_shape[1], gap[2] : gap[2] + new_shape[2]] + new_img2_tensor = img2_tensor[gap2[0] : gap2[0] + new_shape[0], gap2[1] : gap2[1] + new_shape[1], gap2[2] : gap2[2] + new_shape[2]] + new_seg_tensor_all = seg_tensor_all[gap[0] : gap[0] + new_shape[0], gap[1] : gap[1] + new_shape[1], gap[2] : gap[2] + new_shape[2]] + new_seg2_tensor_all = seg2_tensor_all[ + gap2[0] : gap2[0] + new_shape[0], gap2[1] : gap2[1] + new_shape[1], gap2[2] : gap2[2] + new_shape[2] + ] # Add segmentation values to different channels seg_tensor = torch.zeros((1, 5, *new_seg_tensor_all.shape)) for i, value in enumerate([12, 13, 14, 15, 16]): - seg_tensor[0, i] = (new_seg_tensor_all == value) - + seg_tensor[0, i] = new_seg_tensor_all == value + seg2_tensor = torch.zeros((1, 5, *new_seg2_tensor_all.shape)) for i, value in enumerate([50, 45, 44, 43, 42]): - seg2_tensor[0, i] = (new_seg2_tensor_all == value) + seg2_tensor[0, i] = new_seg2_tensor_all == value # Format tensors to match expected input shape (B, C, D, H, W) - img_tensor = torch.cat([new_img_tensor.unsqueeze(0), new_seg_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze(0) # Add batch dimension and second channel - img2_tensor = torch.cat([new_img2_tensor.unsqueeze(0), new_seg2_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze(0) # Add batch dimension and second channel - - # Add batch + img_tensor = torch.cat([new_img_tensor.unsqueeze(0), new_seg_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze( + 0 + ) # Add batch dimension and second channel + img2_tensor = torch.cat([new_img2_tensor.unsqueeze(0), new_seg2_tensor_all.bool().int().unsqueeze(0)], dim=0).unsqueeze( + 0 + ) # Add batch dimension and second channel + + # Add batch img_tensor = torch.cat([img_tensor, img2_tensor], dim=0) seg_tensor = torch.cat([seg_tensor, seg2_tensor], dim=0) @@ -535,10 +596,13 @@ def pad_numpy_array(arr, shape): raise ValueError("NaNs found in augmented image.") if torch.isnan(augmented_seg).any(): raise ValueError("NaNs found in augmented segmentation.") - + + import os + import warnings + import cv2 import numpy as np - import warnings, sys, os + warnings.simplefilter("always") # Convert tensors to numpy arrays @@ -551,25 +615,95 @@ def pad_numpy_array(arr, shape): seg_tensor_np = np.sum(seg_tensor_np, axis=1) augmented_seg_np = np.sum(augmented_seg_np, axis=1) - pad_shape = 2*(np.max(img_tensor_np.shape[2:]),) + pad_shape = 2 * (np.max(img_tensor_np.shape[2:]),) # Combine tensors into single output for visualization - os.makedirs('img', exist_ok=True) - img_line = np.concatenate([normalize(pad_numpy_array(img_tensor_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(img_tensor_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(img_tensor_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - augmented_img_line = np.concatenate([normalize(pad_numpy_array(augmented_img_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - seg_line = np.concatenate([normalize(pad_numpy_array(seg_tensor_np[0, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(seg_tensor_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(seg_tensor_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - augmented_seg_line = np.concatenate([normalize(pad_numpy_array(augmented_seg_np[0, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(augmented_seg_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(augmented_seg_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - not_augmented_channel_line = np.concatenate([normalize(pad_numpy_array(augmented_img_np[0, 1, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[0, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[0, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) + os.makedirs("img", exist_ok=True) + img_line = np.concatenate( + [ + normalize(pad_numpy_array(img_tensor_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(img_tensor_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(img_tensor_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + augmented_img_line = np.concatenate( + [ + normalize(pad_numpy_array(augmented_img_np[0, 0, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[0, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[0, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + seg_line = np.concatenate( + [ + normalize(pad_numpy_array(seg_tensor_np[0, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(seg_tensor_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(seg_tensor_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + augmented_seg_line = np.concatenate( + [ + normalize(pad_numpy_array(augmented_seg_np[0, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_seg_np[0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_seg_np[0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + not_augmented_channel_line = np.concatenate( + [ + normalize(pad_numpy_array(augmented_img_np[0, 1, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[0, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[0, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) combined_img = np.concatenate([img_line, seg_line, augmented_img_line, augmented_seg_line, not_augmented_channel_line], axis=0) - cv2.imwrite('img/combined.png', combined_img*255) - - img_line2 = np.concatenate([normalize(pad_numpy_array(img_tensor_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(img_tensor_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(img_tensor_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - augmented_img_line2 = np.concatenate([normalize(pad_numpy_array(augmented_img_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - seg_line2 = np.concatenate([normalize(pad_numpy_array(seg_tensor_np[1, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(seg_tensor_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(seg_tensor_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - augmented_seg_line2 = np.concatenate([normalize(pad_numpy_array(augmented_seg_np[1, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(augmented_seg_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(augmented_seg_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) - not_augmented_channel_line2 = np.concatenate([normalize(pad_numpy_array(augmented_img_np[1, 1, img_tensor_np.shape[2] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[1, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), normalize(pad_numpy_array(augmented_img_np[1, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape))], axis=1) + cv2.imwrite("img/combined.png", combined_img * 255) + + img_line2 = np.concatenate( + [ + normalize(pad_numpy_array(img_tensor_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(img_tensor_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(img_tensor_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + augmented_img_line2 = np.concatenate( + [ + normalize(pad_numpy_array(augmented_img_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[1, 0, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[1, 0, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + seg_line2 = np.concatenate( + [ + normalize(pad_numpy_array(seg_tensor_np[1, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(seg_tensor_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(seg_tensor_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + augmented_seg_line2 = np.concatenate( + [ + normalize(pad_numpy_array(augmented_seg_np[1, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_seg_np[1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_seg_np[1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) + not_augmented_channel_line2 = np.concatenate( + [ + normalize(pad_numpy_array(augmented_img_np[1, 1, img_tensor_np.shape[2] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[1, 1, :, :, img_tensor_np.shape[4] // 2], pad_shape)), + normalize(pad_numpy_array(augmented_img_np[1, 1, :, img_tensor_np.shape[3] // 2, :], pad_shape)), + ], + axis=1, + ) combined_img2 = np.concatenate([img_line2, seg_line2, augmented_img_line2, augmented_seg_line2, not_augmented_channel_line2], axis=0) - cv2.imwrite('img/combined2.png', combined_img2*255) + cv2.imwrite("img/combined2.png", combined_img2 * 255) # cv2.imwrite('img/orig_img.png', normalize(pad_numpy_array(img_tensor_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape))*255) # cv2.imwrite('img/aug_img.png', normalize(pad_numpy_array(augmented_img_np[1, 0, img_tensor_np.shape[2] // 2], pad_shape))*255) diff --git a/auglab/transforms/gpu/transforms_list.py b/auglab/transforms/gpu/transforms_list.py index 8165c89..930c402 100644 --- a/auglab/transforms/gpu/transforms_list.py +++ b/auglab/transforms/gpu/transforms_list.py @@ -1,29 +1,27 @@ -import os, json +import json +import os +from typing import Any -import torch.nn as nn -import torch import numpy as np +import torch +from torch import Tensor, nn -from auglab.transforms.gpu.base import ImageOnlyTransform -from typing import Any, Dict, Optional, Tuple, Union, List -from kornia.core import Tensor - +from auglab.transforms.gpu.base import AugmentationSequentialCustom, ImageOnlyTransform from auglab.transforms.gpu.contrast import ( - RandomConvTransformGPU, - RandomGaussianNoiseGPU, + RandomBiasFieldGPU, RandomBrightnessGPU, - RandomGammaGPU, + RandomClampGPU, + RandomContrastGPU, + RandomConvTransformGPU, RandomFunctionGPU, + RandomGammaGPU, + RandomGaussianNoiseGPU, RandomHistogramEqualizationGPU, RandomInverseGPU, - RandomBiasFieldGPU, - RandomContrastGPU, ZscoreNormalizationGPU, - RandomClampGPU, ) -from auglab.transforms.gpu.spatial import RandomAffine3DCustom, RandomLowResTransformGPU, RandomFlipTransformGPU, RandomAcqTransformGPU from auglab.transforms.gpu.fromSeg import RandomRedistributeSegGPU -from auglab.transforms.gpu.base import AugmentationSequentialCustom +from auglab.transforms.gpu.spatial import RandomAcqTransformGPU, RandomAffine3DCustom, RandomFlipTransformGPU, RandomLowResTransformGPU class AugTransformsGPURandomOrder(AugmentationSequentialCustom): @@ -34,7 +32,7 @@ class AugTransformsGPURandomOrder(AugmentationSequentialCustom): def __init__(self, json_path: str): # Load transform parameters from JSON config_path = os.path.join(json_path) - with open(config_path, "r") as f: + with open(config_path) as f: config = json.load(f) if "GPU" in config.keys(): @@ -175,17 +173,17 @@ def _build_transforms(self) -> list[nn.Module]: ] function_params = self.transform_params.get("FunctionTransform") if function_params is not None: - for func in func_list: - ta_transforms.append( - RandomFunctionGPU( - func=func, - p=function_params.get("probability", 0), - in_seg=function_params.get("in_seg", 0.0), - out_seg=function_params.get("out_seg", 0.0), - mix_in_out=function_params.get("mix_in_out", False), - retain_stats=function_params.get("retain_stats", False), - ) + ta_transforms.extend( + RandomFunctionGPU( + func=func, + p=function_params.get("probability", 0), + in_seg=function_params.get("in_seg", 0.0), + out_seg=function_params.get("out_seg", 0.0), + mix_in_out=function_params.get("mix_in_out", False), + retain_stats=function_params.get("retain_stats", False), ) + for func in func_list + ) # Bias field artifact bias_field_params = self.transform_params.get("BiasFieldTransform") @@ -333,12 +331,18 @@ def _build_transforms(self) -> list[nn.Module]: choose_x_params = self.transform_params.get("RandomChooseXTransforms") transforms.append( RandomChooseXTransformsGPU( - transforms_list=ta_transforms, num_transforms=len(ta_transforms), p=choose_x_params.get("ta_probability", 1.0), random_order=choose_x_params.get("ta_random_order", True), + transforms_list=ta_transforms, + num_transforms=len(ta_transforms), + p=choose_x_params.get("ta_probability", 1.0), + random_order=choose_x_params.get("ta_random_order", True), ) ) transforms.append( RandomChooseXTransformsGPU( - transforms_list=ge_transforms, num_transforms=len(ge_transforms), p=choose_x_params.get("ge_probability", 1.0), random_order=choose_x_params.get("ge_random_order", True) + transforms_list=ge_transforms, + num_transforms=len(ge_transforms), + p=choose_x_params.get("ge_probability", 1.0), + random_order=choose_x_params.get("ge_random_order", True), ) ) @@ -353,7 +357,7 @@ class AugTransformsGPURandomOrderTA(AugmentationSequentialCustom): def __init__(self, json_path: str): # Load transform parameters from JSON config_path = os.path.join(json_path) - with open(config_path, "r") as f: + with open(config_path) as f: config = json.load(f) if "GPU" in config.keys(): @@ -494,17 +498,17 @@ def _build_transforms(self) -> list[nn.Module]: ] function_params = self.transform_params.get("FunctionTransform") if function_params is not None: - for func in func_list: - ta_transforms.append( - RandomFunctionGPU( - func=func, - p=function_params.get("probability", 0), - in_seg=function_params.get("in_seg", 0.0), - out_seg=function_params.get("out_seg", 0.0), - mix_in_out=function_params.get("mix_in_out", False), - retain_stats=function_params.get("retain_stats", False), - ) + ta_transforms.extend( + RandomFunctionGPU( + func=func, + p=function_params.get("probability", 0), + in_seg=function_params.get("in_seg", 0.0), + out_seg=function_params.get("out_seg", 0.0), + mix_in_out=function_params.get("mix_in_out", False), + retain_stats=function_params.get("retain_stats", False), ) + for func in func_list + ) # Bias field artifact bias_field_params = self.transform_params.get("BiasFieldTransform") @@ -677,7 +681,7 @@ class RandomChooseXTransformsGPU(ImageOnlyTransform): def __init__( self, - transforms_list: List[ImageOnlyTransform], + transforms_list: list[ImageOnlyTransform], num_transforms: int = 1, same_on_batch: bool = False, p: float = 1.0, @@ -692,7 +696,7 @@ def __init__( self.num_transforms = num_transforms self.random_order = random_order - def _apply_mix(self, x: Tensor, seg: Optional[Tensor]) -> Tensor: + def _apply_mix(self, x: Tensor, seg: Tensor | None) -> Tensor: if self.num_transforms == 0 or len(self.transforms_list) == 0: return x @@ -703,7 +707,7 @@ def _apply_mix(self, x: Tensor, seg: Optional[Tensor]) -> Tensor: else: idx = torch.arange(len(self.transforms_list), device=x.device)[:k] - child_params: Dict[str, Tensor] = {} + child_params: dict[str, Tensor] = {} if seg is not None: child_params["seg"] = seg @@ -719,10 +723,8 @@ def _apply_mix(self, x: Tensor, seg: Optional[Tensor]) -> Tensor: return x @torch.no_grad() # disable gradients for efficiency - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: - seg = params.get("seg", None) + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: + seg = params.get("seg") if self.same_on_batch: return self._apply_mix(input, seg) @@ -732,10 +734,7 @@ def apply_transform( for i in range(batch_size): xi = out[i : i + 1] seg_i = None - if seg is not None and isinstance(seg, torch.Tensor) and seg.shape[0] == batch_size: - seg_i = seg[i : i + 1] - else: - seg_i = seg + seg_i = seg[i : i + 1] if seg is not None and isinstance(seg, torch.Tensor) and seg.shape[0] == batch_size else seg xi = self._apply_mix(xi, seg_i) out[i : i + 1] = xi return out @@ -767,7 +766,9 @@ def pad_numpy_array(arr, shape): if __name__ == "__main__": # Example usage import importlib - import auglab.configs as configs + + from auglab import configs + from auglab.transforms.gpu.transforms import AugTransformsGPU from auglab.utils.image import Image, resample_nib configs_path = importlib.resources.files(configs) @@ -856,9 +857,11 @@ def pad_numpy_array(arr, shape): if torch.isnan(augmented_seg).any(): raise ValueError("NaNs found in augmented segmentation.") + import os + import warnings + import cv2 import numpy as np - import warnings, sys, os warnings.simplefilter("always") diff --git a/auglab/transforms/synthseg/__init__.py b/auglab/transforms/synthseg/__init__.py index f947b0b..56091ea 100644 --- a/auglab/transforms/synthseg/__init__.py +++ b/auglab/transforms/synthseg/__init__.py @@ -15,13 +15,13 @@ AugmentationSequentialCustom pipelines). """ +from auglab.transforms.synthseg import functional from auglab.transforms.synthseg.generator import SynthSegGenerator from auglab.transforms.synthseg.transforms import RandomSynthSegGPU, SynthSegTransformsGPU -from auglab.transforms.synthseg import functional __all__ = [ + "RandomSynthSegGPU", "SynthSegGenerator", "SynthSegTransformsGPU", - "RandomSynthSegGPU", "functional", ] diff --git a/auglab/transforms/synthseg/functional.py b/auglab/transforms/synthseg/functional.py index 697308c..b2c4615 100644 --- a/auglab/transforms/synthseg/functional.py +++ b/auglab/transforms/synthseg/functional.py @@ -27,7 +27,8 @@ from __future__ import annotations import math -from typing import List, Optional, Sequence, Tuple, Union +from collections.abc import Sequence +from typing import Union import torch import torch.nn.functional as F @@ -35,22 +36,22 @@ Number = Union[int, float] __all__ = [ - "to_label_map", - "infer_label_values", - "sample_gmm_parameters", - "labels_to_image_gmm", - "sample_affine_matrices", - "random_svf_field", - "warp_volume", "bias_field", - "intensity_augmentation", "blurring_sigma_for_downsampling", - "gaussian_blur_3d", - "sample_resolution", - "mimic_acquisition", + "convert_labels", "em_subdivide_labels", "flip_lr_with_swap", - "convert_labels", + "gaussian_blur_3d", + "infer_label_values", + "intensity_augmentation", + "labels_to_image_gmm", + "mimic_acquisition", + "random_svf_field", + "sample_affine_matrices", + "sample_gmm_parameters", + "sample_resolution", + "to_label_map", + "warp_volume", ] @@ -94,8 +95,8 @@ def infer_label_values(label_map: torch.Tensor) -> torch.Tensor: # + SynthSeg.model_inputs.build_model_inputs) # --------------------------------------------------------------------------- def _draw_value( - prior: Optional[Union[Number, Sequence[Number], torch.Tensor]], - size: Tuple[int, int], + prior: Union[Number, Sequence[Number], torch.Tensor] | None, + size: tuple[int, int], distribution: str, centre: float, default_range: float, @@ -140,9 +141,7 @@ def _draw_value( b = prior_t[1].expand(size).clone() elif prior_t.dim() == 2 and prior_t.shape[0] == 2: if prior_t.shape[1] != n_classes: - raise ValueError( - f"Prior array has {prior_t.shape[1]} classes, expected {n_classes}." - ) + raise ValueError(f"Prior array has {prior_t.shape[1]} classes, expected {n_classes}.") a = prior_t[0].unsqueeze(0).expand(size).clone() b = prior_t[1].unsqueeze(0).expand(size).clone() else: @@ -168,10 +167,10 @@ def sample_gmm_parameters( prior_means=None, prior_stds=None, prior_distributions: str = "uniform", - generation_classes: Optional[Sequence[int]] = None, - background_label_index: Optional[int] = 0, + generation_classes: Sequence[int] | None = None, + background_label_index: int | None = 0, randomise_background: bool = True, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Draw per-label Gaussian means/stds for one minibatch. Mirrors ``SynthSeg.model_inputs.build_model_inputs``. With the default @@ -196,12 +195,8 @@ def sample_gmm_parameters( means = torch.empty(batch, n_classes, n_channels, device=device) stds = torch.empty(batch, n_classes, n_channels, device=device) for ch in range(n_channels): - means[:, :, ch] = _draw_value( - prior_means, (batch, n_classes), prior_distributions, 125.0, 125.0, device, positive_only=True - ) - stds[:, :, ch] = _draw_value( - prior_stds, (batch, n_classes), prior_distributions, 15.0, 15.0, device, positive_only=True - ) + means[:, :, ch] = _draw_value(prior_means, (batch, n_classes), prior_distributions, 125.0, 125.0, device, positive_only=True) + stds[:, :, ch] = _draw_value(prior_stds, (batch, n_classes), prior_distributions, 15.0, 15.0, device, positive_only=True) # Scatter class parameters to per-label parameters. means_lab = means[:, classes, :] @@ -298,6 +293,7 @@ def sample_affine_matrices( translation placed in the last column. Returns ``(B, 4, 4)``. The matrix is applied about the volume centre by :func:`warp_volume`. """ + def draw(bounds, centre): vec = _as_3vec(bounds, device, default=0.0) return centre + (2.0 * torch.rand(batch, 3, device=device) - 1.0) * vec.view(1, 3) @@ -345,7 +341,7 @@ def stack3(rows): return affine -def _identity_grid(shape: Tuple[int, int, int], device: torch.device) -> torch.Tensor: +def _identity_grid(shape: tuple[int, int, int], device: torch.device) -> torch.Tensor: """Voxel-coordinate identity grid in ``(i, j, k)`` order, shape ``(3, D, H, W)``.""" d, h, w = shape zs = torch.arange(d, device=device, dtype=torch.float32) @@ -355,7 +351,7 @@ def _identity_grid(shape: Tuple[int, int, int], device: torch.device) -> torch.T return torch.stack([ii, jj, kk], dim=0) -def _coords_to_grid_sample(coords: torch.Tensor, shape: Tuple[int, int, int]) -> torch.Tensor: +def _coords_to_grid_sample(coords: torch.Tensor, shape: tuple[int, int, int]) -> torch.Tensor: """Convert ``(B, 3, D, H, W)`` voxel coords (i,j,k) to a grid_sample grid. Output ``(B, D, H, W, 3)`` with last-dim order ``(x, y, z) = (k, j, i)`` @@ -373,8 +369,8 @@ def _coords_to_grid_sample(coords: torch.Tensor, shape: Tuple[int, int, int]) -> def warp_volume( volume: torch.Tensor, - affine: Optional[torch.Tensor] = None, - displacement: Optional[torch.Tensor] = None, + affine: torch.Tensor | None = None, + displacement: torch.Tensor | None = None, interp: str = "linear", center: bool = True, padding_mode: str = "zeros", @@ -407,9 +403,7 @@ def warp_volume( if affine is not None: flat = coords.reshape(B, 3, -1) # (B, 3, N) if center: - centre = torch.tensor( - [(D - 1) / 2.0, (H - 1) / 2.0, (W - 1) / 2.0], device=device - ).view(1, 3, 1) + centre = torch.tensor([(D - 1) / 2.0, (H - 1) / 2.0, (W - 1) / 2.0], device=device).view(1, 3, 1) flat = flat - centre linear = affine[:, :3, :3] translation = affine[:, :3, 3:4] @@ -423,9 +417,7 @@ def warp_volume( sample_grid = _coords_to_grid_sample(coords, shape) mode = "nearest" if interp == "nearest" else "bilinear" # 3D 'bilinear' == trilinear - return F.grid_sample( - volume, sample_grid, mode=mode, align_corners=True, padding_mode=padding_mode - ) + return F.grid_sample(volume, sample_grid, mode=mode, align_corners=True, padding_mode=padding_mode) def _integrate_velocity(velocity: torch.Tensor, int_steps: int = 7) -> torch.Tensor: @@ -436,7 +428,7 @@ def _integrate_velocity(velocity: torch.Tensor, int_steps: int = 7) -> torch.Ten yielding a diffeomorphic displacement field. ``velocity`` and the returned displacement are ``(B, 3, D, H, W)`` in voxel units. """ - disp = velocity / (2 ** int_steps) + disp = velocity / (2**int_steps) for _ in range(int_steps): disp = disp + warp_volume(disp, displacement=disp, interp="linear", padding_mode="border") return disp @@ -444,7 +436,7 @@ def _integrate_velocity(velocity: torch.Tensor, int_steps: int = 7) -> torch.Ten def random_svf_field( batch: int, - shape: Tuple[int, int, int], + shape: tuple[int, int, int], device: torch.device, nonlin_std: float = 4.0, nonlin_scale: float = 0.04, @@ -460,7 +452,7 @@ def random_svf_field( if nonlin_std <= 0: return torch.zeros(batch, 3, *shape, device=device) - small = [max(2, int(math.ceil(s * nonlin_scale))) for s in shape] + small = [max(2, math.ceil(s * nonlin_scale)) for s in shape] std = torch.rand(batch, 1, 1, 1, 1, device=device) * nonlin_std velocity = torch.randn(batch, 3, *small, device=device) * std velocity = F.interpolate(velocity, size=shape, mode="trilinear", align_corners=True) @@ -487,7 +479,7 @@ def bias_field( return image B, C, D, H, W = image.shape device = image.device - small = [max(2, int(math.ceil(s * bias_scale))) for s in (D, H, W)] + small = [max(2, math.ceil(s * bias_scale)) for s in (D, H, W)] std = torch.rand(B, 1, 1, 1, 1, device=device) * bias_field_std field = torch.randn(B, C, *small, device=device) * std field = F.interpolate(field, size=(D, H, W), mode="trilinear", align_corners=True) @@ -535,7 +527,7 @@ def intensity_augmentation( def blurring_sigma_for_downsampling( current_res: torch.Tensor, downsample_res: torch.Tensor, - thickness: Optional[torch.Tensor] = None, + thickness: torch.Tensor | None = None, ) -> torch.Tensor: """Per-axis Gaussian blur sigma for a target acquisition resolution. @@ -556,7 +548,7 @@ def blurring_sigma_for_downsampling( def _gaussian_kernel1d(sigma: float, device: torch.device) -> torch.Tensor: if sigma <= 0: return torch.tensor([1.0], device=device) - radius = max(1, int(math.ceil(3.0 * sigma))) + radius = max(1, math.ceil(3.0 * sigma)) x = torch.arange(-radius, radius + 1, device=device, dtype=torch.float32) k = torch.exp(-0.5 * (x / sigma) ** 2) return k / k.sum() @@ -606,7 +598,7 @@ def sample_resolution( max_res_aniso: float = 8.0, prob_iso: float = 0.1, prob_min: float = 0.05, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Sample a random target acquisition resolution and slice thickness. Port of ``lab2im.layers.SampleResolution`` (with ``return_thickness=True``): @@ -641,7 +633,7 @@ def mimic_acquisition( image: torch.Tensor, current_res: torch.Tensor, downsample_res: torch.Tensor, - output_shape: Tuple[int, int, int], + output_shape: tuple[int, int, int], ) -> torch.Tensor: """Downsample to a target resolution, then resample to the output grid. @@ -652,7 +644,7 @@ def mimic_acquisition( B, C, D, H, W = image.shape in_shape = (D, H, W) factor = (current_res / downsample_res).tolist() - down_shape = [max(1, int(round(in_shape[i] * factor[i]))) for i in range(3)] + down_shape = [max(1, round(in_shape[i] * factor[i])) for i in range(3)] x = F.interpolate(image, size=down_shape, mode="nearest") x = F.interpolate(x, size=tuple(output_shape), mode="trilinear", align_corners=True) return x @@ -661,9 +653,7 @@ def mimic_acquisition( # --------------------------------------------------------------------------- # EM label completion for sparse label maps (SynthSeg paper, Sec. 5.4) # --------------------------------------------------------------------------- -def _em_gmm_1d( - x_fit: torch.Tensor, n_components: int, n_iters: int, eps: float -) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: +def _em_gmm_1d(x_fit: torch.Tensor, n_components: int, n_iters: int, eps: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: """Fit a 1D Gaussian mixture by Expectation-Maximization. ``x_fit`` is a 1D tensor of intensities. Returns ``(means, vars, weights)`` @@ -691,8 +681,8 @@ def _em_gmm_1d( - 0.5 * (x - means.view(1, k)) ** 2 / var.view(1, k) ) logp = logp - torch.logsumexp(logp, dim=1, keepdim=True) - resp = logp.exp() # (n, k) - nk = resp.sum(0).clamp_min(eps) # (k,) + resp = logp.exp() # (n, k) + nk = resp.sum(0).clamp_min(eps) # (k,) weights = nk / n means = (resp * x).sum(0) / nk var = (resp * (x - means.view(1, k)) ** 2).sum(0) / nk @@ -701,8 +691,12 @@ def _em_gmm_1d( def _assign_gmm( - x_full: torch.Tensor, means: torch.Tensor, var: torch.Tensor, weights: torch.Tensor, - eps: float, chunk: int = 2_000_000, + x_full: torch.Tensor, + means: torch.Tensor, + var: torch.Tensor, + weights: torch.Tensor, + eps: float, + chunk: int = 2_000_000, ) -> torch.Tensor: """Hard-assign each value in ``x_full`` to its most likely mixture component.""" n = x_full.numel() @@ -713,9 +707,9 @@ def _assign_gmm( m = means.view(1, k) v = var.view(1, k) for s in range(0, n, chunk): - xc = x_full[s:s + chunk].view(-1, 1) + xc = x_full[s : s + chunk].view(-1, 1) logp = logw - half_logvar - 0.5 * (xc - m) ** 2 / v - out[s:s + chunk] = logp.argmax(dim=1) + out[s : s + chunk] = logp.argmax(dim=1) return out @@ -730,7 +724,7 @@ def em_subdivide_labels( channel: int = 0, same_on_batch: bool = False, eps: float = 1e-6, -) -> Tuple[torch.Tensor, List[int], List[int]]: +) -> tuple[torch.Tensor, list[int], list[int]]: """Subdivide each label into intensity-coherent subregions via EM (SynthSeg §5.4). Reproduces SynthSeg's handling of sparse / incomplete label maps: "we enhance @@ -764,11 +758,11 @@ def em_subdivide_labels( """ B = image.shape[0] device = image.device - ref = image[:, channel] # (B, D, H, W) - parents = torch.unique(label_map).long().tolist() # sorted, batch-wide + ref = image[:, channel] # (B, D, H, W) + parents = torch.unique(label_map).long().tolist() # sorted, batch-wide parent_to_idx = {p: i for i, p in enumerate(parents)} lo, hi = int(background_clusters_range[0]), int(background_clusters_range[1]) - mult = max(hi, int(n_foreground_clusters)) + 1 # collision-free encoding + mult = max(hi, int(n_foreground_clusters)) + 1 # collision-free encoding # If the configured background label is absent (e.g. a *complete* one-hot whose # decoding shifted every label by +1, so the real background is no longer 0), @@ -805,10 +799,7 @@ def em_subdivide_labels( else: x_fit = x fit = _em_gmm_1d(x_fit, k, n_iters, eps) - assign = ( - torch.zeros(cnt, dtype=torch.long, device=device) - if fit is None else _assign_gmm(x, *fit, eps=eps) - ) + assign = torch.zeros(cnt, dtype=torch.long, device=device) if fit is None else _assign_gmm(x, *fit, eps=eps) fine[b, 0][mask] = pi * mult + assign gen_values = torch.unique(fine).long().tolist() @@ -822,8 +813,8 @@ def em_subdivide_labels( def flip_lr_with_swap( label_map: torch.Tensor, flip_axis: int, - label_values: Optional[torch.Tensor] = None, - n_neutral_labels: Optional[int] = None, + label_values: torch.Tensor | None = None, + n_neutral_labels: int | None = None, ) -> torch.Tensor: """Flip the label map along ``flip_axis`` and (optionally) swap L/R labels. @@ -850,8 +841,8 @@ def flip_lr_with_swap( return flipped neutral = values[:n_neutral_labels] - left = values[n_neutral_labels:n_neutral_labels + n_sided] - right = values[n_neutral_labels + n_sided:n_neutral_labels + 2 * n_sided] + left = values[n_neutral_labels : n_neutral_labels + n_sided] + right = values[n_neutral_labels + n_sided : n_neutral_labels + 2 * n_sided] source = neutral + left + right dest = neutral + right + left return convert_labels(flipped, source, dest) diff --git a/auglab/transforms/synthseg/generator.py b/auglab/transforms/synthseg/generator.py index b7cba44..1adcafa 100644 --- a/auglab/transforms/synthseg/generator.py +++ b/auglab/transforms/synthseg/generator.py @@ -28,7 +28,8 @@ from __future__ import annotations -from typing import List, Optional, Sequence, Tuple, Union +from collections.abc import Sequence +from typing import Union import torch from torch import nn @@ -93,10 +94,10 @@ class SynthSegGenerator(nn.Module): def __init__( self, - generation_labels: Optional[Sequence[int]] = None, - output_labels: Optional[Sequence[int]] = None, - n_neutral_labels: Optional[int] = None, - generation_classes: Optional[Sequence[int]] = None, + generation_labels: Sequence[int] | None = None, + output_labels: Sequence[int] | None = None, + n_neutral_labels: int | None = None, + generation_classes: Sequence[int] | None = None, n_channels: int = 1, prior_distributions: str = "uniform", prior_means=None, @@ -122,7 +123,7 @@ def __init__( thickness=None, blur_range: float = 1.03, atlas_res: float = 1.0, - output_shape: Optional[Sequence[int]] = None, + output_shape: Sequence[int] | None = None, em_label_completion: bool = False, em_n_foreground_clusters: int = 2, em_background_clusters_range: Sequence[int] = (3, 10), @@ -190,9 +191,7 @@ def __init__( # ------------------------------------------------------------------ @torch.no_grad() - def forward( - self, label_map: torch.Tensor, image: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, torch.Tensor]: + def forward(self, label_map: torch.Tensor, image: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: """Generate an image and its label map from an input label map. Args: @@ -211,9 +210,9 @@ def forward( batch = labels.shape[0] # Label bookkeeping (defaults from config; overridden by EM completion). - gen_labels = self.generation_labels # list[int] or None - out_labels_cfg = self.output_labels # list[int] or None - gen_classes = self.generation_classes # list[int] or None + gen_labels = self.generation_labels # list[int] or None + out_labels_cfg = self.output_labels # list[int] or None + gen_classes = self.generation_classes # list[int] or None n_neutral = self.n_neutral_labels randomise_bg = True @@ -221,13 +220,17 @@ def forward( if self.em_label_completion: if image is None: if not self._warned_em: - print(f"{type(self).__name__}: em_label_completion is enabled but no " - f"image was provided; falling back to plain generation.", flush=True) + print( + f"{type(self).__name__}: em_label_completion is enabled but no " + f"image was provided; falling back to plain generation.", + flush=True, + ) self._warned_em = True else: ref = image if image.dim() == 5 else image.unsqueeze(1) labels, gen_labels, out_labels_cfg = FN.em_subdivide_labels( - ref.float(), labels, + ref.float(), + labels, n_foreground_clusters=self.em_n_foreground_clusters, background_clusters_range=self.em_background_clusters_range, background_label=self.em_background_label, @@ -235,9 +238,9 @@ def forward( max_fit_voxels=self.em_max_fit_voxels, same_on_batch=self.em_same_on_batch, ) - gen_classes = None # each sub-label gets its own Gaussian - n_neutral = None # plain flip (sub-labels carry no L/R structure) - randomise_bg = False # background is now modelled by its clusters + gen_classes = None # each sub-label gets its own Gaussian + n_neutral = None # plain flip (sub-labels carry no L/R structure) + randomise_bg = False # background is now modelled by its clusters # 1. random crop to output_shape (label space) ------------------------ if self.output_shape is not None and tuple(self.output_shape) != tuple(labels.shape[2:]): @@ -247,7 +250,8 @@ def forward( affine = None if self.apply_affine and self._affine_active(): affine = FN.sample_affine_matrices( - batch, device, + batch, + device, scaling_bounds=self.scaling_bounds, rotation_bounds=self.rotation_bounds, shearing_bounds=self.shearing_bounds, @@ -256,33 +260,41 @@ def forward( displacement = None if self.apply_nonlinear and self.nonlin_std and self.nonlin_std > 0: displacement = FN.random_svf_field( - batch, tuple(labels.shape[2:]), device, + batch, + tuple(labels.shape[2:]), + device, nonlin_std=self.nonlin_std, nonlin_scale=self.nonlin_scale, int_steps=self.svf_integration_steps, ) if affine is not None or displacement is not None: - labels = FN.warp_volume( - labels.float(), affine=affine, displacement=displacement, - interp="nearest", padding_mode="zeros", - ).round().long() + labels = ( + FN.warp_volume( + labels.float(), + affine=affine, + displacement=displacement, + interp="nearest", + padding_mode="zeros", + ) + .round() + .long() + ) # 3. left/right flipping (with optional label swap) ------------------- if self.flipping and float(torch.rand((), device=device)) < 0.5: label_values_flip = ( - torch.as_tensor(gen_labels, dtype=torch.long, device=device) - if gen_labels is not None else FN.infer_label_values(labels) + torch.as_tensor(gen_labels, dtype=torch.long, device=device) if gen_labels is not None else FN.infer_label_values(labels) ) labels = FN.flip_lr_with_swap( - labels, self.flip_axis, + labels, + self.flip_axis, label_values=label_values_flip, n_neutral_labels=n_neutral, ) # The generation labels (after potential relabelling) used by the GMM. gen_values = ( - torch.as_tensor(gen_labels, dtype=torch.long, device=device) - if gen_labels is not None else FN.infer_label_values(labels) + torch.as_tensor(gen_labels, dtype=torch.long, device=device) if gen_labels is not None else FN.infer_label_values(labels) ) n_labels = gen_values.numel() bg_index = None @@ -291,7 +303,10 @@ def forward( # 4. GMM intensity sampling ------------------------------------------- means, stds = FN.sample_gmm_parameters( - n_labels, self.n_channels, batch, device, + n_labels, + self.n_channels, + batch, + device, prior_means=self.prior_means, prior_stds=self.prior_stds, prior_distributions=self.prior_distributions, @@ -306,9 +321,7 @@ def forward( # 6. intensity augmentation (clip -> normalise -> gamma) -------------- if self.apply_intensity_augmentation: - synth = FN.intensity_augmentation( - synth, clip=self.clip, gamma_std=self.gamma_std, normalise=self.normalise - ) + synth = FN.intensity_augmentation(synth, clip=self.clip, gamma_std=self.gamma_std, normalise=self.normalise) # 7. resolution randomisation, per channel ---------------------------- if self.apply_resolution: @@ -331,15 +344,14 @@ def _affine_active(self) -> bool: @staticmethod def _random_crop(labels: torch.Tensor, output_shape: Sequence[int]) -> torch.Tensor: B, _, D, H, W = labels.shape - out = [] sizes = (D, H, W) starts = [] for dim, target in zip(sizes, output_shape): - target = min(int(target), dim) - start = int(torch.randint(0, dim - target + 1, (1,))) if dim > target else 0 - starts.append((start, target)) + size = min(int(target), dim) + start = int(torch.randint(0, dim - size + 1, (1,))) if dim > size else 0 + starts.append((start, size)) (sd, td), (sh, th), (sw, tw) = starts - return labels[:, :, sd:sd + td, sh:sh + th, sw:sw + tw] + return labels[:, :, sd : sd + td, sh : sh + th, sw : sw + tw] def _simulate_resolution(self, image: torch.Tensor) -> torch.Tensor: device = image.device @@ -348,11 +360,9 @@ def _simulate_resolution(self, image: torch.Tensor) -> torch.Tensor: channels = [] for c in range(image.shape[1]): - ch = image[:, c:c + 1] + ch = image[:, c : c + 1] if self.randomise_res: - res, thickness = FN.sample_resolution( - atlas_res, self.max_res_iso, self.max_res_aniso - ) + res, thickness = FN.sample_resolution(atlas_res, self.max_res_iso, self.max_res_aniso) else: res = self._fixed_res(c, device) thickness = self._fixed_thickness(c, device, res) @@ -385,23 +395,20 @@ def _fixed_thickness(self, channel: int, device: torch.device, res: torch.Tensor device = "cuda" if torch.cuda.is_available() else "cpu" B, D, H, W = 2, 48, 56, 52 - zz, yy, xx = torch.meshgrid( - torch.arange(D), torch.arange(H), torch.arange(W), indexing="ij" - ) + zz, yy, xx = torch.meshgrid(torch.arange(D), torch.arange(H), torch.arange(W), indexing="ij") centre = torch.tensor([D / 2, H / 2, W / 2]) r = ((zz - centre[0]) ** 2 + (yy - centre[1]) ** 2 + (xx - centre[2]) ** 2).sqrt() vol = torch.zeros(D, H, W, dtype=torch.long) - vol[r < 18] = 1 # "tissue A" - vol[r < 10] = 2 # "tissue B" - vol[(xx > W // 2) & (r < 18)] = 3 # right-side structure + vol[r < 18] = 1 # "tissue A" + vol[r < 10] = 2 # "tissue B" + vol[(xx > W // 2) & (r < 18)] = 3 # right-side structure labels = vol.view(1, 1, D, H, W).repeat(B, 1, 1, 1, 1) gen = SynthSegGenerator(generation_labels=[0, 1, 2, 3], n_channels=1).to(device) image, out_labels = gen(labels.to(device)) print("input labels:", tuple(labels.shape), "values", torch.unique(labels).tolist()) - print("output image :", tuple(image.shape), "range", - (round(float(image.min()), 3), round(float(image.max()), 3))) + print("output image :", tuple(image.shape), "range", (round(float(image.min()), 3), round(float(image.max()), 3))) print("output labels:", tuple(out_labels.shape), "values", torch.unique(out_labels).tolist()) assert image.shape[0] == B and image.shape[1] == 1 assert out_labels.shape[2:] == image.shape[2:] diff --git a/auglab/transforms/synthseg/transforms.py b/auglab/transforms/synthseg/transforms.py index 9c838d2..30c7f23 100644 --- a/auglab/transforms/synthseg/transforms.py +++ b/auglab/transforms/synthseg/transforms.py @@ -23,32 +23,62 @@ import json import os -from typing import Any, Dict, List, Optional +from typing import Any import torch -from torch import nn -from kornia.core import Tensor +from torch import Tensor, nn from auglab.transforms.gpu.base import ImageOnlyTransform from auglab.transforms.synthseg.generator import SynthSegGenerator # Keys understood from the JSON config / kwargs, forwarded to SynthSegGenerator. _GENERATOR_KEYS = { - "generation_labels", "output_labels", "n_neutral_labels", "generation_classes", - "n_channels", "prior_distributions", "prior_means", "prior_stds", - "flipping", "flip_axis", "scaling_bounds", "rotation_bounds", "shearing_bounds", - "translation_bounds", "nonlin_std", "nonlin_scale", "svf_integration_steps", - "bias_field_std", "bias_scale", "gamma_std", "clip", "normalise", - "randomise_res", "max_res_iso", "max_res_aniso", "data_res", "thickness", - "blur_range", "atlas_res", "output_shape", - "em_label_completion", "em_n_foreground_clusters", "em_background_clusters_range", - "em_background_label", "em_n_iters", "em_max_fit_voxels", "em_same_on_batch", - "apply_affine", "apply_nonlinear", "apply_bias_field", - "apply_intensity_augmentation", "apply_resolution", + "generation_labels", + "output_labels", + "n_neutral_labels", + "generation_classes", + "n_channels", + "prior_distributions", + "prior_means", + "prior_stds", + "flipping", + "flip_axis", + "scaling_bounds", + "rotation_bounds", + "shearing_bounds", + "translation_bounds", + "nonlin_std", + "nonlin_scale", + "svf_integration_steps", + "bias_field_std", + "bias_scale", + "gamma_std", + "clip", + "normalise", + "randomise_res", + "max_res_iso", + "max_res_aniso", + "data_res", + "thickness", + "blur_range", + "atlas_res", + "output_shape", + "em_label_completion", + "em_n_foreground_clusters", + "em_background_clusters_range", + "em_background_label", + "em_n_iters", + "em_max_fit_voxels", + "em_same_on_batch", + "apply_affine", + "apply_nonlinear", + "apply_bias_field", + "apply_intensity_augmentation", + "apply_resolution", } -def _filter_generator_kwargs(params: Dict[str, Any]) -> Dict[str, Any]: +def _filter_generator_kwargs(params: dict[str, Any]) -> dict[str, Any]: return {k: v for k, v in params.items() if k in _GENERATOR_KEYS} @@ -71,7 +101,7 @@ class RandomSynthSegGPU(ImageOnlyTransform): def __init__( self, - apply_to_channel: Optional[List[int]] = None, + apply_to_channel: list[int] | None = None, same_on_batch: bool = False, p: float = 0.5, keepdim: bool = True, @@ -86,10 +116,8 @@ def __init__( self.generator = SynthSegGenerator(**gen_kwargs) @torch.no_grad() - def apply_transform( - self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None - ) -> Tensor: - seg = params.get("seg", None) + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: + seg = params.get("seg") if seg is None: return input @@ -124,12 +152,12 @@ class SynthSegTransformsGPU(nn.Module): to ``1.0`` (always synthesise, as in the paper). """ - def __init__(self, json_path: Optional[str] = None, params: Optional[Dict[str, Any]] = None): + def __init__(self, json_path: str | None = None, params: dict[str, Any] | None = None): super().__init__() if params is None: if json_path is None: raise ValueError("Provide either json_path or params.") - with open(os.path.join(json_path), "r") as f: + with open(os.path.join(json_path)) as f: config = json.load(f) else: config = params @@ -183,8 +211,7 @@ def _to_onehot(labels: Tensor, n_channels: int) -> Tensor: # Full end-to-end driver driver = SynthSegTransformsGPU(params={"generation_labels": [0, 1, 2], "n_channels": 1}).to(device) out_img, out_lab = driver(img.to(device), labels.to(device)) - print("driver image", tuple(out_img.shape), "labels", tuple(out_lab.shape), - torch.unique(out_lab).tolist()) + print("driver image", tuple(out_img.shape), "labels", tuple(out_lab.shape), torch.unique(out_lab).tolist()) assert out_img.shape[0] == B and not torch.isnan(out_img).any() # Intensity-only ImageOnlyTransform diff --git a/auglab/utils/image.py b/auglab/utils/image.py index 86e031b..6402d53 100644 --- a/auglab/utils/image.py +++ b/auglab/utils/image.py @@ -1,13 +1,15 @@ +import logging import os -import numpy as np +from copy import deepcopy + import nibabel as nib +import numpy as np from nibabel.processing import resample_from_to -import logging -from copy import deepcopy logger = logging.getLogger(__name__) -class Image(object): + +class Image: """ Compact version of SCT's Image Class (https://github.com/spinalcordtoolbox/spinalcordtoolbox/blob/master/spinalcordtoolbox/image.py#L245) Create an object that behaves similarly to nibabel's image object. Useful additions include: dims, change_orientation and getNonZeroCoordinates. @@ -26,7 +28,7 @@ def __init__(self, param=None, hdr=None, orientation=None, absolutepath=None, di if absolutepath is not None: self._path = os.path.abspath(absolutepath) - + # Case 1: load an image from file if isinstance(param, str): self.loadFromPath(param) @@ -44,19 +46,19 @@ def __init__(self, param=None, hdr=None, orientation=None, absolutepath=None, di self.hdr = hdr.copy() if hdr is not None else nib.Nifti1Header() self.hdr.set_data_shape(self.data.shape) else: - raise TypeError('Image constructor takes at least one argument.') - + raise TypeError("Image constructor takes at least one argument.") + # Fix any mismatch between the array's datatype and the header datatype self.fix_header_dtype() @property def dim(self): return get_dimension(self) - + @property def orientation(self): return get_orientation(self) - + @property def absolutepath(self): """ @@ -74,7 +76,7 @@ def absolutepath(self): the best way to set it. """ return self._path - + @absolutepath.setter def absolutepath(self, value): if value is None: @@ -85,7 +87,7 @@ def absolutepath(self, value): elif not os.path.isabs(value): value = os.path.abspath(value) self._path = value - + @property def header(self): return self.hdr @@ -95,7 +97,13 @@ def header(self, value): self.hdr = value def __deepcopy__(self, memo): - return type(self)(deepcopy(self.data, memo), deepcopy(self.hdr, memo), deepcopy(self.orientation, memo), deepcopy(self.absolutepath, memo), deepcopy(self.dim, memo)) + return type(self)( + deepcopy(self.data, memo), + deepcopy(self.hdr, memo), + deepcopy(self.orientation, memo), + deepcopy(self.absolutepath, memo), + deepcopy(self.dim, memo), + ) def copy(self, image=None): if image is not None: @@ -137,7 +145,7 @@ def change_orientation(self, orientation, inverse=False): """ change_orientation(self, orientation, self, inverse=inverse) return self - + def getNonZeroCoordinates(self, sorting=None, reverse_coord=False): """ This function return all the non-zero coordinates that the image contains. @@ -147,41 +155,38 @@ def getNonZeroCoordinates(self, sorting=None, reverse_coord=False): Removed Coordinate object """ n_dim = 1 - if self.dim[3] == 1: - n_dim = 3 - else: - n_dim = 4 + n_dim = 3 if self.dim[3] == 1 else 4 if self.dim[2] == 1: n_dim = 2 if n_dim == 3: X, Y, Z = (self.data > 0).nonzero() - list_coordinates = [[X[i], Y[i], Z[i], self.data[X[i], Y[i], Z[i]]] for i in range(0, len(X))] + list_coordinates = [[X[i], Y[i], Z[i], self.data[X[i], Y[i], Z[i]]] for i in range(len(X))] elif n_dim == 2: try: X, Y = (self.data > 0).nonzero() - list_coordinates = [[X[i], Y[i], 0, self.data[X[i], Y[i]]] for i in range(0, len(X))] + list_coordinates = [[X[i], Y[i], 0, self.data[X[i], Y[i]]] for i in range(len(X))] except ValueError: X, Y, Z = (self.data > 0).nonzero() - list_coordinates = [[X[i], Y[i], 0, self.data[X[i], Y[i], 0]] for i in range(0, len(X))] + list_coordinates = [[X[i], Y[i], 0, self.data[X[i], Y[i], 0]] for i in range(len(X))] if sorting is not None: if reverse_coord not in [True, False]: - raise ValueError('reverse_coord parameter must be a boolean') + raise ValueError("reverse_coord parameter must be a boolean") - if sorting == 'x': + if sorting == "x": list_coordinates = sorted(list_coordinates, key=lambda el: el[0], reverse=reverse_coord) - elif sorting == 'y': + elif sorting == "y": list_coordinates = sorted(list_coordinates, key=lambda el: el[1], reverse=reverse_coord) - elif sorting == 'z': + elif sorting == "z": list_coordinates = sorted(list_coordinates, key=lambda el: el[2], reverse=reverse_coord) - elif sorting == 'value': + elif sorting == "value": list_coordinates = sorted(list_coordinates, key=lambda el: el[3], reverse=reverse_coord) else: raise ValueError("sorting parameter must be either 'x', 'y', 'z' or 'value'") return list_coordinates - + def change_type(self, dtype): """ Change data type on image. @@ -190,7 +195,7 @@ def change_type(self, dtype): """ change_type(self, dtype, self) return self - + def fix_header_dtype(self): """ Change the header dtype to the match the datatype of the array. @@ -198,15 +203,19 @@ def fix_header_dtype(self): # Using bool for nibabel headers is unsupported, so use uint8 instead: # `nibabel.spatialimages.HeaderDataError: data dtype "bool" not supported` dtype_data = self.data.dtype - if dtype_data == bool: + if dtype_data == bool: # noqa: E721 -- numpy dtype equality against a scalar type, not a type() comparison dtype_data = np.uint8 dtype_header = self.hdr.get_data_dtype() if dtype_header != dtype_data: - logger.warning(f"Image header specifies datatype '{dtype_header}', but array is of type " - f"'{dtype_data}'. Header metadata will be overwritten to use '{dtype_data}'.") + logger.warning( + "Image header specifies datatype '%s', but array is of type '%s'. Header metadata will be overwritten to use '%s'.", + dtype_header, + dtype_data, + dtype_data, + ) self.hdr.set_data_dtype(dtype_data) - + def save(self, path=None, dtype=None, verbose=1, mutable=False): """ Write an image in a nifti file @@ -247,8 +256,7 @@ def save(self, path=None, dtype=None, verbose=1, mutable=False): if self.absolutepath: # Use the original filename, but save to the directory specified by `path` path = os.path.join(os.path.abspath(path), os.path.basename(self.absolutepath)) else: - raise ValueError("Don't know where to save the image (path parameter is dir, but absolutepath is " - "missing)") + raise ValueError("Don't know where to save the image (path parameter is dir, but absolutepath is missing)") # Case 3: `path` points to a file (or a *nonexistent* directory) so use its value as-is # (We're okay with letting nonexistent directories slip through, because it's difficult to distinguish # between nonexistent directories and nonexistent files. Plus, `nibabel` will catch any further errors.) @@ -258,11 +266,11 @@ def save(self, path=None, dtype=None, verbose=1, mutable=False): if os.path.isfile(path) and verbose: logger.warning("File %s already exists. Will overwrite it.", path) if os.path.isabs(path): - logger.debug("Saving image to %s orientation %s shape %s", - path, self.orientation, self.data.shape) + logger.debug("Saving image to %s orientation %s shape %s", path, self.orientation, self.data.shape) else: - logger.debug("Saving image to %s (%s) orientation %s shape %s", - path, os.path.abspath(path), self.orientation, self.data.shape) + logger.debug( + "Saving image to %s (%s) orientation %s shape %s", path, os.path.abspath(path), self.orientation, self.data.shape + ) # Now that `path` has been set and log messages have been written, we can assign it to the image itself self.absolutepath = os.path.abspath(path) @@ -287,7 +295,7 @@ def save(self, path=None, dtype=None, verbose=1, mutable=False): return self -class SlicerOneAxis(object): +class SlicerOneAxis: """ Image slicer to use when you don't care about the 2D slice orientation, and don't want to specify them. @@ -300,7 +308,7 @@ class SlicerOneAxis(object): """ def __init__(self, im, axis="IS"): - opposite_character = {'L': 'R', 'R': 'L', 'A': 'P', 'P': 'A', 'I': 'S', 'S': 'I'} + opposite_character = {"L": "R", "R": "L", "A": "P", "P": "A", "I": "S", "S": "I"} axis_labels = "LRPAIS" if len(axis) != 2: raise ValueError() @@ -339,14 +347,15 @@ def __getitem__(self, idx): raise NotImplementedError() if idx >= self.nb_slices: - raise IndexError("I just have {} slices!".format(self.nb_slices)) + raise IndexError(f"I just have {self.nb_slices} slices!") if self.direction == -1: idx = self.nb_slices - 1 - idx return self.im.data[self._slice(idx)] -def get_dimension(im_file, verbose=1): + +def get_dimension(im_file, verbose=1): # noqa: ARG001 -- verbose kept for API parity with spinalcordtoolbox """ Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ @@ -414,7 +423,7 @@ def change_orientation(im_src, orientation, im_dst=None, inverse=False): # Update data by performing inversions and swaps # axes inversion (flip) - data = im_src_data[::inversion[0], ::inversion[1], ::inversion[2]] + data = im_src_data[:: inversion[0], :: inversion[1], :: inversion[2]] # axes manipulations (transpose) if perm == [1, 0, 2]: @@ -438,9 +447,7 @@ def change_orientation(im_src, orientation, im_dst=None, inverse=False): # Update header im_src_aff = im_src.hdr.get_best_affine() - aff = nib.orientations.inv_ornt_aff( - np.array((perm, inversion)).T, - im_src_data.shape) + aff = nib.orientations.inv_ornt_aff(np.array((perm, inversion)).T, im_src_data.shape) im_dst_aff = np.matmul(im_src_aff, aff) im_dst.header.set_qform(im_dst_aff) @@ -460,7 +467,7 @@ def _get_permutations(im_src_orientation, im_dst_orientation): :return: list of axes permutations and list of inversions to achieve an orientation change """ - opposite_character = {'L': 'R', 'R': 'L', 'A': 'P', 'P': 'A', 'I': 'S', 'S': 'I'} + opposite_character = {"L": "R", "R": "L", "A": "P", "P": "A", "I": "S", "S": "I"} perm = [0, 1, 2] inversion = [1, 1, 1] @@ -491,7 +498,7 @@ def orientation_string_nib2sct(s): :return: SCT reference space code from nibabel one """ - opposite_character = {'L': 'R', 'R': 'L', 'A': 'P', 'P': 'A', 'I': 'S', 'S': 'I'} + opposite_character = {"L": "R", "R": "L", "A": "P", "P": "A", "I": "S", "S": "I"} return "".join([opposite_character[x] for x in s]) @@ -533,12 +540,12 @@ def change_type(im_src, dtype, im_dst=None): max_in = np.nanmax(im_src.data) # find optimum type for the input image - if dtype in ('minimize', 'minimize_int'): + if dtype in ("minimize", "minimize_int"): # warning: does not take intensity resolution into account, neither complex voxels # check if voxel values are real or integer isInteger = True - if dtype == 'minimize': + if dtype == "minimize": for vox in im_src.data.flatten(): if int(vox) != vox: isInteger = False @@ -556,24 +563,22 @@ def change_type(im_src, dtype, im_dst=None): dtype = np.uint64 else: raise ValueError("Maximum value of the image is to big to be represented.") + elif max_in <= np.iinfo(np.int8).max and min_in >= np.iinfo(np.int8).min: + dtype = np.int8 + elif max_in <= np.iinfo(np.int16).max and min_in >= np.iinfo(np.int16).min: + dtype = np.int16 + elif max_in <= np.iinfo(np.int32).max and min_in >= np.iinfo(np.int32).min: + dtype = np.int32 + elif max_in <= np.iinfo(np.int64).max and min_in >= np.iinfo(np.int64).min: + dtype = np.int64 else: - if max_in <= np.iinfo(np.int8).max and min_in >= np.iinfo(np.int8).min: - dtype = np.int8 - elif max_in <= np.iinfo(np.int16).max and min_in >= np.iinfo(np.int16).min: - dtype = np.int16 - elif max_in <= np.iinfo(np.int32).max and min_in >= np.iinfo(np.int32).min: - dtype = np.int32 - elif max_in <= np.iinfo(np.int64).max and min_in >= np.iinfo(np.int64).min: - dtype = np.int64 - else: - raise ValueError("Maximum value of the image is to big to be represented.") - else: - # if max_in <= np.finfo(np.float16).max and min_in >= np.finfo(np.float16).min: - # type = 'np.float16' # not supported by nibabel - if max_in <= np.finfo(np.float32).max and min_in >= np.finfo(np.float32).min: - dtype = np.float32 - elif max_in <= np.finfo(np.float64).max and min_in >= np.finfo(np.float64).min: - dtype = np.float64 + raise ValueError("Maximum value of the image is to big to be represented.") + # if max_in <= np.finfo(np.float16).max and min_in >= np.finfo(np.float16).min: + # type = 'np.float16' # not supported by nibabel + elif max_in <= np.finfo(np.float32).max and min_in >= np.finfo(np.float32).min: + dtype = np.float32 + elif max_in <= np.finfo(np.float64).max and min_in >= np.finfo(np.float64).min: + dtype = np.float64 dtype = to_dtype(dtype) else: @@ -588,7 +593,10 @@ def change_type(im_src, dtype, im_dst=None): if (min_in < min_out) or (max_in > max_out): # This condition is important for binary images since we do not want to scale them - logger.warning(f"To avoid intensity overflow due to convertion to +{dtype.name}+, intensity will be rescaled to the maximum quantization scale") + logger.warning( + "To avoid intensity overflow due to convertion to +%s+, intensity will be rescaled to the maximum quantization scale", + dtype.name, + ) # rescale intensity data_rescaled = im_src.data * (max_out - min_out) / (max_in - min_in) im_dst.data = data_rescaled - (data_rescaled.min() - min_out) @@ -612,15 +620,14 @@ def to_dtype(dtype): if dtype is None: return None - if isinstance(dtype, type): - if isinstance(dtype(0).dtype, np.dtype): - return dtype(0).dtype + if isinstance(dtype, type) and isinstance(dtype(0).dtype, np.dtype): + return dtype(0).dtype if isinstance(dtype, np.dtype): return dtype if isinstance(dtype, str): return np.dtype(dtype) - raise TypeError("data type {}: {} not understood".format(dtype.__class__, dtype)) + raise TypeError(f"data type {dtype.__class__}: {dtype} not understood") def zeros_like(img, dtype=None): @@ -671,10 +678,10 @@ def find_zmin_zmax(im, threshold=0.1): # Make sure image is not empty if not np.any(slicer): - logger.error('Input image is empty') + logger.error("Input image is empty") # Iterate from bottom to top until we find data - for zmin in range(0, len(slicer)): + for zmin in range(len(slicer)): if np.any(slicer[zmin] > threshold): break @@ -686,10 +693,11 @@ def find_zmin_zmax(im, threshold=0.1): return zmin, zmax -def resample_nib(image, new_size=None, new_size_type=None, image_dest=None, interpolation='linear', mode='nearest', - preserve_codes=False, verbose=True): +def resample_nib( + image, new_size=None, new_size_type=None, image_dest=None, interpolation="linear", mode="nearest", preserve_codes=False, verbose=True +): """ - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/blob/master/spinalcordtoolbox/resampling.py + Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/blob/master/spinalcordtoolbox/resampling.py Resample a nibabel or Image object based on a specified resampling factor. Can deal with 2d, 3d or 4d image objects. @@ -715,7 +723,7 @@ def resample_nib(image, new_size=None, new_size_type=None, image_dest=None, inte """ # set interpolation method - dict_interp = {'nn': 0, 'linear': 1, 'spline': 2} + dict_interp = {"nn": 0, "linear": 1, "spline": 2} # If input is an Image object, create nibabel object from it if isinstance(image, nib.nifti1.Nifti1Image): @@ -723,16 +731,17 @@ def resample_nib(image, new_size=None, new_size_type=None, image_dest=None, inte elif isinstance(image, Image): img = nib.nifti1.Nifti1Image(image.data, image.hdr.get_best_affine(), image.hdr) else: - raise TypeError(f'Invalid image type: {type(image)}') + raise TypeError(f"Invalid image type: {type(image)}") # convert to floating point if we're doing arithmetic interpolation - if interpolation != 'nn' and img.get_data_dtype().kind in 'biu': + if interpolation != "nn" and img.get_data_dtype().kind in "biu": original_dtype = img.get_data_dtype() img = nib.nifti1.Nifti1Image(img.get_fdata(), img.header.get_best_affine(), img.header) img.set_data_dtype(img.dataobj.dtype) if verbose: - logger.warning("Converting image from type '%s' to type '%s' for %s interpolation", - original_dtype, img.get_data_dtype(), interpolation) + logger.warning( + "Converting image from type '%s' to type '%s' for %s interpolation", original_dtype, img.get_data_dtype(), interpolation + ) if image_dest is None: # Get dimensions of data @@ -746,15 +755,15 @@ def resample_nib(image, new_size=None, new_size_type=None, image_dest=None, inte ndim_r = 3 # compute new shape based on specific resampling method - if new_size_type == 'vox': + if new_size_type == "vox": shape_r = tuple([int(new_size[i]) for i in range(ndim_r)]) - elif new_size_type == 'factor': + elif new_size_type == "factor": if len(new_size) == 1: # isotropic resampling new_size = tuple([new_size[0] for i in range(ndim_r)]) # compute new shape as: shape_r = shape * f shape_r = tuple([int(np.round(shape[i] * float(new_size[i]))) for i in range(ndim_r)]) - elif new_size_type == 'mm': + elif new_size_type == "mm": if len(new_size) == 1: # isotropic resampling new_size = tuple([new_size[0] for i in range(ndim_r)]) @@ -765,37 +774,37 @@ def resample_nib(image, new_size=None, new_size_type=None, image_dest=None, inte if img.ndim == 4: # Copy over 't' dim (i.e. number of volumes should be unaffected) - shape_r = shape_r + (shape[3],) + shape_r = (*shape_r, shape[3]) # Generate 3d affine transformation: R affine = img.affine[:4, :4] affine[3, :] = np.array([0, 0, 0, 1]) # satisfy to nifti convention. Otherwise it grabs the temporal - logger.debug('Affine matrix: \n' + str(affine)) + logger.debug("Affine matrix: \n%s", affine) R = np.eye(4) for i in range(3): try: R[i, i] = img.shape[i] / float(shape_r[i]) except ZeroDivisionError: - raise ZeroDivisionError("Destination size is zero for dimension {}. You are trying to resample to an " - "unrealistic dimension. Check your NIFTI pixdim values to make sure they are " - "not corrupted.".format(i)) + raise ZeroDivisionError( + f"Destination size is zero for dimension {i}. You are trying to resample to an " + "unrealistic dimension. Check your NIFTI pixdim values to make sure they are " + "not corrupted." + ) from None affine_r = np.dot(affine, R) reference = (shape_r, affine_r) # If reference is provided + elif isinstance(image_dest, nib.nifti1.Nifti1Image): + reference = image_dest + elif isinstance(image_dest, Image): + reference = nib.nifti1.Nifti1Image(image_dest.data, affine=image_dest.hdr.get_best_affine(), header=image_dest.hdr) else: - if isinstance(image_dest, nib.nifti1.Nifti1Image): - reference = image_dest - elif isinstance(image_dest, Image): - reference = nib.nifti1.Nifti1Image(image_dest.data, affine=image_dest.hdr.get_best_affine(), header=image_dest.hdr) - else: - raise TypeError(f'Invalid image type: {type(image_dest)}') + raise TypeError(f"Invalid image type: {type(image_dest)}") if img.ndim == 3: # we use mode 'nearest' to overcome issue #2453 - img_r = resample_from_to( - img, to_vox_map=reference, order=dict_interp[interpolation], mode=mode, cval=0.0, out_class=None) + img_r = resample_from_to(img, to_vox_map=reference, order=dict_interp[interpolation], mode=mode, cval=0.0, out_class=None) elif img.ndim == 4: # TODO: Cover img_dest with 4D volumes @@ -807,22 +816,22 @@ def resample_nib(image, new_size=None, new_size_type=None, image_dest=None, inte data3d = np.asanyarray(img.dataobj)[..., it] nii_tmp = nib.nifti1.Nifti1Image(data3d, affine, dtype=data3d.dtype) img3d_r = resample_from_to( - nii_tmp, to_vox_map=(shape_r[:-1], affine_r), order=dict_interp[interpolation], mode=mode, - cval=0.0, out_class=None) + nii_tmp, to_vox_map=(shape_r[:-1], affine_r), order=dict_interp[interpolation], mode=mode, cval=0.0, out_class=None + ) data4d[..., it] = np.asanyarray(img3d_r.dataobj) # Create 4d nibabel Image img_r = nib.nifti1.Nifti1Image(data4d, affine_r) # Can't be int64 (#4408) # Copy over the TR parameter from original 4D image (otherwise it will be incorrectly set to 1) - img_r.header.set_zooms(list(img_r.header.get_zooms()[0:3]) + [img.header.get_zooms()[3]]) + img_r.header.set_zooms([*list(img_r.header.get_zooms()[0:3]), img.header.get_zooms()[3]]) # preserve the codes from the original image, which will otherwise get overwritten with 0/2 if preserve_codes: - img_r.header['qform_code'] = img.header['qform_code'] - img_r.header['sform_code'] = img.header['sform_code'] + img_r.header["qform_code"] = img.header["qform_code"] + img_r.header["sform_code"] = img.header["sform_code"] # Convert back to proper type if isinstance(image, nib.nifti1.Nifti1Image): return img_r else: assert isinstance(image, Image) # already checked at the start of the function - return Image(np.asanyarray(img_r.dataobj), hdr=img_r.header, orientation=image.orientation, dim=img_r.header.get_data_shape()) \ No newline at end of file + return Image(np.asanyarray(img_r.dataobj), hdr=img_r.header, orientation=image.orientation, dim=img_r.header.get_data_shape()) diff --git a/auglab/utils/utils.py b/auglab/utils/utils.py index 1edc255..319c109 100644 --- a/auglab/utils/utils.py +++ b/auglab/utils/utils.py @@ -1,12 +1,13 @@ -import os -from progress.bar import Bar -import json import argparse +import json +import os + import numpy as np +from progress.bar import Bar -def fetch_image_config(config_data, split='TRAINING'): - ''' +def fetch_image_config(config_data, split="TRAINING"): + """ :param config_data: Config dict where every label used for TRAINING, VALIDATION and/or TESTING has its path specified :param split: Split of the data needed in the config file ('TRAINING', 'VALIDATION', 'TESTING'). :return: out_list: list of dictionary with image and label paths (like monai load_decathlon_datalist) @@ -14,85 +15,89 @@ def fetch_image_config(config_data, split='TRAINING'): {'image': '/workspace/data/chest_19.nii.gz', 'label': '/workspace/data/chest_19_label.nii.gz'}, {'image': '/workspace/data/chest_31.nii.gz', 'label': '/workspace/data/chest_31_label.nii.gz'} ] - ''' + """ # Check config type to ensure that labels paths are specified and not images - if config_data['TYPE'] != 'LABEL': - raise ValueError('TYPE error: Type LABEL not detected') - + if config_data["TYPE"] != "LABEL": + raise ValueError("TYPE error: Type LABEL not detected") + # Get file paths based on split dict_list = config_data[split] - + # Init progression bar - bar = Bar(f'Load {split} data', max=len(dict_list)) - + bar = Bar(f"Load {split} data", max=len(dict_list)) + err = [] out_list = [] for di in dict_list: - input_img_path = os.path.join(config_data['DATASETS_PATH'], di['IMAGE']) - input_seg_path = os.path.join(config_data['DATASETS_PATH'], di['LABEL']) + input_img_path = os.path.join(config_data["DATASETS_PATH"], di["IMAGE"]) + input_seg_path = os.path.join(config_data["DATASETS_PATH"], di["LABEL"]) if not os.path.exists(input_img_path): - err.append([input_img_path, 'path error']) + err.append([input_img_path, "path error"]) else: - out_list.append({'image':os.path.abspath(input_img_path), 'segmentation':os.path.abspath(input_seg_path)}) + out_list.append({"image": os.path.abspath(input_img_path), "segmentation": os.path.abspath(input_seg_path)}) # Plot progress - bar.suffix = f'{dict_list.index(di)+1}/{len(dict_list)}' + bar.suffix = f"{dict_list.index(di) + 1}/{len(dict_list)}" bar.next() bar.finish() return out_list, err + def config2parser(config_path): - ''' - Create a parser object from a json file - ''' + """ + Create a parser object from a json file + """ # Read json file and create a dictionary - with open(config_path, "r") as file: + with open(config_path) as file: config_dict = json.load(file) return argparse.Namespace(**config_dict) def parser2config(args, path_out): - ''' + """ Extract the parameters from an input parser to create a config json file :param args: parser arguments :param path_out: path out of the config file - ''' + """ # Check if path_out exists or create it if not os.path.exists(os.path.dirname(path_out)): os.makedirs(os.path.dirname(path_out)) # Serializing json json_object = json.dumps(vars(args), indent=4) - + # Inform user if os.path.exists(path_out): print(f"The config file {path_out} with all the training parameters was updated") else: print(f"The config file {path_out} with all the training parameters was created") - + # Write json file with open(path_out, "w") as outfile: outfile.write(json_object) + def tuple_type_int(strings): - ''' + """ Copied from https://stackoverflow.com/questions/33564246/passing-a-tuple-as-command-line-argument - ''' + """ strings = strings.replace("(", "").replace(")", "") mapped_int = map(int, strings.split(",")) return tuple(mapped_int) + def tuple_type_float(strings): - ''' + """ Copied from https://stackoverflow.com/questions/33564246/passing-a-tuple-as-command-line-argument - ''' + """ strings = strings.replace("(", "").replace(")", "") mapped_float = map(float, strings.split(",")) return tuple(mapped_float) + def tuple2string(t): - return str(t).replace(' ', '').replace('(','').replace(')','').replace(',','-') + return str(t).replace(" ", "").replace("(", "").replace(")", "").replace(",", "-") def adjust_learning_rate(optimizer, lr, gamma): @@ -102,9 +107,10 @@ def adjust_learning_rate(optimizer, lr, gamma): """ lr *= gamma for param_group in optimizer.param_groups: - param_group['lr'] = lr + param_group["lr"] = lr return lr + def compute_dsc(gt_mask, pred_mask, sigmoid=False): """ :param gt_mask: Ground truth mask used as the reference @@ -115,7 +121,7 @@ def compute_dsc(gt_mask, pred_mask, sigmoid=False): """ if sigmoid: pred_mask = sig_fn(pred_mask) - numerator = 2 * (gt_mask*pred_mask).sum() + numerator = 2 * (gt_mask * pred_mask).sum() denominator = gt_mask.sum() + pred_mask.sum() if denominator == 0: # Both ground truth and prediction are empty @@ -123,8 +129,10 @@ def compute_dsc(gt_mask, pred_mask, sigmoid=False): else: return numerator / denominator + def sig_fn(z): - return 1/(1 + np.exp(-z)) + return 1 / (1 + np.exp(-z)) + def get_validation_image(in_img, target_img, pred_img, sigmoid=False): in_img = in_img.data.cpu().numpy() @@ -143,20 +151,20 @@ def get_validation_image(in_img, target_img, pred_img, sigmoid=False): shape = x.shape # Extract middle slice - x = x[shape[0]//2,:,:] - y = y[shape[0]//2,:,:] - y_pred = y_pred[shape[0]//2,:,:] + x = x[shape[0] // 2, :, :] + y = y[shape[0] // 2, :, :] + y_pred = y_pred[shape[0] // 2, :, :] # Normalize intensity - x = normalize(x)*255 - y = normalize(y)*255 - y_pred = normalize(y_pred)*255 + x = normalize(x) * 255 + y = normalize(y) * 255 + y_pred = normalize(y_pred) * 255 # Regroup batch in_all.append(x) target_all.append(y) pred_all.append(y_pred) - + # Regroup batch into 1 array in_line_arr = np.concatenate(np.array(in_all), axis=1) target_line_arr = np.concatenate(np.array(target_all), axis=1) @@ -164,14 +172,15 @@ def get_validation_image(in_img, target_img, pred_img, sigmoid=False): # Regroup image/target/pred into 1 array img_result = np.concatenate((in_line_arr, target_line_arr, pred_line_arr), axis=0) - + return img_result, target_line_arr, pred_line_arr + def normalize(arr): - ''' + """ Normalize image using percentiles - ''' + """ # Use 10th percentile p10 = np.percentile(arr, 10) p90 = np.percentile(arr, 90) - return ((arr - p10) / (p90 - p10 + 0.00001)) \ No newline at end of file + return (arr - p10) / (p90 - p10 + 0.00001) diff --git a/pyproject.toml b/pyproject.toml index d600529..b406d81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,23 @@ -[project] +[tool.poetry] name = "auglab" -version = "20260109" -requires-python = ">=3.10" +# Placeholder only. The real version is substituted at build time by +# poetry-dynamic-versioning from the git tag -- same arrangement as +# Hendrik-code/TPTBox. Do not bump this by hand. +version = "0.0.0" description = "AugLab investigates the influence of different data augmentation strategies on MRI training performance." readme = "README.md" authors = [ - { name = "Nathan Molinier", email = "nathan.molinier@polymtl.ca"}, - { name = "Hendrik Möller"}, + "Nathan Molinier ", ] +# poetry-core only ever writes the FIRST author into the wheel metadata, so a +# second entry in `authors` would silently vanish from the published package +# (setuptools used to emit both). Listing Hendrik as a maintainer keeps the +# attribution in the metadata as `Maintainer:`. +maintainers = [ + "Hendrik Möller", +] +homepage = "https://github.com/neuropoly/AugLab" +repository = "https://github.com/neuropoly/AugLab" classifiers = [ "Development Status :: 5 - Production/Stable", "Intended Audience :: Science/Research", @@ -17,44 +27,232 @@ classifiers = [ "Topic :: Scientific/Engineering :: Medical Science Apps.", ] keywords = [ - 'deep learning', - 'image segmentation', - 'nnU-Net', - 'nnunet', - 'magnetic resonance imaging', - 'online augmentation', - 'offline augmentation', - 'mri' -] -dependencies = [ - "batchgeneratorsv2", - "kornia", - "torchio" -] - -[project.optional-dependencies] -nnunetv2 = ["nnunetv2"] -all = [ - "monai[all]", - "progress", - "numpy", - "tqdm", - "wandb", + "deep learning", + "image segmentation", + "nnU-Net", + "nnunet", + "magnetic resonance imaging", + "online augmentation", + "offline augmentation", + "mri", ] +# auglab has no __init__.py anywhere, so it is a PEP 420 namespace package. +# poetry-core still walks the tree correctly; the CI build job and +# unit_tests/test_packaging.py assert the wheel really contains every module. +packages = [{ include = "auglab" }] +# The config JSONs are package data, not code, so they need listing explicitly. +# A wheel without them is broken: auglab resolves its default config through +# importlib.resources at runtime. +include = [{ path = "auglab/configs/**/*.json", format = ["sdist", "wheel"] }] -[project.scripts] -auglab_add_nnunettrainer = "auglab.add_trainer:main" +[tool.poetry.dependencies] +python = ">=3.10" +batchgenerators = "*" +batchgeneratorsv2 = "*" +# auglab subclasses kornia's private augmentation internals +# (_AugmentationBase, RigidAffineAugmentationBase3D, augmentation.container.ops, +# _adapted_rsampling, _tuple_range_reader). Those move between minor releases, +# so the range is capped and both ends are exercised by the kornia-compat job +# in .github/workflows/tests.yml. Verified working: 0.7.3 - 0.8.3. +kornia = ">=0.7.3,<0.9" +nibabel = "*" +numpy = "*" +progress = "*" +scipy = "*" +torchio = "*" +torchvision = "*" -[project.urls] -homepage = "https://github.com/neuropoly/AugLab" -repository = "https://github.com/neuropoly/AugLab" +# Optional dependencies are declared as extras rather than poetry groups so +# that plain `pip install -e ".[dev]"` keeps working. Nothing here requires the +# poetry CLI or a poetry.lock -- only the build backend is poetry's. +monai = { version = "*", extras = ["all"], optional = true } +tqdm = { version = "*", optional = true } +wandb = { version = "*", optional = true } +nnunetv2 = { version = "*", optional = true } +build = { version = "*", optional = true } +coverage = { version = ">=7", optional = true } +pre-commit = { version = "*", optional = true } +pytest = { version = ">=8", optional = true } +pytest-cov = { version = "*", optional = true } +# Pinned to the same version as the ruff-pre-commit rev in +# .pre-commit-config.yaml. Ruff adds rules between releases, so an unpinned +# local ruff will disagree with the one CI runs. +ruff = { version = "==0.16.1", optional = true } +twine = { version = "*", optional = true } + +[tool.poetry.extras] +nnunetv2 = ["nnunetv2"] +all = ["monai", "tqdm", "wandb"] +dev = ["build", "coverage", "pre-commit", "pytest", "pytest-cov", "ruff", "twine"] + +[tool.poetry.scripts] +auglab_add_nnunettrainer = "auglab.add_trainer:main" [build-system] -requires = ["pip>=23", "setuptools>=67"] -build-backend = "setuptools.build_meta" +requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"] +build-backend = "poetry_dynamic_versioning.backend" + +[tool.poetry-dynamic-versioning] +enable = true +# Strip an optional r/v/release- prefix; the rest becomes the version. +# The default pattern rejects the "r" prefix already in use (r20260615). +# Handles release and pre-release tags (r20260801, v1.2.3, v1.0.0rc1, +# v2.0.0-beta1). A ".post" tag is not supported and fails the build loudly +# rather than silently dropping the suffix. +pattern = "^(?:[rvV]|release[-_])?(?P\\d+(\\.\\d+)*)(?:[-.]?(?P[a-zA-Z]+)\\.?(?P\\d+)?)?" +# On a tag, the version is exactly the tag. Between tags, bump and mark .dev so +# the build still sorts AFTER the release it follows -- a plain "{base}.devN" +# would sort *before* it. No local version (+g), because PyPI rejects those. +# +# stage/revision must be carried through, or a pre-release tag like v1.0.0rc1 +# would build as plain "1.0.0" and collide with the real 1.0.0 release. +format-jinja = """ +{%- if distance == 0 -%} +{{ serialize_pep440(base, stage=stage, revision=revision) }} +{%- else -%} +{{ serialize_pep440(bump_version(base), stage=stage, revision=revision, dev=distance) }} +{%- endif -%} +""" -[tool.setuptools] -include-package-data = true +[tool.ruff] +line-length = 140 +indent-width = 4 +target-version = "py310" +exclude = [ + "*.ipynb", + ".eggs", + ".venv", + "*.egg-info", + "build", + "dist", + "venv", +] + +[tool.ruff.lint] +select = [ + "A", + "ARG", + "B", + "BLE", + "C4", + "E", + "F", + "FLY", + "FURB", + "G", + "I", + "ICN", + "INT", + "N", + "NPY", + "PERF", + "PGH", + "PIE", + "PL", + "RUF", + "SIM", + "TID", + "TRY", + "UP", + "W", +] -[tool.setuptools.package-data] -'auglab' = ['data/**.json'] +ignore = [ + "A001", # builtin shadowed by a variable (pervasive: `input`, `filter`, `type`) + "A002", # builtin shadowed by an argument (same, and part of the public API) + "ARG002", # unused method argument (required by the batchgenerators/kornia interfaces) + "ARG004", # unused static method argument (same reason) + "B905", # zip() without strict= + "BLE001", # blind `except Exception` + "E501", # line too long (the formatter handles what it can) + "E741", # ambiguous variable name (`l` for label is idiomatic here) + "F811", # redefinition (triggered by the __main__ demo blocks) + "FURB171", # membership test against a single-item container + "N801", # class name not CapWords (transform names mirror nnU-Net's) + "N802", # function name not lowercase + "N803", # argument name not lowercase (tensors are `X`, `Y`) + "N806", # variable in function should be lowercase (same) + "N812", # lowercase imported as non-lowercase (`import torch.nn.functional as F`) + "N999", # invalid module name (`nnUNetTrainerDAExt`) + "NPY002", # legacy np.random (reproducibility of published experiments) + "PERF203", # try/except in a loop + "PGH003", # blanket type: ignore + "PLC0415", # import not at top of file (deliberate, for optional heavy deps) + "PLR0911", # too many return statements + "PLR0912", # too many branches + "PLR0913", # too many arguments (augmentation configs are wide by nature) + "PLR0915", # too many statements + "PLR0917", # too many positional arguments + "PLR2004", # magic value in comparison + "PLW0108", # unnecessary lambda + "RUF002", # ambiguous unicode in docstrings + "RUF003", # ambiguous unicode in comments + "RUF059", # unused unpacked variable + "RUF100", # unused noqa + "SIM105", # contextlib.suppress instead of try/except/pass + "SIM118", # `key in dict` instead of `key in dict.keys()` + "TRY003", # long message outside the exception class + "UP007", # Optional/Union instead of `X | Y` +] + +# Allow fix for all enabled rules (when `--fix` is provided). +fixable = ["ALL"] +unfixable = [] + +# Allow unused variables when underscore-prefixed. +dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" + +[tool.ruff.lint.per-file-ignores] +# Re-exports are the point of an __init__. +"**/__init__.py" = ["F401"] +# Tests are not an importable package and assert against literals. +"unit_tests/**" = ["INP001", "PLR2004"] +# Standalone scripts, not part of the installed package. +"scripts/**" = ["INP001"] +# The trainers mirror nnU-Net's upstream signatures verbatim, including its +# `device: torch.device = torch.device("cuda")` default. Diverging would break +# the drop-in contract with nnUNetTrainer. +"auglab/trainers/**" = ["B008"] + +[tool.ruff.lint.mccabe] +max-complexity = 20 + +[tool.ruff.format] +# Like Black, use double quotes for strings. +quote-style = "double" + +# Like Black, indent with spaces, rather than tabs. +indent-style = "space" + +# Like Black, respect magic trailing commas. +skip-magic-trailing-comma = false + +# Enable reformatting of code snippets in docstrings. +docstring-code-format = true + +# Like Black, automatically detect the appropriate line ending. +line-ending = "auto" + +[tool.pytest.ini_options] +testpaths = ["unit_tests"] +filterwarnings = [ + "ignore::DeprecationWarning", + "ignore::FutureWarning", + "ignore::UserWarning", +] +markers = [ + "slow: marks tests as slow (deselect with '-m \"not slow\"')", +] + +[tool.coverage.run] +source = ["auglab"] + +[tool.coverage.report] +exclude_lines = [ + "pragma: no cover", + "if __name__ == .__main__.:", + "def __repr__", + "def __str__", + "if TYPE_CHECKING:", + "raise NotImplementedError", +] diff --git a/scripts/generate_augmentations.py b/scripts/generate_augmentations.py index c9c0446..39c4f00 100644 --- a/scripts/generate_augmentations.py +++ b/scripts/generate_augmentations.py @@ -1,59 +1,71 @@ -import argparse, textwrap +import argparse import json import multiprocessing as mp +import textwrap +import warnings from functools import partial -from tqdm.contrib.concurrent import process_map from pathlib import Path + import numpy as np import torch -import warnings +from tqdm.contrib.concurrent import process_map -from auglab.utils.utils import fetch_image_config from auglab.transforms.cpu.transforms import AugTransforms from auglab.utils.image import Image, resample_nib, zeros_like +from auglab.utils.utils import fetch_image_config warnings.filterwarnings("ignore") rs = np.random.RandomState() + def main(): # Description and arguments parser = argparse.ArgumentParser( - description=' '.join(f''' - This script processes NIfTI (Neuroimaging Informatics Technology Initiative) image and segmentation files. - It apply transformation on the image and the segmentation to make augmented image. Useful if augmentations cannot be performed on the fly during training. - Based on https://github.com/neuropoly/totalspineseg/blob/b4da40840ad618498be3ec02564d4c5f5fa5c8aa/totalspineseg/utils/augment.py - '''.split()), - formatter_class=argparse.RawTextHelpFormatter + description="This script processes NIfTI (Neuroimaging Informatics Technology Initiative) image and segmentation files. It apply transformation on the image and the segmentation to make augmented image. Useful if augmentations cannot be performed on the fly during training. Based on https://github.com/neuropoly/totalspineseg/blob/b4da40840ad618498be3ec02564d4c5f5fa5c8aa/totalspineseg/utils/augment.py", + formatter_class=argparse.RawTextHelpFormatter, ) parser.add_argument( - '--data', '-d', type=Path, required=True, - help='Data config JSON file containing the IMAGE and LABEL paths of the files used. Only TRAINING will be augmented. See example in auglab.configs.data (required).' + "--data", + "-d", + type=Path, + required=True, + help="Data config JSON file containing the IMAGE and LABEL paths of the files used. Only TRAINING will be augmented. See example in auglab.configs.data (required).", ) parser.add_argument( - '--ofolder', '-o', type=Path, required=True, - help='The folder where output augmented images will be saved with _a1, _a2 etc. suffixes (required).' + "--ofolder", + "-o", + type=Path, + required=True, + help="The folder where output augmented images will be saved with _a1, _a2 etc. suffixes (required).", ) parser.add_argument( - '--transforms', '-t', type=Path, required=True, - help='Transforms config JSON file containing the parameters of the transformations. See example in auglab.configs (required).' + "--transforms", + "-t", + type=Path, + required=True, + help="Transforms config JSON file containing the parameters of the transformations. See example in auglab.configs (required).", ) parser.add_argument( - '--augmentations-per-image', '-n', type=int, default=5, - help='Number of augmentation images to generate. Default is 5.' + "--augmentations-per-image", "-n", type=int, default=5, help="Number of augmentation images to generate. Default is 5." ) parser.add_argument( - '--overwrite', '-r', action="store_true", default=False, - help='If provided, overwrite existing output files, defaults to false (Do not overwrite).' + "--overwrite", + "-r", + action="store_true", + default=False, + help="If provided, overwrite existing output files, defaults to false (Do not overwrite).", ) parser.add_argument( - '--max-workers', '-w', type=int, default=mp.cpu_count(), - help='Max worker to run in parallel proccess, defaults to multiprocessing.cpu_count().' + "--max-workers", + "-w", + type=int, + default=mp.cpu_count(), + help="Max worker to run in parallel proccess, defaults to multiprocessing.cpu_count().", ) parser.add_argument( - '--quiet', '-q', action="store_true", default=False, - help='Do not display inputs and progress bar, defaults to false (display).' + "--quiet", "-q", action="store_true", default=False, help="Do not display inputs and progress bar, defaults to false (display)." ) # Parse the command-line arguments @@ -70,7 +82,8 @@ def main(): # Print the argument values if not quiet if not quiet: - print(textwrap.dedent(f''' + print( + textwrap.dedent(f""" Running {Path(__file__).stem} with the following params: data_json_path = {data_json_path} transforms_json_path = {transforms_json_path} @@ -79,7 +92,8 @@ def main(): overwrite = {overwrite} max_workers = {max_workers} quiet = {quiet} - ''')) + """) + ) augment_mp( data_json_path=data_json_path, @@ -91,35 +105,41 @@ def main(): quiet=quiet, ) + def augment_mp( - data_json_path, - transforms_json_path, - ofolder, - augmentations_per_image=5, - overwrite=False, - max_workers=mp.cpu_count(), - quiet=False, - ): - ''' + data_json_path, + transforms_json_path, + ofolder, + augmentations_per_image=5, + overwrite=False, + max_workers=None, + quiet=False, +): + """ Wrapper function to handle multiprocessing. - ''' + """ + # Resolved here rather than in the signature so the default reflects the + # machine running the call, not the machine that imported the module. + if max_workers is None: + max_workers = mp.cpu_count() + # Convert to Path object data_json_path = Path(data_json_path) transforms_json_path = Path(transforms_json_path) ofolder = Path(ofolder) # Load data config - with open(str(data_json_path), "r") as f: + with open(str(data_json_path)) as f: data_config = json.load(f) - + data_list, _ = fetch_image_config( config_data=data_config, - split='TRAINING', + split="TRAINING", ) # Init transforms if not transforms_json_path.is_file(): - print(f'Error: {str(transforms_json_path)}, Transforms config file not found') + print(f"Error: {transforms_json_path!s}, Transforms config file not found") return process_map( @@ -136,31 +156,32 @@ def augment_mp( disable=quiet, ) + def augment( - data_dict, - augmentations_per_image, - train_transforms_path, - ofolder, - overwrite=False, - ): - ''' + data_dict, + augmentations_per_image, + train_transforms_path, + ofolder, + overwrite=False, +): + """ Augmentation function. - ''' + """ # Load transforms train_transforms = AugTransforms(json_path=str(train_transforms_path)) # Create PATH objects - img_path = Path(data_dict['image']) - seg_path = Path(data_dict['segmentation']) + img_path = Path(data_dict["image"]) + seg_path = Path(data_dict["segmentation"]) # Load images - img = Image(str(img_path)).change_orientation('RPI') # RPI- == LAS+ - seg = Image(str(seg_path)).change_orientation('RPI') + img = Image(str(img_path)).change_orientation("RPI") # RPI- == LAS+ + seg = Image(str(seg_path)).change_orientation("RPI") # Resample to 1mm isotropic pr = 1.0 - img = resample_nib(img, new_size=[pr, pr, pr], new_size_type='mm', interpolation='spline', verbose=False) - seg = resample_nib(seg, new_size=[pr, pr, pr], new_size_type='mm', interpolation='nn', verbose=False) + img = resample_nib(img, new_size=[pr, pr, pr], new_size_type="mm", interpolation="spline", verbose=False) + seg = resample_nib(seg, new_size=[pr, pr, pr], new_size_type="mm", interpolation="nn", verbose=False) # Normalize image using mean and std img.data = (img.data - img.data.mean()) / img.data.std() @@ -172,20 +193,20 @@ def augment( # Create augmentations for i in range(augmentations_per_image): # Create output path - output_image_path = Path(ofolder) / "img" / f"{img_path.name.replace('.nii.gz', '')}_a{i+1}.nii.gz" - output_seg_path = Path(ofolder) / "seg" / f"{seg_path.name.replace('.nii.gz', '')}_a{i+1}.nii.gz" + output_image_path = Path(ofolder) / "img" / f"{img_path.name.replace('.nii.gz', '')}_a{i + 1}.nii.gz" + output_seg_path = Path(ofolder) / "seg" / f"{seg_path.name.replace('.nii.gz', '')}_a{i + 1}.nii.gz" # Generate augmentation if not overwrite and (output_image_path.exists() or output_seg_path.exists()): continue - + # Transform data - tensor_dict = train_transforms({'image': img_tensor.detach().clone(), 'segmentation': seg_tensor.detach().clone()}) - + tensor_dict = train_transforms({"image": img_tensor.detach().clone(), "segmentation": seg_tensor.detach().clone()}) + img_out = zeros_like(img) - img_out.data = tensor_dict['image'].squeeze(0).numpy() + img_out.data = tensor_dict["image"].squeeze(0).numpy() seg_out = zeros_like(seg) - seg_out.data = tensor_dict['segmentation'].squeeze(0).numpy() + seg_out.data = tensor_dict["segmentation"].squeeze(0).numpy() # Save augmented data if not output_image_path.parent.exists(): @@ -195,8 +216,8 @@ def augment( img_out.save(output_image_path) seg_out.save(output_seg_path) -if __name__ == '__main__': +if __name__ == "__main__": main() # augment_mp( # data_json_path="auglab/configs/data/data.json", @@ -206,4 +227,4 @@ def augment( # overwrite=True, # max_workers=mp.cpu_count(), # quiet=False, - # ) \ No newline at end of file + # ) diff --git a/scripts/train_monai.py b/scripts/train_monai.py index 083432c..2a5db85 100644 --- a/scripts/train_monai.py +++ b/scripts/train_monai.py @@ -1,58 +1,94 @@ -''' +""" This script trains a segmentation network using MONAI with augmentations done on the fly during training. -''' +""" -import os -import numpy as np import argparse -import random -import json -import wandb import copy -from tqdm import tqdm import importlib +import json +import os +import random - +import numpy as np import torch -import torch.optim as optim - +import wandb from monai.data import DataLoader, Dataset -from monai.networks.nets import UNet, AttentionUnet, SwinUNETR, UNETR -from monai.losses import DiceCELoss, DiceFocalLoss, DiceLoss +from monai.losses import DiceFocalLoss +from monai.networks.nets import UNETR, AttentionUnet, SwinUNETR from monai.transforms import ( + Compose, + EnsureChannelFirstd, LoadImaged, + NormalizeIntensityd, Orientationd, - EnsureChannelFirstd, - Spacingd, - Compose, RandCropByPosNegLabeld, ResizeWithPadOrCropd, - NormalizeIntensityd, + Spacingd, ) +from torch import optim +from tqdm import tqdm + +from auglab import configs -# Import AugLab custom transforms -from auglab.utils.utils import fetch_image_config, parser2config, tuple_type_float, tuple_type_int, adjust_learning_rate, tuple2string, compute_dsc, get_validation_image -import auglab.configs as configs # Import AugLab GPU transforms 🐞 from auglab.transforms.gpu.transforms import AugTransformsGPU +# Import AugLab custom transforms +from auglab.utils.utils import ( + adjust_learning_rate, + compute_dsc, + fetch_image_config, + get_validation_image, + parser2config, + tuple2string, + tuple_type_float, + tuple_type_int, +) + + def get_parser(): # parse command line arguments - parser = argparse.ArgumentParser(description='Train monai network') - parser.add_argument('--config', required=True, help='Config JSON file where every label used for TRAINING, VALIDATION and TESTING has its path specified ~//config_data.json (Required)') - parser.add_argument('--transforms', default=None, help='Transforms JSON with GPU parameters default="auglab/configs/transform_params_gpu.json"') - parser.add_argument('--model', type=str, default='attunet', choices=['attunet', 'unetr', 'swinunetr'] , help='Model used for training. Options:["attunet", "unetr", "swinunetr"] (default="attunet")') - parser.add_argument('--batch-size', type=int, default=3, help='Training batch size (default=3).') - parser.add_argument('--nb-epochs', type=int, default=300, help='Number of training epochs (default=300).') - parser.add_argument('--start-epoch', type=int, default=0, help='Starting epoch (default=0).') - parser.add_argument('--schedule', type=tuple_type_float, default=tuple([0.3, 0.6, 0.9]), help='Fraction of the max epoch where the learning rate will be reduced of a factor gamma (default=(0.3, 0.6, 0.9)).') - parser.add_argument('--gamma', type=float, default=0.1, help='Factor used to reduce the learning rate (default=0.1)') - parser.add_argument('--channels', type=tuple_type_int, default=(32, 64, 128, 256), help='Channels if attunet selected (default=16,32,64,128,256)') - parser.add_argument('--patch-size', type=tuple_type_int, default=(64, 64, 64), help='Training patch size (default=(64, 64, 64)).') - parser.add_argument('--pixdim', type=tuple_type_float, default=(1, 1, 1), help='Training resolution in RSP orientation (default=(1, 1, 1)).') - parser.add_argument('--lr', default=1e-4, type=float, metavar='LR', help='Initial learning rate (default=1e-4)') - parser.add_argument('--weight-folder', type=str, default=os.path.abspath('weights/'), help='Folder where the weights will be stored and loaded. Will be created if does not exist. (default="src/ply/weights/3DGAN")') - parser.add_argument('--start-weights', type=str, default='', help='Path to the model weights used to start the training.') + parser = argparse.ArgumentParser(description="Train monai network") + parser.add_argument( + "--config", + required=True, + help="Config JSON file where every label used for TRAINING, VALIDATION and TESTING has its path specified ~//config_data.json (Required)", + ) + parser.add_argument( + "--transforms", default=None, help='Transforms JSON with GPU parameters default="auglab/configs/transform_params_gpu.json"' + ) + parser.add_argument( + "--model", + type=str, + default="attunet", + choices=["attunet", "unetr", "swinunetr"], + help='Model used for training. Options:["attunet", "unetr", "swinunetr"] (default="attunet")', + ) + parser.add_argument("--batch-size", type=int, default=3, help="Training batch size (default=3).") + parser.add_argument("--nb-epochs", type=int, default=300, help="Number of training epochs (default=300).") + parser.add_argument("--start-epoch", type=int, default=0, help="Starting epoch (default=0).") + parser.add_argument( + "--schedule", + type=tuple_type_float, + default=(0.3, 0.6, 0.9), + help="Fraction of the max epoch where the learning rate will be reduced of a factor gamma (default=(0.3, 0.6, 0.9)).", + ) + parser.add_argument("--gamma", type=float, default=0.1, help="Factor used to reduce the learning rate (default=0.1)") + parser.add_argument( + "--channels", type=tuple_type_int, default=(32, 64, 128, 256), help="Channels if attunet selected (default=16,32,64,128,256)" + ) + parser.add_argument("--patch-size", type=tuple_type_int, default=(64, 64, 64), help="Training patch size (default=(64, 64, 64)).") + parser.add_argument( + "--pixdim", type=tuple_type_float, default=(1, 1, 1), help="Training resolution in RSP orientation (default=(1, 1, 1))." + ) + parser.add_argument("--lr", default=1e-4, type=float, metavar="LR", help="Initial learning rate (default=1e-4)") + parser.add_argument( + "--weight-folder", + type=str, + default=os.path.abspath("weights/"), + help='Folder where the weights will be stored and loaded. Will be created if does not exist. (default="src/ply/weights/3DGAN")', + ) + parser.add_argument("--start-weights", type=str, default="", help="Path to the model weights used to start the training.") return parser @@ -65,52 +101,49 @@ def main(): ## Set seed seed = 42 - os.environ['PYTHONHASHSEED'] = str(seed) + os.environ["PYTHONHASHSEED"] = str(seed) # Torch RNG torch.manual_seed(seed) - if device.type=='cuda': + if device.type == "cuda": torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) # Python RNG np.random.seed(seed) - random.seed(seed) - + random.seed(seed) + # Load config data # Read json file and create a dictionary - with open(args.config, "r") as file: + with open(args.config) as file: config_data = json.load(file) - + # Load variables weight_folder = args.weight_folder - + # Save training config - model = args.model if args.model != 'attunet' else f'{args.model}{str(args.channels[-1])}' - json_name = f'config_{model}_pixdimRSP_{tuple2string(args.pixdim)}.json' + model = args.model if args.model != "attunet" else f"{args.model}{args.channels[-1]!s}" + json_name = f"config_{model}_pixdimRSP_{tuple2string(args.pixdim)}.json" saved_args = copy.copy(args) parser2config(saved_args, path_out=os.path.join(weight_folder, json_name)) # Create json file with training parameters # Create weights folder to store training weights if not os.path.exists(weight_folder): os.makedirs(weight_folder) - + # Load images for training and validation - print('loading images...') + print("loading images...") train_list, err_train = fetch_image_config( config_data=config_data, - split='TRAINING', + split="TRAINING", ) - + val_list, err_val = fetch_image_config( config_data=config_data, - split='VALIDATION', + split="VALIDATION", ) - + # Load AugLab transform parameters 🐞 configs_path = importlib.resources.files(configs) - if args.transforms is not None: - gpu_transforms_path = args.transforms - else: - gpu_transforms_path = configs_path / "transform_params_gpu.json" + gpu_transforms_path = args.transforms if args.transforms is not None else configs_path / "transform_params_gpu.json" # Compose MONAI and AugLab transforms pixdim = args.pixdim @@ -123,11 +156,19 @@ def main(): Spacingd( keys=["image", "segmentation"], pixdim=pixdim, - mode=(2, 'nearest'), # 2 for spline interpolation + mode=(2, "nearest"), # 2 for spline interpolation ), NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=False), - RandCropByPosNegLabeld(keys=["image", "segmentation"], label_key="segmentation", spatial_size=patch_size, pos=3, neg=1, num_samples=3, allow_smaller=True), - ResizeWithPadOrCropd(keys=["image", "segmentation"], spatial_size=patch_size) + RandCropByPosNegLabeld( + keys=["image", "segmentation"], + label_key="segmentation", + spatial_size=patch_size, + pos=3, + neg=1, + num_samples=3, + allow_smaller=True, + ), + ResizeWithPadOrCropd(keys=["image", "segmentation"], spatial_size=patch_size), ] ) val_transforms = Compose( @@ -138,10 +179,18 @@ def main(): Spacingd( keys=["image", "segmentation"], pixdim=pixdim, - mode=(2, 'nearest'), + mode=(2, "nearest"), ), NormalizeIntensityd(keys=["image"], nonzero=False, channel_wise=False), - RandCropByPosNegLabeld(keys=["image", "segmentation"], label_key="segmentation", spatial_size=patch_size, pos=3, neg=1, num_samples=3, allow_smaller=True), + RandCropByPosNegLabeld( + keys=["image", "segmentation"], + label_key="segmentation", + spatial_size=patch_size, + pos=3, + neg=1, + num_samples=3, + allow_smaller=True, + ), ResizeWithPadOrCropd(keys=["image", "segmentation"], spatial_size=patch_size), ] ) @@ -157,72 +206,49 @@ def main(): ) # Define train and val DataLoader - train_loader = DataLoader( - train_ds, - batch_size=args.batch_size, - shuffle=True, - num_workers=5, - pin_memory=False, - persistent_workers=False - ) + train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=False, persistent_workers=False) - val_loader = DataLoader( - val_ds, - batch_size=args.batch_size, - shuffle=False, - num_workers=5, - pin_memory=False, - persistent_workers=False - ) + val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=5, pin_memory=False, persistent_workers=False) # Load AugLab GPU transforms and set on device 🐞 gpu_transforms = AugTransformsGPU(json_path=gpu_transforms_path).to(device) # Create model - channels=args.channels - if args.model == 'attunet': + channels = args.channels + if args.model == "attunet": model = AttentionUnet( - spatial_dims=3, - in_channels=1, - out_channels=1, - channels=channels, - strides=[2]*(len(channels)-1), - kernel_size=3).to(device) - elif args.model == 'swinunetr': - model = SwinUNETR( - spatial_dims=3, - in_channels=1, - out_channels=1, - img_size=patch_size, - feature_size=24).to(device) - elif args.model == 'unetr': + spatial_dims=3, in_channels=1, out_channels=1, channels=channels, strides=[2] * (len(channels) - 1), kernel_size=3 + ).to(device) + elif args.model == "swinunetr": + model = SwinUNETR(spatial_dims=3, in_channels=1, out_channels=1, img_size=patch_size, feature_size=24).to(device) + elif args.model == "unetr": model = UNETR( - in_channels=1, - out_channels=1, - img_size=patch_size, - feature_size=16, - hidden_size=768, - mlp_dim=3072, - num_heads=12, - pos_embed="perceptron", - norm_name="instance", - res_block=True, - dropout_rate=0.0, - ).to(device) + in_channels=1, + out_channels=1, + img_size=patch_size, + feature_size=16, + hidden_size=768, + mlp_dim=3072, + num_heads=12, + pos_embed="perceptron", + norm_name="instance", + res_block=True, + dropout_rate=0.0, + ).to(device) else: - raise ValueError(f'Specified model {args.model} is unknown') - + raise ValueError(f"Specified model {args.model} is unknown") + # Init weights if weights are specified if args.start_weights: # Check if weights path exists if not os.path.exists(args.start_weights): - raise ValueError(f'Weights path {args.start_weights} does not exist') + raise ValueError(f"Weights path {args.start_weights} does not exist") else: # Load model weights model.load_state_dict(torch.load(args.start_weights, map_location=torch.device(device))["weights"]) - - # Path to the saved weights - weights_path = f'{weight_folder}/{json_name.replace("config_SegVert_","").replace(".json", ".pth")}' + + # Path to the saved weights + weights_path = f"{weight_folder}/{json_name.replace('config_SegVert_', '').replace('.json', '.pth')}" # Init criterion loss_func = DiceFocalLoss(sigmoid=True, smooth_dr=1e-4) @@ -234,13 +260,13 @@ def main(): scaler = torch.amp.GradScaler() # 🐝 Initialize wandb run - wandb.init(project=f'MonaiSeg', config=vars(args)) + wandb.init(project="MonaiSeg", config=vars(args)) # 🐝 Log gen gradients of the models to wandb wandb.watch(model, log_freq=100) - + # 🐝 Add training script as an artifact - artifact_script = wandb.Artifact(name='training', type='file') + artifact_script = wandb.Artifact(name="training", type="file") artifact_script.add_file(local_path=os.path.abspath(__file__), name=os.path.basename(__file__)) wandb.log_artifact(artifact_script) @@ -248,10 +274,10 @@ def main(): val_dsc_best = 0 for epoch in range(args.start_epoch, args.nb_epochs): # Adjust learning rate - if epoch in [int(sch*args.nb_epochs) for sch in args.schedule]: + if epoch in [int(sch * args.nb_epochs) for sch in args.schedule]: lr = adjust_learning_rate(optimizer, lr, gamma=args.gamma) - print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr)) + print(f"\nEpoch: {epoch + 1:d} | LR: {lr:.8f}") # train for one epoch train_loss, train_dsc = train(train_loader, gpu_transforms, model, loss_func, optimizer, scaler, device) @@ -260,20 +286,20 @@ def main(): wandb.log({"Loss_train/epoch": train_loss}) wandb.log({"DSC_train/epoch": train_dsc}) wandb.log({"training_lr/epoch": lr}) - + # evaluate on validation set val_loss, val_dsc = validate(val_loader, model, loss_func, epoch, device) # 🐝 Plot loss and dice similarity coefficient wandb.log({"Loss_val/epoch": val_loss}) wandb.log({"DSC_val/epoch": val_dsc}) - + # remember best acc and save checkpoint if val_dsc > val_dsc_best: val_dsc_best = val_dsc - state = copy.deepcopy({'weights': model.state_dict()}) + state = copy.deepcopy({"weights": model.state_dict()}) torch.save(state, weights_path) - + # 🐝 close wandb run wandb.finish() @@ -288,11 +314,11 @@ def validate(data_loader, model, loss_func, epoch, device): x, y = (batch["image"].to(device), batch["segmentation"].to(device)) # Get output from model - if x.isnan().any(): - print('found a nan in data.') + if x.isnan().any(): + print("found a nan in data.") y_pred = model(x) - if y_pred.isnan().any(): - print('found a nan in output.') + if y_pred.isnan().any(): + print("found a nan in output.") # Compute loss for each element in the batch size loss = loss_func(y_pred, y) @@ -302,18 +328,16 @@ def validate(data_loader, model, loss_func, epoch, device): if dsc > 0: dsc_list.append(dsc) - epoch_iterator.set_description( - "Validation (loss=%2.5f) (DSC=%2.5f)" % (loss.mean().item(), np.mean(dsc_list)) - ) + epoch_iterator.set_description(f"Validation (loss={loss.mean().item():2.5f}) (DSC={np.mean(dsc_list):2.5f})") # Display first image if step == 0: res_img, target_img, pred_img = get_validation_image(x, y, y_pred, sigmoid=True) # 🐝 log visuals for the first validation batch only in wandb - wandb.log({"validation_img/batch_1": wandb.Image(res_img, caption=f'res_{epoch}')}) - wandb.log({"validation_img/groud_truth": wandb.Image(target_img, caption=f'ground_truth_{epoch}')}) - wandb.log({"validation_img/prediction": wandb.Image(pred_img, caption=f'prediction_{epoch}')}) + wandb.log({"validation_img/batch_1": wandb.Image(res_img, caption=f"res_{epoch}")}) + wandb.log({"validation_img/groud_truth": wandb.Image(target_img, caption=f"ground_truth_{epoch}")}) + wandb.log({"validation_img/prediction": wandb.Image(pred_img, caption=f"prediction_{epoch}")}) return loss.mean().item(), np.mean(dsc_list) @@ -322,17 +346,17 @@ def train(data_loader, gpu_transforms, model, loss_func, optimizer, scaler, devi model.train() dsc_list = [0] epoch_iterator = tqdm(data_loader, desc="Training (loss=X.X) (DSC=X.X)", dynamic_ncols=True) - for step, batch in enumerate(epoch_iterator): + for _step, batch in enumerate(epoch_iterator): # Load input and target x, y = batch["image"].to(device), batch["segmentation"].to(device) - - with torch.amp.autocast('cuda'): + + with torch.amp.autocast("cuda"): # Apply GPU transforms 🐞 x_aug, y_aug = gpu_transforms(x, y) - + # Get output from model y_pred = model(x_aug) - + # Compute loss for each element in the batch size loss = 0 loss = loss_func(y_pred, y_aug) @@ -348,11 +372,9 @@ def train(data_loader, gpu_transforms, model, loss_func, optimizer, scaler, devi scaler.step(optimizer) scaler.update() - epoch_iterator.set_description( - "Training (loss=%2.5f) (DSC=%2.5f)" % (loss.mean().item(), np.mean(dsc_list)) - ) + epoch_iterator.set_description(f"Training (loss={loss.mean().item():2.5f}) (DSC={np.mean(dsc_list):2.5f})") return loss.mean().item(), np.mean(dsc_list) - -if __name__=='__main__': - main() \ No newline at end of file + +if __name__ == "__main__": + main() diff --git a/unit_tests/__init__.py b/unit_tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py new file mode 100644 index 0000000..634b6db --- /dev/null +++ b/unit_tests/conftest.py @@ -0,0 +1,90 @@ +"""Shared fixtures for the AugLab test suite. + +Everything here runs on CPU with tiny volumes, so the whole suite stays fast +enough to gate every pull request. No GPU and no image data on disk required. +""" + +from __future__ import annotations + +import importlib.resources +import json +from pathlib import Path + +import pytest +import torch + +# Small enough to be fast, large enough that the spatial transforms (crop, +# low-res simulation, flips) still have something to work with. +VOLUME_SHAPE = (1, 1, 24, 24, 24) +SEED = 1234 + + +@pytest.fixture(autouse=True) +def _seeded(): + """Seed every RNG the transforms reach for, so failures are reproducible.""" + import random + + import numpy as np + + torch.manual_seed(SEED) + random.seed(SEED) + np.random.seed(SEED) + + +@pytest.fixture +def tiny_volume() -> torch.Tensor: + """A [N, C, D, H, W] float image in roughly the range the transforms expect.""" + return torch.rand(*VOLUME_SHAPE, dtype=torch.float32) + + +@pytest.fixture +def tiny_seg() -> torch.Tensor: + """A binary segmentation mask matching `tiny_volume`.""" + seg = torch.zeros(*VOLUME_SHAPE, dtype=torch.float32) + seg[:, :, 6:18, 6:18, 6:18] = 1.0 + return seg + + +def configs_dir() -> Path: + """Locate the packaged config directory. + + Uses importlib.resources rather than a path relative to this file, which is + how auglab.add_trainer resolves its own package data -- so the tests + exercise the same lookup that ships to users. + """ + from auglab import configs + + return Path(str(importlib.resources.files(configs))) + + +def all_config_paths() -> list[Path]: + """Every JSON config shipped in auglab/configs (excluding the data/ examples).""" + return sorted(p for p in configs_dir().glob("*.json")) + + +def gpu_config_paths() -> list[Path]: + """Configs that drive the GPU augmentation pipeline.""" + paths = [p for p in all_config_paths() if p.name.startswith("transform_params_gpu")] + extra = configs_dir() / "transform_params_one-sequence-to-segment-them-all.json" + if extra.is_file(): + paths.append(extra) + return sorted(paths) + + +def requires_external_asset(config_path: Path) -> str | None: + """Return a skip reason if a config needs an asset that is not on this machine. + + RandomDomainTransferGPU loads a precomputed histogram bank from an absolute + path baked into the module, which only exists on the authors' machines. + Rather than fail CI, skip those configs and say why. + """ + from auglab.transforms.gpu.domain_transfer import DEFAULT_BANK_PATH + + params = json.loads(config_path.read_text()) + params = params.get("GPU", params) + if not isinstance(params, dict): + return None + uses_transfer = params.get("RandomDomainTransferGPU") or params.get("DomainTransferTransform") + if uses_transfer and not Path(DEFAULT_BANK_PATH).is_file(): + return f"domain transfer bank not available at {DEFAULT_BANK_PATH}" + return None diff --git a/unit_tests/test_configs.py b/unit_tests/test_configs.py new file mode 100644 index 0000000..9831ba5 --- /dev/null +++ b/unit_tests/test_configs.py @@ -0,0 +1,78 @@ +"""The shipped JSON configs must parse and actually drive the pipeline. + +Every `transform_params_gpu*.json` is built into an `AugTransformsGPU` and run +over a tiny CPU volume. This is what catches a config that references a +transform the code no longer provides, or a parameter that was renamed. +""" + +from __future__ import annotations + +import json + +import pytest +import torch + +from unit_tests.conftest import all_config_paths, gpu_config_paths, requires_external_asset + +ALL_CONFIGS = all_config_paths() +GPU_CONFIGS = gpu_config_paths() + + +def _ids(paths): + return [p.name for p in paths] + + +def test_configs_are_shipped(): + assert ALL_CONFIGS, "no config JSONs found -- package data is missing" + assert GPU_CONFIGS, "no transform_params_gpu*.json found" + + +@pytest.mark.parametrize("config_path", ALL_CONFIGS, ids=_ids(ALL_CONFIGS)) +def test_config_is_valid_json(config_path): + payload = json.loads(config_path.read_text()) + assert isinstance(payload, dict), f"{config_path.name} should hold a JSON object" + + +@pytest.mark.parametrize("config_path", GPU_CONFIGS, ids=_ids(GPU_CONFIGS)) +def test_gpu_config_builds_and_runs(config_path, tiny_volume, tiny_seg): + """Build the pipeline from the config and push one volume through it.""" + skip_reason = requires_external_asset(config_path) + if skip_reason: + pytest.skip(skip_reason) + + from auglab.transforms.gpu.transforms import AugTransformsGPU + + pipeline = AugTransformsGPU(json_path=str(config_path)) + result = pipeline(tiny_volume, tiny_seg) + + image = result[0] if isinstance(result, (list, tuple)) else result + assert image.shape == tiny_volume.shape, f"{config_path.name} changed the volume shape" + assert image.dtype.is_floating_point + assert torch.isfinite(image).all(), f"{config_path.name} produced NaN or Inf" + + +@pytest.mark.parametrize("config_path", GPU_CONFIGS, ids=_ids(GPU_CONFIGS)) +def test_gpu_config_is_deterministic_under_a_seed(config_path, tiny_volume, tiny_seg): + """Same seed, same output -- otherwise published experiments are not reproducible.""" + skip_reason = requires_external_asset(config_path) + if skip_reason: + pytest.skip(skip_reason) + + from auglab.transforms.gpu.transforms import AugTransformsGPU + + def run_once(): + # Some transforms reach for the stdlib/numpy RNGs, not just torch's, + # so all three have to be pinned for the comparison to mean anything. + import random + + import numpy as np + + torch.manual_seed(7) + random.seed(7) + np.random.seed(7) + pipeline = AugTransformsGPU(json_path=str(config_path)) + out = pipeline(tiny_volume.clone(), tiny_seg.clone()) + return out[0] if isinstance(out, (list, tuple)) else out + + first, second = run_once(), run_once() + assert torch.equal(first, second), f"{config_path.name} is not reproducible under a fixed seed" diff --git a/unit_tests/test_imports.py b/unit_tests/test_imports.py new file mode 100644 index 0000000..d20baf8 --- /dev/null +++ b/unit_tests/test_imports.py @@ -0,0 +1,73 @@ +"""Every module in the package must import cleanly. + +This is the cheapest regression net there is. It catches undeclared +dependencies, syntax errors, and undefined names at module scope -- the class +of bug that ruff's F821 found in transforms_list.py. +""" + +from __future__ import annotations + +import importlib +from pathlib import Path + +import pytest + +import auglab + +# Requires the optional `nnunetv2` extra; skipped rather than failed when absent. +OPTIONAL_PREFIXES = ("auglab.trainers", "auglab.add_trainer") + + +def _module_names() -> list[str]: + """Every .py file under auglab/, as a dotted module name. + + Deliberately a filesystem walk rather than pkgutil.walk_packages: several + subdirectories (transforms/, transforms/cpu/, transforms/gpu/, utils/) have + no __init__.py, so auglab resolves as a PEP 420 namespace package and + walk_packages only reaches 4 of the ~24 modules. Walking the tree keeps + this test honest regardless of how the package is laid out. + """ + roots = [Path(p) for p in auglab.__path__] + names = set() + for root in roots: + for path in root.rglob("*.py"): + if "__pycache__" in path.parts: + continue + relative = path.relative_to(root).with_suffix("") + parts = list(relative.parts) + if parts[-1] == "__init__": + parts.pop() + if not parts: + continue + names.add(".".join(["auglab", *parts])) + return sorted(names) + + +MODULES = _module_names() + + +def test_walk_found_modules(): + """Guard against the discovery itself silently returning too little. + + If this trips, either modules were deleted or the package layout changed in + a way that hides them -- both worth noticing. + """ + assert len(MODULES) >= 20, f"expected the full package, discovered only {len(MODULES)}: {MODULES}" + + +@pytest.mark.parametrize("module_name", MODULES) +def test_module_imports(module_name): + if module_name.startswith(OPTIONAL_PREFIXES): + pytest.importorskip("nnunetv2", reason=f"{module_name} needs the nnunetv2 extra") + importlib.import_module(module_name) + + +def test_public_pipeline_entrypoints_are_importable(): + """The classes users actually construct must be reachable from the package.""" + from auglab.transforms.gpu.transforms import AugTransformsGPU + from auglab.transforms.gpu.transforms_list import ( + AugTransformsGPURandomOrder, + AugTransformsGPURandomOrderTA, + ) + + assert all(callable(cls) for cls in (AugTransformsGPU, AugTransformsGPURandomOrder, AugTransformsGPURandomOrderTA)) diff --git a/unit_tests/test_packaging.py b/unit_tests/test_packaging.py new file mode 100644 index 0000000..90726b1 --- /dev/null +++ b/unit_tests/test_packaging.py @@ -0,0 +1,70 @@ +"""The built wheel must actually contain the package. + +auglab has no __init__.py anywhere in its tree, so it is picked up purely by +setuptools' namespace auto-discovery. That works, but it is easy to break +silently -- a wheel that is missing a subpackage installs fine and only fails +at import time for users. This test builds the real artifact and looks inside. +""" + +from __future__ import annotations + +import subprocess +import sys +import zipfile +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent + + +def _source_modules() -> set[str]: + """Every .py file under auglab/, as a wheel-relative path.""" + package_root = REPO_ROOT / "auglab" + return {str(path.relative_to(REPO_ROOT)) for path in package_root.rglob("*.py") if "__pycache__" not in path.parts} + + +@pytest.fixture(scope="module") +def built_wheel(tmp_path_factory) -> Path: + pytest.importorskip("build", reason="the `build` package is needed to test packaging") + out_dir = tmp_path_factory.mktemp("dist") + result = subprocess.run( + [sys.executable, "-m", "build", "--wheel", "--outdir", str(out_dir), str(REPO_ROOT)], + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + pytest.fail(f"wheel build failed:\n{result.stdout}\n{result.stderr}") + wheels = list(out_dir.glob("*.whl")) + assert len(wheels) == 1, f"expected exactly one wheel, got {wheels}" + return wheels[0] + + +@pytest.mark.slow +def test_wheel_contains_every_module(built_wheel): + with zipfile.ZipFile(built_wheel) as archive: + shipped = {name for name in archive.namelist() if name.endswith(".py")} + + missing = _source_modules() - shipped + assert not missing, f"wheel is missing modules: {sorted(missing)}" + + +@pytest.mark.slow +def test_wheel_contains_the_config_data(built_wheel): + """The JSON configs are the package's data; without them nothing runs.""" + with zipfile.ZipFile(built_wheel) as archive: + configs = {name for name in archive.namelist() if name.startswith("auglab/configs/") and name.endswith(".json")} + + assert configs, "wheel ships no config JSONs" + assert any("transform_params_gpu" in name for name in configs), "wheel is missing the default GPU transform config" + + +@pytest.mark.slow +def test_wheel_excludes_scratch_directories(built_wheel): + """Personal scratch configs should not be published to PyPI.""" + with zipfile.ZipFile(built_wheel) as archive: + names = archive.namelist() + + leaked = [name for name in names if "configs_paul" in name] + assert not leaked, f"wheel ships personal scratch configs: {leaked}" diff --git a/unit_tests/test_transforms_gpu.py b/unit_tests/test_transforms_gpu.py new file mode 100644 index 0000000..32e2e10 --- /dev/null +++ b/unit_tests/test_transforms_gpu.py @@ -0,0 +1,134 @@ +"""Each GPU transform, exercised on its own. + +The config-level tests in test_configs.py prove the pipelines people actually +use still work. These prove each transform works in isolation, so a failure +points at one class instead of a whole config. + +Transforms are discovered by introspection rather than listed by hand, so a +newly added transform is covered the moment it lands. +""" + +from __future__ import annotations + +import importlib +import inspect +from pathlib import Path + +import pytest +import torch + +from auglab.transforms.gpu.base import AugmentationSequentialCustom + +TRANSFORM_MODULES = [ + "auglab.transforms.gpu.contrast", + "auglab.transforms.gpu.spatial", + "auglab.transforms.gpu.fromSeg", + "auglab.transforms.gpu.domain_transfer", +] + +# Not augmentations: helper modules that happen to be nn.Module subclasses. +NOT_A_TRANSFORM = {"DifferentiableHistogram3D"} + + +def _discover(): + """Collect transform classes that can be constructed without arguments.""" + found = [] + for module_name in TRANSFORM_MODULES: + module = importlib.import_module(module_name) + for name, obj in vars(module).items(): + if not inspect.isclass(obj) or obj.__module__ != module_name: + continue + if not issubclass(obj, torch.nn.Module) or name in NOT_A_TRANSFORM: + continue + signature = inspect.signature(obj.__init__) + required = [ + param + for param in list(signature.parameters.values())[1:] + if param.default is inspect.Parameter.empty and param.kind not in (param.VAR_POSITIONAL, param.VAR_KEYWORD) + ] + if required: + # Needs caller-supplied configuration; covered via test_configs.py. + continue + found.append((f"{module_name.rsplit('.', 1)[-1]}.{name}", obj, signature)) + return sorted(found, key=lambda item: item[0]) + + +DISCOVERED = _discover() + + +def _build_kwargs(cls, signature) -> dict: + """Construction arguments that make a transform actually do something. + + `p` is forced to 1.0 because most transforms default to a low probability + and would otherwise pass through untouched most of the time. + """ + kwargs = {"p": 1.0} if "p" in signature.parameters else {} + if cls.__name__ == "RandomDomainTransferGPU": + # Every parameter has a default, but the constructor still rejects a + # missing source_label unless it is told to draw from every domain pair. + kwargs["any_source"] = True + return kwargs + + +def _skip_reason(cls) -> str | None: + """Some transforms depend on assets that do not exist on a fresh checkout.""" + if cls.__name__ == "RandomDomainTransferGPU": + from auglab.transforms.gpu.domain_transfer import DEFAULT_BANK_PATH + + if not Path(DEFAULT_BANK_PATH).is_file(): + return f"domain transfer bank not available at {DEFAULT_BANK_PATH}" + return None + + +def test_discovery_found_transforms(): + assert len(DISCOVERED) >= 15, f"expected the bulk of the GPU transforms, found {len(DISCOVERED)}" + + +@pytest.mark.parametrize( + ("cls", "signature"), + [(cls, sig) for _, cls, sig in DISCOVERED], + ids=[name for name, _, _ in DISCOVERED], +) +def test_transform_runs_on_a_tiny_volume(cls, signature, tiny_volume, tiny_seg): + """Drive each transform the way AugTransformsGPU does and check the output.""" + reason = _skip_reason(cls) + if reason: + pytest.skip(reason) + + # Force the transform to actually fire; the default probability is often low. + pipeline = AugmentationSequentialCustom(cls(**_build_kwargs(cls, signature)), data_keys=["input", "mask"], same_on_batch=True) + + result = pipeline(tiny_volume, tiny_seg) + image = result[0] if isinstance(result, (list, tuple)) else result + + assert image.shape == tiny_volume.shape, f"{cls.__name__} changed the volume shape" + assert image.dtype.is_floating_point + assert torch.isfinite(image).all(), f"{cls.__name__} produced NaN or Inf" + + +@pytest.mark.parametrize( + ("cls", "signature"), + [(cls, sig) for _, cls, sig in DISCOVERED], + ids=[name for name, _, _ in DISCOVERED], +) +def test_transform_leaves_the_mask_intact(cls, signature, tiny_volume, tiny_seg): + """Image-only transforms must not silently alter the segmentation labels. + + Spatial transforms legitimately move the mask, so only the label *set* is + checked -- values must stay in {0, 1}, never interpolated into something in + between. + """ + reason = _skip_reason(cls) + if reason: + pytest.skip(reason) + + pipeline = AugmentationSequentialCustom(cls(**_build_kwargs(cls, signature)), data_keys=["input", "mask"], same_on_batch=True) + + result = pipeline(tiny_volume, tiny_seg) + if not isinstance(result, (list, tuple)) or len(result) < 2: + pytest.skip(f"{cls.__name__} does not return a mask") + + mask = result[1] + assert torch.isfinite(mask).all(), f"{cls.__name__} produced a non-finite mask" + unique = torch.unique(mask) + assert unique.numel() <= 2, f"{cls.__name__} interpolated the mask into {unique.numel()} values"