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..6853bab --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,4 @@ +# Owners are requested for review automatically on every pull request. +# See https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners + +* @Hendrik-code @NathanMolinier diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..ba301a9 --- /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. SMAUGLAB_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** +- SmaugLab 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..dd65888 --- /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 SmaugLab 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..92c6829 --- /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)** +- SmaugLab 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..4a23053 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,25 @@ +## 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 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..ef0d7e5 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,117 @@ +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 "smauglab/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 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; smauglab 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..5e260ce --- /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 SmaugLab + run: pip install -e ".[dev]" + + - name: Run tests + run: pytest -v --cov=smauglab --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: + # smauglab 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 SmaugLab 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("smauglab/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 "smauglab/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..60f1785 100644 --- a/.gitignore +++ b/.gitignore @@ -40,6 +40,10 @@ share/python-wheels/ .installed.cfg *.egg MANIFEST +*.json +*.yaml +*.nii +*.nii.gz # PyInstaller # Usually these files are written by a python script from a template @@ -164,6 +168,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..5568e0e --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,116 @@ +# Contributing to SmaugLab + +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/SmaugLab.git +cd SmaugLab + +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 | +| --- | --- | +| `helpers.py` | `SmaugLabTestCase` base class (RNG seeding, test volumes) and config lookup | +| `test_imports.py` | Every module under `smauglab/` 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 | + +Tests are `unittest.TestCase` subclasses, so they run under either runner: + +```bash +pytest # what CI uses +python -m unittest discover -s unit_tests -t . +``` + +Cases that vary over configs or transforms use `subTest`, so one bad config +does not hide the rest and the failure names the offending item — look for +`SUBFAILED(config=...)` in the output. Derive new test classes from +`SmaugLabTestCase` to get seeded RNGs and the shared `tiny_volume()` / +`tiny_seg()` helpers. + +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 `smauglab/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 `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. + +## Dependency pins + +`kornia` is capped at `>=0.7.3,<0.9`. SmaugLab 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. + +`smauglab/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..e2e1853 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,17 @@ [![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/smauglab)](https://pypi.org/project/smauglab/) +[![Python Versions](https://img.shields.io/pypi/pyversions/smauglab)](https://pypi.org/project/smauglab/) +[![tests](https://github.com/neuropoly/SmaugLab/actions/workflows/tests.yml/badge.svg)](https://github.com/neuropoly/SmaugLab/actions/workflows/tests.yml) +[![lint](https://github.com/neuropoly/SmaugLab/actions/workflows/lint.yml/badge.svg)](https://github.com/neuropoly/SmaugLab/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 +# SmaugLab This repository investigates the influence of different data augmentation strategies on MRI training performance. ## Citation -If you use AugLab, please make sure to cite the following paper: +If you use SmaugLab, please make sure to cite the following paper: ``` @article{molinier2026one, @@ -21,9 +25,9 @@ If you use AugLab, please make sure to cite the following paper: ## What is available ? This repository contains: -- A nnUNet [trainer](https://github.com/neuropoly/AugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/auglab/trainers/nnUNetTrainerDAExt.py) with extensive data augmentations -- A basic Monai segmentation [script](https://github.com/neuropoly/AugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/scripts/train_monai.py) incorporating data augmentations -- A [script](https://github.com/neuropoly/AugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/scripts/generate_augmentations.py) generating augmentations from input images and segmentations +- A nnUNet [trainer](https://github.com/neuropoly/SmaugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/smauglab/trainers/nnUNetTrainerDAExt.py) with extensive data augmentations +- A basic Monai segmentation [script](https://github.com/neuropoly/SmaugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/scripts/train_monai.py) incorporating data augmentations +- A [script](https://github.com/neuropoly/SmaugLab/blob/bed6c1b5cf8ec3dbe6165daca507bf431cad65e5/scripts/generate_augmentations.py) generating augmentations from input images and segmentations ## How to install ? @@ -44,11 +48,11 @@ This repository contains: 3. Clone this repository: - Git clone ```bash - git clone git@github.com:neuropoly/AugLab.git - cd AugLab + git clone git@github.com:neuropoly/SmaugLab.git + cd SmaugLab ``` -4. Install AugLab using one of the following commands: +4. Install SmaugLab using one of the following commands: > **Note:** If you pull a new version from GitHub, make sure to rerun this command with the flag `--upgrade` - nnunetv2 only usage (tested with nnunetv2==2.6.2) ```bash @@ -65,29 +69,29 @@ This repository contains: python3 -m pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu118 --upgrade ``` -## Run nnunet training with AugLab trainer +## Run nnunet training with SmaugLab trainer -To use the AugLab trainer with nnUNet, first add the trainer to your nnUNet installation by running: +To use the SmaugLab trainer with nnUNet, first add the trainer to your nnUNet installation by running: ```bash -auglab_add_nnunettrainer --trainer nnUNetTrainerDAExt +smauglab_add_nnunettrainer --trainer nnUNetTrainerDAExt ``` -Then, when you run nnUNet training as usual, specifying the AugLab trainer, for example: +Then, when you run nnUNet training as usual, specifying the SmaugLab trainer, for example: ```bash nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans ``` -You can also specify your data augmentation parameters by providing a JSON file using the environment variable `AUGLAB_PARAMS_GPU_JSON`: -> **Note:** By default [auglab/configs/transform_params_gpu.json](https://github.com/neuropoly/AugLab/blob/main/auglab/configs/transform_params_gpu.json) is used if no file is specified. +You can also specify your data augmentation parameters by providing a JSON file using the environment variable `SMAUGLAB_PARAMS_GPU_JSON`: +> **Note:** By default [smauglab/configs/transform_params_gpu.json](https://github.com/neuropoly/SmaugLab/blob/main/smauglab/configs/transform_params_gpu.json) is used if no file is specified. ```bash -AUGLAB_PARAMS_GPU_JSON=/path/to/your/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans +SMAUGLAB_PARAMS_GPU_JSON=/path/to/your/params.json nnUNetv2_train 100 3d_fullres 0 -tr nnUNetTrainerDAExtGPU -p nnUNetPlans ``` > ⚠️ **Warning** : To avoid any paths issues, please specify an absolute path to your JSON file. -## Run Monai training with AugLab augmentations +## Run Monai training with SmaugLab augmentations -> To use AugLab augmentations in a MONAI training pipeline, refer to the example [training script](https://github.com/neuropoly/AugLab/blob/main/scripts/train_monai.py). Key implementation lines required for proper integration are marked with a 🐞 emoji in the comments. +> To use SmaugLab augmentations in a MONAI training pipeline, refer to the example [training script](https://github.com/neuropoly/SmaugLab/blob/main/scripts/train_monai.py). Key implementation lines required for proper integration are marked with a 🐞 emoji in the comments. To run the Monai training script directly, you need to provide a config JSON (`config.json`) file with paths to the images and labels (ground truth) for TRAINING, VALIDATION and TESTING sets like this: ```json @@ -122,25 +126,39 @@ To run the Monai training script directly, you need to provide a config JSON (`c } ``` -Then run the training script with the following command, specifying the path to your config JSON file and the path to your data augmentation parameters JSON file (if you want to use custom parameters, otherwise the default [transform_params_gpu.json](https://github.com/neuropoly/AugLab/blob/main/auglab/configs/transform_params_gpu.json) is used): +Then run the training script with the following command, specifying the path to your config JSON file and the path to your data augmentation parameters JSON file (if you want to use custom parameters, otherwise the default [transform_params_gpu.json](https://github.com/neuropoly/SmaugLab/blob/main/smauglab/configs/transform_params_gpu.json) is used): ```bash python scripts/train_monai.py --config /config.json --transforms /transform_params_gpu.json ``` Additional parameters can be specified—see `python scripts/train_monai.py -h` for details. If anything is unclear, feel free to open an issue. +## Contributing + +Development setup, the test suite, and the release process are documented in +[CONTRIBUTING.md](CONTRIBUTING.md). The short version: + +```bash +pip install -e ".[dev]" +pre-commit install +pytest +``` + +Pull requests are gated on Ruff (lint + format) and the test suite across +Python 3.10–3.12. + ## How to use my data ? -Scripts developped in this repository use JSON files to specify image and segmentation paths: see this [example](https://github.com/neuropoly/AugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/auglab/configs/data/data.json). +Scripts developped in this repository use JSON files to specify image and segmentation paths: see this [example](https://github.com/neuropoly/SmaugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/smauglab/configs/data/data.json). ## How do I specify my parameters ? -To track parameters used during data augmentation, JSON files are also used: see this [example](https://github.com/neuropoly/AugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/auglab/configs/transform_params.json) +To track parameters used during data augmentation, JSON files are also used: see this [example](https://github.com/neuropoly/SmaugLab/blob/16653a84e031c40e25a72e946c2724494606b21c/smauglab/configs/transform_params.json) ## Citation -If you use AugLab, please make sure to cite the following paper: +If you use SmaugLab, please make sure to cite the following paper: ``` @article{molinier2026one, @@ -149,4 +167,4 @@ If you use AugLab, please make sure to cite the following paper: journal={arXiv preprint arXiv:2605.03098}, year={2026} } -``` \ No newline at end of file +``` diff --git a/auglab/transforms/cpu/artifact.py b/auglab/transforms/cpu/artifact.py deleted file mode 100644 index 83cf246..0000000 --- a/auglab/transforms/cpu/artifact.py +++ /dev/null @@ -1,202 +0,0 @@ -import torch - -from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform - -import torchio as tio -import gc - -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). - 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 - self.spike = spike - self.bias_field = bias_field - self.blur = blur - self.noise = noise - self.swap = swap - self.random_pick = random_pick - - def get_parameters(self, **data_dict) -> dict: - - artifacts = { - "motion": self.motion, - "ghosting": self.ghosting, - "spike": self.spike, - "bias_field": self.bias_field, - "blur": self.blur, - "noise": self.noise, - "swap": self.swap - } - - 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()} - - 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) - return data_dict - - def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor: - if params['motion']: - img, seg = aug_motion(img, seg) - if params['ghosting']: - img, seg = aug_ghosting(img, seg) - if params['spike']: - img, seg = aug_spike(img, seg) - if params['bias_field']: - img, seg = aug_bias_field(img, seg) - if params['blur']: - img, seg = aug_blur(img, seg) - if params['noise']: - img, seg = aug_noise(img, seg) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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) - )) - 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/transforms.py b/auglab/transforms/cpu/transforms.py deleted file mode 100644 index 5508073..0000000 --- a/auglab/transforms/cpu/transforms.py +++ /dev/null @@ -1,359 +0,0 @@ -import os -import json -import torch -import numpy as np -from typing import Union, Tuple - -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.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 auglab.transforms.cpu.artifact import ArtifactTransform -from auglab.transforms.cpu.contrast import ConvTransform, HistogramEqualTransform, FunctionTransform -from auglab.transforms.cpu.fromSeg import RedistributeTransform -from auglab.transforms.cpu.spatial import SpatialCustomTransform, ShapeTransform - -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]): - # Load transform parameters from JSON - config_path = os.path.join(json_path) - with open(config_path, 'r') as f: - config = json.load(f) - - 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 - ) - 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]): - transform_params = self.transform_params - transforms = [] - - # Scharr filter - 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) - )) - - # Apply functions - func_list = [ - lambda x: torch.log(1 + x), - torch.sqrt, - torch.sin, - torch.exp, - lambda x: 1/(1 + torch.exp(-x)), - ] - 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) - )) - - # Histogram manipulations - 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) - )) - - # Redistribute segmentation values - 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) - )) - - # Resolution transforms - 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) - )) - - # Artifacts generation - 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) - )) - - # Spatial transforms - 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) - )) - - # 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') - 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' - ) - ) - - if do_dummy_2d_data_aug: - transforms.append(Convert2DTo3DTransform()) - - # Noise transforms - 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) - )) - - # Gaussian blur - 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) - )) - - # Brightness transforms - 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) - )) - - # Contrast transforms - 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) - )) - - # Simulate low resolution - 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) - )) - - # Gamma transforms - 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) - )) - - # 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() - super().__init__(transforms=self.transforms) - - def _build_transforms(self): - transforms = [] - - # Scharr filter - 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 - )) - - 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.transforms.gpu.transforms import AugTransformsGPU - - 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_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_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) - - # 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 - ) - - augmentor_gpu = AugTransformsGPU(json_path) - - # Apply transforms - 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()}) - - 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) - - nb_img = len(tensor_dict.keys()) - nb_col = 6 - for key in ['image', 'segmentation']: - output = [] - line = [] - aug = [[]] - 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]) - 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]) - 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 diff --git a/auglab/transforms/gpu/transforms.py b/auglab/transforms/gpu/transforms.py deleted file mode 100644 index 956f087..0000000 --- a/auglab/transforms/gpu/transforms.py +++ /dev/null @@ -1,577 +0,0 @@ -import os, json - -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 -from auglab.transforms.gpu.domain_transfer import RandomDomainTransferGPU -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: - config = json.load(f) - - 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 - - def _build_transforms(self) -> list[nn.Module]: - transforms = [] - - # Flipping transforms - 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) - )) - - # Spatial transforms - 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) - )) - - # 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') - 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, - )) - - ## Transfer augmentations (TA) - ######################### - # Replace image with V26_6_2 contrast (K-means + Voronoi + per-label remap) - palette_params = self.transform_params.get("RandomPALETTETransform") - if palette_params is not None: - transforms.append( - RandomPALETTEGPU( - p=palette_params.get("probability", 1.0), - c_choices=palette_params.get("c_choices", [2, 3, 4, 5, 6]), - s_choices=palette_params.get("s_choices", [2, 3, 4, 5, 6, 7, 8, 9, 10]), - blur_sigmas=palette_params.get("blur_sigmas", [0.0, 0.0, 0.0, 0.3, 0.5, 0.8]), - dark_threshold=palette_params.get("dark_threshold", 0.01), - n_kmeans_subsample=palette_params.get("n_kmeans_subsample", 10000), - skip_parcellation_prob=palette_params.get("skip_parcellation_prob", 0.10), - skip_sub_parc_prob=palette_params.get("skip_sub_parc_prob", 0.40), - alpha_magnitude_range=palette_params.get("alpha_magnitude_range", [0.5, 2.0]), - label_remap_prob=palette_params.get("label_remap_prob", 0.5), - min_label_voxels=palette_params.get("min_label_voxels", 4), - label_classes=palette_params.get("label_classes", None), - ) - ) - - # 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') - if domain_params is not None: - transforms.append( - RandomDomainTransferGPU( - bank_path=domain_params.get("bank_path", None), - source_label=domain_params["source_label"], - targets=domain_params.get("targets", None), - include_self=domain_params.get("include_self", False), - any_source=domain_params.get("any_source", False), - sigma=domain_params.get("sigma", 2.0), - apply_to_channel=domain_params.get("apply_to_channel", [0]), - zscore_io=domain_params.get("zscore_io", "auto"), - pct=domain_params.get("pct", 1.0), - blend_targets=domain_params.get("blend_targets", 1), - blend_concentration=domain_params.get("blend_concentration", 1.0), - p_class_mix=domain_params.get("p_class_mix", 0.0), - bias_field_std=domain_params.get("bias_field_std", 0.0), - bias_scale=domain_params.get("bias_scale", 0.03), - p_spatial_mix=domain_params.get("p_spatial_mix", 0.0), - spatial_mix_scale=domain_params.get("spatial_mix_scale", 0.03), - spatial_mix_gain=domain_params.get("spatial_mix_gain", 3.0), - p=domain_params.get("probability", 0.0), - same_on_batch=domain_params.get("same_on_batch", False), - ) - ) - - # Inverse transform (max - pixel_value) - inverse_params = self.transform_params.get('InverseTransform') - if inverse_params is not None: - transforms.append( - RandomInverseGPU( - p=inverse_params.get("probability", 0), - in_seg=inverse_params.get("in_seg", 0.0), - out_seg=inverse_params.get("out_seg", 0.0), - mix_in_out=inverse_params.get("mix_in_out", False), - mix_prob=inverse_params.get("mix_prob", 0.0), - retain_stats=inverse_params.get("retain_stats", False), - ) - ) - - # Histogram manipulations - histo_params = self.transform_params.get('HistogramEqualizationTransform') - if histo_params is not None: - transforms.append( - RandomHistogramEqualizationGPU( - p=histo_params.get("probability", 0), - in_seg=histo_params.get("in_seg", 0.0), - out_seg=histo_params.get("out_seg", 0.0), - mix_in_out=histo_params.get("mix_in_out", False), - mix_prob=histo_params.get("mix_prob", 0.0), - retain_stats=histo_params.get("retain_stats", False), - ) - ) - - # Redistribute segmentation values transform - 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]), - )) - - # Scharr filter - 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), - )) - - # Unsharp masking - 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), - )) - - # RandomConv transform - 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), - )) - - ## General enhancement (GE) - # Clamping transform - 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), - )) - - # Noise transforms - 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), - )) - - # Gaussian blur - 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), - )) - - # Brightness transforms - 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), - )) - - # Gamma transforms - 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') - 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), - )) - - # nnUNetV2 Contrast transforms - 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) - )) - - # Apply functions - func_list = [ - lambda x: torch.log(1 + x), - torch.sqrt, - torch.sin, - torch.exp, - lambda x: 1/(1 + torch.exp(-x)), - ] - function_params = self.transform_params.get('FunctionTransform') - if function_params is not None: - for func in func_list: - 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), - )) - - # Shape transforms (Cropping and Simulating low resolution) - 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) - )) - - 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') - 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) - )) - - # Bias field artifact - 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), - )) - - ## Random Z-score normalization - zscore_params = self.transform_params.get('ZscoreNormalizationTransform') - if zscore_params is not None: - 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). - - Args: - transforms_list: List of initialized ImageOnlyTransform to choose from. - num_transforms: Number of transforms to randomly select and apply. - same_on_batch: apply the same transformation across the batch. - p: probability for applying the X transforms to a batch. This param controls the augmentation - probabilities batch-wise. - keepdim: whether to keep the output shape the same as input ``True`` or broadcast it to the batch - form ``False``. - - """ - - def __init__( - self, - transforms_list: List[ImageOnlyTransform], - num_transforms: int = 1, - same_on_batch: bool = False, - p: float = 1.0, - keepdim: bool = True, - **kwargs, - ) -> None: - super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) - if not isinstance(num_transforms, int) or num_transforms < 0: - raise ValueError(f"num_transforms must be a non-negative int. Got {num_transforms!r}.") - self.transforms_list = nn.ModuleList(transforms_list) - self.num_transforms = num_transforms - - def _apply_mix(self, x: Tensor, seg: Optional[Tensor]) -> Tensor: - if self.num_transforms == 0 or len(self.transforms_list) == 0: - return x - - k = min(self.num_transforms, len(self.transforms_list)) - # sample without replacement - idx = torch.randperm(len(self.transforms_list), device=x.device)[:k] - - child_params: Dict[str, Tensor] = {} - if seg is not None: - child_params["seg"] = seg - - for j in idx.tolist(): - t = self.transforms_list[j] - 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)}" - ) - # 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) - - if self.same_on_batch: - return self._apply_mix(input, seg) - - batch_size = input.shape[0] - out = input - 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 - 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]. - """ - min_val = np.min(arr) - max_val = np.max(arr) - 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) - return padded_arr - -if __name__ == "__main__": - # Example usage - import importlib - import auglab.configs as configs - from auglab.utils.image import Image, resample_nib - - configs_path = importlib.resources.files(configs) - json_path = configs_path / "transform_params_gpu.json" - 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_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_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_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_tensor_all = torch.from_numpy(seg2.data.copy()) - - # Combine two images to same size - new_shape = [] - for dim in range(3): - size1 = img_tensor.shape[dim] - 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) - new_seg2_tensor_all = torch.zeros(new_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]] - - # 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) - - 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) - - # 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([img_tensor, img2_tensor], dim=0) - seg_tensor = torch.cat([seg_tensor, seg2_tensor], dim=0) - - # Move to GPU - img_tensor = img_tensor.cuda(device=7) - seg_tensor = seg_tensor.cuda(device=7) - augmentor = augmentor.cuda(device=7) - - # Apply augmentations - augmented_img, augmented_seg = augmentor(img_tensor.clone(), seg_tensor.clone()) - - if augmented_img.shape != img_tensor.shape: - raise ValueError("Augmented image shape does not match input shape.") - if augmented_seg.shape != seg_tensor.shape: - raise ValueError("Augmented segmentation shape does not match input shape.") - # Check if nans are present - if torch.isnan(augmented_img).any(): - raise ValueError("NaNs found in augmented image.") - if torch.isnan(augmented_seg).any(): - raise ValueError("NaNs found in augmented segmentation.") - - import cv2 - import numpy as np - import warnings, sys, os - warnings.simplefilter("always") - - # Convert tensors to numpy arrays - img_tensor_np = img_tensor.cpu().detach().numpy() - seg_tensor_np = seg_tensor.cpu().detach().numpy() - augmented_img_np = augmented_img.cpu().detach().numpy() - augmented_seg_np = augmented_seg.cpu().detach().numpy() - - # Concatenate segmentation channels for visualization - 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:]),) - - # 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) - 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) - 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/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) - - print(augmentor) diff --git a/pyproject.toml b/pyproject.toml index d600529..cfa8bb1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,28 @@ -[project] -name = "auglab" -version = "20260109" -requires-python = ">=3.10" -description = "AugLab investigates the influence of different data augmentation strategies on MRI training performance." +[tool.poetry] +name = "smauglab" +# Placeholder only. The real version is substituted at build time by +# poetry-dynamic-versioning from the git tag. Do not bump this by hand. +version = "0.0.0" +description = "SmaugLab 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 ", + "Hendrik Möller ", ] +# poetry-core writes only the FIRST entry of `authors` AND only the FIRST entry +# of `maintainers` into the wheel metadata; every later entry is silently +# dropped (setuptools used to emit both authors). So the two lists have to name +# different people or someone disappears from the published package: +# authors[0] -> Author: / Author-email: (Nathan) +# maintainers[0] -> Maintainer: (Hendrik) +# Verified by inspecting the built wheel's METADATA -- do not reorder these +# without checking it again. +maintainers = [ + "Hendrik Möller", + "Nathan Molinier", +] +homepage = "https://github.com/neuropoly/SmaugLab" +repository = "https://github.com/neuropoly/SmaugLab" classifiers = [ "Development Status :: 5 - Production/Stable", "Intended Audience :: Science/Research", @@ -17,44 +32,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", ] +# smauglab 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 = "smauglab" }] +# The config JSONs are package data, not code, so they need listing explicitly. +# A wheel without them is broken: smauglab resolves its default config through +# importlib.resources at runtime. +include = [{ path = "smauglab/configs/**/*.json", format = ["sdist", "wheel"] }] -[project.scripts] -auglab_add_nnunettrainer = "auglab.add_trainer:main" +[tool.poetry.dependencies] +python = ">=3.10" +batchgenerators = "*" +batchgeneratorsv2 = "*" +# smauglab 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] +smauglab_add_nnunettrainer = "smauglab.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.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", +] + +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 = [] -[tool.setuptools] -include-package-data = true +# Allow unused variables when underscore-prefixed. +dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" -[tool.setuptools.package-data] -'auglab' = ['data/**.json'] +[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. +"smauglab/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 = ["smauglab"] + +[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..89dd017 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 smauglab.transforms.cpu.transforms import AugTransforms +from smauglab.utils.image import Image, resample_nib, zeros_like +from smauglab.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 smauglab.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 smauglab.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,15 +216,15 @@ 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", - # transforms_json_path="auglab/configs/transform_params.json", - # ofolder="/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/test-auglab/augmented", + # data_json_path="smauglab/configs/data/data.json", + # transforms_json_path="smauglab/configs/transform_params.json", + # ofolder="/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/test-smauglab/augmented", # augmentations_per_image=2, # 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..4e84e5d 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 smauglab import configs + +# Import SmaugLab GPU transforms 🐞 +from smauglab.transforms.gpu.transforms import AugTransformsGPU + +# Import SmaugLab custom transforms +from smauglab.utils.utils import ( + adjust_learning_rate, + compute_dsc, + fetch_image_config, + get_validation_image, + parser2config, + tuple2string, + tuple_type_float, + tuple_type_int, ) -# 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 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="smauglab/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,54 +101,51 @@ 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 🐞 + + # Load SmaugLab 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 + # Compose MONAI and SmaugLab transforms pixdim = args.pixdim patch_size = args.patch_size train_transforms = Compose( @@ -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 🐞 + # Load SmaugLab 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/auglab/add_trainer.py b/smauglab/add_trainer.py similarity index 82% rename from auglab/add_trainer.py rename to smauglab/add_trainer.py index 996fac1..4134f82 100644 --- a/auglab/add_trainer.py +++ b/smauglab/add_trainer.py @@ -4,23 +4,20 @@ import nnunetv2 -import auglab.trainers as trainers +from smauglab import trainers + def main(): - parser = argparse.ArgumentParser( - description="This script copies an auglab nnUNetTrainer inside the nnunet folder." - ) + parser = argparse.ArgumentParser(description="This script copies a smauglab nnUNetTrainer inside the nnunet folder.") parser.add_argument( - "-t", "--trainer", + "-t", + "--trainer", choices=["nnUNetTrainerDAExt", "nnUNetTrainerTest"], type=str, required=True, help="nnUNetTrainer to be copied. Choices are: nnUNetTrainerDAExt and nnUNetTrainerTest", ) - parser.add_argument( - '--overwrite', action='store_true', - help='Whether to overwrite existing trainer.' - ) + parser.add_argument("--overwrite", action="store_true", help="Whether to overwrite existing trainer.") args = parser.parse_args() # Get trainer name @@ -30,6 +27,7 @@ def main(): # Add trainer add_trainer(trainer_name, overwrite=overwrite) + def add_trainer(trainer_name: str, overwrite: bool = False): # Find trainer path @@ -54,6 +52,7 @@ def add_trainer(trainer_name: str, overwrite: bool = False): print(f"Trainer {trainer_name} was added to {output_path}") else: print(f"Trainer {trainer_name} already exists at {output_path}. Use --overwrite to replace it.") - + + if __name__ == "__main__": main() diff --git a/auglab/configs/__init__.py b/smauglab/configs/__init__.py similarity index 100% rename from auglab/configs/__init__.py rename to smauglab/configs/__init__.py diff --git a/auglab/configs/data/data.json b/smauglab/configs/data/data.json similarity index 99% rename from auglab/configs/data/data.json rename to smauglab/configs/data/data.json index f08343a..b61d982 100644 --- a/auglab/configs/data/data.json +++ b/smauglab/configs/data/data.json @@ -48,4 +48,4 @@ "LABEL": "data-multi-subject/derivatives/labels/sub-tokyoSkyra02/anat/sub-tokyoSkyra02_T2w_label-SC_seg.nii.gz" } ] -} \ No newline at end of file +} diff --git a/auglab/configs/transform_params.json b/smauglab/configs/transform_params.json similarity index 100% rename from auglab/configs/transform_params.json rename to smauglab/configs/transform_params.json diff --git a/auglab/configs/transform_params_gpu.json b/smauglab/configs/transform_params_gpu.json similarity index 99% rename from auglab/configs/transform_params_gpu.json rename to smauglab/configs/transform_params_gpu.json index 6cb8ce9..397eb91 100644 --- a/auglab/configs/transform_params_gpu.json +++ b/smauglab/configs/transform_params_gpu.json @@ -206,4 +206,4 @@ "ZscoreNormalizationTransform": { "probability": 0.00 } -} \ No newline at end of file +} diff --git a/auglab/configs/transform_params_hybrid.json b/smauglab/configs/transform_params_hybrid.json similarity index 99% rename from auglab/configs/transform_params_hybrid.json rename to smauglab/configs/transform_params_hybrid.json index 46503b4..d27cce8 100644 --- a/auglab/configs/transform_params_hybrid.json +++ b/smauglab/configs/transform_params_hybrid.json @@ -44,7 +44,7 @@ "random_pick": true, "probability": 0 }, - + "Comment2": "nnUNet default CPU augmentations", "SpatialTransform": { "patch_center_dist_from_border": 0, @@ -259,4 +259,4 @@ "probability": 0.3 } } -} \ No newline at end of file +} diff --git a/auglab/configs/transform_params_hybrid_TAGE.json b/smauglab/configs/transform_params_hybrid_TAGE.json similarity index 99% rename from auglab/configs/transform_params_hybrid_TAGE.json rename to smauglab/configs/transform_params_hybrid_TAGE.json index 46b5053..9382fc6 100644 --- a/auglab/configs/transform_params_hybrid_TAGE.json +++ b/smauglab/configs/transform_params_hybrid_TAGE.json @@ -261,4 +261,4 @@ "probability": 0 } } -} \ No newline at end of file +} diff --git a/auglab/configs/transform_params_one-sequence-to-segment-them-all.json b/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json similarity index 99% rename from auglab/configs/transform_params_one-sequence-to-segment-them-all.json rename to smauglab/configs/transform_params_one-sequence-to-segment-them-all.json index 19783a5..c037ff6 100755 --- a/auglab/configs/transform_params_one-sequence-to-segment-them-all.json +++ b/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json @@ -162,4 +162,4 @@ "ZscoreNormalizationTransform": { "probability": 0.00 } -} \ No newline at end of file +} diff --git a/auglab/trainers/__init__.py b/smauglab/trainers/__init__.py similarity index 100% rename from auglab/trainers/__init__.py rename to smauglab/trainers/__init__.py diff --git a/auglab/trainers/nnUNetTrainerDAExt.py b/smauglab/trainers/nnUNetTrainerDAExt.py similarity index 59% rename from auglab/trainers/nnUNetTrainerDAExt.py rename to smauglab/trainers/nnUNetTrainerDAExt.py index 206fef0..5e4298f 100644 --- a/auglab/trainers/nnUNetTrainerDAExt.py +++ b/smauglab/trainers/nnUNetTrainerDAExt.py @@ -1,133 +1,124 @@ -from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer -from nnunetv2.training.dataloading.data_loader import nnUNetDataLoader -from nnunetv2.utilities.default_n_proc_DA import get_allowed_n_proc_DA -from nnunetv2.training.dataloading.nnunet_dataset import infer_dataset_class -from nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels -from torch.nn.parallel import DistributedDataParallel as DDP - -from batchgenerators.dataloading.single_threaded_augmenter import SingleThreadedAugmenter -from batchgenerators.dataloading.nondet_multi_threaded_augmenter import NonDetMultiThreadedAugmenter -from typing import Tuple, Union, List +import importlib +import json +import os +import shutil +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 os -import torch -import importlib -from torch import autocast -import json -import shutil - from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer from nnunetv2.utilities.helpers import dummy_context +from torch import autocast + +from smauglab import configs +from smauglab.trainers.utils import DownsampleSegForDSTransformCustom +from smauglab.transforms.cpu.transforms import AugTransforms +from smauglab.transforms.gpu.transforms import AugTransformsGPU -from auglab.transforms.cpu.transforms import AugTransforms -from auglab.transforms.cpu.contrast import ZscoreNormalization -import auglab.configs as configs -from auglab.transforms.gpu.transforms import AugTransformsGPU -from auglab.trainers.utils import DownsampleSegForDSTransformCustom class nnUNetTrainerDAExt(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) - + 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 return rotation_for_DA, do_dummy_2d_data_aug, initial_patch_size, mirror_axes - + @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 = [] ### Adds 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)) + json_path = os.environ.get("SMAUGLAB_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, + ) + ) 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 +127,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 +137,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,27 +148,24 @@ 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 # Load transform parameters from json file configs_path = importlib.resources.files(configs) - json_path = os.environ.get("AUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) + json_path = os.environ.get("SMAUGLAB_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 SmaugLab GPU 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_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 +173,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: + json_path = os.environ.get("SMAUGLAB_PARAMS_GPU_JSON", str(configs_path / "transform_params_gpu.json")) + 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 +221,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 +251,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 +261,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,32 +275,25 @@ 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 ) ) @@ -333,8 +304,8 @@ def get_validation_transforms( 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 +320,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 +344,26 @@ 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()} 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")) + json_path = os.environ.get("SMAUGLAB_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 SmaugLab 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,54 +371,57 @@ 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 = [] ### Adds 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, - )) - + json_path = os.environ.get("SMAUGLAB_PARAMS_HYBRID_JSON", str(configs_path / "transform_params_hybrid.json")) 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 +430,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 +440,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 +467,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 +491,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/smauglab/trainers/nnUNetTrainerTest.py similarity index 62% rename from auglab/trainers/nnUNetTrainerTest.py rename to smauglab/trainers/nnUNetTrainerTest.py index ed732a5..aa0a54d 100644 --- a/auglab/trainers/nnUNetTrainerTest.py +++ b/smauglab/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.transforms.cpu.transforms import AugTransformsTest -from auglab.transforms.gpu.transforms import AugTransformsGPU -from auglab.trainers.utils import DownsampleSegForDSTransformCustom +from smauglab import configs +from smauglab.trainers.utils import DownsampleSegForDSTransformCustom +from smauglab.transforms.cpu.transforms import AugTransformsTest +from smauglab.transforms.gpu.transforms import AugTransformsGPU + 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/smauglab/trainers/utils.py similarity index 61% rename from auglab/trainers/utils.py rename to smauglab/trainers/utils.py index 99da7e4..e950b23 100644 --- a/auglab/trainers/utils.py +++ b/smauglab/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/smauglab/transforms/cpu/artifact.py b/smauglab/transforms/cpu/artifact.py new file mode 100644 index 0000000..24bc978 --- /dev/null +++ b/smauglab/transforms/cpu/artifact.py @@ -0,0 +1,200 @@ +import gc +import random + +import torch +import torchio as tio +from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform + + +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). + 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 + self.spike = spike + self.bias_field = bias_field + self.blur = blur + self.noise = noise + self.swap = swap + self.random_pick = random_pick + + def get_parameters(self, **data_dict) -> dict: + + artifacts = { + "motion": self.motion, + "ghosting": self.ghosting, + "spike": self.spike, + "bias_field": self.bias_field, + "blur": self.blur, + "noise": self.noise, + "swap": self.swap, + } + + 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()} + + 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) + return data_dict + + def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor: + if params["motion"]: + img, seg = aug_motion(img, seg) + if params["ghosting"]: + img, seg = aug_ghosting(img, seg) + if params["spike"]: + img, seg = aug_spike(img, seg) + if params["bias_field"]: + img, seg = aug_bias_field(img, seg) + if params["blur"]: + img, seg = aug_blur(img, seg) + if params["noise"]: + img, seg = aug_noise(img, seg) + 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), + ) + ) + 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))) + 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), + ) + ) + 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))) + 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), + ) + ) + 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))) + 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), + ) + ) + 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))) + 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), + ) + ) + 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))) + 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), + ) + ) + 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))) + 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), + ) + ) + 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))) + img_out, seg_out = subject.image.data, subject.seg.data + del subject + gc.collect() # Force garbage collection + return img_out, seg_out diff --git a/auglab/transforms/cpu/contrast.py b/smauglab/transforms/cpu/contrast.py similarity index 65% rename from auglab/transforms/cpu/contrast.py rename to smauglab/transforms/cpu/contrast.py index e7cfe4f..5918ade 100644 --- a/auglab/transforms/cpu/contrast.py +++ b/smauglab/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/smauglab/transforms/cpu/fromSeg.py similarity index 80% rename from auglab/transforms/cpu/fromSeg.py rename to smauglab/transforms/cpu/fromSeg.py index f0bb11d..cf7c528 100644 --- a/auglab/transforms/cpu/fromSeg.py +++ b/smauglab/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/smauglab/transforms/cpu/spatial.py similarity index 56% rename from auglab/transforms/cpu/spatial.py rename to smauglab/transforms/cpu/spatial.py index e6c58c2..cbcd5a7 100644 --- a/auglab/transforms/cpu/spatial.py +++ b/smauglab/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/smauglab/transforms/cpu/transforms.py b/smauglab/transforms/cpu/transforms.py new file mode 100644 index 0000000..570f811 --- /dev/null +++ b/smauglab/transforms/cpu/transforms.py @@ -0,0 +1,403 @@ +import json +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 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.compose import ComposeTransforms +from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform +from batchgeneratorsv2.transforms.utils.random import RandomTransform + +from smauglab.transforms.cpu.artifact import ArtifactTransform +from smauglab.transforms.cpu.contrast import ConvTransform, FunctionTransform, HistogramEqualTransform +from smauglab.transforms.cpu.fromSeg import RedistributeTransform +from smauglab.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], + ): + # Load transform parameters from JSON + config_path = os.path.join(json_path) + with open(config_path) as f: + config = json.load(f) + + 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 + ) + 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] + ): + transform_params = self.transform_params + transforms = [] + + # Scharr filter + 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), + ) + ) + + # Apply functions + func_list = [ + lambda x: torch.log(1 + x), + torch.sqrt, + torch.sin, + torch.exp, + lambda x: 1 / (1 + torch.exp(-x)), + ] + func_params = transform_params.get("FunctionTransform") + if func_params is not None: + 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") + 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), + ) + ) + + # Redistribute segmentation values + 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), + ) + ) + + # Resolution transforms + 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), + ) + ) + + # Artifacts generation + 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), + ) + ) + + # Spatial transforms + 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), + ) + ) + + # Spatial nnunet transform + if do_dummy_2d_data_aug: + transforms.append(Convert3DTo2DTransform()) + patch_size_spatial = patch_size[1:] + else: + patch_size_spatial = patch_size + + 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", + ) + ) + + if do_dummy_2d_data_aug: + transforms.append(Convert2DTo3DTransform()) + + # Noise transforms + 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), + ) + ) + + # Gaussian blur + 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.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") + 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), + ) + ) + + # Contrast transforms + 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), + ) + ) + + # Simulate low resolution + 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), + ) + ) + + # Gamma transforms + 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), + ) + ) + + # 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() + super().__init__(transforms=self.transforms) + + def _build_transforms(self): + transforms = [] + + # Scharr filter + 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, + ) + ) + + return transforms + + +if __name__ == "__main__": + # Example usage + import importlib + + import cv2 + + from smauglab import configs + from smauglab.transforms.gpu.transforms import AugTransformsGPU + from smauglab.utils.image import Image, resample_nib + from smauglab.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_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_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 + + # 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 + ) + + augmentor_gpu = AugTransformsGPU(json_path) + + # Apply transforms + tensor_dict = {} + gpu = False + for i in range(24): + 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_{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"]: + output = [] + line = [] + aug = [[]] + 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]) + 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]) + 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) diff --git a/auglab/transforms/gpu/base.py b/smauglab/transforms/gpu/base.py similarity index 66% rename from auglab/transforms/gpu/base.py rename to smauglab/transforms/gpu/base.py index b37d1d3..dc97c03 100644 --- a/auglab/transforms/gpu/base.py +++ b/smauglab/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/smauglab/transforms/gpu/contrast.py similarity index 89% rename from auglab/transforms/gpu/contrast.py rename to smauglab/transforms/gpu/contrast.py index 570162b..b09a0a4 100644 --- a/auglab/transforms/gpu/contrast.py +++ b/smauglab/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 torch import Tensor +from torch.nn import functional as F -from typing import Any, Dict, Optional -from kornia.core import Tensor -import random -import math - -from auglab.transforms.gpu.base import ImageOnlyTransform -from typing import Any, Dict, Optional, Tuple, Union, List +from smauglab.transforms.gpu.base import ImageOnlyTransform -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/smauglab/transforms/gpu/domain_transfer.py similarity index 84% rename from auglab/transforms/gpu/domain_transfer.py rename to smauglab/transforms/gpu/domain_transfer.py index 68a02f3..c896237 100644 --- a/auglab/transforms/gpu/domain_transfer.py +++ b/smauglab/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 +from smauglab.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/smauglab/transforms/gpu/fromSeg.py similarity index 84% rename from auglab/transforms/gpu/fromSeg.py rename to smauglab/transforms/gpu/fromSeg.py index 8b0e149..34115f3 100644 --- a/auglab/transforms/gpu/fromSeg.py +++ b/smauglab/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 - +from smauglab.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): @@ -257,7 +265,7 @@ def apply_transform( class RandomPALETTEGPU(ImageOnlyTransform): """ - AugLab GPU augmentation implementing PALETTE synthesis. + SmaugLab GPU augmentation implementing PALETTE synthesis. Pipeline (mirrors src/synthesis/PALETTE_synthesis.py, self-contained): 1. Min-max normalise input to [0, 1] per sample. @@ -269,7 +277,7 @@ class RandomPALETTEGPU(ImageOnlyTransform): label's voxels with a fresh (μ, α). 4. Optional second Gaussian blur, then foreground z-score. - Segmentation is read from params['seg'] (injected by AugLab's pipeline). + Segmentation is read from params['seg'] (injected by SmaugLab's pipeline). Supported formats: one-hot [B, C_seg, D, H, W] or index [B, 1, D, H, W]. Args: @@ -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/smauglab/transforms/gpu/spatial.py similarity index 80% rename from auglab/transforms/gpu/spatial.py rename to smauglab/transforms/gpu/spatial.py index 25708be..97b601f 100644 --- a/auglab/transforms/gpu/spatial.py +++ b/smauglab/transforms/gpu/spatial.py @@ -1,17 +1,21 @@ -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 -from typing import Any, Dict, Optional, Tuple, Union -from auglab.transforms.gpu.base import ImageOnlyTransform +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 smauglab.transforms.gpu.base import ImageOnlyTransform # Affine transform @@ -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/smauglab/transforms/gpu/transforms.py b/smauglab/transforms/gpu/transforms.py new file mode 100644 index 0000000..abc7968 --- /dev/null +++ b/smauglab/transforms/gpu/transforms.py @@ -0,0 +1,711 @@ +import json +import os +from typing import Any + +import numpy as np +import torch +from torch import Tensor, nn + +from smauglab.transforms.gpu.base import AugmentationSequentialCustom, ImageOnlyTransform +from smauglab.transforms.gpu.contrast import ( + RandomBiasFieldGPU, + RandomBrightnessGPU, + RandomClampGPU, + RandomContrastGPU, + RandomConvTransformGPU, + RandomFunctionGPU, + RandomGammaGPU, + RandomGaussianNoiseGPU, + RandomHistogramEqualizationGPU, + RandomInverseGPU, + ZscoreNormalizationGPU, +) +from smauglab.transforms.gpu.domain_transfer import RandomDomainTransferGPU +from smauglab.transforms.gpu.fromSeg import RandomPALETTEGPU, RandomRedistributeSegGPU +from smauglab.transforms.gpu.spatial import ( + RandomAcqTransformGPU, + RandomAffine3DCustom, + RandomCropTransformGPU, + RandomFlipTransformGPU, + RandomLowResTransformGPU, +) +from smauglab.transforms.synthseg.transforms import RandomSynthSegGPU + + +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) as f: + config = json.load(f) + + 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 + + def _build_transforms(self) -> list[nn.Module]: + transforms = [] + + # Flipping transforms + 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), + ) + ) + + # Spatial transforms + 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), + ) + ) + + # 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") + 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, + ) + ) + + ## Transfer augmentations (TA) + ######################### + # Replace image with V26_6_2 contrast (K-means + Voronoi + per-label remap) + palette_params = self.transform_params.get("RandomPALETTETransform") + if palette_params is not None: + transforms.append( + RandomPALETTEGPU( + p=palette_params.get("probability", 1.0), + c_choices=palette_params.get("c_choices", [2, 3, 4, 5, 6]), + s_choices=palette_params.get("s_choices", [2, 3, 4, 5, 6, 7, 8, 9, 10]), + blur_sigmas=palette_params.get("blur_sigmas", [0.0, 0.0, 0.0, 0.3, 0.5, 0.8]), + dark_threshold=palette_params.get("dark_threshold", 0.01), + n_kmeans_subsample=palette_params.get("n_kmeans_subsample", 10000), + skip_parcellation_prob=palette_params.get("skip_parcellation_prob", 0.10), + skip_sub_parc_prob=palette_params.get("skip_sub_parc_prob", 0.40), + alpha_magnitude_range=palette_params.get("alpha_magnitude_range", [0.5, 2.0]), + label_remap_prob=palette_params.get("label_remap_prob", 0.5), + min_label_voxels=palette_params.get("min_label_voxels", 4), + label_classes=palette_params.get("label_classes", None), + ) + ) + + # 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") + if domain_params is not None: + transforms.append( + RandomDomainTransferGPU( + bank_path=domain_params.get("bank_path", None), + source_label=domain_params["source_label"], + targets=domain_params.get("targets", None), + include_self=domain_params.get("include_self", False), + any_source=domain_params.get("any_source", False), + sigma=domain_params.get("sigma", 2.0), + apply_to_channel=domain_params.get("apply_to_channel", [0]), + zscore_io=domain_params.get("zscore_io", "auto"), + pct=domain_params.get("pct", 1.0), + blend_targets=domain_params.get("blend_targets", 1), + blend_concentration=domain_params.get("blend_concentration", 1.0), + p_class_mix=domain_params.get("p_class_mix", 0.0), + bias_field_std=domain_params.get("bias_field_std", 0.0), + bias_scale=domain_params.get("bias_scale", 0.03), + p_spatial_mix=domain_params.get("p_spatial_mix", 0.0), + spatial_mix_scale=domain_params.get("spatial_mix_scale", 0.03), + spatial_mix_gain=domain_params.get("spatial_mix_gain", 3.0), + p=domain_params.get("probability", 0.0), + same_on_batch=domain_params.get("same_on_batch", False), + ) + ) + + # Inverse transform (max - pixel_value) + inverse_params = self.transform_params.get("InverseTransform") + if inverse_params is not None: + transforms.append( + RandomInverseGPU( + p=inverse_params.get("probability", 0), + in_seg=inverse_params.get("in_seg", 0.0), + out_seg=inverse_params.get("out_seg", 0.0), + mix_in_out=inverse_params.get("mix_in_out", False), + mix_prob=inverse_params.get("mix_prob", 0.0), + retain_stats=inverse_params.get("retain_stats", False), + ) + ) + + # Histogram manipulations + histo_params = self.transform_params.get("HistogramEqualizationTransform") + if histo_params is not None: + transforms.append( + RandomHistogramEqualizationGPU( + p=histo_params.get("probability", 0), + in_seg=histo_params.get("in_seg", 0.0), + out_seg=histo_params.get("out_seg", 0.0), + mix_in_out=histo_params.get("mix_in_out", False), + mix_prob=histo_params.get("mix_prob", 0.0), + retain_stats=histo_params.get("retain_stats", False), + ) + ) + + # Redistribute segmentation values transform + 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]), + ) + ) + + # Scharr filter + 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), + ) + ) + + # Unsharp masking + 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), + ) + ) + + # RandomConv transform + 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), + ) + ) + + ## General enhancement (GE) + # Clamping transform + 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), + ) + ) + + # Noise transforms + 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), + ) + ) + + # Gaussian blur + 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), + ) + ) + + # Brightness transforms + 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), + ) + ) + + # Gamma transforms + 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") + 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), + ) + ) + + # nnUNetV2 Contrast transforms + 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), + ) + ) + + # Apply functions + func_list = [ + lambda x: torch.log(1 + x), + torch.sqrt, + torch.sin, + torch.exp, + lambda x: 1 / (1 + torch.exp(-x)), + ] + function_params = self.transform_params.get("FunctionTransform") + if function_params is not None: + 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 + ) + + # Shape transforms (Cropping and Simulating low resolution) + 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), + ) + ) + + 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") + 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), + ) + ) + + # Bias field artifact + 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), + ) + ) + + ## Random Z-score normalization + zscore_params = self.transform_params.get("ZscoreNormalizationTransform") + if zscore_params is not None: + 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). + + Args: + transforms_list: List of initialized ImageOnlyTransform to choose from. + num_transforms: Number of transforms to randomly select and apply. + same_on_batch: apply the same transformation across the batch. + p: probability for applying the X transforms to a batch. This param controls the augmentation + probabilities batch-wise. + keepdim: whether to keep the output shape the same as input ``True`` or broadcast it to the batch + form ``False``. + + """ + + def __init__( + self, + transforms_list: list[ImageOnlyTransform], + num_transforms: int = 1, + same_on_batch: bool = False, + p: float = 1.0, + keepdim: bool = True, + **kwargs, + ) -> None: + super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim) + if not isinstance(num_transforms, int) or num_transforms < 0: + raise ValueError(f"num_transforms must be a non-negative int. Got {num_transforms!r}.") + self.transforms_list = nn.ModuleList(transforms_list) + self.num_transforms = num_transforms + + def _apply_mix(self, x: Tensor, seg: Tensor | None) -> Tensor: + if self.num_transforms == 0 or len(self.transforms_list) == 0: + return x + + k = min(self.num_transforms, len(self.transforms_list)) + # sample without replacement + idx = torch.randperm(len(self.transforms_list), device=x.device)[:k] + + child_params: dict[str, Tensor] = {} + if seg is not None: + child_params["seg"] = seg + + for j in idx.tolist(): + t = self.transforms_list[j] + 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)}") + # 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: Tensor | None = None) -> Tensor: + seg = params.get("seg") + + if self.same_on_batch: + return self._apply_mix(input, seg) + + batch_size = input.shape[0] + out = input + for i in range(batch_size): + xi = out[i : i + 1] + seg_i = None + 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]. + """ + min_val = np.min(arr) + max_val = np.max(arr) + 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) + return padded_arr + + +if __name__ == "__main__": + # Example usage + import importlib + + from smauglab import configs + from smauglab.utils.image import Image, resample_nib + + configs_path = importlib.resources.files(configs) + json_path = configs_path / "transform_params_gpu.json" + 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_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_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_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_tensor_all = torch.from_numpy(seg2.data.copy()) + + # Combine two images to same size + new_shape = [] + for dim in range(3): + size1 = img_tensor.shape[dim] + 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) + new_seg2_tensor_all = torch.zeros(new_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] + ] + + # 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 + + 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 + + # 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([img_tensor, img2_tensor], dim=0) + seg_tensor = torch.cat([seg_tensor, seg2_tensor], dim=0) + + # Move to GPU + img_tensor = img_tensor.cuda(device=7) + seg_tensor = seg_tensor.cuda(device=7) + augmentor = augmentor.cuda(device=7) + + # Apply augmentations + augmented_img, augmented_seg = augmentor(img_tensor.clone(), seg_tensor.clone()) + + if augmented_img.shape != img_tensor.shape: + raise ValueError("Augmented image shape does not match input shape.") + if augmented_seg.shape != seg_tensor.shape: + raise ValueError("Augmented segmentation shape does not match input shape.") + # Check if nans are present + if torch.isnan(augmented_img).any(): + 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 + + warnings.simplefilter("always") + + # Convert tensors to numpy arrays + img_tensor_np = img_tensor.cpu().detach().numpy() + seg_tensor_np = seg_tensor.cpu().detach().numpy() + augmented_img_np = augmented_img.cpu().detach().numpy() + augmented_seg_np = augmented_seg.cpu().detach().numpy() + + # Concatenate segmentation channels for visualization + 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:]),) + + # 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, + ) + 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, + ) + 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/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) + + print(augmentor) diff --git a/auglab/transforms/gpu/transforms_list.py b/smauglab/transforms/gpu/transforms_list.py similarity index 93% rename from auglab/transforms/gpu/transforms_list.py rename to smauglab/transforms/gpu/transforms_list.py index 8165c89..f5a7034 100644 --- a/auglab/transforms/gpu/transforms_list.py +++ b/smauglab/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.contrast import ( - RandomConvTransformGPU, - RandomGaussianNoiseGPU, +from smauglab.transforms.gpu.base import AugmentationSequentialCustom, ImageOnlyTransform +from smauglab.transforms.gpu.contrast import ( + 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 smauglab.transforms.gpu.fromSeg import RandomRedistributeSegGPU +from smauglab.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,8 +766,10 @@ def pad_numpy_array(arr, shape): if __name__ == "__main__": # Example usage import importlib - import auglab.configs as configs - from auglab.utils.image import Image, resample_nib + + from smauglab import configs + from smauglab.transforms.gpu.transforms import AugTransformsGPU + from smauglab.utils.image import Image, resample_nib configs_path = importlib.resources.files(configs) json_path = configs_path / "transform_params_gpu.json" @@ -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/README.md b/smauglab/transforms/synthseg/README.md similarity index 93% rename from auglab/transforms/synthseg/README.md rename to smauglab/transforms/synthseg/README.md index 0d56058..2acabc7 100644 --- a/auglab/transforms/synthseg/README.md +++ b/smauglab/transforms/synthseg/README.md @@ -1,7 +1,7 @@ # SynthSeg generative augmentation A faithful [torch] re-implementation of the **SynthSeg** "brain generator" as an -AugLab augmentation. Unlike every other transform in AugLab — which perturbs a +SmaugLab augmentation. Unlike every other transform in SmaugLab — which perturbs a *real* image — SynthSeg **ignores the input image entirely and synthesises a new image from a label map**, using domain randomisation (a per-label Gaussian mixture model plus random spatial, bias, intensity and resolution corruptions). @@ -63,8 +63,8 @@ Defaults match `BrainGenerator.__init__` (which overrides several ## Implementation notes / deviations -- **3D only** (5D `(B, C, D, H, W)` tensors), matching AugLab's GPU transforms. -- Affine transforms are applied **about the volume centre** (like AugLab's +- **3D only** (5D `(B, C, D, H, W)` tensors), matching SmaugLab's GPU transforms. +- Affine transforms are applied **about the volume centre** (like SmaugLab's `RandomAffine3DCustom`), rather than the corner-origin used by neuron's `affine_to_shift`. This keeps the anatomy in frame and is the standard choice; the visual augmentation is equivalent. @@ -128,8 +128,8 @@ is ignored; `target` is the label map. ```python import importlib, torch -import auglab.configs as configs -from auglab.transforms.synthseg import SynthSegTransformsGPU +import smauglab.configs as configs +from smauglab.transforms.synthseg import SynthSegTransformsGPU cfg = importlib.resources.files(configs) / "synthseg_params.json" synth = SynthSegTransformsGPU(json_path=str(cfg)).to("cuda") @@ -141,7 +141,7 @@ image, label = synth(data, target) # image is fully synthetic, label is deform Or directly with the module API: ```python -from auglab.transforms.synthseg import SynthSegGenerator +from smauglab.transforms.synthseg import SynthSegGenerator gen = SynthSegGenerator(generation_labels=[0, 2, 3, 41, 42, ...], n_neutral_labels=1, n_channels=1).to("cuda") image, label = gen(label_map) # label_map: (B, 1, D, H, W) @@ -150,14 +150,14 @@ image, label = gen(label_map) # label_map: (B, 1, D, H, W) ### 2. As an `ImageOnlyTransform` in an existing GPU pipeline `RandomSynthSegGPU` replaces the image with a GMM synthesis of `params['seg']` -(intensity-only: GMM → bias → intensity → resolution). Put AugLab's geometric +(intensity-only: GMM → bias → intensity → resolution). Put SmaugLab's geometric transforms *before* it so the mask is deformed first and SynthSeg generates from the deformed labels: ```python -from auglab.transforms.gpu.base import AugmentationSequentialCustom -from auglab.transforms.gpu.spatial import RandomAffine3DCustom -from auglab.transforms.synthseg import RandomSynthSegGPU +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from smauglab.transforms.gpu.spatial import RandomAffine3DCustom +from smauglab.transforms.synthseg import RandomSynthSegGPU aug = AugmentationSequentialCustom( RandomAffine3DCustom(degrees=15, scale=[0.8, 1.2], p=1.0), @@ -170,7 +170,7 @@ image, seg = aug(image, seg) > **Note (kornia 0.7.4 quirk):** when an `AugmentationSequentialCustom` is called > with more than one data key, kornia detaches the returned **input image** to the -> CPU (the segmentation/mask stays on the GPU). This affects *every* AugLab GPU +> CPU (the segmentation/mask stays on the GPU). This affects *every* SmaugLab GPU > transform identically, not just SynthSeg — re-`.to(device)` the returned image > if you need it back on the GPU. The full-pipeline `SynthSegTransformsGPU` driver > (section 1) does **not** go through kornia's sequential and is unaffected. @@ -199,6 +199,6 @@ are therefore inert in this path — add an `AffineTransform`/`FlipTransform` bl Both modules are runnable and self-contained (no data files, CPU-friendly): ```bash -python -m auglab.transforms.synthseg.generator -python -m auglab.transforms.synthseg.transforms +python -m smauglab.transforms.synthseg.generator +python -m smauglab.transforms.synthseg.transforms ``` diff --git a/auglab/transforms/synthseg/__init__.py b/smauglab/transforms/synthseg/__init__.py similarity index 71% rename from auglab/transforms/synthseg/__init__.py rename to smauglab/transforms/synthseg/__init__.py index f947b0b..fa13d65 100644 --- a/auglab/transforms/synthseg/__init__.py +++ b/smauglab/transforms/synthseg/__init__.py @@ -1,4 +1,4 @@ -"""SynthSeg generative augmentation for AugLab. +"""SynthSeg generative augmentation for SmaugLab. A faithful torch re-implementation of the SynthSeg "brain generator" (Billot et al., Medical Image Analysis 2023; BBillot/SynthSeg, BBillot/lab2im): @@ -9,19 +9,19 @@ SynthSegGenerator -- the full generative model as an nn.Module (``forward(label_map) -> (image, label)``). SynthSegTransformsGPU -- config-driven driver, ``forward(data, target) -> - (image, target)`` (AugLab calling convention). + (image, target)`` (SmaugLab calling convention). RandomSynthSegGPU -- ImageOnlyTransform that replaces the image with a GMM synthesis of ``params['seg']`` (composes inside AugmentationSequentialCustom pipelines). """ -from auglab.transforms.synthseg.generator import SynthSegGenerator -from auglab.transforms.synthseg.transforms import RandomSynthSegGPU, SynthSegTransformsGPU -from auglab.transforms.synthseg import functional +from smauglab.transforms.synthseg import functional +from smauglab.transforms.synthseg.generator import SynthSegGenerator +from smauglab.transforms.synthseg.transforms import RandomSynthSegGPU, SynthSegTransformsGPU __all__ = [ + "RandomSynthSegGPU", "SynthSegGenerator", "SynthSegTransformsGPU", - "RandomSynthSegGPU", "functional", ] diff --git a/auglab/transforms/synthseg/functional.py b/smauglab/transforms/synthseg/functional.py similarity index 91% rename from auglab/transforms/synthseg/functional.py rename to smauglab/transforms/synthseg/functional.py index 697308c..30d6979 100644 --- a/auglab/transforms/synthseg/functional.py +++ b/smauglab/transforms/synthseg/functional.py @@ -11,7 +11,7 @@ ``BBillot/lab2im``. Each function below cites the corresponding reference layer. Everything operates on 3D volumes stored as ``(B, C, D, H, W)`` torch tensors (label maps as ``(B, 1, D, H, W)`` integer tensors), which is the convention -used throughout AugLab's GPU transforms. +used throughout SmaugLab's GPU transforms. Spatial conventions -------------------- @@ -20,14 +20,15 @@ the ``(x, y, z) = (W, H, D)`` order expected by ``F.grid_sample`` at the very end, with ``align_corners=True`` so that integer voxel indices map exactly. * Affine transforms are applied about the volume centre (standard practice and - matching AugLab's existing ``RandomAffine3DCustom``), so small rotations / + matching SmaugLab's existing ``RandomAffine3DCustom``), so small rotations / scalings keep the anatomy in frame. """ 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", ] @@ -63,7 +64,7 @@ def to_label_map(seg: torch.Tensor) -> torch.Tensor: Accepts: * ``(B, 1, D, H, W)`` -> rounded to integer labels (used directly). * ``(B, C, D, H, W)`` one-hot (C > 1) -> argmax + 1, background (all-zero - across channels) stays 0. This matches AugLab's + across channels) stays 0. This matches SmaugLab's :func:`collapse_onehot_to_index` convention where channel ``c`` encodes label ``c + 1``. * ``(B, D, H, W)`` -> unsqueezed to ``(B, 1, D, H, W)``. @@ -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/smauglab/transforms/synthseg/generator.py similarity index 85% rename from auglab/transforms/synthseg/generator.py rename to smauglab/transforms/synthseg/generator.py index b7cba44..8c22c3a 100644 --- a/auglab/transforms/synthseg/generator.py +++ b/smauglab/transforms/synthseg/generator.py @@ -28,12 +28,13 @@ 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 -from auglab.transforms.synthseg import functional as FN +from smauglab.transforms.synthseg import functional as FN Number = Union[int, float] @@ -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/smauglab/transforms/synthseg/transforms.py similarity index 79% rename from auglab/transforms/synthseg/transforms.py rename to smauglab/transforms/synthseg/transforms.py index 9c838d2..ac31632 100644 --- a/auglab/transforms/synthseg/transforms.py +++ b/smauglab/transforms/synthseg/transforms.py @@ -1,10 +1,10 @@ -"""AugLab-style wrappers around the SynthSeg generative model. +"""SmaugLab-style wrappers around the SynthSeg generative model. Two entry points are provided: * :class:`RandomSynthSegGPU` -- an :class:`ImageOnlyTransform` that *replaces* the image with a GMM-synthesised one derived from ``params['seg']``. It is - intensity-only (no internal spatial deformation), so it composes with AugLab's + intensity-only (no internal spatial deformation), so it composes with SmaugLab's existing geometric transforms (``RandomAffine3DCustom``, ``RandomFlipTransformGPU``, ...) inside an :class:`AugmentationSequentialCustom`: place those *before* it so the mask is deformed first and SynthSeg generates from the deformed labels, @@ -12,7 +12,7 @@ pipeline like any other transform. * :class:`SynthSegTransformsGPU` -- a config-driven top-level module mirroring - :class:`auglab.transforms.gpu.transforms.AugTransformsGPU`. It runs the *full* + :class:`smauglab.transforms.gpu.transforms.AugTransformsGPU`. It runs the *full* SynthSeg pipeline (spatial deform + flip + GMM + bias + intensity + resolution) and returns ``(image, label)`` from ``forward(data, target)`` -- the calling convention used by the nnUNet trainer and ``train_monai.py``. This is the @@ -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 +from smauglab.transforms.gpu.base import ImageOnlyTransform +from smauglab.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} @@ -58,7 +88,7 @@ class RandomSynthSegGPU(ImageOnlyTransform): Intensity-only: GMM sampling -> bias field -> intensity augmentation -> resolution randomisation. Spatial deformation / flipping are disabled so that geometry stays consistent with the segmentation propagated by the surrounding - :class:`AugmentationSequentialCustom` (use AugLab's geometric transforms for + :class:`AugmentationSequentialCustom` (use SmaugLab's geometric transforms for that, placed before this one). Args: @@ -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/smauglab/utils/image.py similarity index 84% rename from auglab/utils/image.py rename to smauglab/utils/image.py index 86e031b..6402d53 100644 --- a/auglab/utils/image.py +++ b/smauglab/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/smauglab/utils/utils.py similarity index 78% rename from auglab/utils/utils.py rename to smauglab/utils/utils.py index 1edc255..319c109 100644 --- a/auglab/utils/utils.py +++ b/smauglab/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/unit_tests/__init__.py b/unit_tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/unit_tests/helpers.py b/unit_tests/helpers.py new file mode 100644 index 0000000..53aacf2 --- /dev/null +++ b/unit_tests/helpers.py @@ -0,0 +1,106 @@ +"""Shared helpers and the base TestCase for the SmaugLab 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 +import random +import unittest +from pathlib import Path + +import numpy as np +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 + + +def seed_everything(seed: int = SEED) -> None: + """Pin every RNG the transforms reach for. + + Some transforms use the stdlib and numpy generators, not just torch's, so + all three have to be set or results are not reproducible between runs. + """ + torch.manual_seed(seed) + random.seed(seed) + np.random.seed(seed) + + +class SmaugLabTestCase(unittest.TestCase): + """Base class that seeds the RNGs and hands out the standard test volumes.""" + + def setUp(self) -> None: + super().setUp() + seed_everything() + + def tiny_volume(self) -> 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) + + def tiny_seg(self) -> 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 assertIsImageLike(self, image: torch.Tensor, reference: torch.Tensor, label: str) -> None: + """Assert `image` is a finite float tensor shaped like `reference`.""" + self.assertEqual(image.shape, reference.shape, f"{label} changed the volume shape") + self.assertTrue(image.dtype.is_floating_point, f"{label} returned a non-float tensor") + self.assertTrue(bool(torch.isfinite(image).all()), f"{label} produced NaN or Inf") + + +def first_output(result): + """Pipelines return either a tensor or a (image, mask) sequence.""" + return result[0] if isinstance(result, (list, tuple)) else result + + +def configs_dir() -> Path: + """Locate the packaged config directory. + + Uses importlib.resources rather than a path relative to this file, which is + how smauglab.add_trainer resolves its own package data -- so the tests + exercise the same lookup that ships to users. + """ + from smauglab import configs + + return Path(str(importlib.resources.files(configs))) + + +def all_config_paths() -> list[Path]: + """Every JSON config shipped in smauglab/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 smauglab.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..e292dee --- /dev/null +++ b/unit_tests/test_configs.py @@ -0,0 +1,82 @@ +"""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 unittest + +import torch + +from unit_tests.helpers import ( + SmaugLabTestCase, + all_config_paths, + first_output, + gpu_config_paths, + requires_external_asset, + seed_everything, +) + +ALL_CONFIGS = all_config_paths() +GPU_CONFIGS = gpu_config_paths() + + +class TestConfigsArePresent(unittest.TestCase): + def test_configs_are_shipped(self): + self.assertTrue(ALL_CONFIGS, "no config JSONs found -- package data is missing") + self.assertTrue(GPU_CONFIGS, "no transform_params_gpu*.json found") + + def test_every_config_is_valid_json(self): + for config_path in ALL_CONFIGS: + with self.subTest(config=config_path.name): + payload = json.loads(config_path.read_text()) + self.assertIsInstance(payload, dict, f"{config_path.name} should hold a JSON object") + + +class TestGpuConfigs(SmaugLabTestCase): + def test_gpu_config_builds_and_runs(self): + """Build the pipeline from each config and push one volume through it.""" + from smauglab.transforms.gpu.transforms import AugTransformsGPU + + for config_path in GPU_CONFIGS: + with self.subTest(config=config_path.name): + skip_reason = requires_external_asset(config_path) + if skip_reason: + self.skipTest(skip_reason) + + volume, seg = self.tiny_volume(), self.tiny_seg() + pipeline = AugTransformsGPU(json_path=str(config_path)) + image = first_output(pipeline(volume, seg)) + + self.assertIsImageLike(image, volume, config_path.name) + + def test_gpu_config_is_deterministic_under_a_seed(self): + """Same seed, same output -- otherwise published experiments are not reproducible.""" + from smauglab.transforms.gpu.transforms import AugTransformsGPU + + def run_once(config_path): + # 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. + seed_everything(7) + pipeline = AugTransformsGPU(json_path=str(config_path)) + return first_output(pipeline(self.tiny_volume(), self.tiny_seg())) + + for config_path in GPU_CONFIGS: + with self.subTest(config=config_path.name): + skip_reason = requires_external_asset(config_path) + if skip_reason: + self.skipTest(skip_reason) + + first, second = run_once(config_path), run_once(config_path) + self.assertTrue( + torch.equal(first, second), + f"{config_path.name} is not reproducible under a fixed seed", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_imports.py b/unit_tests/test_imports.py new file mode 100644 index 0000000..b5b80e2 --- /dev/null +++ b/unit_tests/test_imports.py @@ -0,0 +1,88 @@ +"""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 +import importlib.util +import unittest +from pathlib import Path + +import smauglab + +# Requires the optional `nnunetv2` extra; skipped rather than failed when absent. +OPTIONAL_PREFIXES = ("smauglab.trainers", "smauglab.add_trainer") + + +def module_names() -> list[str]: + """Every .py file under smauglab/, 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 smauglab 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 smauglab.__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(["smauglab", *parts])) + return sorted(names) + + +MODULES = module_names() + + +class TestModuleDiscovery(unittest.TestCase): + def test_walk_found_modules(self): + """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. + """ + self.assertGreaterEqual( + len(MODULES), + 20, + f"expected the full package, discovered only {len(MODULES)}: {MODULES}", + ) + + +class TestModuleImports(unittest.TestCase): + def test_every_module_imports(self): + """Import each module in turn, reporting the module name on failure.""" + have_nnunet = importlib.util.find_spec("nnunetv2") is not None + + for module_name in MODULES: + with self.subTest(module=module_name): + if module_name.startswith(OPTIONAL_PREFIXES) and not have_nnunet: + self.skipTest(f"{module_name} needs the nnunetv2 extra") + importlib.import_module(module_name) + + def test_public_pipeline_entrypoints_are_importable(self): + """The classes users actually construct must be reachable from the package.""" + from smauglab.transforms.gpu.transforms import AugTransformsGPU + from smauglab.transforms.gpu.transforms_list import ( + AugTransformsGPURandomOrder, + AugTransformsGPURandomOrderTA, + ) + + for cls in (AugTransformsGPU, AugTransformsGPURandomOrder, AugTransformsGPURandomOrderTA): + with self.subTest(cls=cls.__name__): + self.assertTrue(callable(cls)) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_packaging.py b/unit_tests/test_packaging.py new file mode 100644 index 0000000..e5767f6 --- /dev/null +++ b/unit_tests/test_packaging.py @@ -0,0 +1,95 @@ +"""The built wheel must actually contain the package. + +smauglab has no __init__.py anywhere, so it is picked up purely by the build +backend's namespace-package handling. 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 importlib.util +import subprocess +import sys +import tempfile +import unittest +import zipfile +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent + + +def source_modules() -> set[str]: + """Every .py file under smauglab/, as a wheel-relative path.""" + package_root = REPO_ROOT / "smauglab" + return {str(path.relative_to(REPO_ROOT)) for path in package_root.rglob("*.py") if "__pycache__" not in path.parts} + + +# Building a wheel is slow relative to the rest of the suite, so this class is +# marked for `pytest -m "not slow"`. The mark is applied at class level, which +# is the form pytest honours on unittest.TestCase subclasses. +@pytest.mark.slow +class TestWheelContents(unittest.TestCase): + """Builds the wheel once for the whole class, then inspects it.""" + + _tmpdir: tempfile.TemporaryDirectory + wheel: Path + + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + if importlib.util.find_spec("build") is None: + raise unittest.SkipTest("the `build` package is needed to test packaging") + + cls._tmpdir = tempfile.TemporaryDirectory() + out_dir = Path(cls._tmpdir.name) + 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: + cls._tmpdir.cleanup() + raise AssertionError(f"wheel build failed:\n{result.stdout}\n{result.stderr}") + + wheels = list(out_dir.glob("*.whl")) + if len(wheels) != 1: + cls._tmpdir.cleanup() + raise AssertionError(f"expected exactly one wheel, got {wheels}") + cls.wheel = wheels[0] + + @classmethod + def tearDownClass(cls) -> None: + cls._tmpdir.cleanup() + super().tearDownClass() + + def wheel_names(self) -> list[str]: + with zipfile.ZipFile(self.wheel) as archive: + return archive.namelist() + + def test_wheel_contains_every_module(self): + shipped = {name for name in self.wheel_names() if name.endswith(".py")} + missing = source_modules() - shipped + self.assertFalse(missing, f"wheel is missing modules: {sorted(missing)}") + + def test_wheel_contains_the_config_data(self): + """The JSON configs are the package's data; without them nothing runs.""" + configs = {n for n in self.wheel_names() if n.startswith("smauglab/configs/") and n.endswith(".json")} + + self.assertTrue(configs, "wheel ships no config JSONs") + self.assertTrue( + any("transform_params_gpu" in name for name in configs), + "wheel is missing the default GPU transform config", + ) + + def test_wheel_excludes_scratch_directories(self): + """Personal scratch configs should not be published to PyPI.""" + leaked = [name for name in self.wheel_names() if "configs_paul" in name] + self.assertFalse(leaked, f"wheel ships personal scratch configs: {leaked}") + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_transforms_gpu.py b/unit_tests/test_transforms_gpu.py new file mode 100644 index 0000000..f391ba7 --- /dev/null +++ b/unit_tests/test_transforms_gpu.py @@ -0,0 +1,142 @@ +"""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 +import unittest +from pathlib import Path + +import torch + +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from unit_tests.helpers import SmaugLabTestCase, first_output + +TRANSFORM_MODULES = [ + "smauglab.transforms.gpu.contrast", + "smauglab.transforms.gpu.spatial", + "smauglab.transforms.gpu.fromSeg", + "smauglab.transforms.gpu.domain_transfer", +] + +# Not augmentations: helper modules that happen to be nn.Module subclasses. +NOT_A_TRANSFORM = {"DifferentiableHistogram3D"} + + +def discover_transforms(): + """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_transforms() + + +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 smauglab.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 + + +class TestTransformDiscovery(unittest.TestCase): + def test_discovery_found_transforms(self): + self.assertGreaterEqual( + len(DISCOVERED), + 15, + f"expected the bulk of the GPU transforms, found {len(DISCOVERED)}", + ) + + +class TestTransformsRunStandalone(SmaugLabTestCase): + def _pipeline(self, cls, signature): + """Drive a single transform the way AugTransformsGPU does.""" + return AugmentationSequentialCustom( + cls(**build_kwargs(cls, signature)), + data_keys=["input", "mask"], + same_on_batch=True, + ) + + def test_transform_runs_on_a_tiny_volume(self): + for label, cls, signature in DISCOVERED: + with self.subTest(transform=label): + reason = skip_reason(cls) + if reason: + self.skipTest(reason) + + volume, seg = self.tiny_volume(), self.tiny_seg() + image = first_output(self._pipeline(cls, signature)(volume, seg)) + + self.assertIsImageLike(image, volume, cls.__name__) + + def test_transform_leaves_the_mask_intact(self): + """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. + """ + for label, cls, signature in DISCOVERED: + with self.subTest(transform=label): + reason = skip_reason(cls) + if reason: + self.skipTest(reason) + + result = self._pipeline(cls, signature)(self.tiny_volume(), self.tiny_seg()) + if not isinstance(result, (list, tuple)) or len(result) < 2: + self.skipTest(f"{cls.__name__} does not return a mask") + + mask = result[1] + self.assertTrue(bool(torch.isfinite(mask).all()), f"{cls.__name__} produced a non-finite mask") + unique = torch.unique(mask) + self.assertLessEqual( + unique.numel(), + 2, + f"{cls.__name__} interpolated the mask into {unique.numel()} values", + ) + + +if __name__ == "__main__": + unittest.main()