diff --git a/.all-contributorsrc b/.all-contributorsrc index d479b0378..1bc81599b 100644 --- a/.all-contributorsrc +++ b/.all-contributorsrc @@ -919,24 +919,6 @@ "contributions": [ "bug" ] - }, - { - "login": "gaoflow", - "name": "Vincent Gao", - "avatar_url": "https://avatars.githubusercontent.com/u/3355843?v=4", - "profile": "https://github.com/gaoflow", - "contributions": [ - "bug", - "code" - }, - { - "login": "Jessy-Ding", - "name": "Mengyuan Ding", - "avatar_url": "https://avatars.githubusercontent.com/u/89423283?v=4", - "profile": "https://github.com/Jessy-Ding", - "contributions": [ - "bug" - ] } ], "contributorsPerLine": 7, diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 55baec499..806c04ff8 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -4,19 +4,12 @@ # https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates version: 2 -multi-ecosystem-groups: - dependencies: - schedule: - interval: "monthly" - updates: - package-ecosystem: "github-actions" directory: "/" - patterns: - - "*" - multi-ecosystem-group: "dependencies" + schedule: + interval: "monthly" - package-ecosystem: "pip" directory: "/" - patterns: - - "*" - multi-ecosystem-group: "dependencies" + schedule: + interval: "monthly" diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 392d5ebae..750ff3fba 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -13,10 +13,10 @@ Fixes #{issue_number}. +CONTRIBUTING (https://github.com/TorchIO-project/torchio/blob/main/CONTRIBUTING.rst) docs. --> -- [ ] I have read the [`CONTRIBUTING`](https://github.com/TorchIO-project/torchio/blob/main/CONTRIBUTING.md) docs and have a developer setup ready +- [ ] I have read the [`CONTRIBUTING`](https://github.com/TorchIO-project/torchio/blob/main/CONTRIBUTING.rst) docs and have a developer setup ready - Changes are - [ ] Non-breaking (would not break existing functionality) - [ ] Breaking (would cause existing functionality to change) diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml index 66849f24e..493182b29 100644 --- a/.github/workflows/code_quality.yml +++ b/.github/workflows/code_quality.yml @@ -29,7 +29,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v7 + uses: actions/checkout@v6 with: fetch-depth: 0 diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index d7f32a78f..aada3e9d7 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -1,32 +1,16 @@ -# This workflow builds the documentation and either: -# - deploys versioned docs to the `gh-pages` branch with mike (on `v1` and -# `main`), or -# - uploads a Smokeshow preview and comments the URL on the PR (other -# branches), using a single "push" trigger so that secrets are available -# even for PRs from forks. -# -# Versioning uses the Material team's fork of mike, integrated with Zensical -# (https://zensical.org/docs/setup/versioning/): -# - `v1` -> version "" with aliases `stable` + `latest` -# - `main` (v2) -> version "" with alias `dev` -# `latest` is the default alias, so docs.torchio.org/ redirects to v1 while v2 -# is a pre-release. Flip with `mike set-default` once v2 becomes stable. -# -# Both `v1` and `main` carry this workflow because Actions runs the version that -# lives on the branch that was pushed. -# -# One-time manual setup for production (cannot be done from CI): -# 1. Let `v1` deploy first so the `latest`/`stable`/default redirect exist -# before the site goes live. -# 2. Settings > Pages > Source = "Deploy from a branch" -> `gh-pages` / root. -# 3. Settings > Pages > Custom domain = `docs.torchio.org`. GitHub writes a -# CNAME file to `gh-pages`, which mike preserves on later deploys. +# This workflow handles both PR doc previews (via Smokeshow) and production +# deployment (via GitHub Pages). We use a single "push" trigger instead of +# splitting into "push" + "pull_request" because: +# - "push" ensures secrets (SMOKESHOW_AUTH_KEY) are always available, +# even for PRs from forks (pull_request events don't expose secrets). +# - A single workflow avoids duplicating the build steps. +# The tradeoff is that sticky-pull-request-comment can't auto-detect the PR +# number from a push event, so we look it up explicitly with `gh pr list`. name: Documentation on: push: - branches: ['**'] # all branches, but not tags (avoids deploying on releases) permissions: contents: write @@ -41,38 +25,49 @@ env: FORCE_COLOR: 1 jobs: - preview: - # Build + Smokeshow preview for branches/PRs that are not deployed versions. - if: github.ref != 'refs/heads/main' && github.ref != 'refs/heads/v1' + build: runs-on: ubuntu-latest steps: - name: Checkout repository - uses: actions/checkout@v7 + uses: actions/checkout@v6 - - name: Install mise-en-place - uses: jdx/mise-action@v4 - with: - github_token: ${{ secrets.GITHUB_TOKEN }} + - name: Install the latest version of uv + uses: astral-sh/setup-uv@v7 + + - name: Install dependencies + run: uv sync --group doc + + - name: Install just + uses: extractions/setup-just@v3 - name: Restore TorchIO cached data id: cache-torchio-data-restore - uses: actions/cache/restore@v6 + uses: actions/cache/restore@v5 with: path: ~/.cache/torchio key: ${{ runner.os }}-torchio-data + - name: Generate plots + run: just generate-plots + + - name: Generate examples gallery + run: just generate-gallery + - name: Build docs - run: mise run docs:build + run: just build-docs - name: Save TorchIO cached data if: steps.cache-torchio-data-restore.outputs.cache-hit != 'true' - uses: actions/cache@v6 + id: cache-torchio-data-save + uses: actions/cache@v5 with: path: ~/.cache/torchio key: ${{ steps.cache-torchio-data-restore.outputs.cache-primary-key }} + # Upload to smokeshow if not on main branch - name: Upload docs to smokeshow id: smokeshow + if: github.ref != 'refs/heads/main' env: SMOKESHOW_AUTH_KEY: ${{ secrets.SMOKESHOW_AUTH_KEY || env.SMOKESHOW_AUTH_KEY }} SMOKESHOW_GITHUB_STATUS_DESCRIPTION: Docs preview @@ -88,10 +83,10 @@ jobs: echo "preview_url=$URL" >> "$GITHUB_OUTPUT" # sticky-pull-request-comment auto-detects the PR number only on - # pull_request events. Since this workflow triggers on push, we look up the - # PR number ourselves. + # pull_request events. Since this workflow triggers on push, we need to + # look up the PR number ourselves. - name: Find PR number - if: steps.smokeshow.outputs.preview_url != '' + if: github.ref != 'refs/heads/main' && steps.smokeshow.outputs.preview_url != '' id: find-pr env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} @@ -100,8 +95,8 @@ jobs: echo "pr_number=$PR_NUMBER" >> "$GITHUB_OUTPUT" - name: Comment PR with preview URL - if: steps.smokeshow.outputs.preview_url != '' && steps.find-pr.outputs.pr_number != '' - uses: marocchino/sticky-pull-request-comment@v3 + if: github.ref != 'refs/heads/main' && steps.smokeshow.outputs.preview_url != '' && steps.find-pr.outputs.pr_number != '' + uses: marocchino/sticky-pull-request-comment@v2 with: header: docs-preview number: ${{ steps.find-pr.outputs.pr_number }} @@ -114,56 +109,32 @@ jobs: Built from ${{ github.sha }} - deploy: - # Deploy versioned docs to the gh-pages branch with mike (production). - if: github.ref == 'refs/heads/main' || github.ref == 'refs/heads/v1' - runs-on: ubuntu-latest - # Serialize deploys across branches so concurrent main/v1 runs don't race on - # the gh-pages branch. - concurrency: - group: docs-deploy - cancel-in-progress: false - steps: - - name: Checkout repository - uses: actions/checkout@v7 + # Upload to GitHub Pages if on main branch + - name: Upload artifacts using actions/upload-pages-artifact + if: github.ref == 'refs/heads/main' + id: deployment + uses: actions/upload-pages-artifact@v4 with: - fetch-depth: 0 + path: ./site - - name: Install mise-en-place - uses: jdx/mise-action@v4 - with: - github_token: ${{ secrets.GITHUB_TOKEN }} + deploy: + if: github.ref == 'refs/heads/main' - - name: Restore TorchIO cached data - id: cache-torchio-data-restore - uses: actions/cache/restore@v6 - with: - path: ~/.cache/torchio - key: ${{ runner.os }}-torchio-data + needs: build - - name: Configure git and fetch gh-pages - run: | - git config user.name "github-actions[bot]" - git config user.email "41898282+github-actions[bot]@users.noreply.github.com" - # Make the existing versioned docs available locally so mike updates - # them instead of recreating the branch from scratch. - git fetch origin gh-pages --depth=1 || true + # Grant GITHUB_TOKEN the permissions required to make a Pages deployment + permissions: + pages: write # to deploy to Pages + id-token: write # to verify the deployment originates from an appropriate source - - name: Deploy versioned docs with mike - run: | - VERSION=$(grep -m1 -E '^version = ' pyproject.toml | sed -E 's/^version = "(.*)"/\1/') - MAJOR_MINOR=$(printf '%s' "$VERSION" | sed -E 's/^([0-9]+\.[0-9]+).*/\1/') - echo "Deploying docs for $VERSION (mike version id: $MAJOR_MINOR)" - if [ "${{ github.ref }}" = "refs/heads/v1" ]; then - mise run docs:deploy -- --push "$MAJOR_MINOR" stable latest - mise run docs:set-default -- --push latest - else - mise run docs:deploy -- --push "$MAJOR_MINOR" dev - fi + runs-on: ubuntu-latest - - name: Save TorchIO cached data - if: steps.cache-torchio-data-restore.outputs.cache-hit != 'true' - uses: actions/cache@v6 - with: - path: ~/.cache/torchio - key: ${{ steps.cache-torchio-data-restore.outputs.cache-primary-key }} + # Deploy to the github-pages environment + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index f0592eee2..2fd1b1a83 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -13,7 +13,7 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v7 + - uses: actions/checkout@v6 with: fetch-depth: 0 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 8bc8f7f74..314f8e04e 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -30,7 +30,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v7 + uses: actions/checkout@v6 with: fetch-depth: 0 @@ -45,26 +45,22 @@ jobs: run: echo "C:/Users/runneradmin/.local/bin" >> $GITHUB_PATH shell: bash - - name: Install ffmpeg for video tests - if: runner.os == 'Linux' - run: sudo apt-get update && sudo apt-get install -y ffmpeg - - name: Setup test suite - run: tox run -v --notest --skip-missing-interpreters false -e test + run: tox run -v --notest --skip-missing-interpreters false -e ${{ matrix.python }} # Run all tests on schedule, but only non-slow tests on push - name: Run pytest run: | if [ "${{ github.event_name }}" == "schedule" ]; then - tox -e test + tox -e pytest else - tox -e test -- -m "not slow" + tox -e pytest -- -m "not slow" fi shell: bash # this wouldn't work on Powershell - name: Upload coverage reports to Codecov if: ${{ matrix.os == 'ubuntu-latest' && matrix.python == '3.14' }} - uses: codecov/codecov-action@v7.0.0 + uses: codecov/codecov-action@v5.5.2 env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} @@ -74,7 +70,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v7 + uses: actions/checkout@v6 with: fetch-depth: 0 @@ -82,27 +78,7 @@ jobs: uses: astral-sh/setup-uv@v7 - name: Test transforms CLI tool - run: uv run -- torchio transform --help + run: uv run -- tiotr --help - name: Test info CLI tool - run: uv run -- torchio info --help - - docs: - name: Documentation snippet tests - runs-on: ubuntu-latest - if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' - - steps: - - name: Checkout repository - uses: actions/checkout@v7 - with: - fetch-depth: 0 - - - name: Install uv - uses: astral-sh/setup-uv@v7 - - - name: Install tox - run: uv tool install --with tox-uv tox - - - name: Run doc snippet tests - run: tox -e docs-test + run: uv run -- tiohd --help diff --git a/.gitignore b/.gitignore index bf2e293ef..9432b2f71 100644 --- a/.gitignore +++ b/.gitignore @@ -115,6 +115,3 @@ uv.lock docs/images/plots/ docs/images/gallery/ docs/examples/*.md - -ideas.md -*.nii* diff --git a/.mise/tasks/bump-python b/.mise/tasks/bump-python deleted file mode 100755 index 8001b7671..000000000 --- a/.mise/tasks/bump-python +++ /dev/null @@ -1,49 +0,0 @@ -#!/usr/bin/env -S uv run --script -#MISE description="Bump minimum Python version in project files" -# /// script -# dependencies = [ -# "packaging", -# ] -# /// -from pathlib import Path - -from packaging.version import Version - -python_version_path = Path(".python-version") -old_version_string = python_version_path.read_text().strip() -old_version = Version(old_version_string) -new_version_string = f"{old_version.major}.{old_version.minor + 1}" -python_version_path.write_text(new_version_string + "\n") - -tests_workflow_path = Path(".github/workflows/tests.yml") -tests_workflow = tests_workflow_path.read_text() -tests_workflow = tests_workflow.replace( - f'"{old_version_string}"]', - f'"{old_version_string}", "{new_version_string}"]', -) -tests_workflow = tests_workflow.replace( - f"matrix.python == '{old_version_string}'", - f"matrix.python == '{new_version_string}'", -) -tests_workflow_path.write_text(tests_workflow) - -scrutinizer_config_path = Path(".scrutinizer.yml") -if scrutinizer_config_path.exists(): - scrutinizer_config = scrutinizer_config_path.read_text() - scrutinizer_config = scrutinizer_config.replace( - old_version_string, - new_version_string, - ) - scrutinizer_config_path.write_text(scrutinizer_config) - -pyproject_path = Path("pyproject.toml") -pyproject_text = pyproject_path.read_text() -old_str = f' "Programming Language :: Python :: {old_version_string}",\n' -new_str = f' "Programming Language :: Python :: {new_version_string}",\n' -pyproject_text = pyproject_text.replace( - old_str, - old_str + new_str, -) -pyproject_path.write_text(pyproject_text) - -print(f"Bumped Python {old_version_string} -> {new_version_string}") diff --git a/.mise/tasks/deprecate-python-files b/.mise/tasks/deprecate-python-files deleted file mode 100755 index ff45f1245..000000000 --- a/.mise/tasks/deprecate-python-files +++ /dev/null @@ -1,53 +0,0 @@ -#!/usr/bin/env -S uv run --script -#MISE description="Update project files for Python version deprecation" -#MISE hide=true -# /// script -# dependencies = [ -# "packaging", -# ] -# /// -from pathlib import Path -from tomllib import load - -from packaging.version import Version - -pyproject_path = Path("pyproject.toml") -with open(pyproject_path, "rb") as f: - pyproject = load(f) - -classifiers = pyproject["project"]["classifiers"] -for classifier in classifiers: - if classifier.startswith("Programming Language :: Python :: 3."): - old_version = Version(classifier.split("::")[-1].strip()) - break - -pyproject_text = pyproject_path.read_text() -to_replace = f' "Programming Language :: Python :: {old_version}",\n' -pyproject_text = pyproject_text.replace(to_replace, "") -new_version = Version(f"{old_version.major}.{old_version.minor + 1}") -pyproject_text = pyproject_text.replace( - f'requires-python = ">={old_version}"', - f'requires-python = ">={new_version}"', -) -pyproject_path.write_text(pyproject_text) - -pre_commit_path = Path(".pre-commit-config.yaml") -if pre_commit_path.exists(): - pre_commit_text = pre_commit_path.read_text() - old_version_pyupgrade = str(old_version).replace(".", "") - new_version_pyupgrade = str(new_version).replace(".", "") - pre_commit_text = pre_commit_text.replace( - f"--py{old_version_pyupgrade}-plus", - f"--py{new_version_pyupgrade}-plus", - ) - pre_commit_path.write_text(pre_commit_text) - -tests_workflow_path = Path(".github/workflows/tests.yml") -tests_workflow = tests_workflow_path.read_text() -tests_workflow = tests_workflow.replace( - f'["{old_version}", ', - "[", -) -tests_workflow_path.write_text(tests_workflow) - -print(f"Deprecated Python {old_version}, new minimum: {new_version}") diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ea891235e..f3312bf45 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -15,7 +15,7 @@ repos: - id: trailing-whitespace # Trims trailing whitespace - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.20 + rev: v0.14.10 hooks: - id: ruff - id: ruff-format @@ -27,6 +27,9 @@ repos: - id: python-check-blanket-type-ignore # enforce that # type: ignore annotations always occur with specific codes - id: python-no-log-warn # check for the deprecated .warn() method of python loggers - id: python-use-type-annotations # enforce that type annotations are used instead of type comments + - id: rst-backticks # detect common mistake of using single backticks when writing rst + - id: rst-directive-colons # detect mistake of rst directive not ending with double colon + - id: rst-inline-touching-normal # detect mistake of inline code touching normal text in rst - repo: https://github.com/asottile/pyupgrade rev: v3.21.2 @@ -45,17 +48,6 @@ repos: require_serial: true types_or: [python, pyi] - - repo: https://github.com/yunojuno/pre-commit-xenon - rev: v0.1 - hooks: - - id: xenon - args: ["--max-average=B", "--max-modules=B", "--max-absolute=B"] - - # - repo: https://github.com/regebro/pyroma - # rev: "5.0.1" - # hooks: - # - id: pyroma - ci: autoupdate_commit_msg: Autoupdate pre-commit hooks autoupdate_schedule: quarterly diff --git a/.python-version b/.python-version index 24ee5b1be..c8cfe3959 100644 --- a/.python-version +++ b/.python-version @@ -1 +1 @@ -3.13 +3.10 diff --git a/CNAME b/CNAME deleted file mode 100644 index a585f638f..000000000 --- a/CNAME +++ /dev/null @@ -1 +0,0 @@ -docs.torchio.org diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md deleted file mode 100644 index 689df18b1..000000000 --- a/CONTRIBUTING.md +++ /dev/null @@ -1,200 +0,0 @@ -# Contributing to TorchIO - -Contributions are welcome and greatly appreciated. -Every little bit helps, and credit will always be given. - -TorchIO development happens on the `main` branch. -Maintenance fixes for TorchIO v1 belong on the `v1` branch. - -## Types of contributions - -### Report bugs - -Report bugs [on GitHub](https://github.com/TorchIO-project/torchio/issues/new?assignees=&labels=bug&template=bug_report.md&title=). - -When reporting a bug, include: - -- Your TorchIO version. -- Any local setup details that might help with troubleshooting. -- Detailed steps to reproduce the bug. -- The full traceback, when there is one. - -You can print your local setup details with: - -```shell -uv run https://raw.githubusercontent.com/TorchIO-project/torchio/refs/heads/main/print_system.py -``` - -### Fix bugs - -Look through the [GitHub issues](https://github.com/TorchIO-project/torchio/issues) for bugs. -Issues tagged with `bug` and `help wanted` are open to whoever wants to implement them. - -### Implement features - -Look through the [GitHub issues](https://github.com/TorchIO-project/torchio/issues) for feature requests. -Issues tagged with `enhancement` and `help wanted` are open to whoever wants to implement them. - -### Write documentation - -TorchIO can always use more documentation, whether in the official docs, docstrings, tutorials, examples, blog posts, or articles. - -Docs are built with [Zensical](https://zensical.org/). -Follow the [Diataxis](https://diataxis.fr/) framework when adding or reorganizing documentation. - -### Submit feedback - -The best way to send feedback is to [file an issue](https://github.com/TorchIO-project/torchio/issues). - -If you are proposing a feature: - -- Explain in detail how it would work. -- Keep the scope as narrow as possible so it is easier to implement. -- Remember that this is a volunteer-driven project and that contributions are welcome. - -## Get started - -### 1. Create or find an issue - -It is good practice to discuss proposed changes before opening a pull request, because the feature might already be implemented or planned. - -### 2. Fork the repository - -[Create a fork](https://github.com/TorchIO-project/torchio/fork) of the -`TorchIO-project/torchio` repository on GitHub. - -### 3. Clone your fork locally - -```bash -git clone git@github.com:your_github_username_here/torchio.git -cd torchio -``` - -### 4. Install the development environment - -TorchIO uses [uv](https://docs.astral.sh/uv/) for Python environments and [mise](https://mise.jdx.dev/) for task automation and tool pinning. - -Install mise, then trust the repository configuration and run the setup task: - -```bash -mise trust -mise run setup -``` - -The setup task installs pinned tools, syncs all dependency groups, and installs the pre-commit hooks through [prek](https://github.com/j178/prek). - -If you do not use mise, the equivalent core setup is: - -```bash -uv sync --all-groups -uvx prek install --install-hooks -``` - -### 5. Create a branch - -Create a branch from `main` for TorchIO v2 changes. If your work addresses an -issue, start the branch name with the issue number: - -```bash -git checkout -b 55-name-of-your-bugfix-or-feature -``` - -### 6. Make your changes - -Follow the existing project style: - -- Use type annotations for new code. -- Prefer existing helpers and patterns over new one-off utilities. -- Add Google-style docstrings to new public classes, functions, and methods. -- Include usage examples in new public docstrings when practical. -- Add or update tests for behavior changes. -- Update documentation when user-facing behavior changes. - -For TorchIO v2 transforms, do not add new `Random*` transform names. -Use the v2 transform names such as `Affine`, `Flip`, and `Noise`. - -### 7. Run checks - -Run the smallest check that covers your change before opening a pull request. - -For unit tests: - -```bash -mise run test -``` - -To run a subset of tests: - -```bash -uv run tox -e test -- tests/data/test_image.py -``` - -For linting, formatting, and type checking: - -```bash -mise run quality -``` - -You can also run individual checks: - -```bash -mise run lint -mise run format-check -mise run types -``` - -To run the pre-commit hooks on all files: - -```bash -mise run prek -``` - -### 8. Check documentation - -If you changed documentation, examples, or docstrings, build the docs and test the code snippets: - -```bash -mise run docs:build -mise run docs:test -``` - -To serve the docs locally while editing: - -```bash -mise run docs:serve -``` - -### 9. Commit and push - -Stage only the files that belong to your change, then commit and push: - -```bash -git add path/to/changed_files -git commit -m "Fix nasty bug" -git push origin 55-name-of-your-bugfix-or-feature -``` - -Write clear commit messages. -These posts have useful guidance: - -- [How to Write a Git Commit Message](https://chris.beams.io/posts/git-commit/) -- [Write Joyous Git Commit Messages](https://medium.com/@joshuatauberer/write-joyous-git-commit-messages-2f98891114c4) - -### 10. Submit a pull request - -Open a pull request on GitHub. -Fill in the template, link the related issue, and mark checklist items that apply to your change. - -## Tips - -Add a fork as a remote with: - -```bash -mise run add-remote your_github_username_here -``` - -Run all tox environments with: - -```bash -mise run tox -``` diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst new file mode 100644 index 000000000..948644f27 --- /dev/null +++ b/CONTRIBUTING.rst @@ -0,0 +1,142 @@ +.. highlight:: shell + +============ +Contributing +============ + +Contributions are welcome, and they are greatly appreciated! Every little bit +helps, and credit will always be given. + +You can contribute in many ways: + +Types of contributions +---------------------- + +Report bugs +~~~~~~~~~~~ + +Report bugs +`on GitHub `_. + +If you are reporting a bug, please include: + +* Your TorchIO version. +* Any details about your local setup that might be helpful in troubleshooting. +* Detailed steps to reproduce the bug. + +Fix bugs +~~~~~~~~ + +Look through the GitHub issues for bugs. Anything tagged with "bug" and "help +wanted" is open to whoever wants to implement it. + +Implement features +~~~~~~~~~~~~~~~~~~ + +Look through the GitHub issues for features. Anything tagged with "enhancement" +and "help wanted" is open to whoever wants to implement it. + +Write documentation +~~~~~~~~~~~~~~~~~~~ + +TorchIO could always use more documentation, whether as part of the +official TorchIO docs, in docstrings, or even on the web in blog posts, +articles, and such. + +Submit feedback +~~~~~~~~~~~~~~~ + +The best way to send feedback is to file an issue at https://github.com/TorchIO-project/torchio/issues. + +If you are proposing a feature: + +* Explain in detail how it would work. +* Keep the scope as narrow as possible, to make it easier to implement. +* Remember that this is a volunteer-driven project, and that contributions + are welcome :) + +Get started! +------------ + +Ready to contribute? Here's how to set up ``torchio`` for local development. + +1) Create an issue on the GitHub repository +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +It's good practice to first discuss the proposed changes as the feature might +already be implemented. + +2) Fork the ``torchio`` repository on GitHub +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Click `here `_ to create your fork. + +3) Clone your fork locally +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: bash + + git clone git@github.com:your_github_username_here/torchio.git + cd torchio + +4) Install your local copy into a virtual environment +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +`uv `_ is recommended for development. +You can use `just `_ to set up the development environment. +This will 1) install ``uv`` if not found and 2) install ``torchio`` and all its +dependencies:: + + just setup + +5) Create a branch for local development using the issue number +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For example, if the issue number is 55:: + + git checkout -b 55-name-of-your-bugfix-or-feature + +Now you can make your changes locally. + +6) Run unit tests +~~~~~~~~~~~~~~~~~ + +When you're done making changes, check that your changes pass the tests +using ``pytest``:: + + uv run pytest -x + +7) Commit your changes and push your branch to GitHub +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +`Here's some great +advice to write good commit +messages `_, and `here's some +more `_):: + + git add . + git commit -m "Fix nasty bug" + git push origin 55-name-of-your-bugfix-or-feature + +8) Check documentation +~~~~~~~~~~~~~~~~~~~~~~ + +If you have modified the documentation or some docstrings, build the docs and +verify that everything looks good:: + + just build-docs + +You can also build, serve and automatically rebuild the docs every +time you modify them and reload them in the browser:: + + just serve-docs + +9) Submit a pull request on GitHub +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Tips +---- + +To run a subset of tests:: + + uv run pytest tests/data/test_image.py diff --git a/README.md b/README.md index 60467b9ca..e05b86c1d 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ > *Tools like TorchIO are a symptom of the maturation of medical AI research using deep learning techniques*. Jack Clark, Policy Director -at [OpenAI](https://openai.com/), Co-Founder and Head of Policy of Anthropic ([link](https://jack-clark.net/2020/03/17/)). +at [OpenAI](https://openai.com/) ([link](https://jack-clark.net/2020/03/17/)). --- @@ -409,10 +409,6 @@ Thanks goes to all these people ([emoji key](https://allcontributors.org/docs/en Anders Dahl Henriksen
Anders Dahl Henriksen
🚧 Eike Petersen
Eike Petersen
🐛 - - Vincent Gao
Vincent Gao
🐛 💻 - Mengyuan Ding
Mengyuan Ding
🐛 - diff --git a/docs/concepts/data-model.md b/docs/concepts/data-model.md deleted file mode 100644 index 91a5e77d3..000000000 --- a/docs/concepts/data-model.md +++ /dev/null @@ -1,312 +0,0 @@ -# Data model - -TorchIO's data model has five core classes: **Image**, **Points**, -**BoundingBoxes**, **Subject**, and **AffineMatrix**. This article explains -what each one does and how they relate. - -## Overview - -```mermaid -classDiagram - class Image { - +data: Tensor (C,I,J,K) - +affine: AffineMatrix - +load() / save() - } - - class ScalarImage { - Intensity data - } - - class LabelMap { - Segmentation labels - } - - class Points { - +data: Tensor (N,3) - +affine: AffineMatrix - +to_world() - } - - class BoundingBoxes { - +data: Tensor (N,6) - +format: BoundingBoxFormat - +labels: Tensor | None - +to_format() - } - - class Subject { - +images - +points - +bounding_boxes - +metadata - } - - class AffineMatrix { - +spacing - +origin - +direction - +orientation - +inverse() - +compose(other) - +apply(points) - } - - Image <|-- ScalarImage - Image <|-- LabelMap - Image --> AffineMatrix : has - Points --> AffineMatrix : has - BoundingBoxes --> AffineMatrix : has - Subject --> Image : contains - Subject --> Points : contains - Subject --> BoundingBoxes : contains -``` - -## Image - -An `Image` represents a single 3D (or multi-channel 3D) medical image. -It stores: - -- A **4D tensor** with shape $(C, I, J, K)$: channels, then three - spatial dimensions. -- An **affine matrix** mapping voxel indices $(i, j, k)$ to world - coordinates $(x, y, z)$ in millimeters. -- Optional **metadata** passed as keyword arguments. - -Images are **lazy**: data is not read from disk until first accessed. -This means you can create thousands of `Image` objects cheaply and -only load what you need. - -Images can be created from multiple sources: - - -```python -image = tio.ScalarImage("t1.nii.gz") # from file (lazy) -image = tio.ScalarImage("s3://bucket/t1.nii.gz") # from cloud (via fsspec) -image = tio.ScalarImage("https://example.com/t1.nii.gz") # from URL -image = tio.ScalarImage(buf, suffix=".nii.gz") # from file-like object -image = tio.ScalarImage(tensor) # from PyTorch tensor -image = tio.ScalarImage(sitk_image) # from SimpleITK -image = tio.ScalarImage(nifti_image) # from NiBabel (lazy) -image = tio.ScalarImage(raw_bytes) # from bytes or BytesIO -image = tio.ScalarImage(zarr_store) # from zarr Store (lazy) -``` - -### Metadata - -Any extra keyword argument is stored as metadata and accessible by -attribute or dict-style lookup: - - -```python -image = tio.ScalarImage("t1.nii.gz", protocol="MPRAGE", te=3.5) -image.protocol # "MPRAGE" -image["te"] # 3.5 -image.metadata # {"protocol": "MPRAGE", "te": 3.5} -``` - -### ScalarImage vs LabelMap - -`ScalarImage` and `LabelMap` are subclasses of `Image`. They carry no -extra data. The distinction is purely semantic: - -- **ScalarImage**: continuous intensity data (MRI signal, CT - Hounsfield units, PET SUV). -- **LabelMap**: discrete segmentation labels (0 = background, 1 = - tumor, etc.). - -Transforms use `isinstance` checks to decide behavior. For example, -spatial transforms use linear interpolation for `ScalarImage` and -nearest-neighbor for `LabelMap`. - -### Tensor layout - -TorchIO uses the convention `(C, I, J, K)`: - -| Axis | Meaning | Example | -|------|---------|---------| -| `C` | Channels | Gradient directions in DWI, components in a vector field | -| `I` | First spatial axis | Left-Right (in RAS) | -| `J` | Second spatial axis | Posterior-Anterior (in RAS) | -| `K` | Third spatial axis | Inferior-Superior (in RAS) | - -Most single-channel images (T1, CT, etc.) have `C = 1`. - -## AffineMatrix - -The `AffineMatrix` class wraps a $4 \times 4$ matrix that maps voxel indices -to world coordinates: - -$$ -\begin{bmatrix} x \\ y \\ z \\ 1 \end{bmatrix} -= -\mathbf{A} -\begin{bmatrix} i \\ j \\ k \\ 1 \end{bmatrix} -$$ - -It provides named access to the components people usually care about: - -- **`spacing`**: voxel size in mm, derived from the column norms of - the rotation-zoom block. -- **`origin`**: world coordinates of the voxel at index $(0, 0, 0)$. -- **`direction`**: $3 \times 3$ rotation matrix (spacing factored out). -- **`orientation`**: anatomical axis codes like `('R', 'A', 'S')`. - -Affines compose via the `@` operator: - - -```python -combined = affine_a @ affine_b -``` - -## Points - -A `Points` object stores an $(N, 3)$ tensor of 3D coordinates in voxel -space, together with an affine for converting to world coordinates: - - -```python -import torch -import torchio as tio - -landmarks = tio.Points( - torch.tensor([[128.0, 100.0, 90.0], [128.0, 130.0, 90.0]]), - affine=image.affine, -) -world = landmarks.to_world() # (N, 3) in mm -``` - -Use cases include anatomical landmarks, fiducial markers, and seed -points. - -## BoundingBoxes - -`BoundingBoxes` stores an $(N, 6)$ tensor of 3D bounding boxes. -Inspired by `torchvision.tv_tensors.BoundingBoxes`, extended to three -dimensions. Two formats are supported: - -The format is parameterised by **axes** and **representation**: - -- **Axes**: any permutation of `IJK` (voxel) or any valid anatomical - triplet like `RAS`, `LPI`, etc. -- **Representation**: *corners* (two opposite corners) or *center_size* - (center + extent along each axis). - -| Predefined | Axes | Representation | -|------------|------|----------------| -| `IJKIJK` | `IJK` | corners: $(i_1, j_1, k_1, i_2, j_2, k_2)$ | -| `IJKWHD` | `IJK` | center + size: $(i_c, j_c, k_c, s_i, s_j, s_k)$ | - -Custom formats are created with -`BoundingBoxFormat("RAS", "corners")`, etc. - -Convert between formats with `to_format()`. This handles -representation changes, axis permutations, and even voxel ↔ anatomical -conversions (using the stored affine). Optionally attach an integer -`labels` tensor to track per-box class IDs. - -## Subject (a.k.a. Study) - -A `Subject` groups images, points, bounding boxes, and metadata -belonging to one imaging session: - - -```python -subject = tio.Subject( - t1=tio.ScalarImage("t1.nii.gz"), - seg=tio.LabelMap("seg.nii.gz"), - landmarks=tio.Points(torch.randn(5, 3)), - tumors=tio.BoundingBoxes( - torch.tensor([[10, 20, 30, 50, 60, 70]]), - format=tio.BoundingBoxFormat.IJKIJK, - ), - age=45, -) -``` - -`Study` is an alias for `Subject`. Both refer to the same class. -In DICOM terminology, a "study" contains "series" (volumes), which -maps directly to this container. Neuroscience users tend to think -in "subjects", radiology users in "studies": - - -```python -study = tio.Study(t1=tio.ScalarImage("t1.nii.gz"), patient_id="abc") -``` - -Contents are classified automatically by type: - -- `Image` instances go to `subject.images` -- `Points` instances go to `subject.points` -- `BoundingBoxes` instances go to `subject.bounding_boxes` -- Everything else is metadata, accessible via `subject.metadata` - -All entries are accessible by name: - - -```python -subject.t1 # the ScalarImage -subject.landmarks # the Points -subject.tumors # the BoundingBoxes -subject.age # 45 -``` - -The `Subject` checks consistency across images. For example, -`subject.spatial_shape` raises an error if the images have different -spatial shapes. - -## How they fit together - -```mermaid -flowchart LR - subgraph Subject - T1[ScalarImage
t1.nii.gz] - SEG[LabelMap
seg.nii.gz] - LM["Points
landmarks"] - BB["BoundingBoxes
tumors"] - META["age: 45"] - end - - T1 -->|".affine"| A1[AffineMatrix
spacing, origin,
direction] - SEG -->|".affine"| A2[AffineMatrix] - LM -->|".affine"| A3[AffineMatrix] - T1 -->|".data"| D1["Tensor (C, I, J, K)"] - SEG -->|".data"| D2["Tensor (C, I, J, K)"] - LM -->|".data"| D3["Tensor (N, 3)"] - BB -->|".data"| D4["Tensor (N, 6)"] -``` - -A typical workflow: - -1. Create `Image` objects from file paths (lazy, no data read). -2. Create `Points` or `BoundingBoxes` from annotations. -3. Group them into a `Subject`. -4. Apply transforms to the `Subject`. This triggers loading and - produces a new `Subject` with transformed data. -5. Access `.data` tensors for training. - -## Batching - -When training a model, you need to stack subjects into batches. -`SubjectsLoader` returns `SubjectsBatch` instances where each -image entry is an `ImagesBatch` with a 5D tensor `(B, C, I, J, K)`: - - -```python -loader = tio.SubjectsLoader(dataset, batch_size=4) -for batch in loader: - batch.t1.data.shape # (4, C, I, J, K) - batch.metadata["age"] # [42, 35, 60, 28] -``` - -Each `ImagesBatch` stores per-sample affine matrices, so subjects -with different spatial properties batch correctly. - -Transforms work directly on batches. By default, transforms that -support it sample independent parameters per batch element (see -[Per-instance augmentation](per-instance-augmentation.md)): - - -```python -augmented = tio.Flip(axes=(0,))(batch) -``` diff --git a/docs/concepts/lazy-loading.md b/docs/concepts/lazy-loading.md deleted file mode 100644 index b5f757286..000000000 --- a/docs/concepts/lazy-loading.md +++ /dev/null @@ -1,243 +0,0 @@ -# Lazy loading and backends - -TorchIO images are lazy by default: creating an `Image` from a file -path, a NiBabel image, or a zarr Store reads nothing from disk. This -article explains when data actually enters memory and how the backend -system works. - -The backends are a lazy **I/O** layer, not a lazy *computation* framework. -They speed up metadata reads (shape, affine, dtype) and region slicing, but -they do not defer arithmetic or transforms. Once you access `.data`, apply a -transform, or build a batch, the full tensor is materialized in memory (see -[When tensors are materialized](#when-tensors-are-materialized)). - -## When is data loaded? - -```mermaid -stateDiagram-v2 - [*] --> Lazy: Image(path) / Image(nifti) / Image(store) - Lazy --> BackendReady: .shape / .affine / .dataobj / [slicing] - Lazy --> Loaded: .data / .load() - BackendReady --> Loaded: .data / .load() - Loaded --> Loaded: .data (cached) - - note right of Lazy: No I/O yet - note right of BackendReady: Header read,
data not in memory - note right of Loaded: Full tensor
in memory -``` - -| Access | What happens | -|--------|-------------| -| `Image(path)` | Nothing. Stores the path. | -| `Image(nifti_image)` | Nothing. Stores a reference to the nibabel object. | -| `Image(zarr_store)` | Nothing. Stores the store reference. | -| `Image(tensor)` | Immediate. The tensor is already in memory. | -| `.shape` | Creates a backend and reads the header. No data loaded. | -| `.spacing`, `.affine` | Same: reads header via backend. | -| `image[slices]` | Reads only the sliced region through the backend. Parent image stays unloaded. | -| `.data` | Loads the full tensor into memory. Cached for subsequent access. | -| `.dataobj` | Returns the raw backend for advanced use. | - -## Backends - -A **backend** is a lazy *I/O adapter*: an object that gives `Image` uniform -access to one image's data, wherever it lives, without loading the whole -volume. Intuitively, a backend is "a thing that can answer five questions", -which is exactly its contract: - -- `shape` -> always `(C, I, J, K)` -- `affine` -> the 4x4 voxel-to-world matrix -- `dtype` -> the on-disk (or in-memory) data type -- `backend[region]` -> read *just* that region, as a 4D tensor -- `to_tensor()` -> materialize the whole volume - -The first three are cheap header reads; the last two are where pixel data is -actually read. Each backend normalizes its storage-specific layout (e.g. a -NIfTI's `(I, J, K)` or `(I, J, K, C)`) into TorchIO's `(C, I, J, K)`. - -Because "backend" is an overloaded word, it helps to say what this is **not**: - -- **Not a compute backend**: it has nothing to do with `torch` devices or - kernels. -- **Not the storage format itself**: it is the *adapter* to a format. Formats - are mapped to backends by a resolver (below). -- **Not a lazy computation framework**: it defers *reads*, not arithmetic. - Transforms, batching, and queues still materialize tensors (see - [When tensors are materialized](#when-tensors-are-materialized)). - -TorchIO does not hard-code the choice in `Image`: it passes a description of the -source (a `BackendRequest`) to a small **resolver**, which consults a registry -of backends in order and returns the first match: - -| Backend | Format | How it works | -|---------|--------|-------------| -| `NibabelBackend` | `.nii`, `.nii.gz`, `nib.Nifti1Image` | Wraps nibabel's `ArrayProxy`. Uncompressed files are memory-mapped; compressed files are read through nibabel's proxy. Also used for NiBabel images passed directly to the constructor. | -| `ZarrBackend` | `.nii.zarr` | Wraps `niizarr.zarr2nii()`. Data is stored in independently compressed chunks. Only the chunks overlapping your slice are read. | -| `NibabelBackend` (via store) | `zarr.Store` | For zarr stores passed to the constructor, `zarr2nii(store)` is called on first access, producing a dask-backed nibabel image. Instantiation is O(1). | -| `TensorBackend` | In-memory | Used for images created from tensors or NumPy arrays. Wraps a PyTorch tensor directly (no numpy round-trip), preserving its device and dtype. | - -For other formats (NRRD, MHA, etc.), there is no lazy backend. Shape -and dtype can still be read from the header via SimpleITK without loading -data, but slicing triggers a full load. You can teach TorchIO about new -formats without modifying `Image`; see -[Extending the backend system](#extending-the-backend-system). - -## Practical impact - -Slicing a lazy image instead of loading it whole avoids allocating and copying -the full tensor. For formats that support random access, it also avoids -reading most of the file. Consider reading a small patch from a large volume: - - -```python -# Full load: reads and allocates the whole volume, then slices -mean_full = tio.ScalarImage("huge_volume.nii").data[:, 100:110, 100:110, 100:110].mean() - -# Lazy slice: reads only the requested region -mean_lazy = tio.ScalarImage("huge_volume.nii")[:, 100:110, 100:110, 100:110].data.mean() -``` - -How much you gain depends strongly on the format: - -- **`.nii` (uncompressed)** is memory-mapped, so the lazy path reads essentially - only the requested bytes. This is true random access and is by far the - fastest, often two orders of magnitude quicker for a small patch. -- **`.nii.zarr` (chunked)** reads only the chunks overlapping the patch, so it - also scales well, especially for remote storage. -- **`.nii.gz` (compressed)** is the subtle case: gzip is a *stream* format, not - a random-access one. To reach the requested region, nibabel must decompress - the stream from the beginning, so the lazy path still does most of the - decompression work. The speedup over a full load is real but modest: it - comes mainly from skipping the full float32 allocation and copy, not from - avoiding decompression. - -At a glance: - -| Format | Partial I/O | Notes | -|---|---|---| -| `.nii` | memory-mapped | true random access, by far the fastest | -| `.nii.zarr` | chunked | reads only overlapping chunks; great for remote storage | -| `.nii.gz` | streamed | gzip is *not* random access; modest speedup | - -## When tensors are materialized - -Laziness applies to *I/O*, not to computation. The full tensor is read into -memory the first time you do any of the following: - -- access `.data` or call `.load()` (the result is then cached); -- apply a transform: transforms operate on materialized tensors; -- build a batch or collate subjects in a `DataLoader`; -- iterate a `Queue` or sampler, which materializes each sampled patch (a - sampler may still use lazy slicing to read only that patch from disk); -- save the image, or call `.numpy()`. - -In other words, lazy reads and slicing speed up *getting at a region of the -data*; they do not turn transforms or batching into deferred operations. - -## The `dataobj` property - -For advanced use, `image.dataobj` gives direct access to the backend: - - -```python -backend = image.dataobj # NibabelBackend, ZarrBackend, or TensorBackend -backend.shape # (C, I, J, K) -backend.affine # 4x4 float64 tensor -patch = backend[:, 50:60, 50:60, 50:60] # torch.Tensor, shape (C, 10, 10, 10) -``` - -Backend slicing follows the same rules as `image[...]`: the result is always a -4D `(C, I, J, K)` `torch.Tensor`, and integer indices keep their axis (so -`backend[0]` has shape `(1, I, J, K)` rather than dropping the channel -dimension). For `TensorBackend`, the slice preserves the tensor's device and -dtype. - -This is useful when you need fine-grained control over what gets read, -or when you want to avoid even the overhead of creating a new `Image` -object. - -## Extending the backend system - -!!! note "Advanced, rarely needed" - - Most users never touch this. The built-in backends already cover NIfTI, - NIfTI-Zarr, zarr stores, NiBabel images, and in-memory tensors. Reach for a - custom backend only when you need lazy access to a format TorchIO does not - support out of the box. - -Backend selection is driven by a registry, so you can support a new format -without editing `Image`. Register a *matcher* (which decides whether a -`BackendRequest` applies) and a *factory* (which builds the backend). - -As a concrete example, here is a lazy backend for plain NumPy `.npy` volumes. -`np.load(..., mmap_mode="r")` memory-maps the file, so reading a small region -only touches the bytes you ask for, exactly like the built-in `.nii` path: - - -```python -import numpy as np -import torch - -from torchio.data import register_backend -from torchio.data.backends import BackendRequest, normalize_index - - -class NpyBackend: - """Lazy, memory-mapped backend for single-channel ``.npy`` volumes.""" - - def __init__(self, path): - self._memmap = np.load(path, mmap_mode="r") # shape (I, J, K), unread - - @property - def shape(self): - i, j, k = self._memmap.shape - return (1, i, j, k) # always (C, I, J, K) - - @property - def affine(self): - return torch.eye(4, dtype=torch.float64) # unknown here: identity - - @property - def dtype(self): - return self._memmap.dtype - - def __getitem__(self, index): - # normalize_index keeps the result 4D and never drops axes. - sc, si, sj, sk = normalize_index(index) - region = np.array(self._memmap[si, sj, sk]) # reads only this block - return torch.from_numpy(region)[None][sc] # add channel axis, then select - - def to_tensor(self): - return torch.from_numpy(np.array(self._memmap))[None] - - -register_backend( - "npy", - lambda request: request.path is not None and request.path.suffix == ".npy", - lambda request: NpyBackend(request.path), -) - -# .npy files are now first-class and lazy: -image = tio.ScalarImage("volume.npy") -print(image.shape) # read from the header, no full load -patch = image[:, 10:20, 10:20, 10:20] # reads just that block -``` - -A matcher can key off anything in the `BackendRequest` (the path, a zarr store, -the reader, and so on), and registered backends are consulted before the -built-ins, so you can also override a built-in for a given source. - -Alternatively, if you already pass a custom `reader` to a specific image, make -it a lazy reader by implementing `create_backend` (see -[Use a custom reader](../how-to/custom-reader.md)). Simple readers that just -return `(tensor, affine)` keep working unchanged: they load eagerly. - -## File format recommendations - -| Use case | Recommended format | -|----------|--------------------| -| Local training with random access | Uncompressed `.nii` (memory-mapped) | -| Storage / archival | `.nii.gz` (compressed) | -| Very large volumes, remote storage | `.nii.zarr` (chunked) | -| Large-scale datasets (100k+ volumes) | `zarr.Store` objects (O(1) instantiation) | -| Interop with non-NIfTI tools | `.nrrd`, `.mha` via SimpleITK | diff --git a/docs/concepts/per-instance-augmentation.md b/docs/concepts/per-instance-augmentation.md deleted file mode 100644 index 1f0681ed0..000000000 --- a/docs/concepts/per-instance-augmentation.md +++ /dev/null @@ -1,127 +0,0 @@ -# Per-instance augmentation - -When a transform runs on a batch, each element can receive its **own** -randomly sampled augmentation. This is the default behavior and mirrors -batched GPU augmentation libraries such as -[BatchAug](https://github.com/halleewong/batchaug) and -[Kornia](https://github.com/kornia/kornia). - -## Why per-instance - -Training augmentation works best when every sample in a mini-batch sees -a different perturbation. If a whole batch shared one rotation angle or -one noise level, the effective augmentation diversity per step would be -much lower. Sampling parameters independently per element restores that -diversity while keeping a single, vectorized call. - -## How it works - -TorchIO converts every input into a -`SubjectsBatch` of 5D tensors -`(B, C, I, J, K)`. A transform's -`make_params` step samples either one -parameter set (shared) or `B` independent sets (per-instance), and -`apply_transform` broadcasts or -loops over the batch dimension accordingly. - -```mermaid -flowchart LR - A["SubjectsBatch (B, C, I, J, K)"] --> B{"per_instance and B > 1
and transform opts in?"} - B -- yes --> C["sample B parameter sets"] - B -- no --> D["sample 1 parameter set"] - C --> E["apply per element"] - D --> F["apply to whole batch"] - E --> G["augmented batch"] - F --> G -``` - -Per-instance sampling activates only for genuine batches -(`batch_size > 1`). A single -`Subject`, `Image`, or tensor always uses the -single-sample path, so existing single-subject code is unchanged. - -## Controlling it - -Every transform inherits a `per_instance` flag (default `True`). Set it -on the individual stochastic transform whose sampling you want to -control: - - -```python -import torchio as tio - -# Independent rotation per batch element (default) -augmented = tio.Affine(degrees=(0, 45))(batch) - -# One rotation applied identically to the whole batch -augmented = tio.Affine(degrees=(0, 45), per_instance=False)(batch) -``` - -### Per-element probability - -When a transform opts into per-element probability and its probability -`p` is below 1, each batch element is gated **independently**: some -elements receive the transform and others are left unchanged. - - -```python -# About half of the batch elements get noise, sampled independently -augmented = tio.Noise(std=(0.05, 0.2), p=0.5)(batch) -``` - -Shape-changing transforms (for example a resampling target), and -transforms that have not opted into per-element probability, keep a -single batch-wide decision: masked and unmasked elements would -otherwise have incompatible shapes. - -### Choosing a different transform per element - -`OneOf` and `SomeOf` also branch per -element: each batch element independently chooses which transform (or -subset of transforms) to apply. - - -```python -# Element 0 might be flipped while element 1 is blurred -augmented = tio.OneOf([tio.Flip(axes=(0,)), tio.Blur(std=(1, 3))])(batch) -``` - -Per-element selection requires shape- and schema-preserving transforms, -so the augmented elements can be re-stacked into one batch. - -## History and inversion - -Each element keeps its own sampled parameters in the transform history. -Calling `unbatch()` returns subjects that -each carry only their own history, and elements gated out by per-element -probability omit that transform entirely. - -Invertible transforms invert each element with its own parameters. -After a per-element `OneOf`/`SomeOf`, -the batch is inverted element by element: - - -```python -restored = augmented.apply_inverse_transform() -``` - -## Capability and roll-out - -Per-instance support is advertised per transform through two -properties: - -- `supports_per_instance_params`: the transform samples parameters - independently per element. -- `supports_per_instance_p`: the transform gates each element - independently with `p`. - -Transforms that do not opt in (for example purely deterministic -preprocessing) keep batch-shared behavior even under the default -`per_instance=True`, so mixing them in a pipeline is always safe. - -!!! note "Stochastic realizations" - Some transforms (such as `Noise` and `BiasField`) draw a full - random field spanning the batch dimension. For these, - `per_instance=False` shares only the sampled *parameters* (e.g. the - noise standard deviation); the per-voxel realization still differs - across elements. diff --git a/docs/concepts/transforms.md b/docs/concepts/transforms.md deleted file mode 100644 index 7dd618c0b..000000000 --- a/docs/concepts/transforms.md +++ /dev/null @@ -1,306 +0,0 @@ -# Transform design - -TorchIO transforms are `torch.nn.Module` subclasses. They accept -Subjects, Images, Tensors, NumPy arrays, SimpleITK images, NiBabel -images, MONAI-style dicts, `ImagesBatch`, or `SubjectsBatch`, -and always return the same type. - -## Unified batch architecture - -Internally, **all inputs are converted to a `SubjectsBatch`** -before the transform runs. A single `Image` becomes a batch of -size 1; a `SubjectsBatch` from a `DataLoader` passes through -directly. This means transform authors write **one batch-oriented -application method** that works identically for single samples and batches -(`apply_transform`), plus `make_params` when parameter construction is -needed: - -```python -from typing import Any - -import torch -import torchio as tio - - -class AddValue(tio.Transform): - """Add a fixed value to every image in a batch.""" - - def __init__(self, value: float) -> None: - super().__init__() - self.value = value - - def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: - """Return the value to add.""" - return {"value": self.value} - - def apply_transform( - self, - batch: tio.SubjectsBatch, - params: dict[str, Any], - ) -> tio.SubjectsBatch: - """Add the value to each 5D image tensor.""" - for image_batch in batch.images.values(): - image_batch.data = image_batch.data + params["value"] - return batch - - -subject = tio.Subject( - image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), - site="A", -) -batch = tio.SubjectsBatch.from_subjects([subject]) -assert subject.image.data.shape == (1, 2, 3, 4) -assert batch.image.data.shape == (1, 1, 2, 3, 4) - -transformed = AddValue(2)(subject) -assert isinstance(transformed, tio.Subject) -assert transformed.image.data.shape == (1, 2, 3, 4) -assert torch.all(transformed.image.data == 2) -``` - -The public call performs the complete round trip: - -```text -Subject -> SubjectsBatch -> apply_transform -> Subject -``` - -An image tensor shaped `(C, I, J, K)` therefore reaches -`apply_transform` as `(B, C, I, J, K)`. For a single `Subject`, -`B` is 1. Negative dimension indices (`-3`, `-2`, `-1`) identify -the spatial axes for both single-element and multi-element batches. - -When a `SubjectsBatch` is passed (e.g., from `SubjectsLoader`), -transforms that support it sample **independent parameters per batch -element** by default, so a single call produces diverse augmentations -(see [Per-instance augmentation](per-instance-augmentation.md)). Pass -`per_instance=False` to share one sampled parameter set across all -elements. Fixed parameters are not sampled and therefore remain shared. -Single inputs are unaffected. - -## The `make_params` / `apply_transform` split - -Every transform has two methods: - -- **`make_params(batch)`**: create or sample parameters for the - `SubjectsBatch`. -- **`apply_transform(batch, params)`**: apply those parameters to the - `SubjectsBatch`. - -This separation (inspired by Torchvision V2) means the same random -parameters are applied consistently to all images in a `Subject`. -Parameters are saved in history for inspection and inversion. - -!!! warning "`apply_transform` is a low-level kernel" - Application code should call the transform itself, for example - `result = transform(subject)`. Calling `apply_transform` directly - bypasses input wrapping, copying, probability handling, history - recording, and output-type restoration. It requires a - `SubjectsBatch`, not a `Subject`. - -### Metadata in a batch - -`Subject.metadata` is a `dict[str, Any]`. After batching, -`batch.metadata` is a `dict[str, list[Any]]`, with one value per batch -element: - -```python -import torch -import torchio as tio - -subjects = [ - tio.Subject( - image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), - site="A", - age=30, - ), - tio.Subject( - image=tio.ScalarImage(torch.ones(1, 2, 3, 4)), - site="B", - age=40, - ), -] -batch = tio.SubjectsBatch.from_subjects(subjects) -assert batch.metadata == {"site": ["A", "B"], "age": [30, 40]} -``` - -The first subject defines the image-name and metadata-key order of the -batch. All subjects must have the same schema, although their local -key order may differ. A custom transform should preserve that shared -schema and keep every metadata list aligned with the batch dimension. - -## Scalar, range, or distribution: one class for both - -Transform parameters accept three forms. No separate -`RandomNoise` class: - - -```python -# Deterministic: always std=0.1 -tio.Noise(std=0.1) - -# Random: sample std ~ U(0.05, 0.2) each call -tio.Noise(std=(0.05, 0.2)) - -# Custom distribution: sample from any torch.distributions.Distribution -from torch.distributions import LogNormal -tio.Noise(std=LogNormal(loc=-2, scale=0.5)) -``` - -This parsing and sampling is handled internally. Any -`torch.distributions.Distribution` can be used -for full control over the sampling strategy. - -!!! note "No arguments means no augmentation" - Augmentation transforms whose strength is sampled from a range - (e.g. `Affine`, `Blur`, `Gamma`) default to a deterministic - **identity** (no-op) when constructed with no arguments, and emit a - warning. Pass a range like `(a, b)` for random augmentation, or a - scalar for a fixed effect. Transforms that draw a random realisation - instead of sampling a scalar parameter (e.g. `Noise`) still apply - with their default parameters. - -## Input flexibility - -Transforms accept multiple input types and return the same type: - - -```python -result = transform(subject) # Subject → Subject -result = transform(image) # Image → Image -result = transform(tensor) # 4D Tensor → 4D Tensor -result = transform(ndarray) # NumPy array → NumPy array -result = transform(sitk_image) # SimpleITK → SimpleITK -result = transform(nifti_image) # NiBabel → NiBabel -result = transform(data_dict) # dict → dict (MONAI-compatible) -``` - -Non-Subject inputs are wrapped in a temporary Subject internally. -Spatial metadata (spacing, affine) is preserved through the -round-trip. - -### MONAI interoperability - -Dict input makes TorchIO transforms usable in MONAI pipelines: - - -```python -# MONAI-style dict -data = {"image": tensor, "label": label_tensor, "age": 42} - -# TorchIO transforms work directly -augmented = tio.Noise(std=0.1)(data) # returns dict -augmented = tio.Flip(axes=(0,))(data) # returns dict -``` - -Tensor values are treated as images; non-tensor values pass through -unchanged. See also [`MonaiAdapter`](../how-to/monai.md) for wrapping -MONAI transforms in TorchIO pipelines. - -## Transform types - -- **`SpatialTransform`**: modifies geometry. Applies to all images - (ScalarImage and LabelMap) and transforms attached Points and - BoundingBoxes. -- **`IntensityTransform`**: modifies voxel values. Applies only to - ScalarImage, leaving LabelMap and annotations untouched. - -## Composition - -**`Compose`** runs transforms sequentially. It deep-copies the input -by default so the original data is preserved: - - -```python -pipeline = tio.Compose([ - tio.Flip(axes=(0,)), - tio.Noise(std=0.05), -]) -result = pipeline(subject) # original unchanged -``` - -**`OneOf`** picks one transform at random (with optional weights): - - -```python -augment = tio.OneOf({ - tio.Noise(std=0.1): 0.7, - tio.Blur(std=1.0): 0.3, -}) -``` - -**`SomeOf`** picks N transforms: - - -```python -augment = tio.SomeOf( - [tio.Noise(std=0.1), tio.Blur(std=(0, 2)), tio.Gamma(log_gamma=(-0.3, 0.3))], - num_transforms=2, -) -``` - -## History, traceability, and replay - -Every transform records an `AppliedTransform` in the Subject's -`applied_transforms` list: - -```python -import torch -import torchio as tio - -subject = tio.Subject(image=tio.ScalarImage(torch.zeros(1, 2, 3, 4))) -result = tio.Noise(std=0.1)(subject) -trace = result.applied_transforms[-1] -assert trace.name == "Noise" -assert trace.params["std"] == 0.1 -``` - -History parameters support inspection and inversion. TorchIO does not -currently expose a public API for applying an arbitrary saved parameter -dictionary to another input. In particular, do not use -`apply_transform(new_subject, params)` for replay: the method requires -an already wrapped `SubjectsBatch` and omits the public-call lifecycle. - -## Hydra configuration - -Transforms can export themselves as Hydra-compatible YAML configs -for reproducible experiment management: - - -```python -pipeline = tio.Compose([ - tio.Flip(axes=(0, 1), p=0.5), - tio.Noise(std=(0.05, 0.2)), -]) -cfg = pipeline.to_hydra() -``` - -```python -{ - "_target_": "torchio.Compose", - "transforms": [ - {"_target_": "torchio.Flip", "p": 0.5, "axes": [0, 1]}, - {"_target_": "torchio.Noise", "std": [0.05, 0.2]}, - ], -} -``` - -Instantiate with `hydra.utils.instantiate(cfg)`. - -## GPU and differentiability - -All transforms are pure PyTorch operations. Spatial transforms use -`torch.nn.functional.grid_sample`, which is differentiable and -GPU-compatible: - - -```python -# Augmentation on GPU or MPS -subject = subject.to("cuda") # or "mps" on Apple Silicon -result = transform(subject) # stays on device - -# Gradients flow through -with torch.enable_grad(): - result = transform(subject) - loss = model(result.t1.data) - loss.backward() -``` diff --git a/docs/data/dataset.md b/docs/data/dataset.md new file mode 100644 index 000000000..a95727a41 --- /dev/null +++ b/docs/data/dataset.md @@ -0,0 +1,5 @@ +# Dataset + +![Training with volumes](../images/diagram_volumes.svg) + +::: torchio.data.SubjectsDataset diff --git a/docs/data/image.md b/docs/data/image.md new file mode 100644 index 000000000..c73914554 --- /dev/null +++ b/docs/data/image.md @@ -0,0 +1,41 @@ +# Image + +The `Image` class, representing one medical image, +stores a 4D tensor, whose voxels encode, e.g., signal intensity or segmentation +labels, and the corresponding affine transform, +typically a rigid (Euclidean) transform, to convert +voxel indices to world coordinates in mm. +Arbitrary fields such as acquisition parameters may also be stored. + +Subclasses are used to indicate specific types of images, +such as `ScalarImage` and `LabelMap`, +which are used to store, e.g., CT scans and segmentations, respectively. + +An instance of `Image` can be created using a filepath, +a PyTorch tensor, or a NumPy array. +This class uses lazy loading, i.e., the data is not loaded from disk at +instantiation time. +Instead, the data is only loaded when needed for an operation +(e.g., if a transform is applied to the image). + +Images can be sliced using the standard NumPy / PyTorch slicing syntax. +This operation updates the coordinates origin in the affine matrix +correspondingly. + +The figure below shows two instances of `Image`. +The instance of `ScalarImage` contains a 4D tensor representing a +diffusion MRI, which contains four 3D volumes (one per gradient direction), +and the associated affine matrix. +Additionally, it stores the strength and direction for each of the four +gradients. +The instance of `LabelMap` contains a brain parcellation of the same +subject, the associated affine matrix, and the name and color of each brain +structure. + +![Data structures](../images/data_structures.png) + +::: torchio.ScalarImage + +::: torchio.LabelMap + +::: torchio.Image diff --git a/docs/data/index.md b/docs/data/index.md new file mode 100644 index 000000000..39395822d --- /dev/null +++ b/docs/data/index.md @@ -0,0 +1,8 @@ +# Data structures + +TorchIO provides data structures for medical images, subjects, and datasets. + +- [Image](image.md) — classes for representing medical images +- [Subject](subject.md) — container for images and metadata associated with a subject +- [Dataset](dataset.md) — dataset for loading and processing collections of subjects +- [Loader](loader.md) — data loader for efficient batching and augmentation diff --git a/docs/data/loader.md b/docs/data/loader.md new file mode 100644 index 000000000..f091947c2 --- /dev/null +++ b/docs/data/loader.md @@ -0,0 +1,3 @@ +# Loader + +::: torchio.data.SubjectsLoader diff --git a/docs/data/subject.md b/docs/data/subject.md new file mode 100644 index 000000000..4fbbf3754 --- /dev/null +++ b/docs/data/subject.md @@ -0,0 +1,15 @@ +# Subject + +The `Subject` is a data structure used to store +images associated with a subject and any other metadata necessary for +processing. + +Subject objects can be sliced using the standard NumPy / PyTorch slicing +syntax, returning a new subject with sliced images. +This is only possible if all images in the subject have the same spatial +shape. + +All transforms applied to a `Subject` are saved +in its `history` attribute. + +::: torchio.Subject diff --git a/docs/datasets.md b/docs/datasets.md index 1ba971a78..3de9feab2 100644 --- a/docs/datasets.md +++ b/docs/datasets.md @@ -1,27 +1,48 @@ -# Built-in datasets +# Medical image datasets -TorchIO provides demo datasets for testing and tutorials. Each -downloads on first use and caches locally. +TorchIO offers tools to easily download publicly available datasets from +different institutions and modalities. -## Synthetic +The interface is similar to `torchvision.datasets`. -::: torchio.datasets.ZonePlate +If you use any of them, please visit the corresponding website (linked in each +description) and make sure you comply with any data usage agreement and you +acknowledge the corresponding authors' publications. -## MNI +If you would like to add a dataset here, please open a +[discussion on the GitHub repository](https://github.com/TorchIO-project/torchio/discussions). -::: torchio.datasets.Colin27 +## CT-RATE + +::: torchio.datasets.CtRate + +## IXI + +::: torchio.datasets.IXI + +::: torchio.datasets.IXITiny + +## EPISURG + +::: torchio.datasets.EPISURG + +## Kaggle datasets + +::: torchio.datasets.RSNAMICCAI + +::: torchio.datasets.RSNACervicalSpineFracture + +## MNI ::: torchio.datasets.ICBM2009CNonlinearSymmetric +::: torchio.datasets.Colin27 + ::: torchio.datasets.Pediatric ::: torchio.datasets.Sheep -## IXI - -::: torchio.datasets.ixi - -::: torchio.datasets.ixi_tiny +::: torchio.datasets.BITE3 ## ITK-SNAP @@ -41,17 +62,18 @@ downloads on first use and caches locally. ## MedMNIST -3D datasets from [MedMNIST v2](https://medmnist.com/). Each function -returns a list of subjects for the requested split. +::: torchio.datasets.OrganMNIST3D -::: torchio.datasets.organ_mnist_3d +::: torchio.datasets.NoduleMNIST3D -::: torchio.datasets.nodule_mnist_3d +::: torchio.datasets.AdrenalMNIST3D -::: torchio.datasets.adrenal_mnist_3d +::: torchio.datasets.FractureMNIST3D -::: torchio.datasets.fracture_mnist_3d +::: torchio.datasets.VesselMNIST3D -::: torchio.datasets.vessel_mnist_3d +::: torchio.datasets.SynapseMNIST3D -::: torchio.datasets.synapse_mnist_3d +## ZonePlate + +::: torchio.datasets.ZonePlate diff --git a/docs/examples/plot_3d_to_2d.py b/docs/examples/plot_3d_to_2d.py index f6516513d..8d04ed25f 100644 --- a/docs/examples/plot_3d_to_2d.py +++ b/docs/examples/plot_3d_to_2d.py @@ -16,7 +16,7 @@ patches_per_volume = 2 subject = tio.datasets.Colin27() -subject.remove_image("head") +subject.remove_image('head') subjects = 50 * [subject] max_side = max(subject.shape) @@ -37,8 +37,8 @@ def plot_batch(sampler): batch = tio.utils.get_first_item(loader) _, axes = plt.subplots(4, 4, figsize=(12, 10)) - for ax, im in zip(axes.flatten(), batch["t1"]["data"], strict=True): - ax.imshow(im.squeeze(), cmap="gray") + for ax, im in zip(axes.flatten(), batch['t1']['data'], strict=True): + ax.imshow(im.squeeze(), cmap='gray') plt.suptitle(sampler.__class__.__name__) plt.tight_layout() @@ -59,7 +59,7 @@ def plot_batch(sampler): # for a :class:`torchio.WeightedSampler`. That way, we ensure that the center # of all patches correspond to brain tissue. -sampler = tio.WeightedSampler(patch_size, probability_map="brain") +sampler = tio.WeightedSampler(patch_size, probability_map='brain') plot_batch(sampler) plt.show() diff --git a/docs/examples/plot_history.py b/docs/examples/plot_history.py index a839cdc57..4281b1bad 100644 --- a/docs/examples/plot_history.py +++ b/docs/examples/plot_history.py @@ -19,7 +19,7 @@ batch_size = 4 subject = tio.datasets.FPG() -subject.remove_image("seg") +subject.remove_image('seg') subjects = 4 * [subject] transform = tio.Compose( @@ -35,12 +35,12 @@ dataset = tio.SubjectsDataset(subjects, transform=transform) transformed = dataset[0] -print("Applied transforms:") -pprint.pprint(transformed.history) -print("\nComposed transform to reproduce history:") -print(transformed.get_composed_history()) -print("\nComposed transform to invert applied transforms when possible:") -print(transformed.get_inverse_transform(ignore_intensity=False)) +print('Applied transforms:') # noqa: T201 +pprint.pprint(transformed.history) # noqa: T203 +print('\nComposed transform to reproduce history:') # noqa: T201 +print(transformed.get_composed_history()) # noqa: T201 +print('\nComposed transform to invert applied transforms when possible:') +print(transformed.get_inverse_transform(ignore_intensity=False)) # noqa: T201 loader = tio.SubjectsLoader( dataset, @@ -49,16 +49,16 @@ ) batch = tio.utils.get_first_item(loader) -print("\nTransforms applied to subjects in batch:") -pprint.pprint(batch[tio.HISTORY]) +print('\nTransforms applied to subjects in batch:') # noqa: T201 +pprint.pprint(batch[tio.HISTORY]) # noqa: T203 for i in range(batch_size): - tensor = batch["t1"][tio.DATA][i] - affine = batch["t1"][tio.AFFINE][i] + tensor = batch['t1'][tio.DATA][i] + affine = batch['t1'][tio.AFFINE][i] image = tio.ScalarImage(tensor=tensor, affine=affine) image.plot(show=False) history = batch[tio.HISTORY][i] - title = ", ".join(t.name for t in history) + title = ', '.join(t.name for t in history) plt.suptitle(title) plt.tight_layout() diff --git a/docs/examples/plot_include_exclude.py b/docs/examples/plot_include_exclude.py index 42f184266..500ccd7f8 100644 --- a/docs/examples/plot_include_exclude.py +++ b/docs/examples/plot_include_exclude.py @@ -16,8 +16,8 @@ subject.plot() transform = tio.Compose( [ - tio.RandomAffine(degrees=(20, 30), exclude=["t1"]), - tio.RandomBlur(std=(3, 4), include=["t2"]), + tio.RandomAffine(degrees=(20, 30), exclude=['t1']), + tio.RandomBlur(std=(3, 4), include=['t2']), ] ) transformed = transform(subject) diff --git a/docs/examples/plot_video.py b/docs/examples/plot_video.py index 89f9170d4..48438e3be 100644 --- a/docs/examples/plot_video.py +++ b/docs/examples/plot_video.py @@ -23,10 +23,10 @@ def read_clip(path, undersample=4): frames = [] for i in range(gif.n_frames): gif.seek(i) - frames.append(np.array(gif.convert("RGB"))) + frames.append(np.array(gif.convert('RGB'))) frames = frames[::undersample] array = np.stack(frames).transpose(3, 1, 2, 0) - delay = gif.info["duration"] + delay = gif.info['duration'] return array, delay @@ -39,19 +39,19 @@ def _update_frame(num): def get_frame(image, i): return image.data[..., i].permute(1, 2, 0).byte() - plt.rcParams["animation.embed_limit"] = 25 + plt.rcParams['animation.embed_limit'] = 25 fig, ax = plt.subplots() im = ax.imshow(get_frame(image, 0)) return animation.FuncAnimation( fig, _update_frame, - repeat_delay=image["delay"], + repeat_delay=image['delay'], frames=image.shape[-1], ) # Source: https://thehigherlearning.wordpress.com/2014/06/25/watching-a-cell-divide-under-an-electron-microscope-is-mesmerizing-gif/ -array, delay = read_clip("nBTu3oi.gif") +array, delay = read_clip('nBTu3oi.gif') plt.imshow(array[..., 0].transpose(1, 2, 0)) plt.plot() image = tio.ScalarImage(tensor=array, delay=delay) diff --git a/docs/gallery.py b/docs/gallery.py index 3db7e3f52..4abc7765c 100644 --- a/docs/gallery.py +++ b/docs/gallery.py @@ -21,17 +21,17 @@ import matplotlib -matplotlib.use("Agg") +matplotlib.use('Agg') -import matplotlib.pyplot as plt +import matplotlib.pyplot as plt # noqa: E402 # --------------------------------------------------------------------------- # Paths # --------------------------------------------------------------------------- _REPO_ROOT = Path(__file__).resolve().parent.parent -_DOCS_DIR = _REPO_ROOT / "docs" -_EXAMPLES_DIR = _DOCS_DIR / "examples" -_IMAGES_DIR = _DOCS_DIR / "images" / "gallery" +_DOCS_DIR = _REPO_ROOT / 'docs' +_EXAMPLES_DIR = _DOCS_DIR / 'examples' +_IMAGES_DIR = _DOCS_DIR / 'images' / 'gallery' # Scripts to skip (animation / external deps) _SKIP: set[str] = set() @@ -39,11 +39,11 @@ # --------------------------------------------------------------------------- # RST → Markdown helpers # --------------------------------------------------------------------------- -_RST_ROLE_RE = re.compile(r":(?:func|class|meth|attr|mod|obj):`~?([^`]+)`") -_RST_REF_RE = re.compile(r":ref:`([^`]+)`") -_RST_LINK_RE = re.compile(r"`([^<]+)<([^>]+)>`_") +_RST_ROLE_RE = re.compile(r':(?:func|class|meth|attr|mod|obj):`~?([^`]+)`') +_RST_REF_RE = re.compile(r':ref:`([^`]+)`') +_RST_LINK_RE = re.compile(r'`([^<]+)<([^>]+)>`_') _RST_HEADING_RE = re.compile(r'^(.+)\n([=\-~^"]+)$', re.MULTILINE) -_RST_DOUBLE_BACKTICK_RE = re.compile(r"``([^`]+)``") +_RST_DOUBLE_BACKTICK_RE = re.compile(r'``([^`]+)``') def _rst_to_md(text: str) -> str: @@ -54,14 +54,14 @@ def _heading_repl(m: re.Match) -> str: title = m.group(1).strip() char = m.group(2)[0] # = is h2, - is h3, ~ is h4 - level = {"=": "##", "-": "##", "~": "###", "^": "####"}.get(char, "##") - return f"{level} {title}" + level = {'=': '##', '-': '##', '~': '###', '^': '####'}.get(char, '##') + return f'{level} {title}' text = _RST_HEADING_RE.sub(_heading_repl, text) - text = _RST_ROLE_RE.sub(r"`\1`", text) - text = _RST_REF_RE.sub(r"\1", text) - text = _RST_LINK_RE.sub(r"[\1](\2)", text) - text = _RST_DOUBLE_BACKTICK_RE.sub(r"`\1`", text) + text = _RST_ROLE_RE.sub(r'`\1`', text) + text = _RST_REF_RE.sub(r'\1', text) + text = _RST_LINK_RE.sub(r'[\1](\2)', text) + text = _RST_DOUBLE_BACKTICK_RE.sub(r'`\1`', text) return text @@ -75,18 +75,18 @@ def _parse_docstring(source: str) -> tuple[str, str, str]: # Match the module docstring m = re.match(r'^("""|\'\'\')(.*?)\1', source, re.DOTALL) if not m: - return ("Untitled", "", source) + return ('Untitled', '', source) doc = m.group(2).strip() - rest = source[m.end() :].lstrip("\n") - lines = doc.split("\n") + rest = source[m.end() :].lstrip('\n') + lines = doc.split('\n') title = lines[0].strip() # Skip the underline (===) desc_lines: list[str] = [] i = 1 - while i < len(lines) and re.match(r"^[=\-~]+$", lines[i].strip()): + while i < len(lines) and re.match(r'^[=\-~]+$', lines[i].strip()): i += 1 desc_lines = [line.strip() for line in lines[i:]] - description = _rst_to_md("\n".join(desc_lines).strip()) + description = _rst_to_md('\n'.join(desc_lines).strip()) return title, description, rest @@ -99,10 +99,10 @@ def _parse_blocks(source: str) -> list[dict]: """ blocks: list[dict] = [] # Split on # %% markers - sections = re.split(r"^# %%\s*$", source, flags=re.MULTILINE) + sections = re.split(r'^# %%\s*$', source, flags=re.MULTILINE) for section in sections: - section = section.strip("\n") + section = section.strip('\n') if not section: continue _parse_section(section, blocks) @@ -117,39 +117,35 @@ def _parse_section(section: str, blocks: list[dict]) -> None: after a previous text block are treated as narrative. Indented comments (inside class/function bodies) stay as code. """ - lines = section.split("\n") + lines = section.split('\n') current_text: list[str] = [] current_code: list[str] = [] in_code = False # once we see a code line, stay in code mode - def flush_code() -> None: - if current_code: - blocks.append({"type": "code", "content": "\n".join(current_code)}) - current_code.clear() - - def flush_text() -> None: - if current_text: - text = _rst_to_md("\n".join(current_text)) - blocks.append({"type": "text", "content": text}) - current_text.clear() - for line in lines: - is_toplevel_comment = not in_code and (line.startswith("# ") or line == "#") + is_toplevel_comment = not in_code and (line.startswith('# ') or line == '#') if is_toplevel_comment: - flush_code() - text_line = line[2:] if line.startswith("# ") else "" + if current_code: + blocks.append({'type': 'code', 'content': '\n'.join(current_code)}) + current_code = [] + text_line = line[2:] if line.startswith('# ') else '' current_text.append(text_line) else: - flush_text() + if current_text: + text = _rst_to_md('\n'.join(current_text)) + blocks.append({'type': 'text', 'content': text}) + current_text = [] current_code.append(line) if line.strip(): in_code = True - flush_text() + if current_text: + text = _rst_to_md('\n'.join(current_text)) + blocks.append({'type': 'text', 'content': text}) if current_code: - code = "\n".join(current_code) + code = '\n'.join(current_code) if code.strip(): - blocks.append({"type": "code", "content": code}) + blocks.append({'type': 'code', 'content': code}) # --------------------------------------------------------------------------- @@ -174,7 +170,7 @@ def _execute_example( fig_counter = 0 # Build combined namespace for execution - namespace: dict = {"__name__": "__main__"} + namespace: dict = {'__name__': '__main__'} # Change cwd to examples dir so relative paths (e.g., GIF) work original_cwd = os.getcwd() @@ -182,51 +178,51 @@ def _execute_example( try: for block in blocks: - if block["type"] != "code": + if block['type'] != 'code': continue - code = block["content"] - # Remove plt.show() calls (we capture figures ourselves) + code = block['content'] + # Remove plt.show() calls — we capture figures ourselves code_exec = re.sub( - r"^\s*plt\.show\(\)\s*$", - "", + r'^\s*plt\.show\(\)\s*$', + '', code, flags=re.MULTILINE, ) - plt.close("all") + plt.close('all') stdout_capture = io.StringIO() try: with contextlib.redirect_stdout(stdout_capture): - exec(code_exec, namespace) + exec(code_exec, namespace) # noqa: S102 except Exception as exc: print( - f" WARNING ({stem}): {type(exc).__name__}: {exc}", + f' WARNING ({stem}): {type(exc).__name__}: {exc}', file=sys.stderr, ) - block["stdout"] = f"{type(exc).__name__}: {exc}" - block["figures"] = [] + block['stdout'] = f'{type(exc).__name__}: {exc}' + block['figures'] = [] continue - block["stdout"] = stdout_capture.getvalue() + block['stdout'] = stdout_capture.getvalue() # Save any animations as GIFs - from matplotlib.animation import FuncAnimation as _FuncAnimation + from matplotlib.animation import FuncAnimation as _FA anim_paths: list[Path] = [] for name, obj in list(namespace.items()): - if isinstance(obj, _FuncAnimation): + if isinstance(obj, _FA): fig_counter += 1 - fname = f"{stem}_{fig_counter:03d}.gif" + fname = f'{stem}_{fig_counter:03d}.gif' fpath = _IMAGES_DIR / fname try: - obj.save(str(fpath), writer="pillow") + obj.save(str(fpath), writer='pillow') anim_paths.append(fpath) all_figures.append(fpath) - print(f" Animation {fpath}", file=sys.stderr) + print(f' Animation {fpath}', file=sys.stderr) except Exception as exc: print( - f" WARNING ({stem}): animation save failed: {exc}", + f' WARNING ({stem}): animation save failed: {exc}', file=sys.stderr, ) # Remove from namespace to avoid re-saving @@ -238,13 +234,13 @@ def _execute_example( if not anim_paths: for fig in figs: fig_counter += 1 - fname = f"{stem}_{fig_counter:03d}.png" + fname = f'{stem}_{fig_counter:03d}.png' fpath = _IMAGES_DIR / fname - fig.savefig(fpath, bbox_inches="tight", dpi=150) + fig.savefig(fpath, bbox_inches='tight', dpi=150) fig_paths.append(fpath) all_figures.append(fpath) - block["figures"] = anim_paths + fig_paths - plt.close("all") + block['figures'] = anim_paths + fig_paths + plt.close('all') finally: os.chdir(original_cwd) @@ -264,49 +260,52 @@ def _generate_page( ) -> str: """Generate the markdown content for a single example page.""" lines: list[str] = [] - lines.append(f"# {title}\n") + lines.append(f'# {title}\n') if description: - lines.append(f"{description}\n") + lines.append(f'{description}\n') for block in blocks: - if block["type"] == "text": - lines.append(block["content"]) - lines.append("") - elif block["type"] == "code": - code = block["content"].rstrip() + if block['type'] == 'text': + lines.append(block['content']) + lines.append('') + elif block['type'] == 'code': + code = block['content'].rstrip() # Remove plt.show() for display code_display = re.sub( - r"\n?\s*plt\.show\(\)\s*$", - "", + r'\n?\s*plt\.show\(\)\s*$', + '', code, ).rstrip() if code_display.strip(): - lines.append(f"```python\n{code_display}\n```\n") + lines.append(f'```python\n{code_display}\n```\n') - stdout = block.get("stdout", "").strip() + stdout = block.get('stdout', '').strip() if stdout: - lines.append(f"```text\n{stdout}\n```\n") + lines.append(f'```text\n{stdout}\n```\n') - for fig_path in block.get("figures", []): + for fig_path in block.get('figures', []): rel = os.path.relpath( str(fig_path.relative_to(_DOCS_DIR)), str(script_path.stem), # page will be at examples/.md - ).replace(os.sep, "/") + ).replace(os.sep, '/') # Relative from docs/examples/ to docs/images/gallery/ rel = os.path.relpath( str(fig_path), str(_EXAMPLES_DIR), - ).replace(os.sep, "/") - lines.append(f"![{title}]({rel})\n") + ).replace(os.sep, '/') + lines.append(f'![{title}]({rel})\n') # Link to source gh_url = ( - f"https://github.com/TorchIO-project/torchio/blob/main/" - f"docs/examples/{script_path.name}" + f'https://github.com/TorchIO-project/torchio/blob/main/' + f'docs/examples/{script_path.name}' + ) + lines.append( + f'\n---\n\n' + f'[:fontawesome-brands-github: View source on GitHub]({gh_url}){{ .md-button }}\n' ) - lines.append(f"\n---\n\n[View source on GitHub]({gh_url})\n") - return "\n".join(lines) + return '\n'.join(lines) def _generate_index( @@ -314,37 +313,37 @@ def _generate_index( ) -> str: """Generate the gallery index page with a card grid.""" lines: list[str] = [] - lines.append("# Examples Gallery\n") + lines.append('# Examples Gallery\n') lines.append( - "Below is a gallery of examples covering several features of TorchIO.\n", + 'Below is a gallery of examples covering several features of TorchIO.\n', ) # Use Material grid cards lines.append('
\n') for ex in examples: - thumb = ex.get("thumbnail") - title = ex["title"] - stem = ex["stem"] - desc = ex.get("description", "") + thumb = ex.get('thumbnail') + title = ex['title'] + stem = ex['stem'] + desc = ex.get('description', '') # First line of description - short_desc = desc.split("\n")[0] if desc else "" + short_desc = desc.split('\n')[0] if desc else '' if thumb: thumb_rel = os.path.relpath( str(thumb), str(_EXAMPLES_DIR), - ).replace(os.sep, "/") - lines.append(f"- [![{title}]({thumb_rel})]({stem}.md)\n") - lines.append(f" **[{title}]({stem}.md)**\n") + ).replace(os.sep, '/') + lines.append(f'- [![{title}]({thumb_rel})]({stem}.md)\n') + lines.append(f' **[{title}]({stem}.md)**\n') else: - lines.append(f"- **[{title}]({stem}.md)**\n") + lines.append(f'- **[{title}]({stem}.md)**\n') if short_desc: - lines.append(f" {short_desc}\n") + lines.append(f' {short_desc}\n') - lines.append("
\n") - return "\n".join(lines) + lines.append('\n') + return '\n'.join(lines) # --------------------------------------------------------------------------- @@ -352,88 +351,88 @@ def _generate_index( # --------------------------------------------------------------------------- -def _make_thumbnail(figures: list[Path]) -> Path | None: - """Create a thumbnail from the first figure, returning its path.""" - if not figures: - return None - first = figures[0] - thumb_path = first.with_name(first.stem + "_thumb" + first.suffix) - from PIL import Image as PILImage - - try: - img = PILImage.open(first) - if getattr(img, "n_frames", 1) > 1: - thumb_path = thumb_path.with_suffix(".png") - img.seek(0) - img_copy = img.copy() - img_copy.thumbnail((400, 300)) - img_copy.save(thumb_path) - else: - img.thumbnail((400, 300)) - img.save(thumb_path) - print(f" Thumbnail {thumb_path}", file=sys.stderr) - return thumb_path - except Exception: - return first - - -def _load_cached_example(script: Path, hash_marker: str) -> dict | None: - """Return cached example metadata if the markdown is up to date.""" - md_path = _EXAMPLES_DIR / f"{script.stem}.md" - if not (md_path.exists() and hash_marker in md_path.read_text()): - return None - print(" Cached (hash match)", file=sys.stderr) - title, description, _ = _parse_docstring(script.read_text()) - thumbs = sorted(_IMAGES_DIR.glob(f"{script.stem}_*_thumb.png")) - return { - "stem": script.stem, - "title": title, - "description": description, - "thumbnail": thumbs[0] if thumbs else None, - } - - -def _process_script(script: Path) -> dict: - """Process a single example script into a metadata dict.""" - print(f"Processing {script.name}...", file=sys.stderr) - - source_hash = hashlib.md5(script.read_bytes()).hexdigest()[:12] - hash_marker = f"" - - cached = _load_cached_example(script, hash_marker) - if cached is not None: - return cached - - blocks, figures, title, description = _execute_example(script) - page = _generate_page(script, blocks, title, description) - page = f"{hash_marker}\n{page}" - md_path = _EXAMPLES_DIR / f"{script.stem}.md" - md_path.write_text(page) - print(f" Generated {md_path}", file=sys.stderr) - - return { - "stem": script.stem, - "title": title, - "description": description, - "thumbnail": _make_thumbnail(figures), - } - - def main() -> None: - scripts = sorted(_EXAMPLES_DIR.glob("plot_*.py")) + scripts = sorted(_EXAMPLES_DIR.glob('plot_*.py')) scripts = [s for s in scripts if s.name not in _SKIP] if not scripts: - print("No example scripts found", file=sys.stderr) + print('No example scripts found', file=sys.stderr) return - examples = [_process_script(script) for script in scripts] + examples: list[dict] = [] + + for script in scripts: + print(f'Processing {script.name}...', file=sys.stderr) + + # Check if we need to regenerate + source_hash = hashlib.md5(script.read_bytes()).hexdigest()[:12] # noqa: S324 + md_path = _EXAMPLES_DIR / f'{script.stem}.md' + hash_marker = f'' + + if md_path.exists() and hash_marker in md_path.read_text(): + print(' Cached (hash match)', file=sys.stderr) + # Still need to extract title for the index + title, description, _ = _parse_docstring(script.read_text()) + # Find existing thumbnail + thumbs = sorted(_IMAGES_DIR.glob(f'{script.stem}_*_thumb.png')) + thumb = thumbs[0] if thumbs else None + examples.append( + { + 'stem': script.stem, + 'title': title, + 'description': description, + 'thumbnail': thumb, + } + ) + continue + blocks, figures, title, description = _execute_example(script) + page = _generate_page(script, blocks, title, description) + page = f'{hash_marker}\n{page}' + md_path.write_text(page) + print(f' Generated {md_path}', file=sys.stderr) + + # Create thumbnail from first figure + thumb = None + if figures: + suffix = figures[0].suffix # .png or .gif + thumb_path = figures[0].with_name( + figures[0].stem + '_thumb' + suffix, + ) + from PIL import Image as PILImage + + try: + img = PILImage.open(figures[0]) + if getattr(img, 'n_frames', 1) > 1: + # For animated GIFs, extract the first frame as PNG thumb + thumb_path = thumb_path.with_suffix('.png') + img.seek(0) + img_copy = img.copy() + img_copy.thumbnail((400, 300)) + img_copy.save(thumb_path) + else: + img.thumbnail((400, 300)) + img.save(thumb_path) + thumb = thumb_path + print(f' Thumbnail {thumb_path}', file=sys.stderr) + except Exception: + thumb = figures[0] + + examples.append( + { + 'stem': script.stem, + 'title': title, + 'description': description, + 'thumbnail': thumb, + } + ) + + # Generate index page index_content = _generate_index(examples) - index_path = _EXAMPLES_DIR / "index.md" + index_path = _EXAMPLES_DIR / 'index.md' index_path.write_text(index_content) - print(f"Generated {index_path}", file=sys.stderr) + print(f'Generated {index_path}', file=sys.stderr) -if __name__ == "__main__": +if __name__ == '__main__': main() diff --git a/docs/get-started/installation.md b/docs/get-started/installation.md deleted file mode 100644 index 0e9efab0a..000000000 --- a/docs/get-started/installation.md +++ /dev/null @@ -1,62 +0,0 @@ -# Installation - -=== "uv" - - ``` - uv add torchio - ``` - -=== "pip" - - ``` - pip install torchio - ``` - -## Optional extras - -For NIfTI-Zarr support (chunked, lazy-loadable volumes): - -=== "uv" - - ``` - uv add torchio --extra zarr - ``` - -=== "pip" - - ``` - pip install "torchio[zarr]" - ``` - -For cloud storage (HTTP/HTTPS URLs work out of the box): - -=== "Azure Blob" - - ``` - pip install "torchio[azure]" - ``` - -=== "S3" - - ``` - pip install "torchio[s3]" - ``` - -=== "Google Cloud" - - ``` - pip install "torchio[gcs]" - ``` - -For an interactive 3D viewer in Jupyter -([NiiVue](https://niivue.com/)): - -``` -pip install "torchio[niivue]" -``` - -## Verify - -```shell -torchio --version -``` diff --git a/docs/get-started/migration.md b/docs/get-started/migration.md deleted file mode 100644 index 99c5a9652..000000000 --- a/docs/get-started/migration.md +++ /dev/null @@ -1,537 +0,0 @@ -# Migrating from v1 to v2 - -This guide covers every breaking change between TorchIO v1 and v2. - -!!! note "Report bugs or request features" - - Hit a snag migrating, or have an idea to improve TorchIO? Please - [start a discussion](https://github.com/TorchIO-project/torchio/discussions) - or [open an issue](https://github.com/TorchIO-project/torchio/issues) on - GitHub to report bugs or request features. - -## Quick checklist - -- Replace `Random*` transform names with their base names (`RandomFlip` → `Flip`) -- Pass explicit ranges for augmentation (renamed transforms are a no-op without arguments): `RandomAffine()` → `Affine(degrees=(-10, 10), scales=(0.9, 1.1))` -- Replace `path=` with positional arg or `source=` in Image constructors -- Replace `.affine` (numpy array) with `.affine.data` where a raw array is needed -- Replace `RescaleIntensity(out_min_max=...)` with `Normalize(out_min=..., out_max=...)` -- Replace `SubjectsDataset` with any `Dataset` passed to `SubjectsLoader` -- Replace `GridAggregator` with `PatchAggregator` -- Rewrite custom transforms to accept a `SubjectsBatch` in - `make_params(batch)` and `apply_transform(batch, params)` - -## Image construction - -**v1:** - - -```python -image = tio.ScalarImage(path="t1.nii.gz") -image = tio.ScalarImage(tensor=tensor, affine=affine_array) -``` - -**v2:** - - -```python -image = tio.ScalarImage("t1.nii.gz") -image = tio.ScalarImage(tensor, affine=tio.AffineMatrix(affine_array)) -``` - -Changes: - -- First positional argument accepts a path, tensor, numpy array, - NiBabel image, SimpleITK image, or `bytes`. The `path` and - `tensor` keyword names are gone: use positional or `source=`. -- `type` parameter removed. Use `ScalarImage` or `LabelMap` directly. -- `affine` accepts `AffineMatrix` objects in addition to arrays. -- New `channels_last` parameter for tensor sources shaped - `(I, J, K, C)`. - -## Affine access - -**v1:** - - -```python -affine_array = image.affine # np.ndarray (4, 4) -spacing = image.spacing # tuple -direction = image.direction # 9-tuple of floats -``` - -**v2:** - - -```python -affine_obj = image.affine # AffineMatrix object -affine_array = image.affine.data # np.ndarray (4, 4) -spacing = image.spacing # tuple (unchanged) -orientation = image.affine.orientation # e.g. ("R", "A", "S") -``` - -The `.affine` property now returns an `AffineMatrix` object. Use -`.affine.data` when you need the raw 4×4 numpy array. - -## Subject construction - -**v1:** - - -```python -subject = tio.Subject({"t1": image, "seg": label}) # dict positional arg -subject = tio.Subject(t1=image, seg=label) -``` - -**v2:** - - -```python -subject = tio.Subject(t1=image, seg=label) # keyword args only -``` - -The positional dictionary form is removed. Use keyword arguments. - -## Transform naming - -v2 removes the `Random*` prefix. Stochasticity is controlled by -parameter type: a scalar is deterministic, a tuple samples uniformly, -and a `Distribution` or `Choice` gives full control. - -!!! warning "Renaming a `Random*` transform changes its default behavior" - In v1, `RandomAffine()` (no arguments) applied random augmentation. - In v2, the renamed `Affine()` (no arguments) is a deterministic - identity (no-op) that emits a warning (randomness is opt-in). Pass a - range like `(a, b)` for random augmentation, or a scalar for a fixed - effect. Transforms that draw a random realisation rather than - sampling a scalar parameter (e.g. `Noise`, `BiasField`, - `ElasticDeformation`, `Swap`) still apply with their default - parameters. - -| v1 | v2 | -|---|---| -| `RandomFlip` | `Flip` | -| `RandomAffine` | `Affine` | -| `RandomElasticDeformation` | `ElasticDeformation` | -| `RandomNoise` | `Noise` | -| `RandomBlur` | `Blur` | -| `RandomMotion` | `Motion` | -| `RandomGhosting` | `Ghosting` | -| `RandomBiasField` | `BiasField` | -| `RandomGamma` | `Gamma` | -| `RandomSpike` | `Spike` | -| `RandomSwap` | `Swap` | -| `RandomAnisotropy` | `Anisotropy` | -| `RandomLabelsToImage` | `LabelsToImage` | -| `RescaleIntensity` | `Normalize` (alias `RescaleIntensity` available) | -| `ZNormalization` | `Standardize` (alias `ZNormalization` available) | - -## Transform parameter changes - -### Flip - -`flip_probability` default changed from **0.5** to **1.0**. If you -relied on the old default, set it explicitly: - - -```python -# v1 (implicit 0.5) -tio.RandomFlip(axes=(0, 1, 2)) - -# v2 (explicit 0.5) -tio.Flip(axes=(0, 1, 2), flip_probability=0.5) -``` - -### Affine - -`scales` and `degrees` now expect explicit ranges instead of -half-widths: - - -```python -# v1: scales=0.1 means range (0.9, 1.1) -tio.RandomAffine(scales=0.1, degrees=10) - -# v2: specify the range directly -tio.Affine(scales=(0.9, 1.1), degrees=(-10, 10)) -``` - -### Normalize (was RescaleIntensity) - -Tuple parameters are split into individual keyword arguments: - - -```python -# v1 -tio.RescaleIntensity( - out_min_max=(0, 1), - percentiles=(0.5, 99.5), -) - -# v2 -tio.Normalize( - out_min=0, - out_max=1, - percentile_low=0.5, - percentile_high=99.5, -) -``` - -Each parameter can independently be a scalar (fixed), a tuple -(uniform range), a `Distribution`, or a `Choice`. - -### HistogramStandardization - -Landmark computation is now a standalone function instead of a -classmethod: - - -```python -# v1 -landmarks = tio.HistogramStandardization.train(paths) -transform = tio.HistogramStandardization({"t1": landmarks}) - -# v2 -from torchio.transforms.intensity.histogram_standardization import ( - compute_histogram_landmarks, -) -landmarks = compute_histogram_landmarks(images) -transform = tio.HistogramStandardization(landmarks, include=["t1"]) -``` - -One instance per modality. For multi-modal subjects, compose: - - -```python -tio.Compose([ - tio.HistogramStandardization(t1_landmarks, include=["t1"]), - tio.HistogramStandardization(t2_landmarks, include=["t2"]), -]) -``` - -## Custom transforms - -### Rewrite the transform hooks - -In v1, custom transforms implemented a subject-level hook: - - -```python -# v1 -def apply_transform(self, subject: tio.Subject) -> tio.Subject: - ... -``` - -In v2, parameter creation and application are separate batch-level -hooks: - - -```python -# v2 -def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: - ... - -def apply_transform( - self, - batch: tio.SubjectsBatch, - params: dict[str, Any], -) -> tio.SubjectsBatch: - ... -``` - -Call the transform normally rather than calling either hook yourself. -For a single subject, TorchIO performs this conversion automatically: - -```text -Subject -> SubjectsBatch -> apply_transform -> Subject -``` - -The following complete transform works for both a single `Subject` and -a `SubjectsBatch`: - -```python -from typing import Any - -import torch -import torchio as tio - - -class AddValue(tio.Transform): - """Add a fixed value to every batched image.""" - - def __init__(self, value: float) -> None: - super().__init__() - self.value = value - self.received_shape: tuple[int, ...] | None = None - - def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: - """Return the value to add.""" - return {"value": self.value} - - def apply_transform( - self, - batch: tio.SubjectsBatch, - params: dict[str, Any], - ) -> tio.SubjectsBatch: - """Add the value to every image tensor.""" - for image_batch in batch.images.values(): - self.received_shape = tuple(image_batch.data.shape) - image_batch.data = image_batch.data + params["value"] - return batch - - -subject = tio.Subject(image=tio.ScalarImage(torch.zeros(1, 2, 3, 4))) -transform = AddValue(2) -result = transform(subject) -assert isinstance(result, tio.Subject) -assert transform.received_shape == (1, 1, 2, 3, 4) -assert result.image.data.shape == (1, 2, 3, 4) -assert torch.all(result.image.data == 2) -``` - -The image is 4D `(C, I, J, K)` before and after the public call, but it -is 5D `(B, C, I, J, K)` inside `apply_transform`. For a single subject, -`B` is 1. - -!!! warning "`apply_transform` is not a public replay method" - It is the low-level batch kernel. Calling it directly bypasses - wrapping, copying, probability handling, history recording, and - output-type restoration. Pass a supported input to the transform - itself instead. - -### Migrate metadata access - -In v1, subject metadata values were scalars or arbitrary objects. In a -v2 batch, each metadata key maps to a list containing one value per -element: - -```python -import torch -import torchio as tio - -subjects = [ - tio.Subject( - image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), - site="A", - age=30, - ), - tio.Subject( - image=tio.ScalarImage(torch.ones(1, 2, 3, 4)), - site="B", - age=40, - ), -] -batch = tio.SubjectsBatch.from_subjects(subjects) -assert batch.metadata == {"site": ["A", "B"], "age": [30, 40]} -``` - -Treat `batch.metadata` as `dict[str, list[Any]]`. Metadata transforms -must keep each list aligned with the batch dimension. Subjects in one -batch must have equivalent image names and metadata keys. The first -subject determines the shared key order; later subjects may use a -different local order, but custom transforms should preserve the batch -schema rather than adding, removing, or renaming keys for only some -elements. - -### Choose deterministic or per-instance behavior - -A fixed scalar is not sampled: transforms such as `Gamma` use that -value for every batch element. For built-in stochastic transforms that -support per-instance sampling, ranges and distributions produce -independent parameters for each batch element by default. Set -`per_instance=False` to share one sampled parameter set: - -```python -import torchio as tio - -independent = tio.Gamma(log_gamma=(-0.3, 0.3)) -shared = tio.Gamma(log_gamma=(-0.3, 0.3), per_instance=False) -deterministic = tio.Gamma(log_gamma=0.2) -``` - -Custom transforms do not gain per-instance sampling automatically. -Unless a transform explicitly implements and advertises that -capability, its parameters remain batch-shared. See -[Per-instance augmentation](../concepts/per-instance-augmentation.md) -for the capability contract and stochastic-realisation caveats. - -### Migrate inherently per-subject logic - -Prefer vectorized operations on 5D tensors or metadata lists. If logic -must call a subject-oriented external API, the current low-level escape -hatch is to unbatch, process every subject without changing its schema, -restack, and adopt the prior history: - -```python -from typing import Any - -import torch -import torchio as tio - - -class StripIdentifier(tio.Transform): - """Strip whitespace from subject identifiers.""" - - def make_params(self, batch: tio.SubjectsBatch) -> dict[str, Any]: - """Return no parameters.""" - return {} - - def apply_transform( - self, - batch: tio.SubjectsBatch, - params: dict[str, Any], - ) -> tio.SubjectsBatch: - """Process metadata one subject at a time.""" - subjects = batch.unbatch() - for subject in subjects: - identifier = subject.metadata["identifier"] - subject.metadata["identifier"] = identifier.strip() - rebuilt = tio.SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(batch, subjects) - return rebuilt - - -subject = tio.Subject( - image=tio.ScalarImage(torch.zeros(1, 2, 3, 4)), - identifier=" sub-01 ", -) -result = StripIdentifier()(subject) -assert result.identifier == "sub-01" -``` - -This pattern is more expensive than vectorized code and requires every -resulting subject to retain a compatible image and metadata schema. A -supported mapping utility is planned, but it is not part of the current -API. - -## New features - -### Choice - -Sample from a discrete set of values: - - -```python -tio.Affine(degrees=tio.Choice([-90, 0, 90, 180])) -``` - -### SomeOf - -Apply a random subset of transforms: - - -```python -tio.SomeOf( - [tio.Flip(axes=(0,)), tio.Noise(std=0.1), tio.Gamma(log_gamma=(-0.3, 0.3))], - num_transforms=(1, 2), -) -``` - -### Operator sugar - - -```python -pipeline = tio.Flip(axes=(0,)) + tio.Noise(std=0.1) # Compose -artifact = tio.Ghosting(intensity=(0.5, 1)) | tio.Spike(intensity=(1, 3)) # OneOf -``` - -### Compose copy control - -`Compose` deep-copies the input once, then all inner transforms -operate in-place. Disable with `copy=False` for nested pipelines: - - -```python -inner = tio.Compose([tio.Flip(axes=(0,))], copy=False) -outer = tio.Compose([inner, tio.Noise(std=0.1)]) -``` - -## Data loading - -### SubjectsDataset removed - -v1 required wrapping subjects in `SubjectsDataset`. -v2 removes this class. Pass any `Dataset` returning `Subject` instances to -`SubjectsLoader`: - - -```python -# v1 -dataset = tio.SubjectsDataset(subjects, transform=augment) -loader = DataLoader(dataset, batch_size=4) - -# v2 -loader = tio.SubjectsLoader(subjects, transform=augment, batch_size=4) -``` - -### Queue - - -```python -# v1 -dataset = tio.SubjectsDataset(subjects) -sampler = tio.UniformSampler(patch_size=96) -queue = tio.Queue(dataset, max_length=300, samples_per_volume=10) - -# v2 -queue = tio.Queue( - subjects, - patch_sampler=tio.UniformSampler(patch_size=96), - max_length=300, - patches_per_volume=10, -) -``` - -### PatchAggregator (was GridAggregator) - - -```python -# v1 -aggregator = tio.GridAggregator(sampler) - -# v2 -aggregator = tio.PatchAggregator(sampler) -``` - -## Transform history - -v2 simplifies the history API: - - -```python -# Both versions -restored = subject.apply_inverse_transform() -inverse = subject.get_inverse_transform() - -# v1 only (removed in v2) -subject.history -subject.get_applied_transforms() -subject.get_composed_history() -``` - -## Imports - -All transforms are available at the top level: - - -```python -# v1 -from torchio.transforms import RandomFlip, RandomAffine -from torchio.transforms.augmentation.intensity import RandomNoise - -# v2 -import torchio as tio -tio.Flip -tio.Affine -tio.Noise -``` - -New exports in v2: - -- `AffineMatrix`: the affine matrix class -- `Points`, `BoundingBoxes`, `BoundingBoxFormat`: annotation types -- `SubjectsBatch`, `ImagesBatch`: batch containers -- `Choice`: discrete parameter sampling utility -- `SomeOf`: random subset composition -- `PatchAggregator`: renamed from `GridAggregator` -- `apply_inverse_transform`: standalone inverse function diff --git a/docs/get-started/quickstart.md b/docs/get-started/quickstart.md deleted file mode 100644 index 0c77088b4..000000000 --- a/docs/get-started/quickstart.md +++ /dev/null @@ -1,257 +0,0 @@ -# Quickstart - -## Loading an image - - -```python -import torchio as tio - -# From a file (lazy, no data read yet) -image = tio.ScalarImage("t1.nii.gz") -print(image.shape) # reads header only: (1, 256, 256, 176) -print(image.spacing) # (1.0, 1.0, 1.0) - -# Data is loaded on first access -tensor = image.data # shape: (1, 256, 256, 176), dtype: float32 - -# From a URL or cloud path -image = tio.ScalarImage("https://example.com/t1.nii.gz") -image = tio.ScalarImage("s3://bucket/t1.nii.gz") - -# From a file-like object -import io -buf = io.BytesIO(open("t1.nii.gz", "rb").read()) -image = tio.ScalarImage(buf, suffix=".nii.gz") -``` - -## Creating from a tensor - - -```python -import torch - -tensor = torch.randn(1, 128, 128, 128) -image = tio.ScalarImage(tensor) -``` - -## Creating from SimpleITK or NiBabel - - -```python -import SimpleITK as sitk -import nibabel as nib - -# From a SimpleITK Image (preserves spacing, origin, direction) -sitk_image = sitk.ReadImage("t1.nii.gz") -image = tio.ScalarImage(sitk_image) - -# From a NiBabel Nifti1Image (lazy, data not loaded yet) -nifti = nib.load("t1.nii.gz") -image = tio.ScalarImage(nifti) -``` - -## Creating from bytes - -If you have raw image bytes (e.g., from an HTTP response or a -database), pass them directly: - - -```python -response = requests.get("https://example.com/brain.nii.gz") -image = tio.ScalarImage(response.content) - -# Or with a BytesIO buffer -import io -buf = io.BytesIO(some_bytes) -image = tio.ScalarImage(buf, suffix=".nii.gz") -``` - -## Creating from a Zarr store - -For large-scale datasets stored as `.nii.zarr`, you can pass a -`zarr.abc.store.Store` directly. Instantiation is O(1): the store -is only accessed when metadata or data is needed: - - -```python -import zarr - -store = zarr.storage.FsspecStore("s3://bucket/brain.nii.zarr", mode="r") -image = tio.ScalarImage(store) # instant, no I/O -print(image.shape) # triggers header read -image.data # triggers full load - -# Select a specific pyramid level -image = tio.ScalarImage(store, reader_kwargs={"level": 1}) -``` - -## Attaching metadata - -Pass any extra keyword arguments to attach metadata to an image: - - -```python -image = tio.ScalarImage("t1.nii.gz", protocol="MPRAGE", te=3.5) -image.protocol # "MPRAGE" (attribute access) -image["te"] # 3.5 (dict-style access) -image.metadata # {"protocol": "MPRAGE", "te": 3.5} -``` - -## Slicing - -Slicing follows the `(C, I, J, K)` layout and keeps things lazy. Only -the requested region is read from disk: - - -```python -image = tio.ScalarImage("big_volume.nii.gz") -patch = image[:, 100:200, 100:200, 50:100] # no full load -patch.data.mean() # reads only this region -``` - -The affine origin is updated automatically so the patch stays in the -correct world coordinates. - -## Grouping data into a Subject - -A `Subject` (also available as `Study`) holds images, spatial -annotations, and metadata: - - -```python -import torch - -subject = tio.Subject( - t1=tio.ScalarImage("t1.nii.gz"), - seg=tio.LabelMap("seg.nii.gz"), - landmarks=tio.Points(torch.tensor([[64.0, 64.0, 32.0]])), - tumors=tio.BoundingBoxes( - torch.tensor([[10, 20, 30, 50, 60, 70]]), - format=tio.BoundingBoxFormat.IJKIJK, - ), - age=45, -) - -subject.t1 # Image access -subject.landmarks # Points access -subject.tumors # BoundingBoxes access -subject["seg"] # dict-style access -subject.age # metadata access (returns 45) -``` - -## Saving - - -```python -# Any format SimpleITK supports -image.save("output.nii.gz") -image.save("output.nrrd") - -# NIfTI-Zarr (chunked, lazy-loadable) -image.save("output.nii.zarr") -``` - -## Batching with a DataLoader - - -```python -from torch.utils.data import Dataset - -class BrainDataset(Dataset): - def __init__(self, paths): - self.subjects = [ - tio.Subject( - t1=tio.ScalarImage(p / "t1.nii.gz"), - seg=tio.LabelMap(p / "seg.nii.gz"), - ) - for p in paths - ] - - def __len__(self): - return len(self.subjects) - - def __getitem__(self, idx): - return self.subjects[idx] - -loader = tio.SubjectsLoader(BrainDataset(paths), batch_size=4) -batch = next(iter(loader)) -batch.t1.data.shape # (4, 1, 256, 256, 176) -``` - -See the [DataLoader how-to guide](../how-to/dataloader.md) for more -details. - -## Applying transforms - -Transforms accept Subjects, Images, Tensors, NumPy arrays, -SimpleITK Images, NiBabel images, or MONAI-style dicts, and return -the same type: - - -```python -# Single deterministic transform -flipped = tio.Flip(axes=(0,))(subject) - -# Random augmentation pipeline -augment = tio.Compose([ - tio.Flip(axes=(0, 1, 2), p=0.5), - tio.Noise(std=(0.01, 0.1)), # random std each call -]) -augmented = augment(subject) - -# Custom distribution for parameters -from torch.distributions import LogNormal -noisy = tio.Noise(std=LogNormal(loc=-2, scale=0.5))(subject) - -# Works directly on tensors too -noisy_tensor = tio.Noise(std=0.05)(tensor) - -# Works with MONAI-style dicts -data = {"image": tensor, "label": label_tensor} -augmented = tio.Noise(std=0.1)(data) # returns dict - -# Works on batches from SubjectsLoader (same params, vectorised) -batch = next(iter(loader)) # SubjectsBatch -augmented_batch = augment(batch) -``` - -See the [transform design concepts](../concepts/transforms.md) for -the full picture. - -## Where to go next - -
- -- **Tutorials** - - --- - - Step-by-step walkthroughs for common workflows. - - → [Tutorials](../tutorials/first-pipeline.md) - -- **How-to guides** - - --- - - Recipes for specific tasks: custom readers, NIfTI-Zarr, etc. - - → [How-to guides](../how-to/dataloader.md) - -- **Concepts** - - --- - - Understand the design: lazy loading, affines, backends. - - → [Concepts](../concepts/data-model.md) - -- **API Reference** - - --- - - Complete reference for all classes and functions. - - → [Reference](../reference/image.md) - -
diff --git a/docs/getting-started.md b/docs/getting-started.md new file mode 100644 index 000000000..5970cdf61 --- /dev/null +++ b/docs/getting-started.md @@ -0,0 +1,140 @@ +# Getting started + +## Installation + +The Python package is hosted on the +[Python Package Index (PyPI)](https://pypi.org/project/torchio/). + +=== "uv" + + ``` + uv add torchio + ``` + +=== "pip" + + ``` + pip install torchio + ``` + +=== "conda" + + ``` + conda install -c conda-forge torchio + ``` + +!!! note "Optional extras" + + TorchIO provides optional extras for additional functionality: + + - **`plot`** – Plotting support ([Matplotlib](https://matplotlib.org/), [colorcet](https://colorcet.holoviz.org/)): `torchio[plot]` + - **`csv`** – CSV/tabular data support ([pandas](https://pandas.pydata.org/)): `torchio[csv]` + - **`video`** – Video export ([ffmpeg-python](https://github.com/kkroening/ffmpeg-python)): `torchio[video]` + - **`sklearn`** – Scikit-learn integration ([scikit-learn](https://scikit-learn.org/)): `torchio[sklearn]` + + Install extras with your package manager, e.g.: + + === "uv" + + ``` + uv add torchio --extra plot --extra csv + ``` + + === "pip" + + ``` + pip install "torchio[plot,csv]" + ``` + +## Hello, World! + +This example shows the basic usage of TorchIO, where an instance of +[`SubjectsDataset`](data/dataset.md#torchio.data.SubjectsDataset) is passed to +a PyTorch [`SubjectsLoader`](data/loader.md#torchio.data.SubjectsLoader) to generate training batches +of 3D images that are loaded, preprocessed and augmented on the fly, +in parallel: + +```python +import torch +import torchio as tio + +# Each instance of tio.Subject is passed arbitrary keyword arguments. +# Typically, these arguments will be instances of tio.Image +subject_a = tio.Subject( + t1=tio.ScalarImage('subject_a.nii.gz'), + label=tio.LabelMap('subject_a.nii'), + diagnosis='positive', + age=36, +) + +# Image files can be in any format supported by SimpleITK or NiBabel, including DICOM +subject_b = tio.Subject( + t1=tio.ScalarImage('subject_b_dicom_folder/'), + label=tio.LabelMap('subject_b_seg.nrrd'), + diagnosis='negative', + age=24, +) + +# Images may also be created using PyTorch tensors or NumPy arrays +tensor_4d = torch.rand(4, 100, 100, 100) +subject_c = tio.Subject( + t1=tio.ScalarImage(tensor=tensor_4d), + label=tio.LabelMap(tensor=(tensor_4d > 0.5)), + diagnosis='negative', + age=19, +) + +subjects_list = [ + subject_a, + subject_b, + subject_c, +] + +# Let's use one preprocessing transform and one augmentation transform +# This transform will be applied only to scalar images: +rescale = tio.RescaleIntensity(out_min_max=(0, 1)) + +# As RandomAffine is faster then RandomElasticDeformation, we choose to +# apply RandomAffine 80% of the times and RandomElasticDeformation the rest +# Also, there is a 25% chance that none of them will be applied +spatial = tio.OneOf({ + tio.RandomAffine(): 0.8, + tio.RandomElasticDeformation(): 0.2, + }, + p=0.75, +) + +# Transforms can be composed as in torchvision.transforms +transforms = [rescale, spatial] +transform = tio.Compose(transforms) + +# SubjectsDataset is a subclass of torch.data.utils.Dataset +subjects_dataset = tio.SubjectsDataset(subjects_list, transform=transform) + +# Images are processed in parallel thanks to a SubjectsLoader +# (which inherits from torch.utils.data.DataLoader) +training_loader = tio.SubjectsLoader( + subjects_dataset, + batch_size=4, + num_workers=4, + shuffle=True, +) + +# Training epoch +for subjects_batch in training_loader: + inputs = subjects_batch['t1'][tio.DATA] + target = subjects_batch['label'][tio.DATA] +``` + +## Tutorials + +[![Google Colab notebook](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/TorchIO-project/torchio/blob/main/tutorials/README.md) + +The best way to quickly understand and try the library is the +[Jupyter Notebooks](https://github.com/TorchIO-project/torchio/blob/main/tutorials/README.md) +hosted on Google Colab. + +They include multiple examples and visualization of most of the classes, +including training of a [3D U-Net](https://github.com/fepegar/unet) for +brain segmentation on $T_1$-weighted MRI with full volumes and +with subvolumes (aka patches or windows). diff --git a/docs/how-to/annotations.md b/docs/how-to/annotations.md deleted file mode 100644 index e1c9d560e..000000000 --- a/docs/how-to/annotations.md +++ /dev/null @@ -1,110 +0,0 @@ -# Add annotations to a Subject - -Use [`Points`](../reference/points.md) and -[`BoundingBoxes`](../reference/bboxes.md) to attach spatial annotations -to a [`Subject`](../reference/subject.md) alongside its images. - -## Landmarks and fiducials - -If you have a set of 3D coordinates (e.g., anatomical landmarks or -fiducial markers), store them as a `Points` object: - -```python -import torch -import torchio as tio - -landmarks = tio.Points( - torch.tensor([ - [128.0, 100.0, 90.0], # anterior commissure - [128.0, 130.0, 90.0], # posterior commissure - ]), -) - -subject = tio.Subject( - t1=tio.ScalarImage("t1.nii.gz"), - landmarks=landmarks, -) -``` - -Coordinates are in **voxel space** (`IJK`) by default. Convert to a -different axis convention: - - -```python -# To world (mm) coordinates via the affine -world_coords = subject.landmarks.to_world() - -# To any axis convention -ras_points = subject.landmarks.to_axes("RAS") -lpi_points = subject.landmarks.to_axes("LPI") -``` - -## Region-of-interest boxes - -For regions of interest such as lesion detections or organ bounding -boxes, use `BoundingBoxes`: - - -```python -detections = tio.BoundingBoxes( - torch.tensor([ - [50, 60, 40, 100, 110, 90], # lesion 1 - [120, 80, 70, 160, 130, 110], # lesion 2 - ]), - format=tio.BoundingBoxFormat.IJKIJK, -) - -subject = tio.Subject( - t1=tio.ScalarImage("t1.nii.gz"), - seg=tio.LabelMap("seg.nii.gz"), - detections=detections, -) -``` - -## Attaching class labels - -Pass an integer tensor of labels to track which class each box belongs -to: - - -```python -detections = tio.BoundingBoxes( - torch.tensor([ - [50, 60, 40, 100, 110, 90], - [120, 80, 70, 160, 130, 110], - ]), - format=tio.BoundingBoxFormat.IJKIJK, - labels=torch.tensor([1, 2]), # e.g., 1=tumor, 2=edema -) -``` - -## Switching box format - -Convert between representations (corners vs center+size) and axis -conventions: - - -```python -# Corners → center + size -whd = detections.to_format(tio.BoundingBoxFormat.IJKWHD) - -# To a different axis convention -from torchio import BoundingBoxFormat -ras_boxes = detections.to_format(BoundingBoxFormat("RAS", "corners")) - -# Custom: KJI center+size -kji_cs = detections.to_format(BoundingBoxFormat("KJI", "center_size")) -``` - -## Accessing annotations from the Subject - -Iterate over specific annotation types: - - -```python -for name, pts in subject.points.items(): - print(f"{name}: {pts.num_points} points") - -for name, bbs in subject.bounding_boxes.items(): - print(f"{name}: {bbs.num_boxes} boxes") -``` diff --git a/docs/how-to/custom-reader.md b/docs/how-to/custom-reader.md deleted file mode 100644 index 5e5c8fccf..000000000 --- a/docs/how-to/custom-reader.md +++ /dev/null @@ -1,122 +0,0 @@ -# Use a custom reader - -TorchIO reads NIfTI files with NiBabel and everything else with -SimpleITK. If your data is in a format neither supports (e.g., a -custom binary format or a NumPy `.npy` file), you can pass a custom -reader. - -## Write a reader function - -A reader is a callable that takes a `Path` and returns a tuple of -`(tensor, affine_array)`: - -```python -from pathlib import Path -import numpy as np -import torch - -def npy_reader(path: Path) -> tuple[torch.Tensor, np.ndarray]: - data = np.load(path) - tensor = torch.from_numpy(data).unsqueeze(0) # add channel dim - affine = np.eye(4) # identity affine (1mm isotropic) - return tensor, affine -``` - -The tensor must be 4D with shape `(C, I, J, K)`. - -## Use it - - -```python -import torchio as tio - -image = tio.ScalarImage("brain.npy", reader=npy_reader) -print(image.shape) # triggers the reader -print(image.spacing) # (1.0, 1.0, 1.0) from the identity affine -``` - -!!! note "Simple readers load eagerly" - - A reader that only returns `(tensor, affine)` cannot read metadata or - regions lazily, so operations like `.shape`, `.dtype`, and slicing trigger - a full load through your reader. This is unchanged behavior. To opt into - lazy access, make your reader a *lazy reader* (below). - -## Lazy custom readers - -!!! note "Advanced, rarely needed" - - Most custom readers can stay simple and eager. Only reach for a lazy reader - when your format can cheaply read metadata or sub-regions and you actually - care about avoiding full loads (for example, very large volumes). - -If your format supports reading the shape, affine, dtype, or sub-regions -without loading everything, implement `create_backend` so TorchIO can access -it lazily. A lazy reader is any object that has a `create_backend` method -returning an object implementing the -[`ImageDataBackend`](../reference/backends.md) protocol. - -Building on the `.npy` reader above, this version is lazy: it returns a -memory-mapped backend, so `.shape` reads only the header and slicing reads only -the requested block. - - -```python -from pathlib import Path - -import numpy as np -import torch -from torchio.data.backends import BackendRequest, ImageDataBackend, normalize_index - - -class NpyBackend: - """Lazy, memory-mapped backend for single-channel ``.npy`` volumes.""" - - def __init__(self, path: Path): - self._memmap = np.load(path, mmap_mode="r") # shape (I, J, K), unread - - @property - def shape(self): - i, j, k = self._memmap.shape - return (1, i, j, k) - - @property - def affine(self): - return torch.eye(4, dtype=torch.float64) - - @property - def dtype(self): - return self._memmap.dtype - - def __getitem__(self, index): - sc, si, sj, sk = normalize_index(index) - return torch.from_numpy(np.array(self._memmap[si, sj, sk]))[None][sc] - - def to_tensor(self): - return torch.from_numpy(np.array(self._memmap))[None] - - -class LazyNpyReader: - """A custom reader for ``.npy`` that supports lazy access.""" - - def __call__(self, path: Path, **kwargs) -> tuple: - # Eager fallback, used only if create_backend is unavailable. - backend = self.create_backend(BackendRequest(path=path)) - return backend.to_tensor(), backend.affine.numpy() - - def create_backend(self, request: BackendRequest) -> ImageDataBackend: - return NpyBackend(request.path) - - -image = tio.ScalarImage("volume.npy", reader=LazyNpyReader()) -print(image.shape) # read from the header, no full load -``` - -With a lazy reader, `.shape`, `.affine`, `.dtype`, and `image.dataobj[...]` -slicing all go through your backend without materializing the full tensor. - -Passing `reader=...` is per image. If instead you want *every* `.npy` file to -use this backend, register it once globally with -[`register_backend`](../reference/backends.md); see -[Lazy loading and backends](../concepts/lazy-loading.md) for the backend -contract and that registry-based alternative. diff --git a/docs/how-to/dataloader.md b/docs/how-to/dataloader.md deleted file mode 100644 index 4a59ee3b6..000000000 --- a/docs/how-to/dataloader.md +++ /dev/null @@ -1,156 +0,0 @@ -# Load subjects with a DataLoader - -Use `SubjectsLoader` to iterate over batches of subjects during -training. It wraps PyTorch's `DataLoader` and returns -`SubjectsBatch` instances with stacked 5D tensors. - -## Basic usage - - -```python -from torch.utils.data import Dataset -import torchio as tio - - -class MyDataset(Dataset): - def __init__(self, paths): - self.subjects = [ - tio.Subject( - image=tio.ScalarImage(p / "image.nii.gz"), - seg=tio.LabelMap(p / "seg.nii.gz"), - ) - for p in paths - ] - - def __len__(self): - return len(self.subjects) - - def __getitem__(self, idx): - return self.subjects[idx] - - -dataset = MyDataset(paths) -loader = tio.SubjectsLoader(dataset, batch_size=4, num_workers=4) - -for batch in loader: - images = batch.image.data # (4, 1, H, W, D) - segs = batch.seg.data # (4, 1, H, W, D) - # ... train your model -``` - -## Accessing metadata in a batch - -Metadata is stored as lists (one value per sample): - - -```python -batch.metadata["age"] # [42, 35, 60, 28] -batch.metadata["name"] # ["sub_0", "sub_1", "sub_2", "sub_3"] -``` - -## Unbatching - -Split a batch back into individual subjects: - - -```python -subjects = batch.unbatch() -for subject in subjects: - print(subject.image.shape) # (1, H, W, D) -``` - -## Using a plain DataLoader - -If you prefer not to use `SubjectsLoader`, pass `collate_subjects` -as the collation function: - - -```python -from torch.utils.data import DataLoader -import torchio as tio - -loader = DataLoader( - dataset, - batch_size=4, - collate_fn=tio.collate_subjects, -) -``` - -## How it works - -Each image's 4D tensor is stacked into a 5D `ImagesBatch` -`(B, C, I, J, K)`. Per-sample affine matrices are stored as a -list. Metadata is collected into lists. - -## Applying transforms to batches - -Transforms work directly on `SubjectsBatch`: - - -```python -batch = next(iter(loader)) -augmented = tio.Flip(axes=(0,), p=0.5)(batch) -augmented.image.data.shape # (4, 1, H, W, D) -``` - -### Per-instance augmentation - -By default, transforms that support it sample **independent -parameters for each element of a batch**, so a single call produces -diverse augmentations (similar to -[BatchAug](https://github.com/halleewong/batchaug) and -[Kornia](https://github.com/kornia/kornia)). For example, a batch -passed through `tio.Affine(degrees=(0, 45))` receives a different -rotation per element, and `tio.Gamma(log_gamma=(-0.3, 0.3))` a -different gamma per element. - -When a transform opts into per-element probability and `p` is below -1, each element is also gated independently: some elements receive -the transform and others are left unchanged. - -To recover the legacy behavior, where one parameter set is sampled -and shared across every element, pass `per_instance=False`: - - -```python -# Same rotation applied to every element in the batch -augmented = tio.Affine(degrees=(0, 45), per_instance=False)(batch) -``` - -Single inputs (a lone `Subject`, `Image`, or tensor) are unaffected -by this flag, since there is only one element to augment. - -Per-instance parameters are recorded per element in the transform -history, so inverting transforms and unbatching back into individual -subjects keep each element's own parameters. - -!!! note - Per-instance support is rolled out per transform. Transforms that - have not been converted yet fall back to batch-shared parameters - even when `per_instance=True`. A transform advertises its support - through the `supports_per_instance_params` and - `supports_per_instance_p` properties. - - -## Loading images without a Subject - -If your dataset returns individual `Image` objects (not `Subject`), -use `ImagesLoader`: - - -```python -class SliceDataset(Dataset): - def __init__(self, paths): - self.images = [tio.ScalarImage(p) for p in paths] - - def __len__(self): - return len(self.images) - - def __getitem__(self, idx): - return self.images[idx] - -loader = tio.ImagesLoader(SliceDataset(paths), batch_size=4) -batch = next(iter(loader)) -batch.data.shape # (4, 1, H, W, D) -batch.affines # list of 4 AffineMatrix instances -``` diff --git a/docs/how-to/monai.md b/docs/how-to/monai.md deleted file mode 100644 index 221916353..000000000 --- a/docs/how-to/monai.md +++ /dev/null @@ -1,92 +0,0 @@ -# MONAI interoperability - -TorchIO and MONAI can work together in both directions. - -## Use MONAI transforms in TorchIO - -Use `MonaiAdapter` to wrap any MONAI transform for use inside TorchIO -pipelines. - -### Installation - -``` -pip install "torchio[monai]" -``` - -### Array transforms - -Array transforms (e.g., `NormalizeIntensity`) are applied to each -`ScalarImage` in the subject individually: - - -```python -from monai.transforms import NormalizeIntensity -import torchio as tio - -adapter = tio.MonaiAdapter(NormalizeIntensity()) -result = adapter(subject) -``` - -Use `include` / `exclude` to control which images are affected: - - -```python -adapter = tio.MonaiAdapter(NormalizeIntensity(), include=["t1"]) -``` - -### Dictionary transforms - -Dictionary transforms (e.g., `NormalizeIntensityd`) operate on the -full subject dictionary. Only the keys specified in the MONAI -transform are modified: - - -```python -from monai.transforms import NormalizeIntensityd - -adapter = tio.MonaiAdapter(NormalizeIntensityd(keys=["t1"])) -result = adapter(subject) -``` - -Spatial dictionary transforms (e.g., `RandSpatialCropd`) propagate -affine changes back to the TorchIO images automatically. - -### Inside a pipeline - -`MonaiAdapter` works in `Compose` like any other transform: - - -```python -pipeline = tio.Compose([ - tio.Flip(axes=(0,), p=0.5), - tio.MonaiAdapter(NormalizeIntensity()), - tio.Noise(std=(0.01, 0.05)), -]) -result = pipeline(subject) -``` - -!!! note - `MonaiAdapter` does **not** record itself in the subject's - transform history, because MONAI transform objects are not - serializable. - -## Use TorchIO transforms in MONAI - -TorchIO transforms accept `dict[str, Tensor]` directly, so they -work in MONAI dict-based pipelines without any adapter: - - -```python -from monai.transforms import Compose as MonaiCompose -from monai.transforms import NormalizeIntensityd -import torchio as tio - -pipeline = MonaiCompose([ - NormalizeIntensityd(keys=["image"]), - tio.Noise(std=0.1), # works on dicts - tio.Flip(axes=(0,), p=0.5), # works on dicts -]) - -data = {"image": tensor, "label": label_tensor} -result = pipeline(data) -``` diff --git a/docs/how-to/patch-inference.md b/docs/how-to/patch-inference.md deleted file mode 100644 index 25e5948f5..000000000 --- a/docs/how-to/patch-inference.md +++ /dev/null @@ -1,125 +0,0 @@ -# Patch-based inference - -3D medical images often don't fit in GPU memory. TorchIO provides -samplers and an aggregator to process volumes in patches. - -## Dense inference with GridSampler - -Extract all patches on a regular grid, run the model on each batch, -and reassemble the output: - - -```python -import torchio as tio -from torch.utils.data import DataLoader -from torchio.loader import SubjectsLoader - -subject = tio.ScalarImage("brain.nii.gz") -subject = tio.Subject(t1=subject) - -# 1. Sample patches -sampler = tio.GridSampler(subject, patch_size=64, patch_overlap=8) - -# 2. Batch with DataLoader -loader = SubjectsLoader(sampler, batch_size=4) - -# 3. Run inference and aggregate -aggregator = tio.PatchAggregator( - spatial_shape=subject.spatial_shape, - overlap_mode="hann", - patch_overlap=8, -) - -for batch in loader: - input_tensor = batch.t1.data - output = model(input_tensor) - locations = [ - batch.metadata["patch_location"][i] - for i in range(batch.batch_size) - ] - aggregator.add_batch(output, locations) - -result = aggregator.get_output() -``` - -## Overlap modes - -| Mode | Best for | How it works | -|------|----------|-------------| -| `"crop"` | Argmax segmentation | Keeps only the non-overlapping center of each patch | -| `"average"` | Probabilistic outputs | Averages all overlapping predictions | -| `"hann"` | Continuous outputs | Weights with a Hann window for smooth blending | - -## Training with random samplers - -For training, use `UniformSampler`, `WeightedSampler`, or -`LabelSampler`. These are `IterableDataset`s that yield patches -on-the-fly: - - -```python -sampler = tio.UniformSampler(subject, patch_size=64, num_patches=200) -loader = SubjectsLoader(sampler, batch_size=8) - -for batch in loader: - output = model(batch.t1.data) - loss = criterion(output, batch.seg.data) - loss.backward() -``` - -### WeightedSampler - -Sample more patches from regions of interest using a probability map: - - -```python -sampler = tio.WeightedSampler( - subject, - patch_size=64, - probability_map="sampling_weights", - num_patches=200, -) -``` - -### LabelSampler - -Center patches on labeled voxels, with optional per-class weights: - - -```python -sampler = tio.LabelSampler( - subject, - patch_size=64, - label_name="seg", - label_probabilities={0: 0.5, 1: 0.5}, - num_patches=200, -) -``` - -## Downsampled outputs - -If your model produces spatially smaller outputs (e.g., a feature -encoder with stride 2), pass `output_shape` to the aggregator: - - -```python -aggregator = tio.PatchAggregator( - spatial_shape=(256, 256, 176), - output_shape=(128, 128, 88), - overlap_mode="average", -) -``` - -## Multiple outputs - -Pass a dict of tensors to aggregate multiple outputs simultaneously: - - -```python -aggregator.add_batch( - {"segmentation": seg_output, "embedding": emb_output}, - locations, -) -seg = aggregator.get_output("segmentation") -emb = aggregator.get_output("embedding") -``` diff --git a/docs/how-to/patch-training.md b/docs/how-to/patch-training.md deleted file mode 100644 index 3f9053286..000000000 --- a/docs/how-to/patch-training.md +++ /dev/null @@ -1,100 +0,0 @@ -# Patch-based training - -For training on large 3D volumes that don't fit in GPU memory, TorchIO -provides a [`Queue`](../reference/patches.md) that loads subjects -in background threads, applies transforms, and extracts random patches -into a buffer. - -## Basic usage - - -```python -import torchio as tio -from torchio.loader import SubjectsLoader - -subjects = [ - tio.Subject(t1=tio.ScalarImage(path)) - for path in training_paths -] - -transform = tio.Compose([ - tio.Flip(axes=(0,), p=0.5), - tio.Noise(std=0.1), -]) - -sampler = tio.UniformSampler(subjects[0], patch_size=64) - -queue = tio.Queue( - subjects, - patch_sampler=sampler, - max_length=300, - patches_per_volume=10, - num_workers=4, - transform=transform, -) - -loader = SubjectsLoader(queue, batch_size=16) - -for epoch in range(num_epochs): - for batch in loader: - inputs = batch.t1.data - outputs = model(inputs) - loss = criterion(outputs, targets) - loss.backward() -``` - -## How it works - -1. **Subjects** are loaded and preprocessed in background threads - (`num_workers` controls the parallelism) -2. **Patches** are extracted via the sampler (up to - `patches_per_volume` per subject) -3. Patches accumulate in a **buffer** (up to `max_length`) -4. When the buffer is full, patches are shuffled and yielded -5. The external `SubjectsLoader` batches them for GPU training - -## Memory estimation - - -```python -print(queue.max_memory_pretty) # e.g., "200.0 MiB" -``` - -## Distributed training - -Pass a `DistributedSampler` as `subject_sampler` so each rank -processes its own subset of subjects: - - -```python -from torch.utils.data.distributed import DistributedSampler - -subject_sampler = DistributedSampler(subjects, shuffle=True) -queue = tio.Queue( - subjects, - patch_sampler=sampler, - subject_sampler=subject_sampler, - shuffle_subjects=False, # sampler controls order -) - -for epoch in range(num_epochs): - subject_sampler.set_epoch(epoch) - for batch in loader: - ... -``` - -Each rank processes only its assigned subjects. The epoch length -is `len(subject_sampler) * patches_per_volume`, not the full -dataset size. - -## Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `max_length` | 300 | Buffer capacity. Larger = more diversity, more RAM | -| `patches_per_volume` | 10 | Max patches per subject (ceiling) | -| `num_workers` | 0 | Background loading threads | -| `shuffle_subjects` | True | Randomize subject order per epoch | -| `shuffle_patches` | True | Randomize patch order in buffer | -| `transform` | None | Applied to each subject before sampling | -| `subject_sampler` | None | For distributed training | diff --git a/docs/how-to/remote-nii-zarr.md b/docs/how-to/remote-nii-zarr.md deleted file mode 100644 index fd870e5d4..000000000 --- a/docs/how-to/remote-nii-zarr.md +++ /dev/null @@ -1,168 +0,0 @@ -# Stream a remote NIfTI-Zarr - -When a `.nii.zarr` volume lives on cloud storage (Azure Blob, S3, GCS), -TorchIO can open it **without downloading the entire file**. Only the -metadata and the chunks you actually read are fetched over the network. - -This is especially useful for large volumes (e.g., 10 GB whole-brain -microscopy) where you only need a small region of interest. - -## Prerequisites - -Install TorchIO with the `zarr` extra and the storage backend you need: - -=== "Azure Blob" - - ``` - pip install "torchio[zarr,azure]" - ``` - -=== "S3" - - ``` - pip install "torchio[zarr,s3]" - ``` - -=== "Google Cloud" - - ``` - pip install "torchio[zarr,gcs]" - ``` - -## How it works - -```mermaid -sequenceDiagram - participant User - participant TorchIO - participant ZarrBackend - participant Cloud as Cloud Storage - - User->>TorchIO: ScalarImage("az://…/brain.nii.zarr") - Note over TorchIO: URI stored, no download - User->>TorchIO: image.shape - TorchIO->>ZarrBackend: open remote store - ZarrBackend->>Cloud: fetch header + metadata - Cloud-->>ZarrBackend: ~KB - ZarrBackend-->>TorchIO: (1, 512, 512, 512) - User->>TorchIO: image[:, 100:200, 100:200, 100:200] - TorchIO->>ZarrBackend: slice - ZarrBackend->>Cloud: fetch overlapping chunks - Cloud-->>ZarrBackend: ~MB - ZarrBackend-->>TorchIO: Tensor (1, 100, 100, 100) -``` - -When you pass a remote `.nii.zarr` URI, TorchIO: - -1. **Stores the URI**: no bytes are downloaded yet. -2. **On first metadata access** (`.shape`, `.affine`, …), opens a remote - zarr store via [fsspec](https://filesystem-spec.readthedocs.io/) and - reads only the header. -3. **On slicing**, fetches only the chunks that overlap with your region - of interest. - -## Authenticate to Azure and crop a region - - -```python -import os - -import torchio as tio - -# Option 1: Authenticate via environment variables (recommended for CI/HPC). -# adlfs picks these up automatically. -os.environ["AZURE_STORAGE_ACCOUNT_NAME"] = "myaccount" -os.environ["AZURE_STORAGE_ACCOUNT_KEY"] = "my-secret-key" # or use SAS, etc. - -image = tio.ScalarImage("az://mycontainer/dataset/brain.nii.zarr") - -# Nothing has been downloaded yet -print(image.shape) # e.g. (1, 512, 512, 512), only metadata fetched -print(image.spacing) # from the NIfTI header stored in the zarr - -# Crop a 100×100×100 ROI, only the overlapping chunks are fetched -roi = image[:, 200:300, 200:300, 200:300] -print(roi.shape) # (1, 100, 100, 100) -print(roi.data.mean()) -``` - - -```python -# Option 2: Pass credentials via reader_kwargs. -# These are forwarded to niizarr → zarr → fsspec → adlfs. -image = tio.ScalarImage( - "az://mycontainer/dataset/brain.nii.zarr", - reader_kwargs={"account_name": "myaccount", "account_key": "my-key"}, -) - -# Apply a TorchIO transform to the ROI -crop = tio.CropOrPad(128) -cropped = crop(roi) -print(cropped.shape) # (1, 128, 128, 128) -``` - -!!! tip "Azure authentication methods" - - The `adlfs` library supports several authentication methods. - Set the appropriate environment variables or pass them via - `reader_kwargs`: - - | Method | Environment variables | - |--------|-----------------------| - | Account key | `AZURE_STORAGE_ACCOUNT_NAME`, `AZURE_STORAGE_ACCOUNT_KEY` | - | SAS token | `AZURE_STORAGE_ACCOUNT_NAME`, `AZURE_STORAGE_SAS_TOKEN` | - | Connection string | `AZURE_STORAGE_CONNECTION_STRING` | - | Default credential (Azure CLI / Managed Identity) | `AZURE_STORAGE_ACCOUNT_NAME` | - - See the [adlfs documentation](https://github.com/fsspec/adlfs) for - the full list. - -## Other cloud providers - -=== "S3" - - - ```python - image = tio.ScalarImage("s3://my-bucket/brain.nii.zarr") - roi = image[:, 100:200, 100:200, 100:200] - ``` - - Authentication is handled by `s3fs`, which reads `~/.aws/credentials` - or the `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` environment - variables. - -=== "Google Cloud" - - - ```python - image = tio.ScalarImage("gs://my-bucket/brain.nii.zarr") - roi = image[:, 100:200, 100:200, 100:200] - ``` - - Authentication is handled by `gcsfs`, which uses Application Default - Credentials or `GOOGLE_APPLICATION_CREDENTIALS`. - -=== "HTTPS" - - - ```python - image = tio.ScalarImage("https://example.com/data/brain.nii.zarr") - roi = image[:, 100:200, 100:200, 100:200] - ``` - - No extra packages needed: `fsspec[http]` is included by default. - -## Comparison with non-Zarr remote files - -For non-Zarr remote files (e.g., `az://…/brain.nii.gz`), TorchIO -downloads the **entire file** to a temporary local path before reading. -This is the expected behavior because formats like `.nii.gz` do not -support partial reads over the network. - -| Source | What happens | -|--------|-------------| -| `az://…/brain.nii.gz` | Full download, then local read | -| `az://…/brain.nii.zarr` | **Streaming**: only metadata + requested chunks | - -If you are working with large remote volumes, converting to `.nii.zarr` -first is strongly recommended. See [Save as NIfTI-Zarr](save-nii-zarr.md). diff --git a/docs/how-to/save-nii-zarr.md b/docs/how-to/save-nii-zarr.md deleted file mode 100644 index 0090c5654..000000000 --- a/docs/how-to/save-nii-zarr.md +++ /dev/null @@ -1,64 +0,0 @@ -# Save as NIfTI-Zarr - -NIfTI-Zarr (`.nii.zarr`) stores image data in independently compressed -chunks, enabling lazy partial reads. This guide shows how to convert -existing volumes. - -## Prerequisites - -Install the `zarr` extra: - -``` -uv add torchio --extra zarr -``` - -## Convert a NIfTI file - - -```python -import torchio as tio - -image = tio.ScalarImage("input.nii.gz") -image.save("output.nii.zarr") -``` - -That is it. The output is a directory (`output.nii.zarr/`) containing -chunked Zarr arrays and the NIfTI header. - -## Convert from a tensor - -```python -import torch -import torchio as tio - -tensor = torch.randn(1, 256, 256, 256) -image = tio.ScalarImage(tensor) -image.save("synthetic.nii.zarr") -``` - -## Verify the result - - -```python -loaded = tio.ScalarImage("output.nii.zarr") -print(loaded.shape) # reads only metadata -print(loaded.spacing) # from the stored affine - -# Lazy slice: reads only the needed chunks -patch = loaded[:, 50:100, 50:100, 50:100] -print(patch.data.mean()) -``` - -## Chunk size - -The default chunk size is 64 voxels per dimension. To customize it, -use `niizarr` directly: - - -```python -import nibabel as nib -from niizarr import nii2zarr - -nii = nib.load("input.nii.gz") -nii2zarr(nii, "output.nii.zarr", chunk=128) -``` diff --git a/docs/how-to/tta.md b/docs/how-to/tta.md deleted file mode 100644 index 18ea49706..000000000 --- a/docs/how-to/tta.md +++ /dev/null @@ -1,127 +0,0 @@ -# Test-time augmentation (TTA) - -Test-time augmentation improves prediction accuracy by averaging -multiple augmented versions of the same input. TorchIO's invertible -transforms make this easy. The pattern mirrors -[v1's TTA workflow](https://docs.torchio.org/transforms/#invertibility). - -## Basic TTA - -The idea: augment the input, predict, copy the history to the -prediction, invert, then average: - - -```python -import torch -import torchio as tio - -model = ... # your trained model -subject = tio.Subject(t1=tio.ScalarImage("t1.nii.gz")) - -augment = tio.Compose([ - tio.Flip(axes=(0, 1, 2), flip_probability=0.5), -]) - -predictions = [] -n_augmentations = 8 - -for _ in range(n_augmentations): - # Augment - augmented = augment(subject) - - # Predict - with torch.no_grad(): - pred = model(augmented.t1.data.unsqueeze(0)) - - # Wrap prediction and copy transform history - pred_subject = tio.Subject( - pred=tio.ScalarImage(pred.squeeze(0)), - ) - pred_subject.applied_transforms = augmented.applied_transforms - - # Invert augmentation on the prediction - restored = pred_subject.apply_inverse_transform( - ignore_intensity=True, - ) - predictions.append(restored.pred.data) - -# Average predictions in the original space -mean_prediction = torch.stack(predictions).mean(0) -``` - -## API - -**On Subject:** - - -```python -# Get a Compose that inverts the history -inverse_transform = subject.get_inverse_transform() - -# Or apply directly -restored = transformed.apply_inverse_transform() - -# Skip intensity transforms (useful for TTA) -restored = transformed.apply_inverse_transform(ignore_intensity=True) -``` - -**Standalone function** (works on any type with history): - - -```python -restored = tio.apply_inverse_transform(data) -``` - -## Which transforms are invertible? - -| Transform | Invertible | Notes | -|-----------|-----------|-------| -| `Flip` | ✅ | Self-inverse (flip twice = identity) | -| `Crop` | ✅ | Inverse is Pad (lost voxels filled with zeros) | -| `Pad` | ✅ | Inverse is Crop | -| `Resample` | ✅ | Restores the original output grid | -| `Affine` | ✅ | Uses the inverse affine matrix | -| `ElasticDeformation` | ✅ | Negates the sampled displacement field | -| `Spatial` | ✅ | Inverts resampling, affine, and elastic parts together | -| `Normalize` | ✅ | Reverses the linear rescaling | -| `Standardize` | ✅ | Multiplies by std and adds mean | -| `BiasField` | ✅ | Divides by the same bias field | -| `Gamma` | ✅ | Applies $1/\gamma$ | -| `OneHot` | ✅ | Takes argmax back to single-channel labels | -| `RemapLabels` | ✅ | Swaps keys and values in the remapping dict | -| `SequentialLabels` | ✅ | Restores original label values | -| `Transpose` | ✅ | Self-inverse (transpose twice = identity) | -| `Anisotropy` | ❌ | Information lost during downsampling | -| `Blur` | ❌ | Skipped silently when `ignore_intensity=True` | -| `Clamp` | ❌ | Skipped silently when `ignore_intensity=True` | -| `Contour` | ❌ | Destructive: interior information is lost | -| `CopyAffine` | ❌ | Metadata-only; not invertible | -| `Ghosting` | ❌ | Skipped silently when `ignore_intensity=True` | -| `HistogramStandardization` | ❌ | Piecewise-linear map is lossy | -| `KeepLargestComponent` | ❌ | Destructive: removed components are lost | -| `LabelsToImage` | ❌ | Generative; not invertible | -| `Lambda` | ❌ | Skipped silently when `ignore_intensity=True` | -| `Mask` | ❌ | Skipped silently when `ignore_intensity=True` | -| `Motion` | ❌ | Skipped silently when `ignore_intensity=True` | -| `Noise` | ❌ | Skipped silently when `ignore_intensity=True` | -| `PCA` | ❌ | Dimensionality reduction is lossy | -| `RemoveLabels` | ❌ | Destructive: removed labels are lost | -| `Resize` | ❌ | Not automatically invertible | -| `Spike` | ❌ | Skipped silently when `ignore_intensity=True` | -| `Swap` | ❌ | Skipped silently when `ignore_intensity=True` | -| `ToReferenceSpace` | ❌ | Metadata-only; not invertible | - -Non-invertible transforms are **skipped with a warning** (not -errored), so TTA works even with mixed pipelines: - - -```python -pipeline = tio.Compose([ - tio.Flip(axes=(0, 1, 2), flip_probability=0.5), - tio.Noise(std=0.1), # skipped during inversion -]) -transformed = pipeline(subject) -restored = transformed.apply_inverse_transform() # only Flip is inverted -``` - -Use `warn=False` to suppress the warnings. diff --git a/docs/how-to/visualization.md b/docs/how-to/visualization.md deleted file mode 100644 index 3f2aed395..000000000 --- a/docs/how-to/visualization.md +++ /dev/null @@ -1,287 +0,0 @@ -# Visualizing images - -TorchIO provides a built-in plotting function that displays three -orthogonal slices (Sagittal, Coronal, Axial) with correct anatomical -orientation and proportions. It works with any image orientation: -RAS, LPS, or anything else. - -!!! note "Optional dependency" - Plotting requires the `plot` extra: - - === "uv" - - ``` - uv add torchio[plot] - ``` - - === "pip" - - ``` - pip install torchio[plot] - ``` - - This installs [matplotlib](https://matplotlib.org/) and - [colorcet](https://colorcet.holoviz.org/) (for categorical - colormaps). - -## Basic usage - - -```python -import torchio as tio - -image = tio.ScalarImage("brain.nii.gz") -image.plot() -``` - -This shows mid-slices through each anatomical plane with: - -- **Correct proportions** between views (from the voxel spacing) -- **Orientation labels** showing both the tensor axis and anatomical - direction, e.g., `J (A ↔ P)` -- **World-coordinate ticks** in mm (derived from the affine matrix) -- **Coloured intersection lines** showing where the other slices are - -## Choosing slices - -By default, the mid-slice along each axis is shown. You can specify -slices by voxel index or by world coordinates in mm: - - -```python -# By voxel index (None = mid-slice) -image.plot(indices=(80, None, 60)) - -# By world coordinates in mm -image.plot(coordinates=(-10.0, 25.5, 30.0)) - -# Mix None and values -image.plot(coordinates=(None, 0.0, None)) -``` - -!!! tip - `indices` and `coordinates` are mutually exclusive. - -## Tick labels - -By default, tick labels show world coordinates in mm. Pass -`voxels=True` to show voxel indices instead: - - -```python -image.plot(voxels=True) -``` - -## Intensity windowing - -For scalar images, the display range is set from the 0.5th to 99.5th -percentile by default. Adjust with: - - -```python -image.plot(percentiles=(1, 99)) -``` - -## Label maps - -Label maps automatically use nearest-neighbour interpolation and a -categorical colormap (from [colorcet](https://colorcet.holoviz.org/) -if installed, otherwise matplotlib's `tab10`): - - -```python -segmentation = tio.LabelMap("seg.nii.gz") -segmentation.plot() -``` - -## Customization - -### Figure size - - -```python -# Scale the default figure size -image.plot(figsize_multiplier=3.0) - -# Or set an exact size -image.plot(figsize=(12, 4)) -``` - -### Colormap and imshow options - - -```python -image.plot(cmap="hot") -image.plot(vmin=0, vmax=1000) # extra kwargs go to ax.imshow() -``` - -### Intersection lines - -The coloured cross-hair lines can be turned off: - - -```python -image.plot(intersections=False) -``` - -### Saving to file - - -```python -image.plot(output_path="slices.png", show=False) -image.plot( - output_path="slices.pdf", - show=False, - savefig_kwargs={"dpi": 300, "bbox_inches": "tight"}, -) -``` - -### Plotting into existing axes - - -```python -import matplotlib.pyplot as plt - -fig, axes = plt.subplots(1, 3, figsize=(15, 5)) -image.plot(axes=axes, show=False) -``` - -## Jupyter notebooks - -In Jupyter, `Image` objects display an inline plot automatically via -`_repr_html_`: - - -```python -image # shows 3 slices + metadata table -``` - -Call `image.plot()` explicitly to get a larger, interactive figure. - -## Subjects - -Plot all images in a subject as a grid: - - -```python -subject = tio.Subject( - t1=tio.ScalarImage("t1.nii.gz"), - seg=tio.LabelMap("seg.nii.gz"), -) -subject.plot() -``` - -Each image gets a row of Sagittal/Coronal/Axial views (or columns -if there are more than 3 images). LabelMaps are automatically -detected and use categorical colormaps. - -### Per-image colormaps - - -```python -subject.plot(cmap_dict={"t1": "hot", "seg": "viridis"}) -``` - -### In Jupyter - -`Subject` objects also display an inline plot via `_repr_html_`: - - -```python -subject # shows grid + metadata tables -``` - -## How it works - -The display always follows the same anatomical convention regardless -of the image orientation: - -| View | Horizontal | Vertical | -|-----------|------------------|-----------------| -| Sagittal | Anterior ↔ Posterior | Inferior ↔ Superior | -| Coronal | Right ↔ Left | Inferior ↔ Superior | -| Axial | Right ↔ Left | Posterior ↔ Anterior | - -The data is flipped and transposed as needed so these directions are -always in the same screen positions. The axis labels show which tensor -axis (I, J, K) maps to each direction, making it easy to relate what -you see to the underlying data layout. - -## Animated GIFs and videos - -You can export an animation sweeping through slices along any -anatomical direction: - - -```python -image.to_gif("brain.gif", seconds=5, direction="I") -image.to_video("brain.mp4", seconds=5, direction="S") -``` - -The `direction` parameter accepts `"I"` (inferior), `"S"` -(superior), `"A"` (anterior), `"P"` (posterior), `"R"` (right), -or `"L"` (left). The image is automatically reoriented so slices -appear in the correct anatomical view. - -### Displaying in Jupyter - -In Jupyter notebooks, calling `to_gif()` or `to_video()` without -a path automatically creates a temporary file and returns an IPython -display object: - - -```python -# Just this: the GIF is displayed inline -image.to_gif() - -# Same for video -image.to_video() -``` - -You can also pass an explicit path if you want to keep the file: - - -```python -image.to_gif("/tmp/image.gif") -image.to_video("/tmp/image.mp4") -``` - -Outside Jupyter, a path is required: - - -```python -image.to_gif("brain.gif") -image.to_video("brain.mp4") -``` - -### From the command line - - -```bash -torchio animate brain.nii.gz brain.gif -torchio animate brain.nii.gz brain.mp4 --seconds 10 --direction S -``` - -!!! note "Optional dependencies" - GIFs require `Pillow` (included in the `[plot]` extra). - Videos require `ffmpeg-python` (`pip install torchio[video]`) - and a working `ffmpeg` installation. - -## Interactive 3D viewer - -For an interactive viewer in Jupyter notebooks, use -`plot_interactive()`. It uses -[NiiVue](https://niivue.com/) via -[ipyniivue](https://github.com/niivue/ipyniivue) and supports -scrolling through slices, zooming, and crosshair navigation: - - -```python -image.plot_interactive() -``` - -The viewer uses **radiological convention** (left hemisphere on the -right side of the screen). - -!!! note "Optional dependency" - Requires `ipyniivue`: `pip install torchio[niivue]` diff --git a/docs/index.md b/docs/index.md index 4c54bec6d..a0bc7009c 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,79 +1,93 @@ # TorchIO -TorchIO is an open-source Python library for efficient loading, -preprocessing, augmentation, and patch-based sampling of 3D medical images -in deep learning, following the design of PyTorch. - -> *Tools like TorchIO are a symptom of the maturation of medical AI research -> using deep learning techniques*. -> -> Jack Clark, Policy Director at [OpenAI](https://openai.com/), Co-Founder and -> Head of Policy of Anthropic ([link](https://jack-clark.net/2020/03/17/)) - -!!! warning "TorchIO v2 is an experimental pre-release" - - This site documents **TorchIO v2**, which is under active development and - may contain bugs and breaking changes. The current stable release is **v1** - (`pip install torchio`). To try v2, run `pip install --pre torchio`. Use the - version selector in the header to switch between versions. - -## Quick example - -Augment a whole batch of subjects on the GPU in a few lines: - - -```python -import torchio as tio - -# Build subjects (lazy: only headers are read until .data is accessed) -dirs = ["sub-01", "sub-02", "sub-03", "sub-04"] -subjects = [ - tio.Subject( - t1=tio.ScalarImage(f"{dir}/t1.nii.gz"), - seg=tio.LabelMap(f"{dir}/seg.nii.gz"), - ) - for dir in dirs -] - -# Random augmentation pipeline -transform = tio.Compose([ - tio.Flip(), - tio.Affine(degrees=(-15, 15)), - tio.Standardize(), - tio.Noise(std=(0, 0.1)), -]) - -# Stack into a batch and augment all subjects on the GPU in one call -batch = tio.SubjectsBatch.from_subjects(subjects).to("cuda") # or "mps" -augmented = transform(batch) - -print(augmented.t1.data.shape) # (4, 1, 256, 256, 176) -print(augmented.t1.data.device) # cuda:0 -``` +[![PyPI downloads](https://img.shields.io/pypi/dm/torchio.svg?label=PyPI%20downloads&logo=python&logoColor=white)](https://pypi.org/project/torchio/) +[![PyPI version](https://img.shields.io/pypi/v/torchio?label=PyPI%20version&logo=python&logoColor=white)](https://pypi.org/project/torchio/) +[![Conda version](https://img.shields.io/conda/v/conda-forge/torchio.svg?label=conda-forge&logo=conda-forge)](https://anaconda.org/conda-forge/torchio) +[![Google Colab notebooks](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/TorchIO-project/torchio/blob/main/tutorials/README.md) +[![Documentation status](https://github.com/TorchIO-project/torchio/actions/workflows/docs.yml/badge.svg)](https://github.com/TorchIO-project/torchio/actions/workflows/docs.yml) +[![Tests status](https://github.com/TorchIO-project/torchio/actions/workflows/tests.yml/badge.svg)](https://github.com/TorchIO-project/torchio/actions/workflows/tests.yml) +[![Code style: Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v1.json)](https://docs.astral.sh/ruff/) +[![Coverage status](https://codecov.io/gh/TorchIO-project/torchio/branch/main/graphs/badge.svg)](https://app.codecov.io/github/TorchIO-project/torchio) +[![Code quality](https://img.shields.io/scrutinizer/g/TorchIO-project/torchio.svg?label=Code%20quality&logo=scrutinizer)](https://scrutinizer-ci.com/g/TorchIO-project/torchio/?branch=main) +[![YouTube](https://img.shields.io/youtube/views/UEUVSw5-M9M?label=watch&style=social)](https://www.youtube.com/watch?v=UEUVSw5-M9M) + +TorchIO is an open-source Python library for efficient loading, preprocessing, +augmentation and patch-based sampling of 3D medical images in deep learning, +following the design of PyTorch. + +It includes multiple intensity and spatial transforms for data augmentation and +preprocessing. +These transforms include typical computer vision operations +such as random affine transformations and also domain-specific ones such as +simulation of intensity artifacts due to +[MRI magnetic field inhomogeneity (bias)](https://mriquestions.com/why-homogeneity.html) +or [k-space motion artifacts](http://proceedings.mlr.press/v102/shaw19a.html). + +TorchIO is part of the official [PyTorch Ecosystem](https://pytorch.org/ecosystem/), +and was featured at +the [PyTorch Ecosystem Day 2021](https://pytorch.org/blog/ecosystem_day_2021) and +the [PyTorch Developer Day 2021](https://pytorch.org/blog/pytorch-developer-day-2021). + +Many groups have used TorchIO for their research. +The complete list of citations is available on [Google Scholar](https://scholar.google.co.uk/scholar?cites=8711392719159421861&sciodt=0,5&hl=en), and the +[dependents list](https://github.com/TorchIO-project/torchio/network/dependents) is +available on GitHub. + +The code is available on [GitHub](https://github.com/TorchIO-project/torchio). +If you like TorchIO, please go to the repository and star it! + +Star + + +See [Getting started](getting-started.md) for installation instructions and a +usage overview. + +Contributions are more than welcome. +Please check our [contributing guide](https://github.com/TorchIO-project/torchio/blob/main/CONTRIBUTING.rst) +if you would like to contribute. + +If you have questions, feel free to ask in the +[discussions tab](https://github.com/TorchIO-project/torchio/discussions). + +Discuss + +If you found a bug or have a feature request, please +[open an issue](https://github.com/TorchIO-project/torchio/issues). + +Issue -## Where to go next +## Credits -New to TorchIO? Start with the [quickstart](get-started/quickstart.md). -Upgrading from v1? See the [migration guide](get-started/migration.md). +If you use this library for your research, +please cite our paper: + +[F. Pérez-García, R. Sparks, and S. Ourselin. TorchIO: a Python library for +efficient loading, preprocessing, augmentation and patch-based sampling of +medical images in deep learning. Computer Methods and Programs in Biomedicine +(June 2021), p. 106236. ISSN: +0169-2607. doi:10.1016/j.cmpb.2021.106236.](https://doi.org/10.1016/j.cmpb.2021.106236) + +BibTeX: + +```bibtex +@article{perez-garcia_torchio_2021, + title = {{TorchIO}: a {Python} library for efficient loading, preprocessing, augmentation and patch-based sampling of medical images in deep learning}, + journal = {Computer Methods and Programs in Biomedicine}, + pages = {106236}, + year = {2021}, + issn = {0169-2607}, + doi = {https://doi.org/10.1016/j.cmpb.2021.106236}, + url = {https://www.sciencedirect.com/science/article/pii/S0169260721003102}, + author = {P{\'e}rez-Garc{\'i}a, Fernando and Sparks, Rachel and Ourselin, S{\'e}bastien}, +} +``` -## Credits +This project was originally supported by the following institutions: + +- [Engineering and Physical Sciences Research Council (EPSRC) & UK Research and Innovation (UKRI)](https://epsrc.ukri.org/) +- [EPSRC Centre for Doctoral Training in Intelligent, Integrated Imaging In Healthcare (i4health)](https://www.ucl.ac.uk/intelligent-imaging-healthcare/) (University College London) +- [Wellcome / EPSRC Centre for Interventional and Surgical Sciences (WEISS)](https://www.ucl.ac.uk/interventional-surgical-sciences/) (University College London) +- [School of Biomedical Engineering & Imaging Sciences (BMEIS)](https://www.kcl.ac.uk/bmeis) (King's College London) -If you use this library for your research, please cite our paper: - -> F. Perez-Garcia, R. Sparks, and S. Ourselin. -> *TorchIO: a Python library for efficient loading, preprocessing, -> augmentation and patch-based sampling of medical images in deep learning.* -> Computer Methods and Programs in Biomedicine (June 2021), p. 106236. -> [doi:10.1016/j.cmpb.2021.106236](https://doi.org/10.1016/j.cmpb.2021.106236). - -## Related projects - -- [MONAI](https://monai.readthedocs.io) -- [Cornucopia](https://cornucopia.readthedocs.io/) -- [batchgenerators](https://github.com/MIC-DKFZ/batchgenerators) ([v2](https://github.com/MIC-DKFZ/batchgeneratorsv2)) -- [BatchAug](https://github.com/halleewong/batchaug) -- [volumentations](https://github.com/ZFTurbo/volumentations) (low activity) -- [Rising](https://rising.readthedocs.io) (archived) -- [pymia](https://pymia.readthedocs.io) (low activity) -- [MedicalTorch](https://medicaltorch.readthedocs.io) (abandoned) -- [Eisen](https://github.com/eisen-ai/eisen-core) (abandoned) +This library has been greatly inspired by +[NiftyNet](https://github.com/NifTK/NiftyNet), which is no longer maintained. diff --git a/docs/interfaces/cli.md b/docs/interfaces/cli.md new file mode 100644 index 000000000..1d42124c5 --- /dev/null +++ b/docs/interfaces/cli.md @@ -0,0 +1,25 @@ +# Command-line tools + +## `tiotr` + +A transform can be quickly applied to an image file using the command-line +tool `tiotr`, which is automatically installed by `pip` +during installation of TorchIO: + +``` +$ tiotr input.nii RandomAffine output.nii.gz --kwargs "degrees=(0,0,10) scales=0.1" --seed 42 +``` + +For more information, run `tiotr --help`. + +## `tiohd` + +To print some image metadata, `tiohd` can be used. Adding the `--plot` +argument will plot the image using Matplotlib: + +``` +$ tiohd ~/.cache/torchio/mni_colin27_1998_nifti/colin27_t1_tal_lin.nii +ScalarImage(shape: (1, 181, 217, 181); spacing: (1.00, 1.00, 1.00); orientation: RAS+; dtype: torch.FloatTensor; memory: 27.1 MiB) +``` + +For more information, run `tiohd --help`. diff --git a/docs/interfaces/slicer.md b/docs/interfaces/slicer.md new file mode 100644 index 000000000..ece2654c5 --- /dev/null +++ b/docs/interfaces/slicer.md @@ -0,0 +1,59 @@ +# 3D Slicer GUI + +[3D Slicer](https://www.slicer.org/) is an open-source software platform for +medical image informatics, image processing, +and three-dimensional visualization. + +TorchIO provides a 3D Slicer extension for quick experimentation and +visualization of the package features without any coding. + +The TorchIO extension can be easily installed using the +[Extensions Manager](https://slicer.readthedocs.io/en/latest/user_guide/extensions_manager.html). + +The code and installation instructions are available on +[GitHub](https://github.com/fepegar/SlicerTorchIO). + +!!! note + The Preview version (built nightly) is recommended. You can download + and install Slicer from [their download website](https://download.slicer.org/) + or, if you are on macOS, using [Homebrew](https://docs.brew.sh/): + + ``` + brew tap homebrew/cask-versions && brew cask install slicer-preview + ``` + +## TorchIO Transforms + +This module can be used to quickly visualize the effect of each transform +parameter. +That way, users can have an intuitive feeling of what the output +of a transform looks like without any coding at all. + +![TorchIO Transforms module for 3D Slicer](https://raw.githubusercontent.com/fepegar/SlicerTorchIO/master/Screenshots/TorchIO.png) + +### Usage example + +Go to the `Sample Data` module to get an image we can use: + +![Go to Sample Data module](https://raw.githubusercontent.com/fepegar/SlicerTorchIO/master/Screenshots/usage_1.png) + +Click on an image to download, for example MRHead[^1], +and go to the `TorchIO Transforms` module: + +[^1]: All the data in `Sample Data` can be downloaded and used in the TorchIO + Python library using the `torchio.datasets.slicer.Slicer` class. + +![Download MRHead and go to TorchIO Transforms module](https://raw.githubusercontent.com/fepegar/SlicerTorchIO/master/Screenshots/usage_2.png) + +Select the input and output volume nodes: + +![Select volume nodes](https://raw.githubusercontent.com/fepegar/SlicerTorchIO/master/Screenshots/usage_3.png) + +Modify the transform parameters and click on `Apply transform`. +Hovering the mouse over the transforms will show tooltips extracted from the +TorchIO documentation. + +![Apply transform](https://raw.githubusercontent.com/fepegar/SlicerTorchIO/master/Screenshots/usage_4.png) + +You can click on the `Toggle volumes` button to switch between input and +output volumes. diff --git a/docs/patches/index.md b/docs/patches/index.md new file mode 100644 index 000000000..6ffacb673 --- /dev/null +++ b/docs/patches/index.md @@ -0,0 +1,52 @@ +# Patch-based pipelines + +The number of pixels in 2D images used in deep learning +is rarely larger than one million. +For example, the input size of several popular image classification +models is 224 × 224 × 3 = 150 528 pixels +(588 KiB if 32 bits per pixel are used). +In contrast, 3D medical images often contain hundreds of +millions of voxels, and downsampling might not be acceptable when small details +should be preserved. +For example, the size of a high-resolution lung CT-scan used for quantifying +chronic obstructive pulmonary disease damage in a research setting, +with spacing 0.66 × 0.66 × 0.30 mm, +is 512 × 512 × 1069 = 280 231 936 voxels +(1.04 GiB if 32 bits per voxel are used). + +In computer vision applications, images used for training are grouped in +batches whose size is often in the order of +hundreds +or even thousands of training instances, +depending on the available GPU memory. +In medical image applications, batches rarely contain more than one or two +training instances due to their larger memory footprint compared to natural +images. +This reduces the utility of techniques +such as batch normalization, +which rely on batches being large enough to estimate +dataset variance appropriately. +Moreover, large image size and small batches result in longer training time, +hindering the experimental cycle that is necessary for hyperparameter +optimization. +In cases where GPU memory is limited and the network architecture is +large, it is possible that not even the entirety of a single volume can be +processed during a single training iteration. +To overcome this challenge, it is common in medical imaging to train using +subsets of the image, or image *patches*, +randomly extracted from the volumes. + +Networks can be trained with 2D slices extracted from 3D volumes, +aggregating the inference results to generate a 3D volume. +This can be seen as a specific case of patch-based training, +where the size of the patches along a dimension is one. +Other methods extract volumetric patches for training, +that are often cubes, +if the voxel spacing is isotropic, +or cuboids adapted to the +anisotropic spacing of the training images. + +![Training with patches](../images/diagram_patches.svg) + +- [Training](training.md) — patch samplers and queue for training +- [Inference](inference.md) — grid sampler and aggregator for dense inference diff --git a/docs/patches/inference.md b/docs/patches/inference.md new file mode 100644 index 000000000..585283c1f --- /dev/null +++ b/docs/patches/inference.md @@ -0,0 +1,36 @@ +# Inference + +Here is an example that uses a grid sampler and aggregator to perform dense +inference across a 3D image using patches: + +```python +>>> import torch +>>> import torch.nn as nn +>>> import torchio as tio +>>> patch_overlap = 4, 4, 4 # or just 4 +>>> patch_size = 88, 88, 60 +>>> subject = tio.datasets.Colin27() +>>> subject +Colin27(Keys: ('t1', 'head', 'brain'); images: 3) +>>> grid_sampler = tio.inference.GridSampler( +... subject, +... patch_size, +... patch_overlap, +... ) +>>> patch_loader = tio.SubjectsLoader(grid_sampler, batch_size=4) +>>> aggregator = tio.inference.GridAggregator(grid_sampler) +>>> model = nn.Identity().eval() +>>> with torch.no_grad(): +... for patches_batch in patch_loader: +... input_tensor = patches_batch['t1'][tio.DATA] +... locations = patches_batch[tio.LOCATION] +... logits = model(input_tensor) +... labels = logits.argmax(dim=tio.CHANNELS_DIMENSION, keepdim=True) +... outputs = labels +... aggregator.add_batch(outputs, locations) +>>> output_tensor = aggregator.get_output_tensor() +``` + +::: torchio.data.GridSampler + +::: torchio.data.GridAggregator diff --git a/docs/patches/training.md b/docs/patches/training.md new file mode 100644 index 000000000..b7390d86f --- /dev/null +++ b/docs/patches/training.md @@ -0,0 +1,23 @@ +# Training + +## Patch samplers + +Samplers are used to randomly extract patches from volumes. +They are called with a sample generated by a +[`SubjectsDataset`](../data/dataset.md#torchio.data.SubjectsDataset) and return a Python generator that yields +cropped versions of the sample. + +For more information about patch-based training, see +[this NiftyNet tutorial](https://niftynet.readthedocs.io/en/dev/window_sizes.html). + +::: torchio.data.UniformSampler + +::: torchio.data.WeightedSampler + +::: torchio.data.LabelSampler + +::: torchio.data.PatchSampler + +::: torchio.data.GridSampler + +::: torchio.data.Queue diff --git a/docs/plot_directive.py b/docs/plot_directive.py index edbcd8bb8..bb6c76498 100644 --- a/docs/plot_directive.py +++ b/docs/plot_directive.py @@ -19,15 +19,15 @@ import matplotlib -matplotlib.use("Agg") +matplotlib.use('Agg') -import matplotlib.pyplot as plt +import matplotlib.pyplot as plt # noqa: E402 -_DOCS_DIR = Path("docs") -_PLOT_DIR = _DOCS_DIR / "images" / "plots" +_DOCS_DIR = Path('docs') +_PLOT_DIR = _DOCS_DIR / 'images' / 'plots' _FENCE_RE = re.compile( - r"^(\s*)(`{3,})python\s+plot\s*$\n(.*?)^\1\2\s*$", + r'^(\s*)(`{3,})python\s+plot\s*$\n(.*?)^\1\2\s*$', re.MULTILINE | re.DOTALL, ) @@ -35,7 +35,7 @@ def process_file(md_path: Path) -> bool: """Process a single markdown file. Returns True if modified.""" text = md_path.read_text() - if "python plot" not in text: + if 'python plot' not in text: return False page_dir = str(md_path.parent.relative_to(_DOCS_DIR)) @@ -44,41 +44,41 @@ def process_file(md_path: Path) -> bool: def _replace(match: re.Match) -> str: indent = match.group(1) code = textwrap.dedent(match.group(3)) - digest = hashlib.md5(code.encode()).hexdigest()[:12] - fname = f"plot_{digest}.png" + digest = hashlib.md5(code.encode()).hexdigest()[:12] # noqa: S324 + fname = f'plot_{digest}.png' fpath = _PLOT_DIR / fname if not fpath.exists(): - plt.close("all") + plt.close('all') try: - exec(code, {"__name__": "__main__"}) + exec(code, {'__name__': '__main__'}) # noqa: S102 except Exception as exc: - msg = f"{type(exc).__name__}: {exc}" - print(f" WARNING: Plot failed: {msg}", file=sys.stderr) + msg = f'{type(exc).__name__}: {exc}' + print(f' WARNING: Plot failed: {msg}', file=sys.stderr) return ( f'{indent}!!! warning "Plot generation failed"\n{indent} {msg}\n' ) fig = plt.gcf() - fig.savefig(fpath, bbox_inches="tight", dpi=150) - plt.close("all") - print(f" Generated {fpath}", file=sys.stderr) + fig.savefig(fpath, bbox_inches='tight', dpi=150) + plt.close('all') + print(f' Generated {fpath}', file=sys.stderr) else: - print(f" Cached {fpath}", file=sys.stderr) + print(f' Cached {fpath}', file=sys.stderr) img_rel = os.path.relpath( str(fpath.relative_to(_DOCS_DIR)), page_dir, - ).replace(os.sep, "/") + ).replace(os.sep, '/') - bt = "`" * 3 - indented_code = textwrap.indent(code.rstrip(), indent + " ") + bt = '`' * 3 + indented_code = textwrap.indent(code.rstrip(), indent + ' ') return ( - f"{indent}![plot]({img_rel})\n" - f"\n" + f'{indent}![plot]({img_rel})\n' + f'\n' f'{indent}??? note "Source code"\n' - f"{indent} {bt}python\n" - f"{indented_code}\n" - f"{indent} {bt}\n" + f'{indent} {bt}python\n' + f'{indented_code}\n' + f'{indent} {bt}\n' ) new_text = _FENCE_RE.sub(_replace, text) @@ -89,14 +89,14 @@ def _replace(match: re.Match) -> str: def main() -> None: - md_files = sorted(_DOCS_DIR.rglob("*.md")) + md_files = sorted(_DOCS_DIR.rglob('*.md')) modified = 0 for md_path in md_files: - print(f"Processing {md_path}", file=sys.stderr) + print(f'Processing {md_path}', file=sys.stderr) if process_file(md_path): modified += 1 - print(f"Modified {modified} file(s)", file=sys.stderr) + print(f'Modified {modified} file(s)', file=sys.stderr) -if __name__ == "__main__": +if __name__ == '__main__': main() diff --git a/docs/reference/affine.md b/docs/reference/affine.md deleted file mode 100644 index 021a80f6a..000000000 --- a/docs/reference/affine.md +++ /dev/null @@ -1,3 +0,0 @@ -# AffineMatrix - -::: torchio.AffineMatrix diff --git a/docs/reference/axes.md b/docs/reference/axes.md deleted file mode 100644 index ccf3dc5dd..000000000 --- a/docs/reference/axes.md +++ /dev/null @@ -1,9 +0,0 @@ -# Axes - -::: torchio.data.axes.AxesType - -::: torchio.data.axes.validate_axes - -::: torchio.data.axes.axes_type - -::: torchio.data.axes.get_axis_mapping diff --git a/docs/reference/backends.md b/docs/reference/backends.md deleted file mode 100644 index 8fccccd8a..000000000 --- a/docs/reference/backends.md +++ /dev/null @@ -1,28 +0,0 @@ -# Backends - -Lazy image data backends and the registry used to select them. See -[Lazy loading and backends](../concepts/lazy-loading.md) for an overview. - -## Protocols - -::: torchio.data.backends.ImageDataBackend - -::: torchio.data.backends.LazyReader - -## Registration - -::: torchio.data.backends.BackendRequest - -::: torchio.data.backends.register_backend - -::: torchio.data.backends.unregister_backend - -::: torchio.data.backends.resolve_backend - -## Built-in backends - -::: torchio.data.backends.NibabelBackend - -::: torchio.data.backends.ZarrBackend - -::: torchio.data.backends.TensorBackend diff --git a/docs/reference/bboxes.md b/docs/reference/bboxes.md deleted file mode 100644 index e117932c5..000000000 --- a/docs/reference/bboxes.md +++ /dev/null @@ -1,7 +0,0 @@ -# BoundingBoxes - -::: torchio.BoundingBoxes - -::: torchio.BoundingBoxFormat - -::: torchio.Representation diff --git a/docs/reference/cli.md b/docs/reference/cli.md deleted file mode 100644 index 88b6d6a13..000000000 --- a/docs/reference/cli.md +++ /dev/null @@ -1,56 +0,0 @@ -# Command-line interface - -TorchIO includes a `torchio` command with several subcommands. - -## Usage - -``` -torchio [plot|animate|info|convert|transform|cache] [options] -``` - -Pass `--version` to print the installed TorchIO version and exit. - -## Reference - -The CLI is built with [tyro](https://brentyi.github.io/tyro/). -Help text for each subcommand is generated from the source code. - -::: torchio.cli.Plot - options: - show_root_heading: true - heading_level: 3 - -::: torchio.cli.Animate - options: - show_root_heading: true - heading_level: 3 - -::: torchio.cli.Info - options: - show_root_heading: true - heading_level: 3 - -::: torchio.cli.Convert - options: - show_root_heading: true - heading_level: 3 - -::: torchio.cli.Transform - options: - show_root_heading: true - heading_level: 3 - -::: torchio.cli.Cache - options: - show_root_heading: true - heading_level: 3 - -::: torchio.cli.Dir - options: - show_root_heading: true - heading_level: 4 - -::: torchio.cli.Clean - options: - show_root_heading: true - heading_level: 4 diff --git a/docs/reference/image.md b/docs/reference/image.md deleted file mode 100644 index 0ec2cf4a0..000000000 --- a/docs/reference/image.md +++ /dev/null @@ -1,7 +0,0 @@ -# Image - -::: torchio.ScalarImage - -::: torchio.LabelMap - -::: torchio.Image diff --git a/docs/reference/loader.md b/docs/reference/loader.md deleted file mode 100644 index d4fdad170..000000000 --- a/docs/reference/loader.md +++ /dev/null @@ -1,19 +0,0 @@ -# Data loading - -## Loaders - -::: torchio.SubjectsLoader - -::: torchio.ImagesLoader - -## Collation functions - -::: torchio.collate_subjects - -::: torchio.collate_images - -## Batch containers - -::: torchio.SubjectsBatch - -::: torchio.ImagesBatch diff --git a/docs/reference/patches.md b/docs/reference/patches.md deleted file mode 100644 index 0b1dd3f6e..000000000 --- a/docs/reference/patches.md +++ /dev/null @@ -1,36 +0,0 @@ -# Patch-based pipelines - -::: torchio.data.GridSampler - options: - show_root_heading: true - heading_level: 2 - -::: torchio.data.UniformSampler - options: - show_root_heading: true - heading_level: 2 - -::: torchio.data.WeightedSampler - options: - show_root_heading: true - heading_level: 2 - -::: torchio.data.LabelSampler - options: - show_root_heading: true - heading_level: 2 - -::: torchio.data.PatchAggregator - options: - show_root_heading: true - heading_level: 2 - -::: torchio.data.Queue - options: - show_root_heading: true - heading_level: 2 - -::: torchio.data.PatchLocation - options: - show_root_heading: true - heading_level: 2 diff --git a/docs/reference/points.md b/docs/reference/points.md deleted file mode 100644 index ac7585546..000000000 --- a/docs/reference/points.md +++ /dev/null @@ -1,3 +0,0 @@ -# Points - -::: torchio.Points diff --git a/docs/reference/random_parameters.md b/docs/reference/random_parameters.md deleted file mode 100644 index c0f8cf1c8..000000000 --- a/docs/reference/random_parameters.md +++ /dev/null @@ -1,16 +0,0 @@ -# Random parameters - -Many transforms accept a value, a range, or a distribution for their -randomizable parameters (for example `degrees`, `scales`, or `std`). You pass -the specification directly and the transform samples from it at apply time: - -- a scalar is deterministic, e.g. `degrees=10`; -- a 2-tuple `(a, b)` samples uniformly, e.g. `degrees=(-10, 10)`; -- a 3-tuple sets per-axis values and a 6-tuple per-axis ranges (for spatial - parameters); -- a [`Choice`](#choice) samples from a discrete set; -- any `torch.distributions.Distribution` samples from that distribution. - -## Choice - -::: torchio.Choice diff --git a/docs/reference/subject.md b/docs/reference/subject.md deleted file mode 100644 index 4b894ff8a..000000000 --- a/docs/reference/subject.md +++ /dev/null @@ -1,8 +0,0 @@ -# Subject / Study - -::: torchio.Subject - -!!! note - `tio.Study` is an alias for `tio.Subject`. Both refer to the - same class. Use whichever fits your domain: "subject" for - neuroscience, "study" for radiology. diff --git a/docs/reference/transforms.md b/docs/reference/transforms.md deleted file mode 100644 index d5eaedb59..000000000 --- a/docs/reference/transforms.md +++ /dev/null @@ -1,13 +0,0 @@ -# Transforms - -::: torchio.Transform - -::: torchio.SpatialTransform - -::: torchio.IntensityTransform - -::: torchio.AppliedTransform - -::: torchio.To - -::: torchio.MonaiAdapter diff --git a/docs/reference/transforms/affine_transform.md b/docs/reference/transforms/affine_transform.md deleted file mode 100644 index 96652b257..000000000 --- a/docs/reference/transforms/affine_transform.md +++ /dev/null @@ -1,5 +0,0 @@ -# Affine - -::: torchio.Affine - options: - show_root_heading: false diff --git a/docs/reference/transforms/anisotropy.md b/docs/reference/transforms/anisotropy.md deleted file mode 100644 index 642016895..000000000 --- a/docs/reference/transforms/anisotropy.md +++ /dev/null @@ -1,5 +0,0 @@ -# Anisotropy - -::: torchio.Anisotropy - options: - show_root_heading: false diff --git a/docs/reference/transforms/bias_field.md b/docs/reference/transforms/bias_field.md deleted file mode 100644 index cab083173..000000000 --- a/docs/reference/transforms/bias_field.md +++ /dev/null @@ -1,5 +0,0 @@ -# BiasField - -::: torchio.BiasField - options: - show_root_heading: false diff --git a/docs/reference/transforms/blur.md b/docs/reference/transforms/blur.md deleted file mode 100644 index d98a48765..000000000 --- a/docs/reference/transforms/blur.md +++ /dev/null @@ -1,5 +0,0 @@ -# Blur - -::: torchio.Blur - options: - show_root_heading: false diff --git a/docs/reference/transforms/clamp.md b/docs/reference/transforms/clamp.md deleted file mode 100644 index ecec00dfc..000000000 --- a/docs/reference/transforms/clamp.md +++ /dev/null @@ -1,5 +0,0 @@ -# Clamp - -::: torchio.Clamp - options: - show_root_heading: false diff --git a/docs/reference/transforms/compose.md b/docs/reference/transforms/compose.md deleted file mode 100644 index 79fcad9f7..000000000 --- a/docs/reference/transforms/compose.md +++ /dev/null @@ -1,5 +0,0 @@ -# Compose - -::: torchio.Compose - options: - show_root_heading: false diff --git a/docs/reference/transforms/contour.md b/docs/reference/transforms/contour.md deleted file mode 100644 index 61696739e..000000000 --- a/docs/reference/transforms/contour.md +++ /dev/null @@ -1,5 +0,0 @@ -# Contour - -::: torchio.Contour - options: - show_root_heading: false diff --git a/docs/reference/transforms/copy_affine.md b/docs/reference/transforms/copy_affine.md deleted file mode 100644 index 4954b757f..000000000 --- a/docs/reference/transforms/copy_affine.md +++ /dev/null @@ -1,5 +0,0 @@ -# CopyAffine - -::: torchio.CopyAffine - options: - show_root_heading: false diff --git a/docs/reference/transforms/cornucopia_adapter.md b/docs/reference/transforms/cornucopia_adapter.md deleted file mode 100644 index 8bd8fbede..000000000 --- a/docs/reference/transforms/cornucopia_adapter.md +++ /dev/null @@ -1,5 +0,0 @@ -# CornucopiaAdapter - -::: torchio.CornucopiaAdapter - options: - show_root_heading: false diff --git a/docs/reference/transforms/crop.md b/docs/reference/transforms/crop.md deleted file mode 100644 index a23902b39..000000000 --- a/docs/reference/transforms/crop.md +++ /dev/null @@ -1,5 +0,0 @@ -# Crop - -::: torchio.Crop - options: - show_root_heading: false diff --git a/docs/reference/transforms/crop_or_pad.md b/docs/reference/transforms/crop_or_pad.md deleted file mode 100644 index f99c199e5..000000000 --- a/docs/reference/transforms/crop_or_pad.md +++ /dev/null @@ -1,5 +0,0 @@ -# CropOrPad - -::: torchio.CropOrPad - options: - show_root_heading: false diff --git a/docs/reference/transforms/elastic_deformation.md b/docs/reference/transforms/elastic_deformation.md deleted file mode 100644 index 2fec7fdd0..000000000 --- a/docs/reference/transforms/elastic_deformation.md +++ /dev/null @@ -1,5 +0,0 @@ -# ElasticDeformation - -::: torchio.ElasticDeformation - options: - show_root_heading: false diff --git a/docs/reference/transforms/ensure_shape_multiple.md b/docs/reference/transforms/ensure_shape_multiple.md deleted file mode 100644 index 5ef52e820..000000000 --- a/docs/reference/transforms/ensure_shape_multiple.md +++ /dev/null @@ -1,5 +0,0 @@ -# EnsureShapeMultiple - -::: torchio.EnsureShapeMultiple - options: - show_root_heading: false diff --git a/docs/reference/transforms/flip.md b/docs/reference/transforms/flip.md deleted file mode 100644 index 6770e75f3..000000000 --- a/docs/reference/transforms/flip.md +++ /dev/null @@ -1,5 +0,0 @@ -# Flip - -::: torchio.Flip - options: - show_root_heading: false diff --git a/docs/reference/transforms/gamma.md b/docs/reference/transforms/gamma.md deleted file mode 100644 index b9edca689..000000000 --- a/docs/reference/transforms/gamma.md +++ /dev/null @@ -1,5 +0,0 @@ -# Gamma - -::: torchio.Gamma - options: - show_root_heading: false diff --git a/docs/reference/transforms/ghosting.md b/docs/reference/transforms/ghosting.md deleted file mode 100644 index ab71e08cc..000000000 --- a/docs/reference/transforms/ghosting.md +++ /dev/null @@ -1,5 +0,0 @@ -# Ghosting - -::: torchio.Ghosting - options: - show_root_heading: false diff --git a/docs/reference/transforms/histogram_standardization.md b/docs/reference/transforms/histogram_standardization.md deleted file mode 100644 index 0935c1f09..000000000 --- a/docs/reference/transforms/histogram_standardization.md +++ /dev/null @@ -1,11 +0,0 @@ -# HistogramStandardization - -::: torchio.HistogramStandardization - options: - show_root_heading: false - -## Landmark computation - -::: torchio.transforms.intensity.histogram_standardization.compute_histogram_landmarks - options: - show_root_heading: true diff --git a/docs/reference/transforms/keep_largest.md b/docs/reference/transforms/keep_largest.md deleted file mode 100644 index e57f73ea8..000000000 --- a/docs/reference/transforms/keep_largest.md +++ /dev/null @@ -1,5 +0,0 @@ -# KeepLargestComponent - -::: torchio.KeepLargestComponent - options: - show_root_heading: false diff --git a/docs/reference/transforms/labels_to_image.md b/docs/reference/transforms/labels_to_image.md deleted file mode 100644 index 6a75051ae..000000000 --- a/docs/reference/transforms/labels_to_image.md +++ /dev/null @@ -1,5 +0,0 @@ -# LabelsToImage - -::: torchio.LabelsToImage - options: - show_root_heading: false diff --git a/docs/reference/transforms/lambda_transform.md b/docs/reference/transforms/lambda_transform.md deleted file mode 100644 index a3550e74c..000000000 --- a/docs/reference/transforms/lambda_transform.md +++ /dev/null @@ -1,5 +0,0 @@ -# Lambda - -::: torchio.Lambda - options: - show_root_heading: false diff --git a/docs/reference/transforms/mask.md b/docs/reference/transforms/mask.md deleted file mode 100644 index 71b9f765f..000000000 --- a/docs/reference/transforms/mask.md +++ /dev/null @@ -1,5 +0,0 @@ -# Mask - -::: torchio.Mask - options: - show_root_heading: false diff --git a/docs/reference/transforms/monai_adapter.md b/docs/reference/transforms/monai_adapter.md deleted file mode 100644 index 8f134b82a..000000000 --- a/docs/reference/transforms/monai_adapter.md +++ /dev/null @@ -1,5 +0,0 @@ -# MonaiAdapter - -::: torchio.MonaiAdapter - options: - show_root_heading: false diff --git a/docs/reference/transforms/motion.md b/docs/reference/transforms/motion.md deleted file mode 100644 index 107081415..000000000 --- a/docs/reference/transforms/motion.md +++ /dev/null @@ -1,5 +0,0 @@ -# Motion - -::: torchio.Motion - options: - show_root_heading: false diff --git a/docs/reference/transforms/noise.md b/docs/reference/transforms/noise.md deleted file mode 100644 index 48b2ab95a..000000000 --- a/docs/reference/transforms/noise.md +++ /dev/null @@ -1,5 +0,0 @@ -# Noise - -::: torchio.Noise - options: - show_root_heading: false diff --git a/docs/reference/transforms/normalize.md b/docs/reference/transforms/normalize.md deleted file mode 100644 index 9844c0ca2..000000000 --- a/docs/reference/transforms/normalize.md +++ /dev/null @@ -1,5 +0,0 @@ -# Normalize - -::: torchio.Normalize - options: - show_root_heading: false diff --git a/docs/reference/transforms/one_hot.md b/docs/reference/transforms/one_hot.md deleted file mode 100644 index 169b571fd..000000000 --- a/docs/reference/transforms/one_hot.md +++ /dev/null @@ -1,5 +0,0 @@ -# OneHot - -::: torchio.OneHot - options: - show_root_heading: false diff --git a/docs/reference/transforms/one_of.md b/docs/reference/transforms/one_of.md deleted file mode 100644 index bcd8f55e7..000000000 --- a/docs/reference/transforms/one_of.md +++ /dev/null @@ -1,5 +0,0 @@ -# OneOf - -::: torchio.OneOf - options: - show_root_heading: false diff --git a/docs/reference/transforms/pad.md b/docs/reference/transforms/pad.md deleted file mode 100644 index aefb3327b..000000000 --- a/docs/reference/transforms/pad.md +++ /dev/null @@ -1,3 +0,0 @@ -# Pad - -::: torchio.Pad diff --git a/docs/reference/transforms/pca.md b/docs/reference/transforms/pca.md deleted file mode 100644 index cd35f59c5..000000000 --- a/docs/reference/transforms/pca.md +++ /dev/null @@ -1,5 +0,0 @@ -# PCA - -::: torchio.PCA - options: - show_root_heading: false diff --git a/docs/reference/transforms/remap_labels.md b/docs/reference/transforms/remap_labels.md deleted file mode 100644 index 4b0d4afcd..000000000 --- a/docs/reference/transforms/remap_labels.md +++ /dev/null @@ -1,5 +0,0 @@ -# RemapLabels - -::: torchio.RemapLabels - options: - show_root_heading: false diff --git a/docs/reference/transforms/remove_labels.md b/docs/reference/transforms/remove_labels.md deleted file mode 100644 index 091bc92c9..000000000 --- a/docs/reference/transforms/remove_labels.md +++ /dev/null @@ -1,5 +0,0 @@ -# RemoveLabels - -::: torchio.RemoveLabels - options: - show_root_heading: false diff --git a/docs/reference/transforms/reorient.md b/docs/reference/transforms/reorient.md deleted file mode 100644 index 612a84959..000000000 --- a/docs/reference/transforms/reorient.md +++ /dev/null @@ -1,5 +0,0 @@ -# Reorient - -::: torchio.Reorient - options: - show_root_heading: false diff --git a/docs/reference/transforms/resample.md b/docs/reference/transforms/resample.md deleted file mode 100644 index e8c935161..000000000 --- a/docs/reference/transforms/resample.md +++ /dev/null @@ -1,5 +0,0 @@ -# Resample - -::: torchio.Resample - options: - show_root_heading: false diff --git a/docs/reference/transforms/resize.md b/docs/reference/transforms/resize.md deleted file mode 100644 index 7483d791a..000000000 --- a/docs/reference/transforms/resize.md +++ /dev/null @@ -1,5 +0,0 @@ -# Resize - -::: torchio.Resize - options: - show_root_heading: false diff --git a/docs/reference/transforms/sequential_labels.md b/docs/reference/transforms/sequential_labels.md deleted file mode 100644 index 359f3bfca..000000000 --- a/docs/reference/transforms/sequential_labels.md +++ /dev/null @@ -1,5 +0,0 @@ -# SequentialLabels - -::: torchio.SequentialLabels - options: - show_root_heading: false diff --git a/docs/reference/transforms/some_of.md b/docs/reference/transforms/some_of.md deleted file mode 100644 index ca7ef0eee..000000000 --- a/docs/reference/transforms/some_of.md +++ /dev/null @@ -1,5 +0,0 @@ -# SomeOf - -::: torchio.SomeOf - options: - show_root_heading: false diff --git a/docs/reference/transforms/spatial.md b/docs/reference/transforms/spatial.md deleted file mode 100644 index cbb2ffd61..000000000 --- a/docs/reference/transforms/spatial.md +++ /dev/null @@ -1,5 +0,0 @@ -# Spatial - -::: torchio.Spatial - options: - show_root_heading: false diff --git a/docs/reference/transforms/spike.md b/docs/reference/transforms/spike.md deleted file mode 100644 index c2f35bd45..000000000 --- a/docs/reference/transforms/spike.md +++ /dev/null @@ -1,5 +0,0 @@ -# Spike - -::: torchio.Spike - options: - show_root_heading: false diff --git a/docs/reference/transforms/standardize.md b/docs/reference/transforms/standardize.md deleted file mode 100644 index af86eb0b7..000000000 --- a/docs/reference/transforms/standardize.md +++ /dev/null @@ -1,5 +0,0 @@ -# Standardize - -::: torchio.Standardize - options: - show_root_heading: false diff --git a/docs/reference/transforms/swap.md b/docs/reference/transforms/swap.md deleted file mode 100644 index 783b28517..000000000 --- a/docs/reference/transforms/swap.md +++ /dev/null @@ -1,5 +0,0 @@ -# Swap - -::: torchio.Swap - options: - show_root_heading: false diff --git a/docs/reference/transforms/to.md b/docs/reference/transforms/to.md deleted file mode 100644 index 24b08a5b8..000000000 --- a/docs/reference/transforms/to.md +++ /dev/null @@ -1,5 +0,0 @@ -# To - -::: torchio.To - options: - show_root_heading: false diff --git a/docs/reference/transforms/to_reference_space.md b/docs/reference/transforms/to_reference_space.md deleted file mode 100644 index 91acd2f8b..000000000 --- a/docs/reference/transforms/to_reference_space.md +++ /dev/null @@ -1,5 +0,0 @@ -# ToReferenceSpace - -::: torchio.ToReferenceSpace - options: - show_root_heading: false diff --git a/docs/reference/transforms/transpose.md b/docs/reference/transforms/transpose.md deleted file mode 100644 index 590c51506..000000000 --- a/docs/reference/transforms/transpose.md +++ /dev/null @@ -1,5 +0,0 @@ -# Transpose - -::: torchio.Transpose - options: - show_root_heading: false diff --git a/docs/reference/visualization.md b/docs/reference/visualization.md deleted file mode 100644 index fb917dfd8..000000000 --- a/docs/reference/visualization.md +++ /dev/null @@ -1,9 +0,0 @@ -# Visualization - -::: torchio.visualization.plot_image - options: - show_root_heading: true - -::: torchio.visualization.plot_subject - options: - show_root_heading: true diff --git a/docs/transforms/augmentation/Compose.md b/docs/transforms/augmentation/Compose.md new file mode 100644 index 000000000..9305c3beb --- /dev/null +++ b/docs/transforms/augmentation/Compose.md @@ -0,0 +1,3 @@ +# Compose + +::: torchio.transforms.Compose diff --git a/docs/transforms/augmentation/OneOf.md b/docs/transforms/augmentation/OneOf.md new file mode 100644 index 000000000..dcc338a24 --- /dev/null +++ b/docs/transforms/augmentation/OneOf.md @@ -0,0 +1,3 @@ +# OneOf + +::: torchio.transforms.OneOf diff --git a/docs/transforms/augmentation/RandomAffine.md b/docs/transforms/augmentation/RandomAffine.md new file mode 100644 index 000000000..58621fff4 --- /dev/null +++ b/docs/transforms/augmentation/RandomAffine.md @@ -0,0 +1,13 @@ +# RandomAffine + +::: torchio.transforms.RandomAffine + +```python plot +import torchio as tio +subject = tio.datasets.Slicer('CTChest') +ct = subject.CT_chest +transform = tio.RandomAffine() +ct_transformed = transform(ct) +subject.add_image(ct_transformed, 'Transformed') +subject.plot() +``` diff --git a/docs/transforms/augmentation/RandomAffineElasticDeformation.md b/docs/transforms/augmentation/RandomAffineElasticDeformation.md new file mode 100644 index 000000000..9bed21560 --- /dev/null +++ b/docs/transforms/augmentation/RandomAffineElasticDeformation.md @@ -0,0 +1,14 @@ +# RandomAffineElasticDeformation + +::: torchio.transforms.RandomAffineElasticDeformation + +```python plot +import torchio as tio +subject = tio.datasets.Slicer('CTChest') +ct = subject.CT_chest +elastic_kwargs = {'max_displacement': (17, 12, 2)} +transform = tio.RandomAffineElasticDeformation(elastic_kwargs=elastic_kwargs) +ct_transformed = transform(ct) +subject.add_image(ct_transformed, 'Transformed') +subject.plot() +``` diff --git a/docs/transforms/augmentation/RandomAnisotropy.md b/docs/transforms/augmentation/RandomAnisotropy.md new file mode 100644 index 000000000..eba313261 --- /dev/null +++ b/docs/transforms/augmentation/RandomAnisotropy.md @@ -0,0 +1,5 @@ +# RandomAnisotropy + +![Simulation of an image with highly anisotropic spacing](../../images/random_anisotropy.png) + +::: torchio.transforms.RandomAnisotropy diff --git a/docs/transforms/augmentation/RandomBiasField.md b/docs/transforms/augmentation/RandomBiasField.md new file mode 100644 index 000000000..6d18a12aa --- /dev/null +++ b/docs/transforms/augmentation/RandomBiasField.md @@ -0,0 +1,5 @@ +# RandomBiasField + +![MRI bias field artifact](../../images/random_bias_field.gif) + +::: torchio.transforms.RandomBiasField diff --git a/docs/transforms/augmentation/RandomBlur.md b/docs/transforms/augmentation/RandomBlur.md new file mode 100644 index 000000000..3291be200 --- /dev/null +++ b/docs/transforms/augmentation/RandomBlur.md @@ -0,0 +1,3 @@ +# RandomBlur + +::: torchio.transforms.RandomBlur diff --git a/docs/transforms/augmentation/RandomElasticDeformation.md b/docs/transforms/augmentation/RandomElasticDeformation.md new file mode 100644 index 000000000..ce2fa7f76 --- /dev/null +++ b/docs/transforms/augmentation/RandomElasticDeformation.md @@ -0,0 +1,5 @@ +# RandomElasticDeformation + +![Random elastic deformation](../../images/random_elastic_deformation.gif) + +::: torchio.transforms.RandomElasticDeformation diff --git a/docs/transforms/augmentation/RandomFlip.md b/docs/transforms/augmentation/RandomFlip.md new file mode 100644 index 000000000..75830e63d --- /dev/null +++ b/docs/transforms/augmentation/RandomFlip.md @@ -0,0 +1,3 @@ +# RandomFlip + +::: torchio.transforms.RandomFlip diff --git a/docs/transforms/augmentation/RandomGamma.md b/docs/transforms/augmentation/RandomGamma.md new file mode 100644 index 000000000..5becf06d3 --- /dev/null +++ b/docs/transforms/augmentation/RandomGamma.md @@ -0,0 +1,16 @@ +# RandomGamma + +::: torchio.transforms.RandomGamma + +```python plot +import torchio as tio +subject = tio.datasets.FPG() +subject.remove_image('seg') +transform = tio.RandomGamma(log_gamma=(-0.3, -0.3)) +transformed = transform(subject) +subject.add_image(transformed.t1, 'log -0.3') +transform = tio.RandomGamma(log_gamma=(0.3, 0.3)) +transformed = transform(subject) +subject.add_image(transformed.t1, 'log 0.3') +subject.plot() +``` diff --git a/docs/transforms/augmentation/RandomGhosting.md b/docs/transforms/augmentation/RandomGhosting.md new file mode 100644 index 000000000..9702e33c8 --- /dev/null +++ b/docs/transforms/augmentation/RandomGhosting.md @@ -0,0 +1,5 @@ +# RandomGhosting + +![MRI k-space ghosting artifacts](../../images/random_ghosting.gif) + +::: torchio.transforms.RandomGhosting diff --git a/docs/transforms/augmentation/RandomLabelsToImage.md b/docs/transforms/augmentation/RandomLabelsToImage.md new file mode 100644 index 000000000..c1948245d --- /dev/null +++ b/docs/transforms/augmentation/RandomLabelsToImage.md @@ -0,0 +1,27 @@ +# RandomLabelsToImage + +::: torchio.transforms.RandomLabelsToImage + +```python plot +import torch +import torchio as tio +torch.manual_seed(42) +colin = tio.datasets.Colin27(2008) +label_map = colin.cls +colin.remove_image('t1') +colin.remove_image('t2') +colin.remove_image('pd') +downsample = tio.Resample(1) +blurring_transform = tio.RandomBlur(std=0.6) +create_synthetic_image = tio.RandomLabelsToImage( + image_key='synthetic', + ignore_background=True, +) +transform = tio.Compose(( + downsample, + create_synthetic_image, + blurring_transform, +)) +colin_synth = transform(colin) +colin_synth.plot() +``` diff --git a/docs/transforms/augmentation/RandomMotion.md b/docs/transforms/augmentation/RandomMotion.md new file mode 100644 index 000000000..eef639ff0 --- /dev/null +++ b/docs/transforms/augmentation/RandomMotion.md @@ -0,0 +1,5 @@ +# RandomMotion + +![MRI k-space motion artifacts](../../images/random_motion.gif) + +::: torchio.transforms.RandomMotion diff --git a/docs/transforms/augmentation/RandomNoise.md b/docs/transforms/augmentation/RandomNoise.md new file mode 100644 index 000000000..f846a045a --- /dev/null +++ b/docs/transforms/augmentation/RandomNoise.md @@ -0,0 +1,5 @@ +# RandomNoise + +![Random Gaussian noise](../../images/random_noise.gif) + +::: torchio.transforms.RandomNoise diff --git a/docs/transforms/augmentation/RandomSpike.md b/docs/transforms/augmentation/RandomSpike.md new file mode 100644 index 000000000..5b059ea40 --- /dev/null +++ b/docs/transforms/augmentation/RandomSpike.md @@ -0,0 +1,5 @@ +# RandomSpike + +![MRI k-space spike artifacts](../../images/random_spike.gif) + +::: torchio.transforms.RandomSpike diff --git a/docs/transforms/augmentation/RandomSwap.md b/docs/transforms/augmentation/RandomSwap.md new file mode 100644 index 000000000..86d90c97a --- /dev/null +++ b/docs/transforms/augmentation/RandomSwap.md @@ -0,0 +1,5 @@ +# RandomSwap + +![Random patches swapping](../../images/random_swap.jpg) + +::: torchio.transforms.RandomSwap diff --git a/docs/transforms/augmentation/index.md b/docs/transforms/augmentation/index.md new file mode 100644 index 000000000..98d2bc92a --- /dev/null +++ b/docs/transforms/augmentation/index.md @@ -0,0 +1,42 @@ +# Augmentation + +Augmentation transforms generate different results every time they are called. + +![Augmented image](../../images/augmentation.gif) + +## Base class + +::: torchio.transforms.augmentation.RandomTransform + options: + show_root_heading: true + +## Composition + +| Transform | Description | +|-----------|-------------| +| [`Compose`](Compose.md) | Compose several transforms together | +| [`OneOf`](OneOf.md) | Apply one of the given transforms | + +## Spatial + +| Transform | Description | +|-----------|-------------| +| [`RandomFlip`](RandomFlip.md) | Randomly reverse the order of elements in an image along the given axes | +| [`RandomAffine`](RandomAffine.md) | Apply a random affine transformation | +| [`RandomElasticDeformation`](RandomElasticDeformation.md) | Apply a random elastic deformation | +| [`RandomAffineElasticDeformation`](RandomAffineElasticDeformation.md) | Apply random affine and elastic deformation | +| [`RandomAnisotropy`](RandomAnisotropy.md) | Downsample an image along an axis and upsample back | + +## Intensity + +| Transform | Description | +|-----------|-------------| +| [`RandomMotion`](RandomMotion.md) | Simulate MRI motion artifacts | +| [`RandomGhosting`](RandomGhosting.md) | Simulate MRI ghosting artifacts | +| [`RandomSpike`](RandomSpike.md) | Simulate MRI spike artifacts | +| [`RandomBiasField`](RandomBiasField.md) | Simulate MRI bias field artifacts | +| [`RandomBlur`](RandomBlur.md) | Blur an image using a random-sized Gaussian filter | +| [`RandomNoise`](RandomNoise.md) | Add Gaussian noise | +| [`RandomSwap`](RandomSwap.md) | Randomly swap patches in an image | +| [`RandomLabelsToImage`](RandomLabelsToImage.md) | Generate an image from a segmentation | +| [`RandomGamma`](RandomGamma.md) | Randomly change contrast of an image | diff --git a/docs/transforms/index.md b/docs/transforms/index.md new file mode 100644 index 000000000..645f6d688 --- /dev/null +++ b/docs/transforms/index.md @@ -0,0 +1,197 @@ +# Transforms + +![Augmented image](../images/fpg_progressive.gif) + +TorchIO transforms take as input instances of +[`Subject`][torchio.Subject] or +[`Image`][torchio.Image] (and its subclasses), 4D PyTorch tensors, +4D NumPy arrays, SimpleITK images, NiBabel images, or Python dictionaries +(see [`Transform`][torchio.transforms.Transform]). + +For example: + +```python +>>> import torch +>>> import numpy as np +>>> import torchio as tio +>>> affine_transform = tio.RandomAffine() +>>> tensor = torch.rand(1, 256, 256, 159) +>>> transformed_tensor = affine_transform(tensor) +>>> type(transformed_tensor) + +>>> array = np.random.rand(1, 256, 256, 159) +>>> transformed_array = affine_transform(array) +>>> type(transformed_array) + +>>> subject = tio.datasets.Colin27() +>>> transformed_subject = affine_transform(subject) +>>> transformed_subject +Subject(Keys: ('t1', 'head', 'brain'); images: 3) +``` + +Transforms can also be applied from the command line using +`torchio-transform`. + +All transforms inherit from [`Transform`][torchio.transforms.Transform]: + +::: torchio.transforms.Transform + options: + members: + - apply_transform + - __call__ + +## Composability + +Transforms can be composed to create directed acyclic graphs defining the +probability that each transform will be applied. + +For example, to obtain the following graph: + +![Composed transform](../images/composed.png) + +We can type: + +```python +>>> import torchio as tio +>>> spatial_transforms = { +... tio.RandomElasticDeformation(): 0.2, +... tio.RandomAffine(): 0.8, +... } +>>> transform = tio.Compose([ +... tio.OneOf(spatial_transforms, p=0.5), +... tio.RescaleIntensity(out_min_max=(0, 1)), +... ]) +``` + +## Interoperability with MONAI + +[MONAI](https://monai.io/) dictionary transforms can be used inside TorchIO +pipelines via [`MonaiAdapter`][torchio.transforms.MonaiAdapter]. The adapter +handles conversion between TorchIO's [`Subject`][torchio.Subject] and MONAI's +expected dictionary format, including affine propagation for spatial transforms. + +```python +>>> import torchio as tio +>>> from monai.transforms import NormalizeIntensityd, RandSpatialCropd +>>> pipeline = tio.Compose([ +... tio.ToCanonical(), +... tio.MonaiAdapter(NormalizeIntensityd(keys=["t1"])), +... tio.RandomFlip(), +... tio.MonaiAdapter( +... RandSpatialCropd(keys=["t1", "seg"], roi_size=[64, 64, 64]), +... ), +... ]) +``` + +## Reproducibility + +When transforms are instantiated, we typically need to pass values that will be +used to sample the transform parameters when the +[`__call__`][torchio.transforms.Transform.__call__] method of the transform is +called, i.e., when the transform instance is called. + +All random transforms have a corresponding deterministic class, that can be +applied again to obtain exactly the same result. +The [`Subject`][torchio.Subject] class contains some convenient methods to +reproduce transforms: + +```python +>>> import torchio as tio +>>> subject = tio.datasets.FPG() +>>> transforms = ( +... tio.CropOrPad((100, 200, 300)), +... tio.RandomFlip(axes=['LR', 'AP', 'IS']), +... tio.OneOf([tio.RandomAnisotropy(), tio.RandomElasticDeformation()]), +... ) +>>> transform = tio.Compose(transforms) +>>> transformed = transform(subject) +>>> reproduce_transform = transformed.get_composed_history() +>>> reproduced = reproduce_transform(subject) +``` + +## Invertibility + +Inverting transforms can be especially useful in scenarios in which one needs +to apply some transformation, infer a segmentation on the transformed data and +apply the inverse transform to the inference in order to bring it back to the +original space. + +This is particularly useful, for example, for +[test-time augmentation](https://www.nature.com/articles/s41598-020-61808-3) +or [aleatoric uncertainty estimation](https://www.sciencedirect.com/science/article/pii/S0925231219301961): + +```python +>>> import torchio as tio +>>> # Mock a segmentation CNN +>>> def model(x): +... return x +... +>>> subject = tio.datasets.Colin27() +>>> transform = tio.RandomAffine() +>>> segmentations = [] +>>> num_segmentations = 10 +>>> for _ in range(num_segmentations): +... transform = tio.RandomAffine(image_interpolation='bspline') +... transformed = transform(subject) +... segmentation = model(transformed) +... transformed_native_space = segmentation.apply_inverse_transform(image_interpolation='linear') +... segmentations.append(transformed_native_space) +... +``` + +Transforms can be classified in three types, according to their degree of +invertibility: + +- **Lossless**: transforms that can be inverted with no loss of information, + such as [`RandomFlip`][torchio.transforms.RandomFlip], + [`Pad`][torchio.transforms.Pad], + or [`RandomNoise`][torchio.transforms.RandomNoise]. + +- **Lossy**: transforms that can be inverted with some loss of information, such + as [`RandomAffine`][torchio.transforms.RandomAffine], + or [`Crop`][torchio.transforms.Crop]. + +- **Impossible**: transforms that cannot be inverted, such as + [`RandomBlur`][torchio.transforms.RandomBlur]. + +Non-invertible transforms will be ignored by the +[`apply_inverse_transform`][torchio.Subject.apply_inverse_transform] method of +[`Subject`][torchio.Subject]. + +## Interpolation + +Some transforms such as +[`RandomAffine`][torchio.transforms.RandomAffine] or +[`RandomMotion`][torchio.transforms.RandomMotion] +need to interpolate intensity values during resampling. + +The available interpolation strategies can be inferred from the elements of +[`Interpolation`][torchio.transforms.interpolation.Interpolation]. + +`'linear'` interpolation, the default in TorchIO for scalar images, +is usually a good compromise between image quality and speed. +It is therefore a good choice for data augmentation during training. + +Methods such as `'bspline'` or `'lanczos'` generate +high-quality results, but are generally slower. They can be used to obtain +optimal resampling results during offline data preprocessing. + +`'nearest'` can be used for quick experimentation as it is very +fast, but produces relatively poor results for scalar images. +It is the default interpolation type for label maps, as categorical values for +the different labels need to preserved after interpolation. + +When instantiating transforms, it is possible to specify independently the +interpolation type for label maps and scalar images, as shown in the +documentation for, e.g., [`Resample`][torchio.transforms.Resample]. + +Visit the +[SimpleITK docs](https://simpleitk.org/doxygen/latest/html/namespaceitk_1_1simple.html#a7cb1ef8bd02c669c02ea2f9f5aa374e5) +for technical documentation and +[Cambridge in Colour](https://www.cambridgeincolour.com/tutorials/image-interpolation.htm) +for some further general explanations of digital image interpolation. + +::: torchio.transforms.interpolation.Interpolation + options: + members: true + show_root_heading: true diff --git a/docs/transforms/others/Lambda.md b/docs/transforms/others/Lambda.md new file mode 100644 index 000000000..f043b5720 --- /dev/null +++ b/docs/transforms/others/Lambda.md @@ -0,0 +1,3 @@ +# Lambda + +::: torchio.transforms.Lambda diff --git a/docs/transforms/others/MonaiAdapter.md b/docs/transforms/others/MonaiAdapter.md new file mode 100644 index 000000000..5c418e4fe --- /dev/null +++ b/docs/transforms/others/MonaiAdapter.md @@ -0,0 +1,3 @@ +# MonaiAdapter + +::: torchio.transforms.MonaiAdapter diff --git a/docs/transforms/preprocessing/Clamp.md b/docs/transforms/preprocessing/Clamp.md new file mode 100644 index 000000000..ac0c1bec4 --- /dev/null +++ b/docs/transforms/preprocessing/Clamp.md @@ -0,0 +1,14 @@ +# Clamp + +::: torchio.transforms.Clamp + +```python plot +import torchio as tio +subject = tio.datasets.Slicer('CTChest') +ct = subject.CT_chest +HOUNSFIELD_AIR, HOUNSFIELD_BONE = -1000, 1000 +clamp = tio.Clamp(out_min=HOUNSFIELD_AIR, out_max=HOUNSFIELD_BONE) +ct_clamped = clamp(ct) +subject.add_image(ct_clamped, 'Clamped') +subject.plot() +``` diff --git a/docs/transforms/preprocessing/Contour.md b/docs/transforms/preprocessing/Contour.md new file mode 100644 index 000000000..214943dac --- /dev/null +++ b/docs/transforms/preprocessing/Contour.md @@ -0,0 +1,3 @@ +# Contour + +::: torchio.transforms.Contour diff --git a/docs/transforms/preprocessing/CopyAffine.md b/docs/transforms/preprocessing/CopyAffine.md new file mode 100644 index 000000000..e1e5db2a3 --- /dev/null +++ b/docs/transforms/preprocessing/CopyAffine.md @@ -0,0 +1,3 @@ +# CopyAffine + +::: torchio.transforms.CopyAffine diff --git a/docs/transforms/preprocessing/Crop.md b/docs/transforms/preprocessing/Crop.md new file mode 100644 index 000000000..b749386c4 --- /dev/null +++ b/docs/transforms/preprocessing/Crop.md @@ -0,0 +1,3 @@ +# Crop + +::: torchio.transforms.Crop diff --git a/docs/transforms/preprocessing/CropOrPad.md b/docs/transforms/preprocessing/CropOrPad.md new file mode 100644 index 000000000..d55846e56 --- /dev/null +++ b/docs/transforms/preprocessing/CropOrPad.md @@ -0,0 +1,15 @@ +# CropOrPad + +::: torchio.transforms.CropOrPad + options: + members: + - _get_six_bounds_parameters + +```python plot +import torchio as tio +t1 = tio.datasets.Colin27().t1 +crop_pad = tio.CropOrPad((512, 512, 32)) +t1_pad_crop = crop_pad(t1) +subject = tio.Subject(t1=t1, crop_pad=t1_pad_crop) +subject.plot() +``` diff --git a/docs/transforms/preprocessing/EnsureShapeMultiple.md b/docs/transforms/preprocessing/EnsureShapeMultiple.md new file mode 100644 index 000000000..852d844dd --- /dev/null +++ b/docs/transforms/preprocessing/EnsureShapeMultiple.md @@ -0,0 +1,3 @@ +# EnsureShapeMultiple + +::: torchio.transforms.EnsureShapeMultiple diff --git a/docs/transforms/preprocessing/HistogramStandardization.md b/docs/transforms/preprocessing/HistogramStandardization.md new file mode 100644 index 000000000..b7b1c96a4 --- /dev/null +++ b/docs/transforms/preprocessing/HistogramStandardization.md @@ -0,0 +1,7 @@ +# HistogramStandardization + +![Histogram standardization](../../images/histogram_standardization.png) + +::: torchio.transforms.HistogramStandardization + options: + members: true diff --git a/docs/transforms/preprocessing/KeepLargestComponent.md b/docs/transforms/preprocessing/KeepLargestComponent.md new file mode 100644 index 000000000..d4dca660c --- /dev/null +++ b/docs/transforms/preprocessing/KeepLargestComponent.md @@ -0,0 +1,3 @@ +# KeepLargestComponent + +::: torchio.transforms.KeepLargestComponent diff --git a/docs/transforms/preprocessing/Mask.md b/docs/transforms/preprocessing/Mask.md new file mode 100644 index 000000000..588a953ce --- /dev/null +++ b/docs/transforms/preprocessing/Mask.md @@ -0,0 +1,13 @@ +# Mask + +::: torchio.transforms.Mask + +```python plot +import torchio as tio +subject = tio.datasets.Colin27() +subject.remove_image('head') +mask = tio.Mask('brain') +masked = mask(subject) +subject.add_image(masked.t1, 'Masked') +subject.plot() +``` diff --git a/docs/transforms/preprocessing/OneHot.md b/docs/transforms/preprocessing/OneHot.md new file mode 100644 index 000000000..f08c3215c --- /dev/null +++ b/docs/transforms/preprocessing/OneHot.md @@ -0,0 +1,3 @@ +# OneHot + +::: torchio.transforms.OneHot diff --git a/docs/transforms/preprocessing/PCA.md b/docs/transforms/preprocessing/PCA.md new file mode 100644 index 000000000..8544d7283 --- /dev/null +++ b/docs/transforms/preprocessing/PCA.md @@ -0,0 +1,3 @@ +# PCA + +::: torchio.transforms.PCA diff --git a/docs/transforms/preprocessing/Pad.md b/docs/transforms/preprocessing/Pad.md new file mode 100644 index 000000000..df5c30c0d --- /dev/null +++ b/docs/transforms/preprocessing/Pad.md @@ -0,0 +1,3 @@ +# Pad + +::: torchio.transforms.Pad diff --git a/docs/transforms/preprocessing/RemapLabels.md b/docs/transforms/preprocessing/RemapLabels.md new file mode 100644 index 000000000..265f2ec11 --- /dev/null +++ b/docs/transforms/preprocessing/RemapLabels.md @@ -0,0 +1,51 @@ +# RemapLabels + +::: torchio.transforms.RemapLabels + +```python plot +import torchio as tio + +subject = tio.datasets.FPG() +subject.remove_image('t1') + +background_labels = (0, 1, 2, 3, 4) +csf_labels = (5, 12, 16, 47, 52, 53) +white_matter_labels = ( + 45, 46, + 66, 67, + 81, 82, + 83, 84, + 85, 86, + 87, + 89, 90, + 91, 92, + 93, 94, +) + +not_gray_matter_labels = ( + background_labels + + csf_labels + + white_matter_labels +) + +gray_matter_labels = [ + label for label in subject.GIF_COLORS + if label not in not_gray_matter_labels +] + +labels_groups = ( + background_labels, + gray_matter_labels, + white_matter_labels, + csf_labels, +) +remapping = {} +for target, labels in enumerate(labels_groups): + for label in labels: + remapping[label] = target + +parcellation_to_tissues = tio.RemapLabels(remapping) +tissues = parcellation_to_tissues(subject).seg +subject.add_image(tissues, 'remapped') +subject.plot() +``` diff --git a/docs/transforms/preprocessing/RemoveLabels.md b/docs/transforms/preprocessing/RemoveLabels.md new file mode 100644 index 000000000..d3219e744 --- /dev/null +++ b/docs/transforms/preprocessing/RemoveLabels.md @@ -0,0 +1,39 @@ +# RemoveLabels + +::: torchio.transforms.RemoveLabels + +```python plot +import torchio as tio +colin = tio.datasets.Colin27(2008) +label_map = colin.cls +colin.remove_image('t2') +colin.remove_image('pd') +names_to_remove = ( + 'Fat', + 'Muscles', + 'Skin and Muscles', + 'Skull', + 'Fat 2', + 'Dura', + 'Marrow' +) +labels = [colin.NAME_TO_LABEL[name] for name in names_to_remove] +skull_stripping = tio.RemoveLabels(labels) +only_brain = skull_stripping(label_map) +colin.add_image(only_brain, 'brain') +colors = { + 0: (0, 0, 0), + 1: (127, 255, 212), + 2: (96, 204, 96), + 3: (240, 230, 140), + 4: (176, 48, 96), + 5: (48, 176, 96), + 6: (220, 247, 164), + 7: (103, 255, 255), + 9: (205, 62, 78), + 10: (238, 186, 243), + 11: (119, 159, 176), + 12: (220, 216, 20), +} +colin.plot(cmap_dict={'cls': colors, 'brain': colors}) +``` diff --git a/docs/transforms/preprocessing/Resample.md b/docs/transforms/preprocessing/Resample.md new file mode 100644 index 000000000..d3a4ba540 --- /dev/null +++ b/docs/transforms/preprocessing/Resample.md @@ -0,0 +1,13 @@ +# Resample + +::: torchio.transforms.Resample + +```python plot +import torchio as tio +subject = tio.datasets.FPG() +subject.remove_image('seg') +resample = tio.Resample(8) +t1_resampled = resample(subject.t1) +subject.add_image(t1_resampled, 'Downsampled') +subject.plot() +``` diff --git a/docs/transforms/preprocessing/RescaleIntensity.md b/docs/transforms/preprocessing/RescaleIntensity.md new file mode 100644 index 000000000..f3bdcf766 --- /dev/null +++ b/docs/transforms/preprocessing/RescaleIntensity.md @@ -0,0 +1,3 @@ +# RescaleIntensity + +::: torchio.transforms.RescaleIntensity diff --git a/docs/transforms/preprocessing/Resize.md b/docs/transforms/preprocessing/Resize.md new file mode 100644 index 000000000..47ece61a0 --- /dev/null +++ b/docs/transforms/preprocessing/Resize.md @@ -0,0 +1,3 @@ +# Resize + +::: torchio.transforms.Resize diff --git a/docs/transforms/preprocessing/SequentialLabels.md b/docs/transforms/preprocessing/SequentialLabels.md new file mode 100644 index 000000000..dba07e7bc --- /dev/null +++ b/docs/transforms/preprocessing/SequentialLabels.md @@ -0,0 +1,3 @@ +# SequentialLabels + +::: torchio.transforms.SequentialLabels diff --git a/docs/transforms/preprocessing/To.md b/docs/transforms/preprocessing/To.md new file mode 100644 index 000000000..d349f6bb7 --- /dev/null +++ b/docs/transforms/preprocessing/To.md @@ -0,0 +1,3 @@ +# To + +::: torchio.transforms.To diff --git a/docs/transforms/preprocessing/ToCanonical.md b/docs/transforms/preprocessing/ToCanonical.md new file mode 100644 index 000000000..ede7229cd --- /dev/null +++ b/docs/transforms/preprocessing/ToCanonical.md @@ -0,0 +1,3 @@ +# ToCanonical + +::: torchio.transforms.ToCanonical diff --git a/docs/transforms/preprocessing/ToOrientation.md b/docs/transforms/preprocessing/ToOrientation.md new file mode 100644 index 000000000..d7c4d593a --- /dev/null +++ b/docs/transforms/preprocessing/ToOrientation.md @@ -0,0 +1,3 @@ +# ToOrientation + +::: torchio.transforms.ToOrientation diff --git a/docs/transforms/preprocessing/ToReferenceSpace.md b/docs/transforms/preprocessing/ToReferenceSpace.md new file mode 100644 index 000000000..f34b5981e --- /dev/null +++ b/docs/transforms/preprocessing/ToReferenceSpace.md @@ -0,0 +1,3 @@ +# ToReferenceSpace + +::: torchio.transforms.ToReferenceSpace diff --git a/docs/transforms/preprocessing/Transpose.md b/docs/transforms/preprocessing/Transpose.md new file mode 100644 index 000000000..d2a9e1a56 --- /dev/null +++ b/docs/transforms/preprocessing/Transpose.md @@ -0,0 +1,3 @@ +# Transpose + +::: torchio.transforms.Transpose diff --git a/docs/transforms/preprocessing/ZNormalization.md b/docs/transforms/preprocessing/ZNormalization.md new file mode 100644 index 000000000..9b14b8e3a --- /dev/null +++ b/docs/transforms/preprocessing/ZNormalization.md @@ -0,0 +1,3 @@ +# ZNormalization + +::: torchio.transforms.ZNormalization diff --git a/docs/transforms/preprocessing/index.md b/docs/transforms/preprocessing/index.md new file mode 100644 index 000000000..5a24e28df --- /dev/null +++ b/docs/transforms/preprocessing/index.md @@ -0,0 +1,44 @@ +# Preprocessing + +## Intensity + +| Transform | Description | +|-----------|-------------| +| [`RescaleIntensity`](RescaleIntensity.md) | Rescale intensity values to a certain range | +| [`ZNormalization`](ZNormalization.md) | Subtract mean and divide by standard deviation | +| [`HistogramStandardization`](HistogramStandardization.md) | Standardize histogram of foreground intensities | +| [`Mask`](Mask.md) | Mask an image using a label map | +| [`Clamp`](Clamp.md) | Clamp intensity values into a range | +| [`PCA`](PCA.md) | Reduce the number of channels using PCA | +| [`To`](To.md) | Change the dtype or device of image data | + +::: torchio.transforms.preprocessing.intensity.NormalizationTransform + options: + show_root_heading: true + +## Spatial + +| Transform | Description | +|-----------|-------------| +| [`CropOrPad`](CropOrPad.md) | Crop or pad an image to a target shape | +| [`Crop`](Crop.md) | Crop an image | +| [`Pad`](Pad.md) | Pad an image | +| [`Resize`](Resize.md) | Resize an image to a target shape | +| [`Resample`](Resample.md) | Resample an image to a different voxel spacing | +| [`ToCanonical`](ToCanonical.md) | Reorder data to canonical orientation | +| [`ToOrientation`](ToOrientation.md) | Reorder data to a given orientation | +| [`ToReferenceSpace`](ToReferenceSpace.md) | Resample to a reference image space | +| [`Transpose`](Transpose.md) | Transpose spatial dimensions | +| [`EnsureShapeMultiple`](EnsureShapeMultiple.md) | Pad to ensure shape is a multiple of a value | +| [`CopyAffine`](CopyAffine.md) | Copy the affine matrix from one image to another | + +## Label + +| Transform | Description | +|-----------|-------------| +| [`RemapLabels`](RemapLabels.md) | Remap integer labels in a segmentation | +| [`RemoveLabels`](RemoveLabels.md) | Remove labels from a segmentation | +| [`SequentialLabels`](SequentialLabels.md) | Map labels to sequential integers | +| [`OneHot`](OneHot.md) | Convert a label map to one-hot encoding | +| [`Contour`](Contour.md) | Create a binary image with contour of each label | +| [`KeepLargestComponent`](KeepLargestComponent.md) | Keep the largest connected component | diff --git a/docs/tutorials/annotations.md b/docs/tutorials/annotations.md deleted file mode 100644 index b7b1b79ef..000000000 --- a/docs/tutorials/annotations.md +++ /dev/null @@ -1,217 +0,0 @@ -# Working with annotations - -This tutorial introduces `Points` and `BoundingBoxes`, two data -structures for storing spatial annotations alongside medical images. - -## Prerequisites - -``` -uv add torchio -``` - -## Why annotations? - -Medical image analysis often involves more than just voxel intensities. -You might need to track: - -- **Anatomical landmarks** (fiducial points placed by a clinician) -- **Lesion detections** (bounding boxes from a detection model) -- **Seed points** for region growing algorithms - -TorchIO's `Points` and `BoundingBoxes` classes keep these annotations -together with the images they refer to, inside a single `Subject`. - -## Step 1: Create some points - -```python -import torch -import torchio as tio - -# Three landmarks in voxel coordinates (IJK) -landmarks = tio.Points( - torch.tensor([ - [64.0, 64.0, 32.0], - [100.0, 80.0, 50.0], - [30.0, 120.0, 45.0], - ]), -) -print(landmarks) # Points(num_points=3, axes='IJK') -print(landmarks.num_points) # 3 -print(landmarks.axes) # 'IJK' -``` - -Points are stored as an $(N, 3)$ tensor: - - -```python -print(landmarks.data.shape) # torch.Size([3, 3]) -``` - -## Step 2: Convert between axis conventions - -Points default to `IJK` (voxel indices). Convert to any other -convention: - - -```python -# With identity affine, IJK == RAS numerically -ras = landmarks.to_axes("RAS") -print(ras.axes) # 'RAS' - -# Or get world coordinates directly -world = landmarks.to_world() -``` - -When the points come from a real image, pass its affine: - - -```python -image = tio.ScalarImage("t1.nii.gz") -landmarks = tio.Points( - torch.tensor([[64.0, 64.0, 32.0]]), - affine=image.affine, -) -ras = landmarks.to_axes("RAS") # now in mm -``` - -Supported axis conventions: - -| Type | Examples | -|------|----------| -| Voxel | `IJK`, `KJI`, `JIK`, ... (any permutation) | -| Anatomical | `RAS`, `LPI`, `AIR`, ... (one from each pair {R,L}, {A,P}, {S,I}) | - -## Step 3: Create bounding boxes - -Bounding boxes are 6-element vectors. The format is defined by axes -and representation: - -| Predefined | Meaning | -|------------|---------| -| `IJKIJK` | Corners: $(i_1, j_1, k_1, i_2, j_2, k_2)$ | -| `IJKWHD` | Center + size: $(i_c, j_c, k_c, s_i, s_j, s_k)$ | - - -```python -boxes = tio.BoundingBoxes( - torch.tensor([ - [10, 20, 30, 50, 60, 70], - [80, 90, 40, 120, 130, 80], - ]), - format=tio.BoundingBoxFormat.IJKIJK, -) -print(boxes) # BoundingBoxes(num_boxes=2, ...) -print(boxes.num_boxes) # 2 -``` - -Convert between formats: - - -```python -whd = boxes.to_format(tio.BoundingBoxFormat.IJKWHD) -print(whd.data[0]) # center=(30, 40, 50), size=(40, 40, 40) -``` - -Custom formats work too: - - -```python -from torchio import BoundingBoxFormat - -ras_corners = BoundingBoxFormat("RAS", "corners") -ras_boxes = boxes.to_format(ras_corners) -``` - -## Step 4: Attach labels - -Each box can carry a class label: - - -```python -boxes = tio.BoundingBoxes( - torch.tensor([[10, 20, 30, 50, 60, 70]]), - format=tio.BoundingBoxFormat.IJKIJK, - labels=torch.tensor([1]), -) -print(boxes.labels) # tensor([1]) -``` - -## Step 5: Build a Subject with everything - -A `Subject` automatically sorts its contents: - - -```python -subject = tio.Subject( - t1=tio.ScalarImage(torch.randn(1, 64, 64, 64)), - seg=tio.LabelMap((torch.randn(1, 64, 64, 64) > 0).float()), - landmarks=landmarks, - tumors=boxes, - age=45, -) -``` - -Access each type: - - -```python -subject.t1 # Image -subject.landmarks # Points -subject.tumors # BoundingBoxes -subject.age # metadata (45) -``` - -Inspect what the Subject contains: - - -```python -print(subject) -# Subject(images: ('t1', 'seg'); points: ('landmarks',); bboxes: ('tumors',)) -``` - -List all entries of a given type: - - -```python -subject.images() # {'t1': ..., 'seg': ...} -subject.points() # {'landmarks': ...} -subject.bounding_boxes() # {'tumors': ...} -``` - -## Summary - -```mermaid -classDiagram - class Subject { - images - points - bounding_boxes - metadata - } - class Image { - data: (C, I, J, K) - affine - } - class Points { - data: (N, 3) - axes: str - to_axes(target) - } - class BoundingBoxes { - data: (N, 6) - format: BoundingBoxFormat - labels - to_format(target) - } - Subject --> Image - Subject --> Points - Subject --> BoundingBoxes -``` - -You now know how to: - -- Store landmark coordinates with `Points` -- Define regions of interest with `BoundingBoxes` -- Convert between axis conventions (`IJK`, `RAS`, `LPI`, etc.) -- Attach labels to bounding boxes -- Group everything in a `Subject` diff --git a/docs/tutorials/augmentation.md b/docs/tutorials/augmentation.md deleted file mode 100644 index 5e259a865..000000000 --- a/docs/tutorials/augmentation.md +++ /dev/null @@ -1,220 +0,0 @@ -# Augmentation pipelines - -This tutorial shows how to build data augmentation pipelines with -TorchIO transforms. You will learn how to apply spatial, intensity, -and artifact transforms, compose them into pipelines, and use -`Choice` for discrete parameter sampling. - -## A simple pipeline - - -```python -import torchio as tio - -augmentation = tio.Compose([ - tio.Flip(axes=(0, 1, 2), flip_probability=0.5), - tio.Affine(degrees=10, translation=5), - tio.Noise(std=(0.01, 0.1)), - tio.Gamma(log_gamma=(-0.3, 0.3)), -]) - -subject = tio.Subject( - t1=tio.ScalarImage("t1.nii.gz"), - seg=tio.LabelMap("seg.nii.gz"), -) -augmented = augmentation(subject) -``` - -Every call to `augmentation(subject)` produces a different result. -Spatial transforms (Flip, Affine) are applied consistently to all -images in the subject: the T1 and segmentation are transformed -together. - -## Deterministic vs random - -The same class handles both cases. A scalar gives a fixed value; -a tuple gives a uniform range: - - -```python -# Always rotate 90° -tio.Affine(degrees=90) - -# Rotate uniformly between -15° and 15° -tio.Affine(degrees=15) - -# Rotate uniformly between 5° and 20° -tio.Affine(degrees=(5, 20)) -``` - -For discrete choices, use `Choice`: - - -```python -# Rotate by exactly -90, 0, 90, or 180 degrees -tio.Affine(degrees=tio.Choice([-90, 0, 90, 180])) -``` - -You can mix `Choice`, ranges, and fixed values per axis: - - -```python -# Fixed along I, random along J, discrete along K -tio.Affine(degrees=(0, (-10, 10), tio.Choice([-90, 0, 90]))) -``` - -## Probability control - -Every transform has a `p` parameter (probability of being applied): - - -```python -# 50% chance of adding noise -tio.Noise(std=0.1, p=0.5) -``` - -## Composition strategies - -### Compose: apply all in sequence - - -```python -pipeline = tio.Compose([ - tio.Affine(degrees=10), - tio.Noise(std=0.05), - tio.Gamma(log_gamma=0.3), -]) -``` - -### OneOf: pick one at random - - -```python -artifact = tio.OneOf({ - tio.Ghosting(intensity=0.5): 0.4, - tio.Spike(intensity=1.0): 0.3, - tio.Motion(degrees=5): 0.3, -}) -``` - -### SomeOf: pick N at random - - -```python -augment = tio.SomeOf( - [ - tio.Flip(axes=(0, 1, 2)), - tio.Blur(std=1.0), - tio.Noise(std=0.05), - tio.Gamma(log_gamma=0.3), - ], - num_transforms=(1, 3), # apply 1 to 3 of the 4 -) -``` - -### Operator sugar - -You can use `+` for Compose and `|` for OneOf: - - -```python -pipeline = tio.Flip(axes=(0,)) + tio.Noise(std=0.05) + tio.Gamma(log_gamma=(-0.3, 0.3)) -artifact = ( - tio.Ghosting(intensity=(0.5, 1)) | tio.Spike(intensity=(1, 3)) | tio.Motion() -) -``` - -## Preprocessing + augmentation - -A common pattern separates preprocessing (applied once) from -augmentation (applied each epoch): - - -```python -preprocessing = tio.Compose([ - tio.Resample(target=1.0), # isotropic 1mm - tio.CropOrPad(target_shape=128), # fixed shape - tio.Normalize(out_min=-1, out_max=1), # rescale to [-1, 1] -]) - -augmentation = tio.Compose([ - tio.Flip(axes=(0, 1, 2), flip_probability=0.5), - tio.Affine(degrees=15, translation=5, p=0.8), - tio.OneOf({ - tio.Noise(std=(0.01, 0.1)): 0.5, - tio.Blur(std=(0.5, 2.0)): 0.3, - tio.BiasField(coefficients=0.5): 0.2, - }), - tio.Gamma(log_gamma=(-0.3, 0.3), p=0.5), -]) - -full_pipeline = preprocessing + augmentation -``` - -## MRI artifact simulation - -TorchIO includes several MRI-specific artifact transforms: - - -```python -artifacts = tio.Compose([ - tio.Motion(degrees=5, num_transforms=2, p=0.3), - tio.Ghosting(num_ghosts=5, intensity=0.5, p=0.3), - tio.Spike(num_spikes=1, intensity=1.5, p=0.2), - tio.Anisotropy(downsampling=3, p=0.3), - tio.BiasField(coefficients=0.5, p=0.3), -]) -``` - -These are useful for training models that are robust to common MRI -acquisition artifacts. - -## Label-aware transforms - -Some transforms only affect label maps: - - -```python -label_pipeline = tio.Compose([ - tio.SequentialLabels(), # renumber to 0, 1, 2, ... - tio.RemoveLabels([4, 5]), # drop unwanted labels - tio.KeepLargestComponent(labels=[1]), # clean up label 1 -]) -``` - -## SynthSeg-style synthesis - -Generate synthetic images from label maps: - - -```python -synth = tio.Compose([ - tio.LabelsToImage(label_key="seg"), - tio.Blur(std=(0.5, 1.5)), - tio.BiasField(coefficients=0.5), - tio.Gamma(log_gamma=(-0.3, 0.3)), - tio.Noise(std=(0.01, 0.05)), -]) -``` - -## Summary - -| Goal | Transform(s) | -|------|-------------| -| Random flipping | `Flip` | -| Rotation / scaling / shearing | `Affine` | -| Elastic deformation | `ElasticDeformation` | -| Change resolution | `Resample`, `Resize` | -| Fixed shape | `CropOrPad` | -| Intensity normalization | `Normalize`, `Standardize`, `HistogramStandardization` | -| Gaussian noise | `Noise` | -| Gaussian blur | `Blur` | -| Gamma correction | `Gamma` | -| Bias field | `BiasField` | -| MRI motion | `Motion` | -| MRI ghosting | `Ghosting` | -| K-space spikes | `Spike` | -| Simulate low-res axis | `Anisotropy` | -| Label cleanup | `RemoveLabels`, `KeepLargestComponent`, `SequentialLabels` | -| Synthetic images | `LabelsToImage` | -| Self-supervised | `Swap` | diff --git a/docs/tutorials/first-pipeline.md b/docs/tutorials/first-pipeline.md deleted file mode 100644 index 09204a8bd..000000000 --- a/docs/tutorials/first-pipeline.md +++ /dev/null @@ -1,122 +0,0 @@ -# Your first pipeline - -This tutorial walks you through loading medical images, grouping them -into a Subject, and saving results. By the end you will understand -the core data structures. - -## Prerequisites - -``` -uv add torchio -``` - -You need at least one NIfTI file. If you do not have one, create a -synthetic image: - -```python -import torch -import torchio as tio - -image = tio.ScalarImage(torch.randn(1, 64, 64, 64)) -image.save("/tmp/synthetic.nii.gz") -``` - -## Step 1: Load an image - -```python -import torchio as tio - -t1 = tio.ScalarImage("/tmp/synthetic.nii.gz") -``` - -At this point, **no data has been read from disk**. TorchIO uses lazy -loading. The file is only read when you access `.data`, `.spacing`, -or apply a transform. - - -```python -print(t1.shape) # (1, 64, 64, 64), reads only the header -print(t1.is_loaded) # False -print(t1.spacing) # (1.0, 1.0, 1.0) -``` - -## Step 2: Access the data - - -```python -tensor = t1.data # triggers the load -print(t1.is_loaded) # True -print(tensor.shape) # torch.Size([1, 64, 64, 64]) -print(tensor.dtype) # torch.float32 -``` - -The data tensor has shape `(C, I, J, K)` where `C` is the number of -channels and `I, J, K` are the spatial dimensions. - -## Step 3: Create a Subject - -A `Subject` groups related images, annotations, and metadata: - - -```python -import torch - -seg_tensor = (torch.randn(1, 64, 64, 64) > 0).float() -seg = tio.LabelMap(seg_tensor) - -subject = tio.Subject( - t1=t1, - seg=seg, - age=30, -) -``` - -Access images and metadata by name: - - -```python -subject.t1 # the ScalarImage -subject.seg # the LabelMap -subject.age # 30 -subject.spatial_shape # (64, 64, 64), checked across all images -``` - -!!! tip "Annotations" - - You can also add `Points` and `BoundingBoxes` to a Subject. - See the [annotations tutorial](annotations.md) for details. - -## Step 4: Slice a region - - -```python -patch = subject.t1[:, 10:30, 10:30, 10:30] -print(patch.shape) # (1, 20, 20, 20) -print(patch.origin) # shifted by 10 voxels in each direction -``` - -## Step 5: Save the result - - -```python -patch.save("/tmp/patch.nii.gz") -``` - -## Summary - -```mermaid -flowchart LR - A[NIfTI file] -->|"ScalarImage(path)"| B[Lazy image] - B -->|".data"| C[Loaded tensor] - B -->|"[slicing]"| D[Cropped image] - D -->|".save()"| E[Output file] - B --> F[Subject] - G[LabelMap] --> F -``` - -You now know how to: - -- Load images lazily with `ScalarImage` and `LabelMap` -- Group them into a `Subject` -- Slice regions without loading the full volume -- Save results to disk diff --git a/docs/tutorials/large-volumes.md b/docs/tutorials/large-volumes.md deleted file mode 100644 index b4a8781d0..000000000 --- a/docs/tutorials/large-volumes.md +++ /dev/null @@ -1,120 +0,0 @@ -# Working with large volumes - -Medical images can be very large. A single high-resolution MRI might -occupy several gigabytes. This tutorial shows how to work with such -volumes efficiently using TorchIO's lazy loading and backend system. - -## The problem - -Loading a 724 x 868 x 724 float32 volume allocates ~1.8 GB of RAM (and -can take many seconds to read and decompress). If you only need a small -region, that is wasteful. - -## Lazy slicing - -TorchIO images are lazy by default. Slicing a lazy image reads **only -the requested region** through the backend, without loading the full -volume into memory: - - -```python -import torchio as tio - -image = tio.ScalarImage("huge_volume.nii.gz") - -# This is fast: reads only a 10x10x10 patch -patch = image[:, 100:110, 100:110, 100:110] -print(patch.data.mean()) - -# The original image was never fully loaded -print(image.is_loaded) # False -``` - -For uncompressed `.nii` (memory-mapped) and `.nii.zarr` (chunked), this can be -orders of magnitude faster than a full load. For `.nii.gz` the gain is more -modest, because gzip must be decompressed from the start (see below). - -## File format comparison - -Not all formats are equally efficient for partial reads: - -| Format | Extension | Partial I/O | Notes | -|--------|-----------|------------|-------| -| Uncompressed NIfTI | `.nii` | Memory-mapped | Best for local random access | -| Compressed NIfTI | `.nii.gz` | Buffered by nibabel | Modest speedup for small regions; gzip has no true random access | -| NIfTI-Zarr | `.nii.zarr` | Chunked reads | Best for very large volumes and remote storage | - -## NIfTI-Zarr - -NIfTI-Zarr stores data in independently compressed chunks. Reading one -chunk does not require decompressing any other. This is ideal for: - -- Volumes too large to fit in memory -- Remote storage (S3, GCS) where you want to fetch only what you need -- Collaborative workflows where different users need different regions - -### Converting to NIfTI-Zarr - - -```python -image = tio.ScalarImage("volume.nii.gz") -image.save("volume.nii.zarr") -``` - -### Reading from NIfTI-Zarr - - -```python -image = tio.ScalarImage("volume.nii.zarr") -print(image.shape) # reads only metadata - -# Read a small region: only the overlapping chunks are decompressed -patch = image[:, 50:60, 50:60, 50:60] -``` - -!!! note - - NIfTI-Zarr support requires the `zarr` extra: - - ``` - uv add torchio --extra zarr - ``` - -## The backend system - -TorchIO uses a pluggable backend system to support different storage -formats. The choice is not hard-coded in `Image`: a resolver consults a -registry of backends, so new formats can be added without modifying `Image` -(see [Lazy loading and backends](../concepts/lazy-loading.md)). - -```mermaid -flowchart TD - I["Image(path)"] --> check{Resolver} - check -->|".nii"| NB[NibabelBackend
memory-mapped] - check -->|".nii.gz"| NB - check -->|".nii.zarr"| ZB[ZarrBackend
chunked reads] - check -->|other| SITK[SimpleITK reader
full load] - - NB -->|"image.dataobj[slices]"| partial[Partial read] - ZB -->|"image.dataobj[slices]"| partial - NB -->|"image.data"| full[Full tensor] - ZB -->|"image.data"| full - SITK -->|"image.data"| full -``` - -- **`image.data`**: materializes the full tensor (triggers load if needed) -- **`image.dataobj`**: returns the lazy backend for advanced slicing; - `image.dataobj[slices]` returns a 4D `(C, I, J, K)` `torch.Tensor` -- **`image[slices]`**: uses the backend automatically, returns a new image - -## Tips - -1. **Use uncompressed `.nii` for local training**: memory-mapping gives - true random access with zero decompression cost. -2. **Use `.nii.zarr` for large shared volumes**: chunked storage means - each worker reads only what it needs. -3. **Avoid `.nii.gz` for random access**: gzip is a stream format. - Nibabel handles it well for small slices, but for repeated random - access, convert to `.nii` or `.nii.zarr`. -4. **Slice before loading**: `image[:, 100:200].data` is much cheaper - than `image.data[:, 100:200]`. diff --git a/justfile b/justfile new file mode 100644 index 000000000..5ec14d302 --- /dev/null +++ b/justfile @@ -0,0 +1,162 @@ +default: + @just --list + +clean: + rm -rf .mypy_cache + rm -rf .pytest_cache + rm -rf .tox + rm -rf .venv + rm -rf dist + rm -rf **/__pycache__ + rm -rf src/*.egg-info + rm -f .coverage + rm -f coverage.* + +@install_uv: + if ! command -v uv >/dev/null 2>&1; then \ + echo "uv is not installed. Installing..."; \ + curl -LsSf https://astral.sh/uv/install.sh | sh; \ + fi + +setup: install_uv + uv sync --all-extras --all-groups + uv run prek install + +bump part="patch": + uv run bump-my-version bump {{part}} --verbose + +bump-dry part="patch": + uv run bump-my-version bump {{part}} --dry-run --verbose --allow-dirty + +bump-python: + #!/usr/bin/env -S uv run --script + # /// script + # dependencies = [ + # "packaging", + # ] + # /// + from pathlib import Path + from packaging.version import Version + python_version_path = Path(".python-version") + old_version_string = python_version_path.read_text().strip() + old_version = Version(old_version_string) + new_version_string = f"{old_version.major}.{old_version.minor + 1}" + python_version_path.write_text(new_version_string + "\n") + + tests_workflow_path = Path(".github/workflows/tests.yml") + tests_workflow = tests_workflow_path.read_text() + tests_workflow = tests_workflow.replace( + f'"{old_version_string}"]', + f'"{old_version_string}", "{new_version_string}"]', + ) + tests_workflow = tests_workflow.replace( + f"matrix.python == '{old_version_string}'", + f"matrix.python == '{new_version_string}'", + ) + tests_workflow_path.write_text(tests_workflow) + + scrutinizer_config_path = Path(".scrutinizer.yml") + scrutinizer_config = scrutinizer_config_path.read_text() + scrutinizer_config = scrutinizer_config.replace( + old_version_string, + new_version_string, + ) + scrutinizer_config_path.write_text(scrutinizer_config) + + pyproject_path = Path("pyproject.toml") + pyproject_text = pyproject_path.read_text() + old_str = f' "Programming Language :: Python :: {old_version_string}",\n' + new_str = f' "Programming Language :: Python :: {new_version_string}",\n' + pyproject_text = pyproject_text.replace( + old_str, + old_str + new_str, + ) + pyproject_path.write_text(pyproject_text) + +deprecate-python-in-files: + #!/usr/bin/env -S uv run --script + # /// script + # dependencies = [ + # "packaging", + # ] + # /// + from pathlib import Path + from packaging.version import Version + from tomllib import load + + pyproject_path = Path("pyproject.toml") + with open(pyproject_path, "rb") as f: + pyproject = load(f) + + classifiers = pyproject["project"]["classifiers"] + for classifier in classifiers: + if classifier.startswith("Programming Language :: Python :: 3."): + old_version = Version(classifier.split("::")[-1].strip()) + break + pyproject_text = pyproject_path.read_text() + to_replace = f' "Programming Language :: Python :: {old_version}",\n' + pyproject_text = pyproject_text.replace(to_replace, "") + new_version = Version(f"{old_version.major}.{old_version.minor + 1}") + pyproject_text = pyproject_text.replace( + f'requires-python = ">={old_version}"', + f'requires-python = ">={new_version}"', + ) + pyproject_path.write_text(pyproject_text) + + pre_commit_path = Path(".pre-commit-config.yaml") + pre_commit_text = pre_commit_path.read_text() + old_version_pyupgrade = str(old_version).replace(".", "") + new_version_pyupgrade = str(new_version).replace(".", "") + pre_commit_text = pre_commit_text.replace( + f"--py{old_version_pyupgrade}-plus", + f"--py{new_version_pyupgrade}-plus", + ) + pre_commit_path.write_text(pre_commit_text) + + tests_workflow_path = Path(".github/workflows/tests.yml") + tests_workflow = tests_workflow_path.read_text() + tests_workflow = tests_workflow.replace( + f'["{old_version}", ', + f'[', + ) + tests_workflow_path.write_text(tests_workflow) + +quality_cmd := "uv run --group quality" + +deprecate-python: deprecate-python-in-files + uv run pre-commit run --all-files pyupgrade + {{quality_cmd}} -- ruff check --fix src docs tests + +push: + git push && git push --tags + +types: + {{quality_cmd}} -- tox -e types + +lint: + {{quality_cmd}} -- ruff check + +format: + {{quality_cmd}} -- ruff format --diff + +test: + uv run --group test -- tox -e pytest + +add-remote remote: + git remote add {{remote}} git@github.com:{{remote}}/torchio.git + +docs_cmd := "uv run --group doc" + +# Generate/update plot images from ```python plot blocks in docs +generate-plots: + uv run --group doc -- python docs/plot_directive.py + +# Generate examples gallery pages from Python scripts in docs/examples +generate-gallery: + uv run --group doc -- python docs/gallery.py + +build-docs: + uv run --group doc -- zensical build + +serve-docs: + uv run --group doc -- zensical serve diff --git a/mise.toml b/mise.toml deleted file mode 100644 index 52fa4e5bf..000000000 --- a/mise.toml +++ /dev/null @@ -1,111 +0,0 @@ -[tools] -prek = "latest" -uv = "latest" - -[tasks.setup] -description = "Install all dependencies and prek hooks" -run = """ -uv sync --all-groups -mise trust -mise install -prek install --install-hooks -""" - -[tasks.test] -description = "Run tests with pytest" -run = "uv run tox -e test" - -[tasks.lint] -description = "Lint with ruff" -run = "uv run tox -e lint" - -[tasks.format] -description = "Format with ruff" -run = "uv run --group quality -- ruff format src/ tests/" - -[tasks.format-check] -description = "Check formatting with ruff" -run = "uv run tox -e format" - -[tasks.types] -description = "Type check with ty" -run = "uv run tox -e types" - -[tasks.quality] -description = "Run all quality checks (lint, format, types)" -depends = ["lint", "format-check", "types"] - -[tasks.prek] -description = "Run prek hooks on all files" -run = "uv run --group quality -- prek run --all-files" - -[tasks.tox] -description = "Run all tox environments" -run = "uv run --group quality -- tox" - -[tasks.bump] -description = "Bump version (usage: mise run bump -- pre_n|pre_l|patch|minor|major)" -run = "uv run -- bump-my-version bump --verbose" - -[tasks.bump-dry] -description = "Dry-run version bump" -run = "uv run -- bump-my-version bump --dry-run --verbose --allow-dirty" - -[tasks.deprecate-python] -description = "Remove oldest Python, update configs, run pyupgrade and ruff" -depends = ["deprecate-python-files"] -run = """ -uv run --group quality -- prek run --all-files pyupgrade -uv run --group quality -- ruff check --fix src docs tests -""" - -[tasks.push] -description = "Push commits and tags" -run = "git push && git push --tags" - -[tasks.add-remote] -description = "Add a fork as a Git remote (usage: mise run add-remote username)" -run = 'git remote add "$1" "git@github.com:$1/torchio.git"' - -[tasks."docs:serve"] -description = "Serve documentation locally" -run = "uv run --group docs -- zensical serve --open" - -[tasks."docs:build"] -description = "Build documentation" -run = "uv run --group docs -- zensical build" - -[tasks."docs:deploy"] -description = "Deploy a versioned docs build with mike (usage: mise run docs:deploy -- [mike-args] [alias...])" -run = 'uv run --group docs-deploy -- mike deploy --update-aliases "$@"' - -[tasks."docs:set-default"] -description = "Set the docs version users are redirected to (usage: mise run docs:set-default -- [mike-args] )" -run = 'uv run --group docs-deploy -- mike set-default "$@"' - -[tasks."docs:test"] -description = "Test documentation code snippets" -run = "uv run tox -e docs-test" - -[tasks."docs:plots"] -description = "Generate plot images from docs" -run = "uv run --group docs -- python docs/plot_directive.py" - -[tasks."docs:gallery"] -description = "Generate examples gallery pages" -run = "uv run --group docs -- python docs/gallery.py" - -[tasks.clean] -description = "Remove build artifacts and caches" -run = """ -rm -rf .mypy_cache -rm -rf .pytest_cache -rm -rf .ruff_cache -rm -rf .tox -rm -rf dist -rm -rf site -rm -rf **/__pycache__ -rm -rf src/*.egg-info -rm -f .coverage -rm -f coverage.* -""" diff --git a/print_system.py b/print_system.py index 5637a5f07..089f8432d 100644 --- a/print_system.py +++ b/print_system.py @@ -8,11 +8,11 @@ import torchio as tio -sitk_version = re.findall("SimpleITK Version: (.*?)\n", str(sitk.Version()))[0] +sitk_version = re.findall('SimpleITK Version: (.*?)\n', str(sitk.Version()))[0] -print("Platform: ", platform.platform()) -print("TorchIO: ", tio.__version__) -print("PyTorch: ", torch.__version__) -print("SimpleITK: ", sitk_version) -print("NumPy: ", numpy.__version__) -print("Python: ", sys.version) +print('Platform: ', platform.platform()) +print('TorchIO: ', tio.__version__) +print('PyTorch: ', torch.__version__) +print('SimpleITK: ', sitk_version) +print('NumPy: ', numpy.__version__) +print('Python: ', sys.version) diff --git a/pyproject.toml b/pyproject.toml index 33e75d266..6b4e01be2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,21 +1,17 @@ +[build-system] +requires = ["uv_build"] +build-backend = "uv_build" + [project] name = "torchio" -version = "2.0.0a2" -description = "Medical image preprocessing, augmentation, and patch-based training." -requires-python = ">=3.10" +version = "0.23.1" +description = "Tools for medical image processing with PyTorch" license = "Apache-2.0" readme = "README.md" -authors = [{ name = "TorchIO contributors", email = "fepegar@gmail.com" }] -maintainers = [{ name = "Fernando Pérez-García", email = "fepegar@gmail.com" }] -keywords = [ - "medical image computing", - "pytorch", - "preprocessing", - "augmentation", - "mri", -] +authors = [{ name = "TorchIO contributors" }] +maintainers = [{ name = "Fernando Perez-Garcia", email = "fepegar@gmail.com" }] classifiers = [ - "Development Status :: 5 - Production/Stable", + "Development Status :: 4 - Beta", "Environment :: Console", "Intended Audience :: Science/Research", "Natural Language :: English", @@ -34,23 +30,43 @@ classifiers = [ "Topic :: Scientific/Engineering :: Medical Science Apps.", "Typing :: Typed", ] +keywords = ["medical", "image processing", "pytorch", "augmentation", "mri"] +requires-python = ">=3.10" dependencies = [ - "einops", - "fsspec[http]", - "humanize", - "jaxtyping", - "loguru", - "nibabel", - "numpy", - "platformdirs", - "rich", - "simpleitk", - "torch", - "torch-interpol>=0.3.0", - "typing_extensions>=4.15.0", - "tyro>=1.0.15", + "deprecated>=1.2", + "einops>=0.3", + "humanize>=0.1", + "jaxtyping>=0.2", + "nibabel>=3", + "numpy>=1.20", + "packaging>=20", + "rich>=10", + "scipy>=1.7", + "simpleitk >=1.3, !=2.0.*, !=2.1.1.1", + "torch>=1.9", + "tqdm>=4.40", + "typer>=0.1", +] + +[project.optional-dependencies] +csv = ["pandas>=1"] +plot = ["colorcet", "matplotlib>=3.4"] +video = ["ffmpeg-python>=0.2.0"] +sklearn = ["scikit-learn>=1.6.1"] +monai = ["monai>=1.4"] +all = [ + "torchio[csv]", + "torchio[monai]", + "torchio[plot]", + "torchio[sklearn]", + "torchio[video]", ] +[project.scripts] +tiohd = "torchio.cli.print_info:app" +tiotr = "torchio.cli.apply_transform:app" +torchio-transform = "torchio.cli.apply_transform:app" + [project.urls] Homepage = "https://torchio.org" Source = "https://github.com/TorchIO-project/torchio" @@ -58,28 +74,60 @@ Source = "https://github.com/TorchIO-project/torchio" Documentation = "https://docs.torchio.org" "Release notes" = "https://github.com/TorchIO-project/torchio/releases" -[project.scripts] -torchio = "torchio.cli:main" +[dependency-groups] +dev = [ + { include-group = "doc" }, + { include-group = "maintain" }, + { include-group = "quality" }, + { include-group = "test" }, + "ipykernel", + "ipywidgets", + "prek", +] +doc = [ + "colorcet>=3.1.0", + "distinctipy", + "einops", + "matplotlib", + "mkdocstrings-python>=2.0.2", + "zensical>=0.0.23", +] +maintain = ["bump-my-version"] +quality = ["ruff"] +types = [ + "matplotlib", + "ty", + "pandas-stubs", + "pillow", + "pip", + "types-deprecated", + "types-tqdm", +] +test = [ + "coverage", + "matplotlib", + "monai>=1.4", + "parameterized", + "pillow", + "pytest>=9", + "pytest-sugar", + "tox-uv", +] -[project.optional-dependencies] -cornucopia = ["cornucopia"] -plot = ["matplotlib", "colorcet"] -niivue = ["ipyniivue"] -video = ["ffmpeg-python"] -monai = ["monai"] -s3 = ["s3fs"] -azure = ["adlfs"] -gcs = ["gcsfs"] -# nifti-zarr (1.0.0rc6) calls isinstance(x, zarr.storage.StoreLike), but in -# zarr >=3.2 StoreLike is a PEP 695 type alias (not isinstance-able), which -# raises TypeError. zarr 3.2 also requires Python >=3.12, so this only bites on -# 3.12+. Cap zarr until nifti-zarr supports zarr >=3.2. See copilot_log.md. -zarr = ["nifti-zarr", "zarr<3.2"] +[tool.bumpversion] +current_version = "0.23.1" +commit = true +tag = true +[[tool.bumpversion.files]] +filename = "src/torchio/__init__.py" +search = "__version__ = '{current_version}'" +replace = "__version__ = '{new_version}'" -[build-system] -requires = ["uv_build>=0.11.26,<0.12.0"] -build-backend = "uv_build" +[[tool.bumpversion.files]] +filename = "pyproject.toml" +search = 'version = "{current_version}"' +replace = 'version = "{new_version}"' [tool.pytest.ini_options] addopts = ["-ra", "--strict-config", "--strict-markers"] @@ -89,123 +137,26 @@ filterwarnings = [ "ignore:Casting complex values to real discards the imaginary part", # Raised by SimpleITK on CI "ignore:invalid escape sequence", - # Tests deliberately construct no-op transforms (identity tests, - # deterministic building blocks); the user-facing no-op warning is - # asserted explicitly in test_identity_warning.py. - "ignore:.* is a no-op with the given parameters:UserWarning", ] log_level = "INFO" markers = [ "slow: marks tests as slow (deselect with '-m \"not slow\"')", "serial", ] -minversion = "8" +minversion = "9" testpaths = ["tests"] xfail_strict = true -[tool.coverage.run] -source = ["torchio"] - -[tool.coverage.report] -fail_under = 100 -show_missing = true -skip_empty = true - -[tool.ruff] -src = ["src"] +[tool.ruff.format] +quote-style = 'single' [tool.ruff.lint] -select = ["A", "B", "C4", "E", "F", "I", "N", "PT", "RUF", "SIM", "UP", "W"] -ignore = [ - "N813", # https://docs.astral.sh/ruff/rules/camelcase-imported-as-lowercase/ -] +preview = true +select = ["B", "E", "F", "I", "W"] +ignore = ["E203", "E501", "N813"] [tool.ruff.lint.isort] force-single-line = true -[tool.ruff.lint.per-file-ignores] -"tests/**" = [ - "B018", # https://docs.astral.sh/ruff/rules/useless-expression/ -] -"src/torchio/data/bboxes.py" = [ - "A002", # `format` parameter shadows builtin: canonical name for this concept -] - -[tool.bumpversion] -current_version = "2.0.0a2" -commit = true -tag = true -parse = '''(?x) - (?P0|[1-9]\d*)\. - (?P0|[1-9]\d*)\. - (?P0|[1-9]\d*) - (?: - (?Pa|b|rc) - (?P0|[1-9]\d*) - )? -''' -serialize = [ - "{major}.{minor}.{patch}{pre_l}{pre_n}", - "{major}.{minor}.{patch}", -] - -[tool.bumpversion.parts.pre_l] -values = ["a", "b", "rc", "final"] -optional_value = "final" -first_value = "final" - -[tool.bumpversion.parts.pre_n] -first_value = "1" - -[[tool.bumpversion.files]] -filename = "pyproject.toml" -search = 'version = "{current_version}"' -replace = 'version = "{new_version}"' - -[dependency-groups] -dev = [ - "bump-my-version", - "ipykernel>=7.3.0", - "ipywidgets>=8.1.8", - "tox-uv", -] -test = [ - "colorcet", - "cornucopia", - "ffmpeg-python", - "matplotlib", - "monai", - "nifti-zarr", - "zarr<3.2", # see [project.optional-dependencies].zarr for why zarr is capped - # TODO: drop the <9 cap once pytest-codeblocks supports pytest 9. - # The latest release (0.17.0) registers `pytest_collect_file(path, ...)` - # using the `py.path.local` `path` argument removed in pytest 9, so the - # plugin fails to register and breaks the whole test suite, not just the - # docs snippets. See https://github.com/nschloe/pytest-codeblocks. - "pytest<10", - "pytest-codeblocks>=0.18.0", - "pytest-cov", - "pytest-sugar", -] -docs = [ - "mkdocstrings-python", - "nbconvert", - "zensical", -] -docs-deploy = [ - { include-group = "docs" }, - "mike", -] -quality = ["ruff"] -types = [ - "ipykernel", - "matplotlib", - "ty", - "zarr<3.2", # see [project.optional-dependencies].zarr for why zarr is capped -] - -# Zensical integrates with the Material team's fork of mike for versioned docs. -# The fork is not published on PyPI, so it is pinned directly from GitHub. -# See https://zensical.org/docs/setup/versioning/ -[tool.uv.sources] -mike = { git = "https://github.com/squidfunk/mike.git" } +[tool.coverage.run] +omit = ["**/_remote_module_non_scriptable.py"] diff --git a/src/torchio/__init__.py b/src/torchio/__init__.py index ecc5e78c8..21089d92b 100644 --- a/src/torchio/__init__.py +++ b/src/torchio/__init__.py @@ -1,193 +1,48 @@ -"""TorchIO: Medical image preprocessing, augmentation, and patch-based training.""" +"""Top-level package for torchio.""" + +__author__ = """Fernando Perez-Garcia""" +__email__ = 'fepegar@gmail.com' +__version__ = '0.23.1' -from importlib.metadata import version from . import datasets -from .data.affine import AffineMatrix -from .data.aggregator import PatchAggregator -from .data.batch import ImagesBatch -from .data.batch import StudiesBatch -from .data.batch import SubjectsBatch -from .data.bboxes import BoundingBoxes -from .data.bboxes import BoundingBoxFormat -from .data.bboxes import Representation -from .data.image import Image -from .data.image import LabelMap -from .data.image import ScalarImage -from .data.patch import PatchLocation -from .data.points import Points -from .data.queue import Queue -from .data.sampler import GridSampler -from .data.sampler import LabelSampler -from .data.sampler import PatchSampler -from .data.sampler import UniformSampler -from .data.sampler import WeightedSampler -from .data.subject import Study -from .data.subject import Subject -from .io import read_matrix -from .io import write_matrix -from .loader import ImagesLoader -from .loader import StudiesLoader -from .loader import SubjectsLoader -from .loader import collate_images -from .loader import collate_studies -from .loader import collate_subjects -from .logging import enable_logging -from .transforms import PCA -from .transforms import Affine -from .transforms import Anisotropy -from .transforms import AppliedTransform -from .transforms import BiasField -from .transforms import Blur -from .transforms import Choice -from .transforms import Clamp -from .transforms import Compose -from .transforms import Contour -from .transforms import CopyAffine -from .transforms import CornucopiaAdapter -from .transforms import Crop -from .transforms import CropOrPad -from .transforms import ElasticDeformation -from .transforms import EnsureShapeMultiple -from .transforms import Flip -from .transforms import Gamma -from .transforms import Ghosting -from .transforms import HistogramStandardization -from .transforms import IntensityTransform -from .transforms import KeepLargestComponent -from .transforms import LabelsToImage -from .transforms import Lambda -from .transforms import Mask -from .transforms import MonaiAdapter -from .transforms import Motion -from .transforms import Noise -from .transforms import Normalize -from .transforms import OneHot -from .transforms import OneOf -from .transforms import Pad -from .transforms import RemapLabels -from .transforms import RemoveLabels -from .transforms import Reorient -from .transforms import Resample -from .transforms import RescaleIntensity -from .transforms import Resize -from .transforms import SequentialLabels -from .transforms import SomeOf -from .transforms import Spatial -from .transforms import SpatialTransform -from .transforms import Spike -from .transforms import Standardize -from .transforms import Swap -from .transforms import To -from .transforms import ToReferenceSpace -from .transforms import Transform -from .transforms import Transpose -from .transforms import ZNormalization -from .transforms.inverse import apply_inverse_transform -from .types import TypeAffineMatrix -from .types import TypeDirection -from .types import TypeImageData -from .types import TypeOrientationCodes -from .types import TypeOrigin -from .types import TypePath -from .types import TypeSpacing -from .types import TypeSpatialShape -from .types import TypeTensorShape -from .types import TypeWorldPoints +from . import reference +from . import utils +from .constants import * # noqa: F401, F403 +from .data import GridAggregator +from .data import GridSampler +from .data import Image +from .data import LabelMap +from .data import LabelSampler +from .data import Queue +from .data import ScalarImage +from .data import Subject +from .data import SubjectsDataset +from .data import SubjectsLoader +from .data import UniformSampler +from .data import WeightedSampler +from .data import inference +from .data import io +from .data import sampler +from .transforms import * # noqa: F401, F403 __all__ = [ - "PCA", - "Affine", - "AffineMatrix", - "Anisotropy", - "AppliedTransform", - "BiasField", - "Blur", - "BoundingBoxFormat", - "BoundingBoxes", - "Choice", - "Clamp", - "Compose", - "Contour", - "CopyAffine", - "CornucopiaAdapter", - "Crop", - "CropOrPad", - "ElasticDeformation", - "EnsureShapeMultiple", - "Flip", - "Gamma", - "Ghosting", - "GridSampler", - "HistogramStandardization", - "Image", - "ImagesBatch", - "ImagesLoader", - "IntensityTransform", - "KeepLargestComponent", - "LabelMap", - "LabelSampler", - "LabelsToImage", - "Lambda", - "Mask", - "MonaiAdapter", - "Motion", - "Noise", - "Normalize", - "OneHot", - "OneOf", - "Pad", - "PatchAggregator", - "PatchLocation", - "PatchSampler", - "Points", - "Queue", - "RemapLabels", - "RemoveLabels", - "Reorient", - "Representation", - "Resample", - "RescaleIntensity", - "Resize", - "ScalarImage", - "SequentialLabels", - "SomeOf", - "Spatial", - "SpatialTransform", - "Spike", - "Standardize", - "StudiesBatch", - "StudiesLoader", - "Study", - "Subject", - "SubjectsBatch", - "SubjectsLoader", - "Swap", - "To", - "ToReferenceSpace", - "Transform", - "Transpose", - "TypeAffineMatrix", - "TypeDirection", - "TypeImageData", - "TypeOrientationCodes", - "TypeOrigin", - "TypePath", - "TypeSpacing", - "TypeSpatialShape", - "TypeTensorShape", - "TypeWorldPoints", - "UniformSampler", - "WeightedSampler", - "ZNormalization", - "apply_inverse_transform", - "collate_images", - "collate_studies", - "collate_subjects", - "datasets", - "enable_logging", - "read_matrix", - "write_matrix", + 'utils', + 'io', + 'sampler', + 'inference', + 'SubjectsDataset', + 'SubjectsLoader', + 'Image', + 'ScalarImage', + 'LabelMap', + 'Queue', + 'Subject', + 'datasets', + 'reference', + 'WeightedSampler', + 'UniformSampler', + 'LabelSampler', + 'GridSampler', + 'GridAggregator', ] - -__version__ = version(__name__) diff --git a/src/torchio/cli.py b/src/torchio/cli.py deleted file mode 100644 index 593476ec8..000000000 --- a/src/torchio/cli.py +++ /dev/null @@ -1,268 +0,0 @@ -"""TorchIO command-line interface.""" - -from __future__ import annotations - -import ast -import shutil -import sys -from dataclasses import dataclass -from dataclasses import field -from pathlib import Path -from typing import Annotated -from typing import Union - -import tyro - -import torchio as tio -from torchio.download import get_torchio_cache_dir - -# --------------------------------------------------------------------------- -# Subcommands -# --------------------------------------------------------------------------- - - -@dataclass -class Plot: - """Plot 3 orthogonal slices of an image.""" - - path: Annotated[Path, tyro.conf.Positional] - """Path to the image file.""" - - channel: int = 0 - """Channel index to display.""" - - output: Path | None = None - """Save the figure to a file instead of displaying.""" - - indices: tuple[int, int, int] | None = None - """Slice indices (i, j, k). Defaults to mid-slices.""" - - def run(self) -> None: - image = tio.ScalarImage(self.path) - show = self.output is None - image.plot( - channel=self.channel, - indices=self.indices, - output_path=self.output, - show=show, - ) - - -@dataclass -class Animate: - """Create an animated GIF or MP4 sweeping through slices. - - The output format is inferred from the file extension: - `.gif` produces an animated GIF, `.mp4` produces a video. - - Examples:: - - torchio animate brain.nii.gz brain.gif - torchio animate brain.nii.gz brain.mp4 --seconds 10 --direction S - """ - - path: Annotated[Path, tyro.conf.Positional] - """Path to the input image.""" - - output: Annotated[Path, tyro.conf.Positional] - """Output path (.gif or .mp4).""" - - seconds: float = 5.0 - """Duration of the animation in seconds.""" - - direction: str = "I" - """Anatomical sweep direction (I, S, A, P, R, or L).""" - - def run(self) -> None: - image = tio.ScalarImage(self.path) - suffix = self.output.suffix.lower() - if suffix == ".gif": - image.to_gif( - self.output, - seconds=self.seconds, - direction=self.direction, - ) - elif suffix == ".mp4": - image.to_video( - self.output, - seconds=self.seconds, - direction=self.direction, - ) - else: - msg = f"Unsupported output format {self.output.suffix!r}. Use .gif or .mp4." - print(msg, file=sys.stderr) - sys.exit(1) - print(f"Created {self.output}") - - -@dataclass -class Info: - """Print image metadata to stdout.""" - - path: Annotated[Path, tyro.conf.Positional] - """Path to the image file.""" - - def run(self) -> None: - image = tio.ScalarImage(self.path) - print(repr(image)) - - -@dataclass -class Convert: - """Convert an image between formats. - - Supports all SimpleITK formats plus NIfTI-Zarr (.nii.zarr). - The output format is inferred from the file extension. - """ - - input: Annotated[Path, tyro.conf.Positional] - """Path to the input image.""" - - output: Annotated[Path, tyro.conf.Positional] - """Path for the output image.""" - - def run(self) -> None: - image = tio.ScalarImage(self.input) - output_str = str(self.output) - if output_str.endswith(".nii.zarr"): - image.to_nifti_zarr(self.output) - else: - image.save(self.output) - - -@dataclass -class Transform: - """Apply a transform to an image. - - Extra arguments are passed as key=value pairs to the transform. - - Examples:: - - torchio transform brain.nii.gz noisy.nii.gz Noise std=0.1 - torchio transform brain.nii.gz cropped.nii.gz CropOrPad target_shape=128 - """ - - input: Annotated[Path, tyro.conf.Positional] - """Path to the input image.""" - - output: Annotated[Path, tyro.conf.Positional] - """Path for the output image.""" - - name: Annotated[str, tyro.conf.Positional] - """Transform class name (e.g., Noise, Flip, CropOrPad).""" - - device: str = "cpu" - """Device to run the transform on (e.g., "cpu", "cuda", "cuda:0" or "mps").""" - - args: Annotated[list[str], tyro.conf.Positional] = field( - default_factory=list, - ) - """Extra arguments as key=value pairs (e.g., std=0.1).""" - - def run(self) -> None: - transform_cls = _get_transform_class(self.name) - kwargs = _parse_kwargs(self.args) - transform = transform_cls(**kwargs) - image = tio.ScalarImage(self.input).to(self.device) - result = transform(image) - result.save(self.output) - - -@dataclass -class Dir: - """Print the cache directory path.""" - - def run(self) -> None: - print(get_torchio_cache_dir()) - - -@dataclass -class Clean: - """Clear cached data.""" - - dataset: str | None = None - """Dataset name to clear (e.g., 'colin27', 'fpg'). Clears all if omitted.""" - - def run(self) -> None: - cache_dir = get_torchio_cache_dir() - if self.dataset is not None: - target = cache_dir / self.dataset - if not target.exists(): - print(f"No cached data for {self.dataset!r}") - return - shutil.rmtree(target) - print(f"Cleared cache for {self.dataset!r}") - else: - if not cache_dir.exists(): - print("Cache is already empty") - return - shutil.rmtree(cache_dir) - print(f"Cleared all cached data from {cache_dir}") - - -@dataclass -class Cache: - """Manage the TorchIO data cache.""" - - command: tyro.conf.OmitSubcommandPrefixes[Dir | Clean] - - def run(self) -> None: - self.command.run() - - -Command = Union[Plot, Animate, Info, Convert, Transform, Cache] # noqa: UP007 (tyro needs Union) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _get_transform_class(name: str) -> type: - """Look up a transform class by name.""" - from torchio.transforms.transform import _TRANSFORM_REGISTRY - - if name not in _TRANSFORM_REGISTRY: - available = sorted(_TRANSFORM_REGISTRY.keys()) - print(f"Unknown transform {name!r}.", file=sys.stderr) - print(f"Available: {', '.join(available)}", file=sys.stderr) - sys.exit(1) - return _TRANSFORM_REGISTRY[name] - - -def _parse_kwargs(args: list[str]) -> dict[str, object]: - """Parse key=value pairs into a dict, inferring Python types.""" - kwargs: dict[str, object] = {} - for arg in args: - if "=" not in arg: - print(f"Invalid argument {arg!r}, expected key=value", file=sys.stderr) - sys.exit(1) - key, value_str = arg.split("=", 1) - kwargs[key] = _parse_value(value_str) - return kwargs - - -def _parse_value(value_str: str) -> object: - """Parse a string value into a Python literal (int, float, bool, tuple, etc.).""" - try: - return ast.literal_eval(value_str) - except (ValueError, SyntaxError): - return value_str - - -# --------------------------------------------------------------------------- -# Entry point -# --------------------------------------------------------------------------- - - -def main() -> None: - """TorchIO CLI: tools for medical image processing.""" - if "--version" in sys.argv[1:]: - print(f"torchio {tio.__version__}") - raise SystemExit(0) - cmd = tyro.cli(Command) - cmd.run() - - -if __name__ == "__main__": - main() diff --git a/src/torchio/cli/__init__.py b/src/torchio/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/torchio/cli/apply_transform.py b/src/torchio/cli/apply_transform.py new file mode 100644 index 000000000..dc45c602a --- /dev/null +++ b/src/torchio/cli/apply_transform.py @@ -0,0 +1,116 @@ +# pylint: disable=import-outside-toplevel + +from pathlib import Path + +import typer +from rich.progress import Progress +from rich.progress import SpinnerColumn +from rich.progress import TextColumn + +app = typer.Typer() + + +@app.command() +def main( + input_path: Path = typer.Argument( # noqa: B008 + ..., + exists=True, + file_okay=True, + dir_okay=True, + readable=True, + ), + transform_name: str = typer.Argument(...), # noqa: B008 + output_path: Path = typer.Argument( # noqa: B008 + ..., + file_okay=True, + dir_okay=False, + writable=True, + ), + kwargs: str = typer.Option( # noqa: B008 + None, + '--kwargs', + '-k', + help='String of kwargs, e.g. "degrees=(-5,15) num_transforms=3".', + ), + imclass: str = typer.Option( # noqa: B008 + 'ScalarImage', + '--imclass', + '-c', + help=( + 'Name of the subclass of torchio.Image' + ' that will be used to instantiate the image.' + ), + ), + seed: int = typer.Option( # noqa: B008 + None, + '--seed', + '-s', + help='Seed for PyTorch random number generator.', + ), + verbose: bool = typer.Option( # noqa: B008 + False, + help='Print random transform parameters.', + ), + show_progress: bool = typer.Option( # noqa: B008 + True, + '--show-progress/--hide-progress', + '-p/-P', + help='Show animations indicating progress.', + ), +): + """Apply transform to an image. + + Examples: + $ tiotr input.nrrd RandomMotion output.nii "degrees=(-5,15) num_transforms=3" -v + """ + # Imports are placed here so that the tool loads faster if not being run + import torch + + import torchio.transforms as transforms + from torchio.utils import apply_transform_to_file + + try: + transform_class = getattr(transforms, transform_name) + except AttributeError as error: + message = f'Transform "{transform_name}" not found in torchio' + raise ValueError(message) from error + + params_dict = get_params_dict_from_kwargs(kwargs) + transform = transform_class(**params_dict) + if seed is not None: + torch.manual_seed(seed) + with Progress( + SpinnerColumn(), + TextColumn('[progress.description]{task.description}'), + transient=True, + disable=not show_progress, + ) as progress: + progress.add_task('Applying transform', total=1) + apply_transform_to_file( + input_path, + transform, + output_path, + verbose=verbose, + class_=imclass, + ) + + +def get_params_dict_from_kwargs(kwargs): + from torchio.utils import guess_type + + params_dict = {} + if kwargs is not None: + for substring in kwargs.split(): + try: + key, value_string = substring.split('=') + except ValueError as error: + message = f'Arguments string "{kwargs}" not valid' + raise ValueError(message) from error + + value = guess_type(value_string) + params_dict[key] = value + return params_dict + + +if __name__ == '__main__': + app() diff --git a/src/torchio/cli/print_info.py b/src/torchio/cli/print_info.py new file mode 100644 index 000000000..e3fe167c6 --- /dev/null +++ b/src/torchio/cli/print_info.py @@ -0,0 +1,64 @@ +# pylint: disable=import-outside-toplevel +from pathlib import Path + +import typer + +app = typer.Typer() + + +@app.command() +def main( + input_path: Path = typer.Argument( # noqa: B008 + ..., + exists=True, + file_okay=True, + dir_okay=True, + readable=True, + ), + plot: bool = typer.Option( # noqa: B008 + False, + '--plot/--no-plot', + '-p/-P', + help='Plot the image using Matplotlib or Pillow.', + ), + show: bool = typer.Option( # noqa: B008 + False, + '--show/--no-show', + '-s/-S', + help='Show the image using specialized visualisation software.', + ), + label: bool = typer.Option( # noqa: B008 + False, + '--label/--scalar', + '-l/-s', + help='Use torchio.LabelMap to instantiate the image.', + ), + load: bool = typer.Option( # noqa: B008 + False, + help=( + 'Load the image from disk so that information about data type and memory ' + 'can be displayed. Slower, especially for large images.' + ), + ), +) -> None: + """Print information about an image and, optionally, show it. + + Examples: + $ tiohd input.nii.gz + """ + # Imports are placed here so that the tool loads faster if not being run + import torchio as tio + + class_ = tio.LabelMap if label else tio.ScalarImage + image = class_(input_path) + if load: + image.load() + print(image) # noqa: T201 + if plot: + image.plot() + if show: + image.show() + + +if __name__ == '__main__': + app() diff --git a/src/torchio/constants.py b/src/torchio/constants.py new file mode 100644 index 000000000..37b80d562 --- /dev/null +++ b/src/torchio/constants.py @@ -0,0 +1,36 @@ +import torch + +# Image types +INTENSITY = 'intensity' +LABEL = 'label' +SAMPLING_MAP = 'sampling_map' + +# Keys for dataset samples +PATH = 'path' +TYPE = 'type' +STEM = 'stem' +DATA = 'data' +AFFINE = 'affine' +TENSOR = 'tensor' + +# For aggregator +IMAGE = 'image' +LOCATION = 'location' + +# For special collate function +HISTORY = 'history' + +# In PyTorch convention +CHANNELS_DIMENSION = 1 + +# Code repository +REPO_URL = 'https://github.com/TorchIO-project/torchio/' + +# Data repository +DATA_REPO = 'https://github.com/TorchIO-project/torchio-data/raw/main/data/' + +# Floating point error +MIN_FLOAT_32 = torch.finfo(torch.float32).eps + +# For the queue +NUM_SAMPLES = 'num_samples' diff --git a/src/torchio/data/__init__.py b/src/torchio/data/__init__.py index 3d1a5b644..d3250c983 100644 --- a/src/torchio/data/__init__.py +++ b/src/torchio/data/__init__.py @@ -1,57 +1,29 @@ -"""Data classes for TorchIO.""" - -from .aggregator import PatchAggregator -from .backends import BackendRequest -from .backends import ImageDataBackend -from .backends import LazyReader -from .backends import register_backend -from .backends import resolve_backend -from .backends import unregister_backend -from .batch import ImagesBatch -from .batch import StudiesBatch -from .batch import SubjectsBatch -from .bboxes import BoundingBoxes -from .bboxes import BoundingBoxFormat -from .bboxes import Representation +from .dataset import SubjectsDataset from .image import Image from .image import LabelMap from .image import ScalarImage -from .patch import PatchLocation -from .points import Points +from .inference import GridAggregator +from .loader import SubjectsLoader from .queue import Queue from .sampler import GridSampler from .sampler import LabelSampler from .sampler import PatchSampler from .sampler import UniformSampler from .sampler import WeightedSampler -from .subject import Study from .subject import Subject __all__ = [ - "BackendRequest", - "BoundingBoxFormat", - "BoundingBoxes", - "GridSampler", - "Image", - "ImageDataBackend", - "ImagesBatch", - "LabelMap", - "LabelSampler", - "LazyReader", - "PatchAggregator", - "PatchLocation", - "PatchSampler", - "Points", - "Queue", - "Representation", - "ScalarImage", - "StudiesBatch", - "Study", - "Subject", - "SubjectsBatch", - "UniformSampler", - "WeightedSampler", - "register_backend", - "resolve_backend", - "unregister_backend", + 'Queue', + 'Subject', + 'SubjectsDataset', + 'SubjectsLoader', + 'Image', + 'ScalarImage', + 'LabelMap', + 'GridSampler', + 'GridAggregator', + 'PatchSampler', + 'LabelSampler', + 'WeightedSampler', + 'UniformSampler', ] diff --git a/src/torchio/data/affine.py b/src/torchio/data/affine.py deleted file mode 100644 index b5762cc62..000000000 --- a/src/torchio/data/affine.py +++ /dev/null @@ -1,248 +0,0 @@ -"""Affine matrix class.""" - -from __future__ import annotations - -import contextlib -from typing import Any - -import nibabel as nib -import numpy as np -import numpy.typing as npt -import torch -from torch import Tensor - -from ..types import TypeDirection -from ..types import TypeOrientationCodes -from ..types import TypeOrigin -from ..types import TypeSpacing - - -class AffineMatrix: - r"""$4 \times 4$ affine matrix mapping voxel indices to world coordinates. - - Stores the matrix as a `torch.Tensor` so it can live on the same - device as the image data. Named properties expose spacing, origin, - direction, and orientation. Composition uses the `@` operator. - - Args: - matrix: $4 \times 4$ array-like, `torch.Tensor`, or `None` - (identity). NumPy arrays are converted to tensors. - - Examples: - >>> import torchio as tio - >>> affine = tio.AffineMatrix() - >>> affine.spacing - (1.0, 1.0, 1.0) - >>> affine.orientation - ('R', 'A', 'S') - """ - - __slots__ = ("_matrix",) - - def __init__( - self, - matrix: AffineMatrix | Tensor | npt.ArrayLike | None = None, - ) -> None: - if matrix is None: - self._matrix = torch.eye(4, dtype=torch.float64) - elif isinstance(matrix, AffineMatrix): - self._matrix = matrix._matrix.clone() - elif isinstance(matrix, Tensor): - if matrix.shape != (4, 4): - msg = f"AffineMatrix must be 4x4, got {tuple(matrix.shape)}" - raise ValueError(msg) - self._matrix = matrix.to(torch.float64).clone() - else: - m = np.asarray(matrix, dtype=np.float64) - if m.shape != (4, 4): - msg = f"AffineMatrix must be 4x4, got {m.shape}" - raise ValueError(msg) - self._matrix = torch.as_tensor(m.copy(), dtype=torch.float64) - - # --- Construction helpers --- - - @classmethod - def from_spacing( - cls, - spacing: TypeSpacing, - *, - origin: TypeOrigin = (0.0, 0.0, 0.0), - direction: npt.ArrayLike | Tensor | None = None, - ) -> AffineMatrix: - """Create an affine from spacing, origin, and direction. - - Args: - spacing: Voxel size in mm along each axis. - origin: World coordinates of the first voxel center. - direction: 3x3 rotation/direction matrix. Identity if not given. - """ - matrix = torch.eye(4, dtype=torch.float64) - if direction is not None: - if isinstance(direction, Tensor): - matrix[:3, :3] = direction.to(torch.float64) - else: - matrix[:3, :3] = torch.as_tensor( - np.asarray(direction, dtype=np.float64), - ) - sp = torch.as_tensor(spacing, dtype=torch.float64) - matrix[:3, :3] *= sp - matrix[:3, 3] = torch.as_tensor(origin, dtype=torch.float64) - return cls(matrix) - - # --- Properties --- - - @property - def data(self) -> Tensor: - """The underlying 4x4 tensor.""" - return self._matrix - - @property - def device(self) -> torch.device: - """Device the affine matrix resides on.""" - return self._matrix.device - - @property - def spacing(self) -> TypeSpacing: - """Voxel spacing in mm, derived from the rotation-zoom block.""" - rz = self._matrix[:3, :3] - sp = torch.sqrt(torch.sum(rz**2, dim=0)) - return (float(sp[0]), float(sp[1]), float(sp[2])) - - @property - def origin(self) -> TypeOrigin: - """World coordinates of the first voxel center.""" - o = self._matrix[:3, 3] - return (float(o[0]), float(o[1]), float(o[2])) - - @property - def direction(self) -> TypeDirection: - """3x3 direction (rotation) matrix, with spacing factored out.""" - rz = self._matrix[:3, :3] - sp = torch.sqrt(torch.sum(rz**2, dim=0)) - return rz / sp - - @property - def orientation(self) -> TypeOrientationCodes: - """Anatomical orientation codes (e.g., `('R', 'A', 'S')`).""" - codes = nib.orientations.aff2axcodes(self._matrix.cpu().numpy()) - return (codes[0], codes[1], codes[2]) - - @property - def euler_angles(self) -> tuple[float, float, float]: - """Euler angles in degrees (XYZ intrinsic convention). - - Computed from the direction (rotation) matrix. All zeros means - the image axes are perfectly aligned with the scanner axes; non-zero - values indicate an oblique acquisition. - """ - r = self.direction - # XYZ intrinsic = ZYX extrinsic - # r is a (3, 3) Tensor (from the direction property) - sy = torch.sqrt(r[0, 0] ** 2 + r[1, 0] ** 2) - singular = float(sy) < 1e-6 - if not singular: - x = torch.atan2(r[2, 1], r[2, 2]) - y = torch.atan2(-r[2, 0], sy) - z = torch.atan2(r[1, 0], r[0, 0]) - else: - x = torch.atan2(-r[1, 2], r[1, 1]) - y = torch.atan2(-r[2, 0], sy) - z = torch.zeros(1, dtype=r.dtype, device=r.device) - deg = 180.0 / torch.pi - return (float(x * deg), float(y * deg), float(z * deg)) - - # --- Methods --- - - def to(self, *args: Any, **kwargs: Any) -> AffineMatrix: - """Move the affine to a device. - - The affine always stays in float64 for precision. On devices - that don't support float64 (e.g., MPS), it remains on CPU. - - Returns: - `self` (modified in-place). - """ - with contextlib.suppress(TypeError): - # MPS doesn't support float64, so keep on CPU - self._matrix = self._matrix.to(*args, **kwargs).to(torch.float64) - return self - - def clone(self) -> AffineMatrix: - """Return a deep copy.""" - return AffineMatrix(self._matrix.clone()) - - def inverse(self) -> AffineMatrix: - """Return the inverse affine.""" - return AffineMatrix(torch.linalg.inv(self._matrix)) - - def compose(self, other: AffineMatrix) -> AffineMatrix: - """Return `self @ other` as a new `AffineMatrix`. - - Equivalent to using the `@` operator. - """ - return AffineMatrix(self._matrix @ other._matrix) - - def apply(self, points: Tensor | npt.ArrayLike) -> Tensor: - """Apply the affine to an (N, 3) set of points. - - Args: - points: Tensor or array of shape (N, 3). - - Returns: - Transformed points as a tensor, shape (N, 3). - """ - if not isinstance(points, Tensor): - pts = torch.as_tensor( - np.asarray(points, dtype=np.float64), - dtype=torch.float64, - ) - else: - pts = points.to(torch.float64) - pts = pts.to(self._matrix.device) - ones = torch.ones(pts.shape[0], 1, dtype=torch.float64, device=pts.device) - homogeneous = torch.cat([pts, ones], dim=1) - transformed = (self._matrix @ homogeneous.T).T - return transformed[:, :3] - - def numpy(self) -> npt.NDArray[np.float64]: - """Return the underlying 4x4 matrix as a numpy array.""" - return self._matrix.cpu().numpy() - - # --- Dunder methods --- - - def __matmul__(self, other: object) -> AffineMatrix: - """Compose two affines via the `@` operator.""" - if not isinstance(other, AffineMatrix): - return NotImplemented - return self.compose(other) - - def __array__( - self, - dtype: npt.DTypeLike | None = None, - copy: bool | None = None, - ) -> npt.NDArray[np.float64]: - arr = self._matrix.cpu().numpy() - if dtype is not None: - return np.array(arr, dtype=dtype, copy=copy) - if copy: - return arr.copy() - return arr - - def __repr__(self) -> str: - sp = ", ".join(f"{s:.2f}" for s in self.spacing) - ori = "".join(self.orientation) - o = ", ".join(f"{v:.2f}" for v in self.origin) - return f"AffineMatrix(spacing=({sp}), origin=({o}), orientation={ori}+)" - - def __eq__(self, other: object) -> bool: - if not isinstance(other, AffineMatrix): - return NotImplemented - return torch.equal(self._matrix, other._matrix) - - def __copy__(self) -> AffineMatrix: - return self.clone() - - def __deepcopy__(self, memo: dict) -> AffineMatrix: - new = self.clone() - memo[id(self)] = new - return new diff --git a/src/torchio/data/aggregator.py b/src/torchio/data/aggregator.py deleted file mode 100644 index 29dced1e3..000000000 --- a/src/torchio/data/aggregator.py +++ /dev/null @@ -1,245 +0,0 @@ -"""Patch aggregator for dense inference.""" - -from __future__ import annotations - -import torch -from torch import Tensor - -from ..types import TypeThreeInts -from .patch import PatchLocation - - -class PatchAggregator: - """Reassemble patches into a full volume. - - Handles overlapping patches with configurable blending modes. - Supports outputs of different spatial sizes than the input - patches (e.g., downsampled feature maps or embeddings). - - Args: - spatial_shape: Output volume spatial shape `(I, J, K)`. - overlap_mode: How to handle overlapping regions: - `'crop'` keeps only non-overlapping centers (fast, - best for argmax segmentation); - `'average'` averages overlapping values (best for - probabilistic outputs); - `'hann'` uses Hann-window weighting (smoothest, - best for continuous outputs). - patch_overlap: The overlap used during sampling, needed - for `'crop'` mode to compute how much to trim. - output_shape: If the model output is spatially smaller - than the input patch (e.g., due to strided - convolutions), specify the output volume shape here. - Patch locations will be scaled accordingly. - - Examples: - >>> aggregator = tio.PatchAggregator( - ... spatial_shape=(256, 256, 176), - ... overlap_mode="hann", - ... ) - >>> for batch in loader: - ... outputs = model(batch.t1.data) - ... aggregator.add_batch(outputs, locations) - >>> volume = aggregator.get_output() - """ - - def __init__( - self, - spatial_shape: TypeThreeInts, - overlap_mode: str = "crop", - patch_overlap: int | TypeThreeInts = 0, - output_shape: TypeThreeInts | None = None, - ) -> None: - _validate_overlap_mode(overlap_mode) - self.input_spatial_shape = spatial_shape - self.overlap_mode = overlap_mode - - if isinstance(patch_overlap, int): - patch_overlap = (patch_overlap, patch_overlap, patch_overlap) - self.patch_overlap: TypeThreeInts = patch_overlap - - if output_shape is not None: - self.spatial_shape = output_shape - self._scale = ( - output_shape[0] / spatial_shape[0], - output_shape[1] / spatial_shape[1], - output_shape[2] / spatial_shape[2], - ) - else: - self.spatial_shape = spatial_shape - self._scale = (1.0, 1.0, 1.0) - - self._outputs: dict[str, Tensor] = {} - self._counts: dict[str, Tensor] = {} - self._hann_cache: dict[TypeThreeInts, Tensor] = {} - - def add_batch( - self, - batch: Tensor | dict[str, Tensor], - locations: list[PatchLocation], - ) -> None: - """Add a batch of model outputs to the aggregation buffer. - - Args: - batch: 5D tensor `(B, C, I, J, K)` or dict of such - tensors keyed by name. - locations: List of `PatchLocation` for each item in - the batch. - """ - tensors: dict[str, Tensor] = ( - {"__default__": batch} if isinstance(batch, Tensor) else batch - ) - - for key, tensor in tensors.items(): - tensor = tensor.cpu() - for idx, loc in enumerate(locations): - patch = tensor[idx] - if self._scale != (1.0, 1.0, 1.0): - loc = loc.scaled(self._scale) - self._add_patch(key, patch, loc) - - def get_output(self, key: str | None = None) -> Tensor: - """Get the aggregated output volume. - - Args: - key: Name of the output to retrieve. If `None` and - only a single (unnamed) output was added, return it. - - Returns: - The aggregated tensor with shape `(C, I, J, K)`. - """ - resolve_key = key if key is not None else "__default__" - if resolve_key not in self._outputs: - available = [k for k in self._outputs if k != "__default__"] - msg = f"No output for key {key!r}. Available: {available}" - raise KeyError(msg) - - output = self._outputs[resolve_key] - - if self.overlap_mode in ("average", "hann"): - counts = self._counts[resolve_key] - counts = counts.clamp(min=1) - output = output / counts - - return output - - def _add_patch( - self, - key: str, - patch: Tensor, - location: PatchLocation, - ) -> None: - self._ensure_buffer(key, patch) - match self.overlap_mode: - case "crop": - self._add_crop(key, patch, location) - case "average": - self._add_average(key, patch, location) - case "hann": - self._add_hann(key, patch, location) - - def _ensure_buffer(self, key: str, patch: Tensor) -> None: - if key in self._outputs: - return - num_channels = patch.shape[0] - self._outputs[key] = torch.zeros( - num_channels, - *self.spatial_shape, - dtype=patch.dtype, - ) - if self.overlap_mode in ("average", "hann"): - self._counts[key] = torch.zeros( - num_channels, - *self.spatial_shape, - dtype=patch.dtype, - ) - - def _add_crop( - self, - key: str, - patch: Tensor, - location: PatchLocation, - ) -> None: - """Place only the non-overlapping center of the patch.""" - scaled_overlap = ( - round(self.patch_overlap[0] * self._scale[0]), - round(self.patch_overlap[1] * self._scale[1]), - round(self.patch_overlap[2] * self._scale[2]), - ) - half = [o // 2 for o in scaled_overlap] - ini = list(location.index_ini) - fin = list(location.index_fin) - crop_ini = [0, 0, 0] - crop_fin = list(location.size) - - for d in range(3): - if ini[d] > 0: - ini[d] += half[d] - crop_ini[d] += half[d] - if fin[d] < self.spatial_shape[d]: - fin[d] -= half[d] - crop_fin[d] -= half[d] - - cropped = patch[ - :, - crop_ini[0] : crop_fin[0], - crop_ini[1] : crop_fin[1], - crop_ini[2] : crop_fin[2], - ] - self._outputs[key][ - :, - ini[0] : fin[0], - ini[1] : fin[1], - ini[2] : fin[2], - ] = cropped - - def _add_average( - self, - key: str, - patch: Tensor, - location: PatchLocation, - ) -> None: - si, sj, sk = location.to_slices() - self._outputs[key][:, si, sj, sk] += patch - self._counts[key][:, si, sj, sk] += 1 - - def _add_hann( - self, - key: str, - patch: Tensor, - location: PatchLocation, - ) -> None: - patch_shape = ( - patch.shape[-3], - patch.shape[-2], - patch.shape[-1], - ) - window = self._get_hann_window(patch_shape) - si, sj, sk = location.to_slices() - self._outputs[key][:, si, sj, sk] += patch * window - self._counts[key][:, si, sj, sk] += window - - def _get_hann_window(self, patch_size: TypeThreeInts) -> Tensor: - if patch_size in self._hann_cache: - return self._hann_cache[patch_size] - window = _build_hann_3d(patch_size) - self._hann_cache[patch_size] = window - return window - - -def _validate_overlap_mode(mode: str) -> None: - valid = ("crop", "average", "hann") - if mode not in valid: - msg = f"overlap_mode must be one of {valid}, got {mode!r}" - raise ValueError(msg) - - -def _build_hann_3d(patch_size: TypeThreeInts) -> Tensor: - """Build a 3D Hann window for smooth patch blending.""" - window = torch.ones(1) - for dim, size in enumerate(patch_size): - shape = [1, 1, 1] - shape[dim] = size - w = torch.hann_window(size + 2, periodic=False)[1:-1] - window = window * w.reshape(shape) - return window diff --git a/src/torchio/data/axes.py b/src/torchio/data/axes.py deleted file mode 100644 index b24a3d85d..000000000 --- a/src/torchio/data/axes.py +++ /dev/null @@ -1,170 +0,0 @@ -"""Axis validation and conversion utilities. - -Axis strings are 3-character uppercase strings describing the ordering and -orientation of coordinate axes. Two families are supported: - -- **Voxel**: permutations of `"IJK"` (6 options). -- **Anatomical**: one letter from each pair {R, L}, {A, P}, {S, I} - in any order (48 options). - -The string `"IJK"` is always interpreted as voxel (not -Inferior-J-K), because J and K have no anatomical meaning. -""" - -from __future__ import annotations - -from enum import Enum - -_VOXEL_LETTERS = frozenset("IJK") - -# Each tuple is a pair of opposite anatomical directions. -ANATOMICAL_PAIRS: tuple[tuple[str, str], ...] = ( - ("R", "L"), - ("A", "P"), - ("S", "I"), -) - -# Flat set of all anatomical letters. -_ANATOMICAL_LETTERS = frozenset(letter for pair in ANATOMICAL_PAIRS for letter in pair) - -# Map each anatomical letter to its opposite. -_OPPOSITE: dict[str, str] = {} -for _a, _b in ANATOMICAL_PAIRS: - _OPPOSITE[_a] = _b - _OPPOSITE[_b] = _a - -# Map each anatomical letter to its canonical pair (sorted tuple). -_LETTER_TO_PAIR: dict[str, tuple[str, str]] = {} -for _pair in ANATOMICAL_PAIRS: - for _letter in _pair: - _LETTER_TO_PAIR[_letter] = _pair - - -class AxesType(Enum): - """Whether an axis string describes voxel or anatomical coordinates.""" - - VOXEL = "voxel" - ANATOMICAL = "anatomical" - - -def validate_axes(axes: str) -> str: - """Validate a 3-character axis string. - - Args: - axes: Axis string to validate. - - Returns: - The validated string (unchanged). - - Raises: - ValueError: If the string is not a valid axis specification. - """ - if len(axes) != 3: - msg = f"Axis string must be 3 characters, got {len(axes)}: {axes!r}" - raise ValueError(msg) - if _is_voxel(axes) or _is_anatomical(axes): - return axes - msg = ( - f"Invalid axis string {axes!r}. Must be a permutation of 'IJK'" - " (voxel) or one letter from each anatomical pair" - " {R,L}, {A,P}, {S,I}." - ) - raise ValueError(msg) - - -def axes_type(axes: str) -> AxesType: - """Return whether *axes* is a voxel or anatomical axis string. - - The string must already be valid (call - [`validate_axes`][torchio.data.axes.validate_axes] first). - """ - if _is_voxel(axes): - return AxesType.VOXEL - return AxesType.ANATOMICAL - - -def get_axis_mapping( - src: str, - tgt: str, -) -> tuple[tuple[int, int, int], tuple[bool, bool, bool]]: - """Compute the column permutation and flips to go from *src* to *tgt*. - - Both strings must be the same type (both voxel or both anatomical). - - Args: - src: Source axis string. - tgt: Target axis string. - - Returns: - A tuple `(permutation, flips)` where *permutation* gives the - source column index for each target column and *flips* indicates - whether each target column should be negated. - - Raises: - ValueError: If the two strings are not the same type. - """ - src_type = axes_type(src) - tgt_type = axes_type(tgt) - if src_type != tgt_type: - msg = ( - f"Cannot compute axis mapping between different types:" - f" {src!r} ({src_type.value}) and {tgt!r} ({tgt_type.value})." - " Use the affine to convert between voxel and anatomical." - " Both must be the same type." - ) - raise ValueError(msg) - - if src_type == AxesType.VOXEL: - return _voxel_mapping(src, tgt) - return _anatomical_mapping(src, tgt) - - -# --- Private helpers --- - - -def _is_voxel(axes: str) -> bool: - return set(axes) == _VOXEL_LETTERS and len(axes) == 3 - - -def _is_anatomical(axes: str) -> bool: - if len(axes) != 3: - return False - if not all(c in _ANATOMICAL_LETTERS for c in axes): - return False - # Each pair must be represented exactly once. - seen_pairs: set[tuple[str, str]] = set() - for c in axes: - pair = _LETTER_TO_PAIR[c] - if pair in seen_pairs: - return False - seen_pairs.add(pair) - return len(seen_pairs) == 3 - - -def _voxel_mapping( - src: str, - tgt: str, -) -> tuple[tuple[int, int, int], tuple[bool, bool, bool]]: - perm = tuple(src.index(c) for c in tgt) - flips = (False, False, False) - assert len(perm) == 3 - return (perm[0], perm[1], perm[2]), flips - - -def _anatomical_mapping( - src: str, - tgt: str, -) -> tuple[tuple[int, int, int], tuple[bool, bool, bool]]: - perm: list[int] = [] - flips: list[bool] = [] - for tgt_letter in tgt: - tgt_pair = _LETTER_TO_PAIR[tgt_letter] - # Find which source column belongs to the same pair. - for src_idx, src_letter in enumerate(src): - if _LETTER_TO_PAIR[src_letter] == tgt_pair: - perm.append(src_idx) - flips.append(src_letter != tgt_letter) - break - assert len(perm) == 3 - assert len(flips) == 3 - return (perm[0], perm[1], perm[2]), (flips[0], flips[1], flips[2]) diff --git a/src/torchio/data/backends.py b/src/torchio/data/backends.py deleted file mode 100644 index 18012b37f..000000000 --- a/src/torchio/data/backends.py +++ /dev/null @@ -1,577 +0,0 @@ -"""Lazy image data backends. - -Backends provide a uniform interface for accessing image data without -requiring full materialization into memory. The `Image` class uses backends -internally. Users interact with images via `.data` (materialized tensor) -and `.dataobj` (lazy backend for advanced use like slicing). -""" - -from __future__ import annotations - -import types -from collections.abc import Callable -from collections.abc import Mapping -from dataclasses import dataclass -from dataclasses import field -from pathlib import Path -from typing import Any -from typing import Protocol -from typing import runtime_checkable - -import nibabel as nib -import nibabel.spatialimages -import numpy as np -import torch -from einops import rearrange -from torch import Tensor - -from ..types import SliceIndex -from ..types import TypeAffineMatrix -from ..types import TypeTensorShape - - -def _expand_ellipsis( - items: tuple[int | slice | types.EllipsisType, ...], - *, - ndim: int, -) -> tuple[int | slice, ...]: - """Replace a single `Ellipsis` with enough full slices to reach *ndim*.""" - n_ellipsis = sum(1 for s in items if s is Ellipsis) - if n_ellipsis == 0: - return tuple(s for s in items if s is not Ellipsis) - if n_ellipsis > 1: - msg = "Only one ellipsis is allowed" - raise IndexError(msg) - idx = items.index(Ellipsis) - n_explicit = len(items) - 1 - n_fill = max(ndim - n_explicit, 0) - expanded = items[:idx] + (slice(None),) * n_fill + items[idx + 1 :] - return tuple(s for s in expanded if s is not Ellipsis) - - -def normalize_index(item: SliceIndex, *, ndim: int = 4) -> tuple[slice, ...]: - """Normalize an index into a tuple of exactly *ndim* slices. - - Integer indices are converted to size-1 slices so that dimensions are - never dropped, a single `Ellipsis` expands to full slices, and missing - trailing indices are right-padded with full slices. This guarantees that - direct backend slicing (`image.dataobj[...]`) follows the same shape rules - as [`Image.__getitem__`][torchio.Image.__getitem__] and always returns - `(C, I, J, K)` data. - - Args: - item: Integer, slice, `Ellipsis`, or a tuple of those. - ndim: Number of dimensions to produce (4 for `(C, I, J, K)`). - - Returns: - A tuple of exactly *ndim* `slice` objects. - - Raises: - IndexError: If more than *ndim* indices or several ellipses are given. - TypeError: If an index element is not an `int`, `slice`, or `Ellipsis`. - """ - match item: - case int() | slice(): - items: tuple[int | slice | types.EllipsisType, ...] = (item,) - case tuple(): - items = item - case _ if item is Ellipsis: - items = (item,) - case _: - msg = f"Index type {type(item).__name__} not understood" - raise TypeError(msg) - - items = _expand_ellipsis(items, ndim=ndim) - - if len(items) > ndim: - msg = ( - f"Too many indices: expected at most {ndim} (C, I, J, K), got {len(items)}" - ) - raise IndexError(msg) - - parsed: list[slice] = [] - for s in items: - match s: - case int(): - # Keep the axis (size 1) instead of dropping it. ``slice(-1, - # 0)`` would be empty, so the last element needs a special case. - parsed.append(slice(s, None) if s == -1 else slice(s, s + 1)) - case slice(): - parsed.append(s) - case _: - msg = f"Index type {type(s).__name__} not understood" - raise TypeError(msg) - while len(parsed) < ndim: - parsed.append(slice(None)) - return tuple(parsed) - - -@runtime_checkable -class ImageDataBackend(Protocol): - """Protocol for lazy image data access. - - Implementations wrap different storage formats (in-memory tensors, NIfTI - files via nibabel, NIfTI-Zarr via dask) behind a uniform, lazy I/O - interface. This is an I/O adapter layer, not a lazy computation framework: - it speeds up metadata reads and region slicing, but does not defer - arithmetic or transforms. - - Contract: - - `shape` is always 4D `(C, I, J, K)`, even for 3D NIfTI (channel - dimension 1). - - `affine` is a `float64` `torch.Tensor` of shape `(4, 4)`. - - `dtype` reports the on-disk (or in-memory) `numpy` dtype. - - `__getitem__` accepts the same indexing as - [`Image.__getitem__`][torchio.Image.__getitem__], always returns a - 4D `torch.Tensor` in `(C, I, J, K)` layout, and never drops axes - (integer indices keep a size-1 dimension). - - `to_tensor` materializes the full volume as a `torch.Tensor` - preserving the on-disk dtype where PyTorch supports it. - """ - - @property - def shape(self) -> TypeTensorShape: - """Shape as (C, I, J, K).""" - ... - - @property - def affine(self) -> TypeAffineMatrix: - """$4 \\times 4$ affine matrix as a float64 tensor.""" - ... - - @property - def dtype(self) -> np.dtype: - """Data type of the image on disk.""" - ... - - def __getitem__(self, slices: SliceIndex) -> Tensor: - """Slice the data, returning a 4D `(C, I, J, K)` tensor. - - Integer indices keep a size-1 dimension (axes are never dropped). - Tensor-backed images preserve their device and dtype; lazy backends - read only the requested region and convert it to a tensor. - """ - ... - - def to_tensor(self) -> Tensor: - """Materialize the full data as a tensor preserving the on-disk dtype.""" - ... - - -class TensorBackend: - """Backend wrapping an in-memory PyTorch tensor. - - Used for images created from tensors or NumPy arrays (NumPy arrays are - converted to tensors first). - - Args: - data: 4D tensor with shape (C, I, J, K). - affine: $4 \\times 4$ affine tensor. Identity if not given. - """ - - __slots__ = ("_affine", "_data") - - def __init__( - self, - data: Tensor, - affine: TypeAffineMatrix | None = None, - ) -> None: - self._data = data - if affine is not None: - self._affine = affine - else: - self._affine = torch.eye(4, dtype=torch.float64) - - @property - def shape(self) -> TypeTensorShape: - s = self._data.shape - return (int(s[0]), int(s[1]), int(s[2]), int(s[3])) - - @property - def affine(self) -> TypeAffineMatrix: - return self._affine - - @property - def dtype(self) -> np.dtype: - # Map torch dtype to numpy for protocol compatibility - return torch.empty(0, dtype=self._data.dtype).numpy().dtype - - def __getitem__(self, slices: SliceIndex) -> Tensor: - """Slice the tensor, preserving device, dtype, and 4D layout.""" - return self._data[normalize_index(slices)] - - def to_tensor(self) -> Tensor: - return self._data.clone() - - -class NibabelBackend: - """Backend wrapping a nibabel image for lazy NIfTI access. - - Data is accessed through nibabel's `dataobj` proxy, which supports - memory-mapped reads. Shape and affine are read from the header - without loading data. - - This backend also works with NIfTI-Zarr files loaded via `niizarr`, - since `zarr2nii` returns a standard `nibabel.Nifti1Image` whose - `dataobj` is a dask array. - - Args: - nii: A nibabel image (typically from `nib.load()` or `zarr2nii()`). - affine: Optional $4 \\times 4$ affine that overrides the affine stored - in the NIfTI header. Used when the user passes an explicit affine - to `Image`, so that `image.affine` and `image.dataobj.affine` agree. - """ - - __slots__ = ("_affine_override", "_nii", "_shape") - - def __init__( - self, - nii: nib.spatialimages.SpatialImage, - affine: TypeAffineMatrix | None = None, - ) -> None: - self._nii = nii - self._affine_override = ( - torch.as_tensor(affine, dtype=torch.float64) if affine is not None else None - ) - header_shape = nii.header.get_data_shape() - ndim = len(header_shape) - if ndim == 3: - si, sj, sk = header_shape - self._shape: TypeTensorShape = (1, int(si), int(sj), int(sk)) - elif ndim == 4: - si, sj, sk, c = header_shape - self._shape = (int(c), int(si), int(sj), int(sk)) - elif ndim == 5 and header_shape[3] == 1: - # 5D vector NIfTI written by SimpleITK: (I, J, K, 1, C) - si, sj, sk, _, c = header_shape - self._shape = (int(c), int(si), int(sj), int(sk)) - else: - msg = f"Expected 3D or 4D NIfTI, got {ndim}D" - raise ValueError(msg) - - @property - def shape(self) -> TypeTensorShape: - return self._shape - - @property - def affine(self) -> TypeAffineMatrix: - if self._affine_override is not None: - return self._affine_override - return torch.as_tensor( - self._nii.header.get_best_affine(), - dtype=torch.float64, - ) - - @property - def dtype(self) -> np.dtype: - return self._nii.header.get_data_dtype() - - def __getitem__(self, slices: SliceIndex) -> Tensor: - """Slice in (C, I, J, K) space, returning a 4D tensor. - - The (C, I, J, K) index is translated to the on-disk layout, which is - (I, J, K) for 3D, (I, J, K, C) for 4D, and (I, J, K, 1, C) for 5D - vector NIfTI. Integer indices keep their axis (size 1), so the result - is always 4D. - """ - sc, si, sj, sk = normalize_index(slices) - ndim_on_disk = len(self._nii.header.get_data_shape()) - if ndim_on_disk == 3: - # On disk (I, J, K); channel axis is synthetic (size 1). - data = np.asarray(self._nii.dataobj[si, sj, sk]) - array = rearrange(data, "i j k -> 1 i j k")[sc] - elif ndim_on_disk == 4: - # On disk (I, J, K, C). - data = np.asarray(self._nii.dataobj[si, sj, sk, sc]) - array = rearrange(data, "i j k c -> c i j k") - elif ndim_on_disk == 5: - # 5D vector NIfTI written by SimpleITK: (I, J, K, 1, C). - data = np.asarray(self._nii.dataobj[si, sj, sk, :, sc]) - array = rearrange(data, "i j k 1 c -> c i j k") - else: - msg = f"Expected 3D, 4D, or 5D NIfTI, got {ndim_on_disk}D" - raise ValueError(msg) - from .io import _numpy_to_tensor - - array = np.ascontiguousarray(array) - if not array.flags.writeable: - # Proxy reads (e.g. memory-mapped NIfTI) can be read-only, which - # PyTorch does not support; copy so the resulting tensor is safe - # to mutate. - array = array.copy() - return _numpy_to_tensor(array) - - def to_tensor(self) -> Tensor: - """Materialize the full image preserving the on-disk dtype.""" - data = np.asarray(self._nii.dataobj) - ndim = data.ndim - if ndim == 3: - data = rearrange(data, "i j k -> 1 i j k") - elif ndim == 4: - data = rearrange(data, "i j k c -> c i j k") - elif ndim == 5 and data.shape[3] == 1: - # 5D vector NIfTI written by SimpleITK: (I, J, K, 1, C) - data = rearrange(data, "i j k 1 c -> c i j k") - else: - msg = f"Expected 3D or 4D data, got {ndim}D" - raise ValueError(msg) - from .io import _numpy_to_tensor - - return _numpy_to_tensor(np.ascontiguousarray(data)) - - -class ZarrBackend: - """Backend wrapping a NIfTI-Zarr file for chunked lazy access. - - NIfTI-Zarr files are loaded via `niizarr.zarr2nii()`, which returns - a nibabel image with a dask array as its `dataobj`. This backend - delegates to `NibabelBackend` for the actual data access. - - Requires the `nifti-zarr` package. - - Args: - path: Path to a `.nii.zarr` directory or a remote `.nii.zarr` URI. - affine: Optional $4 \\times 4$ affine that overrides the affine stored - in the NIfTI-Zarr metadata. - **kwargs: Extra keyword arguments forwarded to `niizarr.zarr2nii()`. - """ - - __slots__ = ("_nibabel_backend",) - - def __init__( - self, - path: str | object, - affine: TypeAffineMatrix | None = None, - **kwargs: Any, - ) -> None: - from ..external.imports import get_niizarr - - niizarr = get_niizarr() - nii = niizarr.zarr2nii(str(path), **kwargs) - self._nibabel_backend = NibabelBackend(nii, affine=affine) - - @property - def shape(self) -> TypeTensorShape: - return self._nibabel_backend.shape - - @property - def affine(self) -> TypeAffineMatrix: - return self._nibabel_backend.affine - - @property - def dtype(self) -> np.dtype: - return self._nibabel_backend.dtype - - def __getitem__(self, slices: SliceIndex) -> Tensor: - return self._nibabel_backend[slices] - - def to_tensor(self) -> Tensor: - return self._nibabel_backend.to_tensor() - - -# ── Backend resolution and registration ────────────────────────────── - - -@dataclass(frozen=True) -class BackendRequest: - """Description of an image source used to resolve a lazy backend. - - A request decouples backend selection from the `Image` class: resolvers - and custom backends receive a `BackendRequest` instead of an `Image`. - - Attributes: - path: Resolved filesystem (or fsspec) path to the image, if any. - remote_zarr_uri: Remote NIfTI-Zarr URI, if the source is remote. - zarr_store: An open Zarr store object, if the source is a store. - affine: Optional $4 \\times 4$ affine override to apply to the backend. - reader_kwargs: Extra keyword arguments forwarded to the loader. - reader: The reader configured on the `Image`. Custom readers that - implement [`LazyReader`][torchio.data.backends.LazyReader] can - build a lazy backend instead of loading the whole volume. - """ - - path: Path | None = None - remote_zarr_uri: str | None = None - zarr_store: Any = None - affine: TypeAffineMatrix | None = None - reader_kwargs: Mapping[str, Any] = field(default_factory=dict) - reader: Any = None - - -@runtime_checkable -class LazyReader(Protocol): - """A custom reader that can build a lazy backend. - - Readers passed to `Image` are normally simple callables returning - `(tensor, affine)`, which always load the whole volume. A reader that also - implements `create_backend` opts in to lazy access: `Image.shape`, - `affine`, `dtype`, and slicing then go through the returned backend without - materializing the full tensor. - """ - - def create_backend(self, request: BackendRequest) -> ImageDataBackend: - """Build a lazy backend for *request*.""" - ... - - -BackendMatcher = Callable[[BackendRequest], bool] -"""Predicate deciding whether a backend can handle a `BackendRequest`.""" - -BackendFactory = Callable[[BackendRequest], ImageDataBackend] -"""Callable that builds an `ImageDataBackend` from a `BackendRequest`.""" - - -@dataclass(frozen=True) -class _BackendEntry: - name: str - matcher: BackendMatcher - factory: BackendFactory - - -_BACKEND_REGISTRY: list[_BackendEntry] = [] - - -def register_backend( - name: str, - matcher: BackendMatcher, - factory: BackendFactory, - *, - prepend: bool = True, -) -> None: - """Register a lazy image data backend. - - Registered backends are consulted by - [`resolve_backend`][torchio.data.backends.resolve_backend] in order. This is - the extension point for supporting new formats without editing the `Image` - class. - - Args: - name: Identifier for the backend, used to unregister it later. - Registering a new backend with an existing name replaces it. - matcher: Predicate returning `True` if this backend can handle the - given [`BackendRequest`][torchio.data.backends.BackendRequest]. - factory: Callable that builds the backend from the request. - prepend: If `True` (default), the backend is consulted before existing - registrations, so custom backends take priority over the built-ins. - """ - unregister_backend(name) - entry = _BackendEntry(name=name, matcher=matcher, factory=factory) - if prepend: - _BACKEND_REGISTRY.insert(0, entry) - else: - _BACKEND_REGISTRY.append(entry) - - -def unregister_backend(name: str) -> None: - """Remove a previously registered backend by name (no-op if absent). - - Args: - name: The name passed to - [`register_backend`][torchio.data.backends.register_backend]. - """ - _BACKEND_REGISTRY[:] = [e for e in _BACKEND_REGISTRY if e.name != name] - - -def resolve_backend(request: BackendRequest) -> ImageDataBackend | None: - """Resolve a lazy backend for *request*. - - Args: - request: The source description. - - Returns: - The first matching backend, or `None` if no registered backend can - handle the request (for example a non-NIfTI file path, where the - caller falls back to a full read). - """ - for entry in _BACKEND_REGISTRY: - if entry.matcher(request): - return entry.factory(request) - return None - - -# -- Built-in backends ---------------------------------------------------- - - -def _match_custom_reader(request: BackendRequest) -> bool: - return request.reader is not None and isinstance(request.reader, LazyReader) - - -def _build_custom_reader(request: BackendRequest) -> ImageDataBackend: - reader: LazyReader = request.reader - return reader.create_backend(request) - - -def _match_remote_zarr(request: BackendRequest) -> bool: - return request.remote_zarr_uri is not None - - -def _build_remote_zarr(request: BackendRequest) -> ImageDataBackend: - assert request.remote_zarr_uri is not None - return ZarrBackend( - request.remote_zarr_uri, - affine=request.affine, - **dict(request.reader_kwargs), - ) - - -def _match_zarr_store(request: BackendRequest) -> bool: - return request.zarr_store is not None - - -def _build_zarr_store(request: BackendRequest) -> ImageDataBackend: - from ..external.imports import get_niizarr - - niizarr = get_niizarr() - nii = niizarr.zarr2nii(request.zarr_store, **dict(request.reader_kwargs)) - return NibabelBackend(nii, affine=request.affine) - - -def _match_nifti_zarr_path(request: BackendRequest) -> bool: - from .io import is_nifti_zarr - - return request.path is not None and is_nifti_zarr(request.path) - - -def _build_nifti_zarr_path(request: BackendRequest) -> ImageDataBackend: - assert request.path is not None - return ZarrBackend( - request.path, - affine=request.affine, - **dict(request.reader_kwargs), - ) - - -def _match_nifti_path(request: BackendRequest) -> bool: - from .io import is_nifti - - return request.path is not None and is_nifti(request.path) - - -def _build_nifti_path(request: BackendRequest) -> ImageDataBackend: - assert request.path is not None - nii = nib.load(request.path, **dict(request.reader_kwargs)) - assert isinstance(nii, nib.spatialimages.SpatialImage) - return NibabelBackend(nii, affine=request.affine) - - -def _register_builtin_backends() -> None: - """Register the backends shipped with TorchIO (NIfTI and NIfTI-Zarr). - - A custom-reader entry is registered first so that a - [`LazyReader`][torchio.data.backends.LazyReader] takes priority over the - format-based built-ins. - """ - register_backend( - "custom-reader", _match_custom_reader, _build_custom_reader, prepend=False - ) - register_backend( - "remote-nifti-zarr", _match_remote_zarr, _build_remote_zarr, prepend=False - ) - register_backend("zarr-store", _match_zarr_store, _build_zarr_store, prepend=False) - register_backend( - "nifti-zarr", _match_nifti_zarr_path, _build_nifti_zarr_path, prepend=False - ) - register_backend("nifti", _match_nifti_path, _build_nifti_path, prepend=False) - - -_register_builtin_backends() diff --git a/src/torchio/data/batch.py b/src/torchio/data/batch.py deleted file mode 100644 index 13b058eb4..000000000 --- a/src/torchio/data/batch.py +++ /dev/null @@ -1,399 +0,0 @@ -"""Batch containers for stacked images and subjects.""" - -from __future__ import annotations - -import dataclasses -from typing import Any - -import torch -from torch import Tensor -from typing_extensions import Self - -from .affine import AffineMatrix -from .image import Image -from .image import ScalarImage -from .invertible import Invertible - -#: Reserved param keys used for per-instance history bookkeeping. -_BATCH_META_KEYS = ("_batch_size", "_batched_keys", "_keep") - - -class ImagesBatch(Invertible): - """A batch of images with per-sample affines. - - Wraps a 5D tensor `(B, C, I, J, K)` and a list of `AffineMatrix` - matrices (one per sample). Created by stacking multiple `Image` - objects or directly from a 5D tensor. - - Args: - data: 5D tensor with shape `(B, C, I, J, K)`. - affines: List of affine matrices, one per sample. - image_class: The `Image` subclass to use when unbatching. - """ - - def __init__( - self, - data: Tensor, - affines: list[AffineMatrix], - *, - image_class: type[Image] = ScalarImage, - ) -> None: - if data.ndim != 5: - msg = f"Expected 5D tensor (B, C, I, J, K), got {data.ndim}D" - raise ValueError(msg) - if len(affines) != data.shape[0]: - msg = f"Expected {data.shape[0]} affines, got {len(affines)}" - raise ValueError(msg) - self._data = data - self._affines = affines - self._image_class = image_class - self.applied_transforms: list[Any] = [] - - @classmethod - def from_images(cls, images: list[Image]) -> Self: - """Stack a list of images into a batch. - - All images must have the same shape. - - Args: - images: List of `Image` instances to stack. - """ - if not images: - msg = "Cannot create batch from empty list" - raise ValueError(msg) - tensors = [img.data for img in images] - stacked = torch.stack(tensors) - affines = [img.affine.clone() for img in images] - image_class = type(images[0]) - return cls(stacked, affines, image_class=image_class) - - @property - def data(self) -> Tensor: - """5D tensor with shape `(B, C, I, J, K)`.""" - return self._data - - @data.setter - def data(self, value: Tensor) -> None: - if value.ndim != 5: - msg = f"Expected 5D tensor, got {value.ndim}D" - raise ValueError(msg) - self._data = value - - @property - def affines(self) -> list[AffineMatrix]: - """List of affine matrices, one per sample.""" - return self._affines - - @property - def batch_size(self) -> int: - """Number of samples in the batch.""" - return self._data.shape[0] - - @property - def device(self) -> torch.device: - """Device the batch data resides on.""" - return self._data.device - - def to(self, *args: Any, **kwargs: Any) -> Self: - """Move batch data to a device and/or cast dtype.""" - self._data = self._data.to(*args, **kwargs) - for affine in self._affines: - affine.to(*args, **kwargs) - return self - - def __getitem__(self, index: int) -> Image: - """Get a single image from the batch by index.""" - return self._image_class( - self._data[index], - affine=self._affines[index].clone(), - ) - - def __len__(self) -> int: - return self.batch_size - - def unbatch(self) -> list[Image]: - """Split the batch into individual images.""" - return [self[i] for i in range(self.batch_size)] - - def __repr__(self) -> str: - b, c, i, j, k = self._data.shape - cls = self._image_class.__name__ - return f"ImagesBatch({cls}, batch_size={b}, shape=({c}, {i}, {j}, {k}))" - - -class SubjectsBatch(Invertible): - """A batch of subjects with stacked image data. - - Each named image entry becomes an `ImagesBatch`. Metadata is - stored as lists (one value per sample). - - Created by `SubjectsLoader` or `SubjectsBatch.from_subjects()`. - """ - - def __init__( - self, - images: dict[str, ImagesBatch], - *, - metadata: dict[str, list[Any]] | None = None, - ) -> None: - self._images = images - self._metadata: dict[str, list[Any]] = metadata or {} - self.applied_transforms: list[Any] = [] - # When per-element branching occurs (e.g. per-instance OneOf), - # this stores the frozen per-element history prefix. Transforms - # applied afterwards still append to `applied_transforms`, and - # `unbatch()` merges the prefix with the sliced suffix. - self._per_element_history: list[list[Any]] | None = None - - def set_per_element_history(self, histories: list[list[Any]]) -> None: - """Freeze a distinct transform history for each batch element. - - Used when different elements receive different transforms (for - example per-instance [`OneOf`][torchio.OneOf]). Resets the shared - `applied_transforms` so that subsequent transforms accumulate as - a common suffix. - - Args: - histories: One history list per batch element. - """ - if len(histories) != self.batch_size: - msg = ( - f"Expected {self.batch_size} per-element histories," - f" got {len(histories)}" - ) - raise ValueError(msg) - self._per_element_history = [list(history) for history in histories] - self.applied_transforms = [] - - @classmethod - def from_subjects(cls, subjects: list[Any]) -> Self: - """Stack a list of subjects into a batch. - - Args: - subjects: List of `Subject` instances. - """ - from .subject import Subject - - if not subjects: - msg = "Cannot create batch from empty list" - raise ValueError(msg) - - # Collect image names and types from the first subject - first: Subject = subjects[0] - image_names = list(first.images.keys()) - - # Stack images - images: dict[str, ImagesBatch] = {} - for name in image_names: - img_list = [sub.images[name] for sub in subjects] - images[name] = ImagesBatch.from_images(img_list) - - # Collect metadata (non-image, non-annotation entries) - metadata: dict[str, list[Any]] = {} - for key in first.metadata: - metadata[key] = [sub.metadata[key] for sub in subjects] - - return cls(images, metadata=metadata) - - @property - def batch_size(self) -> int: - """Number of samples in the batch.""" - first = next(iter(self._images.values())) - return first.batch_size - - @property - def images(self) -> dict[str, ImagesBatch]: - """Dict of named image batches.""" - return self._images - - @property - def metadata(self) -> dict[str, list[Any]]: - """Metadata lists (one value per sample).""" - return self._metadata - - @property - def device(self) -> torch.device: - """Device of the batch data.""" - first = next(iter(self._images.values())) - return first.device - - def to(self, *args: Any, **kwargs: Any) -> Self: - """Move all data to a device and/or cast dtype.""" - for batch in self._images.values(): - batch.to(*args, **kwargs) - return self - - def __getitem__(self, key: str) -> ImagesBatch: - """Get a named image batch.""" - return self._images[key] - - def __getattr__(self, name: str) -> ImagesBatch: - """Attribute-style access to image batches.""" - if name.startswith("_"): - raise AttributeError(name) - if name in self._images: - return self._images[name] - msg = f"SubjectsBatch has no attribute {name!r}" - raise AttributeError(msg) - - def unbatch(self) -> list[Any]: - """Split the batch back into individual Subjects. - - Per-instance transform history is sliced so that each subject - receives only its own sampled parameters; transforms that were - gated out for an element (per-element probability) are omitted - from that subject's history. - """ - from .subject import Subject - - n = self.batch_size - subjects = [] - for i in range(n): - kwargs: dict[str, Any] = {} - for name, img_batch in self._images.items(): - kwargs[name] = img_batch[i] - for key, values in self._metadata.items(): - kwargs[key] = values[i] - sub = Subject(**kwargs) - suffix = _slice_history(self.applied_transforms, i) - if self._per_element_history is not None: - sub.applied_transforms = list(self._per_element_history[i]) + suffix - else: - sub.applied_transforms = suffix - subjects.append(sub) - return subjects - - def __len__(self) -> int: - return self.batch_size - - def adopt_history(self, source: SubjectsBatch, subjects: list[Any]) -> None: - """Carry transform history from *source* after rebuilding the batch. - - Used by code that unbatches, processes, and re-stacks subjects - (for example the MONAI and Cornucopia adapters). Preserves a - per-element history if *source* had one, otherwise copies the - shared history. - - Args: - source: The batch the subjects were unbatched from. - subjects: The processed subjects, in batch order. - """ - if source._per_element_history is not None: - self.set_per_element_history([s.applied_transforms for s in subjects]) - else: - self.applied_transforms = list(source.applied_transforms) - - def clear_history(self) -> None: - """Remove all applied transform records, including per-element ones.""" - self.applied_transforms = [] - self._per_element_history = None - - def get_inverse_transform(self, **kwargs: Any) -> Any: - """Build a transform that inverts the recorded history. - - Raises: - RuntimeError: If the batch carries per-element histories (from - a per-instance `OneOf`/`SomeOf`), since a single batch - inverse is ambiguous. Call `apply_inverse_transform` - (which inverts each element) or `unbatch()` and invert - each subject. - """ - if self._per_element_history is not None: - msg = ( - "This batch has per-element transform histories from a" - " per-instance OneOf/SomeOf, so a single batch inverse is" - " ambiguous. Call apply_inverse_transform() (which inverts" - " each element) or unbatch() and invert each subject." - ) - raise RuntimeError(msg) - return super().get_inverse_transform(**kwargs) - - def apply_inverse_transform(self, **kwargs: Any) -> SubjectsBatch: - """Apply the inverse of the recorded history. - - When the batch carries per-element histories, each element is - inverted independently and the results are re-stacked. - - Args: - **kwargs: Forwarded to `get_inverse_transform`. - - Returns: - A batch with the transforms undone. - """ - if self._per_element_history is not None: - inverted = [s.apply_inverse_transform(**kwargs) for s in self.unbatch()] - return type(self).from_subjects(inverted) - return super().apply_inverse_transform(**kwargs) - - def __repr__(self) -> str: - names = ", ".join(self._images.keys()) - return f"SubjectsBatch(batch_size={self.batch_size}, images=[{names}])" - - -# Alias for radiology users (see Subject/Study note in subject.py). -StudiesBatch = SubjectsBatch - - -def _slice_params( - params: dict[str, Any], - index: int, - batched_keys: list[str], -) -> dict[str, Any]: - """Slice a per-instance params dict down to a single element. - - Args: - params: The batch-level parameter dict. - index: The batch element to extract. - batched_keys: Names of the params that hold one value per - element. - - Returns: - A new params dict with per-element values resolved and the - internal bookkeeping keys removed. - """ - sliced: dict[str, Any] = {} - for key, value in params.items(): - if key in _BATCH_META_KEYS: - continue - if key in batched_keys and isinstance(value, list): - sliced[key] = value[index] - else: - sliced[key] = value - return sliced - - -def _slice_history(history: list[Any], index: int) -> list[Any]: - """Build the per-subject transform history for batch element *index*. - - Batch-shared traces are copied unchanged. Per-instance traces are - sliced to the element's own parameters, and traces whose per-element - keep mask excludes this element are dropped. - - Args: - history: The batch-level list of `AppliedTransform` records. - index: The batch element whose history to build. - - Returns: - The list of `AppliedTransform` records for the element. - """ - sliced: list[Any] = [] - for trace in history: - params = getattr(trace, "params", None) - if not isinstance(params, dict) or "_batched_keys" not in params: - sliced.append(trace) - continue - expected_size = params.get("_batch_size") - if expected_size is not None and not 0 <= index < expected_size: - msg = ( - f"Cannot extract per-instance history for element {index}:" - f" the transform was recorded for a batch of size" - f" {expected_size}" - ) - raise IndexError(msg) - keep = params.get("_keep") - if keep is not None and not keep[index]: - continue - batched_keys = params["_batched_keys"] - new_params = _slice_params(params, index, batched_keys) - sliced.append(dataclasses.replace(trace, params=new_params)) - return sliced diff --git a/src/torchio/data/bboxes.py b/src/torchio/data/bboxes.py deleted file mode 100644 index 4228fe790..000000000 --- a/src/torchio/data/bboxes.py +++ /dev/null @@ -1,476 +0,0 @@ -"""3D bounding boxes with flexible axis conventions. - -Inspired by `torchvision.tv_tensors.BoundingBoxes`, extended to 3D with -support for arbitrary voxel and anatomical axis orderings. -""" - -from __future__ import annotations - -from enum import Enum -from typing import Any - -import numpy as np -import numpy.typing as npt -import torch -from torch import Tensor -from typing_extensions import Self - -from .affine import AffineMatrix -from .axes import AxesType -from .axes import axes_type -from .axes import get_axis_mapping -from .axes import validate_axes - - -class Representation(Enum): - """How the six columns of a bounding box are interpreted. - - Attributes: - CORNERS: Two corners: $(a_1, b_1, c_1, a_2, b_2, c_2)$. - CENTER_SIZE: Center + size: $(a_c, b_c, c_c, s_a, s_b, s_c)$. - """ - - CORNERS = "corners" - CENTER_SIZE = "center_size" - - -class BoundingBoxFormat: - """Format specification for 3D bounding boxes. - - A format is defined by two components: - - - **axes**: a 3-character string specifying the coordinate system. - Voxel axes are permutations of `"IJK"`. - Anatomical axes use one letter from each pair - `{R, L}`, `{A, P}`, `{S, I}` (e.g., `"RAS"`, `"LPI"`). - - **representation**: either *corners* (two opposite corners) or - *center_size* (center point + extent along each axis). - - Args: - axes: 3-character axis string. - representation: How the 6 values encode the box. - - Examples: - >>> from torchio.data.bboxes import BoundingBoxFormat, Representation - >>> BoundingBoxFormat("IJK", Representation.CORNERS) - BoundingBoxFormat(axes='IJK', representation='corners') - >>> BoundingBoxFormat("RAS", "center_size") - BoundingBoxFormat(axes='RAS', representation='center_size') - """ - - # Predefined convenience formats, set after the class body. - IJKIJK: BoundingBoxFormat - IJKWHD: BoundingBoxFormat - - __slots__ = ("_axes", "_representation") - - def __init__( - self, - axes: str, - representation: Representation | str = Representation.CORNERS, - ) -> None: - self._axes = validate_axes(axes) - if isinstance(representation, str): - representation = Representation(representation) - self._representation = representation - - @property - def axes(self) -> str: - """3-character axis string (e.g., `'IJK'`, `'RAS'`).""" - return self._axes - - @property - def representation(self) -> Representation: - """Corners or center-size.""" - return self._representation - - def __eq__(self, other: object) -> bool: - if not isinstance(other, BoundingBoxFormat): - return NotImplemented - return ( - self._axes == other._axes and self._representation == other._representation - ) - - def __hash__(self) -> int: - return hash((self._axes, self._representation)) - - def __repr__(self) -> str: - return ( - f"BoundingBoxFormat(axes={self._axes!r}," - f" representation={self._representation.value!r})" - ) - - -# --- Predefined formats --- -BoundingBoxFormat.IJKIJK = BoundingBoxFormat("IJK", Representation.CORNERS) -BoundingBoxFormat.IJKWHD = BoundingBoxFormat("IJK", Representation.CENTER_SIZE) - - -# --- Representation conversion helpers --- - - -def _corners_to_center_size(data: Tensor) -> Tensor: - a1, b1, c1, a2, b2, c2 = data.unbind(-1) - ac = (a1 + a2) / 2 - bc = (b1 + b2) / 2 - cc = (c1 + c2) / 2 - sa = a2 - a1 - sb = b2 - b1 - sc = c2 - c1 - return torch.stack([ac, bc, cc, sa, sb, sc], dim=-1) - - -def _center_size_to_corners(data: Tensor) -> Tensor: - ac, bc, cc, sa, sb, sc = data.unbind(-1) - a1 = ac - sa / 2 - b1 = bc - sb / 2 - c1 = cc - sc / 2 - a2 = ac + sa / 2 - b2 = bc + sb / 2 - c2 = cc + sc / 2 - return torch.stack([a1, b1, c1, a2, b2, c2], dim=-1) - - -# --- Axis reordering helpers --- - - -def _permute_corners( - data: Tensor, - perm: tuple[int, int, int], - flips: tuple[bool, bool, bool], -) -> Tensor: - """Permute and flip columns of a corners-format (N, 6) tensor.""" - p0, p1, p2 = perm - # Separate the two triplets. - corner1 = data[:, :3][:, [p0, p1, p2]] - corner2 = data[:, 3:][:, [p0, p1, p2]] - # Apply flips (negation). - for col, flip in enumerate(flips): - if flip: - c1 = -corner1[:, col].clone() - c2 = -corner2[:, col].clone() - # After negation the min/max may swap, so ensure corner1 < corner2. - corner1[:, col] = torch.min(c1, c2) - corner2[:, col] = torch.max(c1, c2) - return torch.cat([corner1, corner2], dim=-1) - - -def _permute_center_size( - data: Tensor, - perm: tuple[int, int, int], - flips: tuple[bool, bool, bool], -) -> Tensor: - """Permute and flip columns of a center-size-format (N, 6) tensor.""" - p0, p1, p2 = perm - center = data[:, :3][:, [p0, p1, p2]] - size = data[:, 3:][:, [p0, p1, p2]] - # Flips negate the center; sizes stay positive. - for col, flip in enumerate(flips): - if flip: - center[:, col] = -center[:, col] - return torch.cat([center, size], dim=-1) - - -# --- Voxel ↔ anatomical helpers --- - - -def _ijk_corners_to_world( - data: Tensor, - affine: AffineMatrix, -) -> Tensor: - """Convert (N, 6) corners from IJK voxel to world (RAS) coordinates.""" - c1 = data[:, :3] - c2 = data[:, 3:] - w1 = affine.apply(c1).to(torch.float32) - w2 = affine.apply(c2).to(torch.float32) - # After affine, min/max might swap, so normalize. - lo = torch.min(w1, w2) - hi = torch.max(w1, w2) - return torch.cat([lo, hi], dim=-1) - - -def _world_corners_to_ijk( - data: Tensor, - affine: AffineMatrix, -) -> Tensor: - """Convert (N, 6) corners from world (RAS) to IJK voxel coordinates.""" - inv = affine.inverse() - c1 = data[:, :3] - c2 = data[:, 3:] - v1 = inv.apply(c1).to(torch.float32) - v2 = inv.apply(c2).to(torch.float32) - lo = torch.min(v1, v2) - hi = torch.max(v1, v2) - return torch.cat([lo, hi], dim=-1) - - -class BoundingBoxes: - r"""3D bounding boxes with flexible axis conventions. - - Inspired by `torchvision.tv_tensors.BoundingBoxes`, extended to 3D. - One instance holds $N$ boxes, each a 6-element vector whose meaning - is determined by the - [`format`][torchio.data.bboxes.BoundingBoxFormat]. - - Args: - data: $(N, 6)$ tensor or array. - format: Interpretation of the 6 columns. - labels: Optional $(N,)$ integer tensor of class labels per box. - affine: $4 \times 4$ affine matrix. Identity if not given. - metadata: Arbitrary metadata dict. - - Examples: - >>> import torch, torchio as tio - >>> boxes = tio.BoundingBoxes( - ... torch.tensor([[10, 20, 30, 50, 60, 70]]), - ... format=tio.BoundingBoxFormat.IJKIJK, - ... ) - >>> boxes.num_boxes - 1 - """ - - def __init__( - self, - data: Tensor | npt.ArrayLike, - *, - format: BoundingBoxFormat, - labels: Tensor | None = None, - affine: AffineMatrix | npt.ArrayLike | None = None, - metadata: dict[str, Any] | None = None, - ) -> None: - self._data = self._parse_data(data) - self._format = format - self._labels = self._parse_labels(labels, self._data.shape[0]) - self._affine = self._parse_affine(affine) - self._metadata: dict[str, Any] = dict(metadata) if metadata else {} - - # --- Parsing --- - - @staticmethod - def _parse_data(data: Tensor | npt.ArrayLike) -> Tensor: - if not isinstance(data, Tensor): - data = torch.as_tensor(np.asarray(data), dtype=torch.float32) - if data.ndim != 2 or data.shape[1] != 6: - msg = f"BoundingBoxes must have shape (N, 6), got {tuple(data.shape)}" - raise ValueError(msg) - return data - - @staticmethod - def _parse_labels(labels: Tensor | None, n: int) -> Tensor | None: - if labels is None: - return None - if labels.shape[0] != n: - msg = f"Expected {n} labels, got {labels.shape[0]}" - raise ValueError(msg) - return labels - - @staticmethod - def _parse_affine(affine: AffineMatrix | npt.ArrayLike | None) -> AffineMatrix: - if affine is None: - return AffineMatrix() - if isinstance(affine, AffineMatrix): - return affine - return AffineMatrix(affine) - - # --- Properties --- - - @property - def data(self) -> Tensor: - """$(N, 6)$ tensor of bounding box coordinates.""" - return self._data - - @property - def format(self) -> BoundingBoxFormat: - """Interpretation of the 6 columns.""" - return self._format - - @property - def labels(self) -> Tensor | None: - """$(N,)$ integer labels, or `None`.""" - return self._labels - - @property - def affine(self) -> AffineMatrix: - r"""$4 \times 4$ affine mapping voxel to world coordinates.""" - return self._affine - - @property - def metadata(self) -> dict[str, Any]: - """Arbitrary metadata dict.""" - return self._metadata - - @property - def num_boxes(self) -> int: - """Number of bounding boxes.""" - return self._data.shape[0] - - @property - def device(self) -> torch.device: - """Device the bounding box data resides on.""" - return self._data.device - - def to(self, *args: Any, **kwargs: Any) -> Self: - """Move bounding box data to a device and/or cast to a dtype. - - Returns: - `self` (modified in-place). - """ - self._data = self._data.to(*args, **kwargs) - if self._labels is not None: - self._labels = self._labels.to(*args, **kwargs) - return self - - # --- Methods --- - - def to_format(self, format: BoundingBoxFormat) -> Self: - """Convert to a different bounding box format. - - Handles representation changes (corners ↔ center-size), axis - permutations within the same type, and voxel ↔ anatomical - conversions (using the stored affine). - - Args: - format: Target format. - - Returns: - New `BoundingBoxes` in the target format. - """ - if format == self._format: - return self._clone(format=format) - - src_axes = self._format.axes - tgt_axes = format.axes - src_repr = self._format.representation - tgt_repr = format.representation - - src_type = axes_type(src_axes) - tgt_type = axes_type(tgt_axes) - - # Step 1: normalise to corners in source axes. - data = self._data - if src_repr == Representation.CENTER_SIZE: - data = _center_size_to_corners(data) - - # Step 2: axis conversion (now in corners). - if src_axes != tgt_axes: - if src_type == tgt_type: - # Same family: permute + flip. - perm, flips = get_axis_mapping(src_axes, tgt_axes) - data = _permute_corners(data, perm, flips) - else: - # Cross-type: go through world coordinates. - data = self._cross_type_corners( - data, - src_axes, - src_type, - tgt_axes, - tgt_type, - ) - - # Step 3: convert to target representation. - if tgt_repr == Representation.CENTER_SIZE: - data = _corners_to_center_size(data) - - return self._clone(data=data, format=format) - - def new_like( - self, - *, - data: Tensor | npt.ArrayLike, - labels: Tensor | None = None, - affine: AffineMatrix | npt.ArrayLike | None = None, - ) -> Self: - """Create new BoundingBoxes with the same format and metadata. - - Args: - data: New $(N, 6)$ coordinates. - labels: New labels. If `None`, no labels. - affine: New affine. If `None`, uses `self.affine`. - """ - new_affine = ( - self._parse_affine(affine) if affine is not None else self._affine.clone() - ) - return type(self)( - data, - format=self._format, - labels=labels, - affine=new_affine, - metadata=dict(self._metadata), - ) - - # --- Internal --- - - def _clone( - self, - *, - data: Tensor | None = None, - format: BoundingBoxFormat | None = None, - ) -> Self: - return type(self)( - data if data is not None else self._data.clone(), - format=format if format is not None else self._format, - labels=self._labels.clone() if self._labels is not None else None, - affine=self._affine.clone(), - metadata=dict(self._metadata), - ) - - def _cross_type_corners( - self, - data: Tensor, - src_axes: str, - src_type: AxesType, - tgt_axes: str, - tgt_type: AxesType, - ) -> Tensor: - """Convert corners between voxel and anatomical coordinate systems.""" - if src_type == AxesType.VOXEL: - # Voxel → RAS via affine, then optionally reorder/flip within - # anatomical. - # First normalise voxel order to IJK. - if src_axes != "IJK": - perm, _ = get_axis_mapping(src_axes, "IJK") - data = _permute_corners(data, perm, (False, False, False)) - # Apply affine to get RAS. - data = _ijk_corners_to_world(data, self._affine) - # The affine's orientation tells us what "world" actually is. - world_axes = "".join(self._affine.orientation) - if world_axes != tgt_axes: - perm, flips = get_axis_mapping(world_axes, tgt_axes) - data = _permute_corners(data, perm, flips) - else: - # Anatomical → voxel. - # First normalise to the affine's world system. - world_axes = "".join(self._affine.orientation) - if src_axes != world_axes: - perm, flips = get_axis_mapping(src_axes, world_axes) - data = _permute_corners(data, perm, flips) - # World → IJK via inverse affine. - data = _world_corners_to_ijk(data, self._affine) - # Reorder to target voxel axes if needed. - if tgt_axes != "IJK": - perm, _ = get_axis_mapping("IJK", tgt_axes) - data = _permute_corners(data, perm, (False, False, False)) - return data - - # --- Dunder --- - - def __len__(self) -> int: - return self.num_boxes - - def __repr__(self) -> str: - return ( - f"BoundingBoxes(num_boxes={self.num_boxes}," - f" axes={self._format.axes!r}," - f" representation={self._format.representation.value!r})" - ) - - def __deepcopy__(self, memo: dict) -> Self: - new = type(self)( - self._data.clone(), - format=self._format, - labels=self._labels.clone() if self._labels is not None else None, - affine=self._affine.clone(), - metadata=dict(self._metadata), - ) - memo[id(self)] = new - return new diff --git a/src/torchio/data/dataset.py b/src/torchio/data/dataset.py new file mode 100644 index 000000000..00f804cc7 --- /dev/null +++ b/src/torchio/data/dataset.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import copy +from collections.abc import Callable +from collections.abc import Iterable +from collections.abc import Sequence + +from torch.utils.data import Dataset + +from ..utils import get_subjects_from_batch +from .subject import Subject + + +class SubjectsDataset(Dataset[Subject]): + """Base TorchIO dataset. + + Reader of 3D medical images that directly inherits from the PyTorch + [`Dataset`][torch.utils.data.Dataset]. It can be used with a + [`SubjectsLoader`][torchio.SubjectsLoader] for efficient loading and + augmentation. It receives a list of instances of [`Subject`][torchio.Subject] + and an optional transform applied to the volumes after loading. + + Args: + subjects: List of instances of [`Subject`][torchio.Subject]. + transform: An instance of [`Transform`][torchio.transforms.Transform] + that will be applied to each subject. + load_getitem: Load all subject images before returning it in + `__getitem__()`. Set it to `False` if some of the images will + not be needed during training. + + Examples: + >>> import torchio as tio + >>> subject_a = tio.Subject( + ... t1=tio.ScalarImage('t1.nrrd',), + ... t2=tio.ScalarImage('t2.mha',), + ... label=tio.LabelMap('t1_seg.nii.gz'), + ... age=31, + ... name='Fernando Perez', + ... ) + >>> subject_b = tio.Subject( + ... t1=tio.ScalarImage('colin27_t1_tal_lin.minc',), + ... t2=tio.ScalarImage('colin27_t2_tal_lin_dicom/',), + ... label=tio.LabelMap('colin27_seg1.nii.gz'), + ... age=56, + ... name='Colin Holmes', + ... ) + >>> subjects_list = [subject_a, subject_b] + >>> transforms = [ + ... tio.RescaleIntensity(out_min_max=(0, 1)), + ... tio.RandomAffine(), + ... ] + >>> transform = tio.Compose(transforms) + >>> subjects_dataset = tio.SubjectsDataset(subjects_list, transform=transform) + >>> subject = subjects_dataset[0] + + !!! tip "Iterating without loading" + To quickly iterate over the subjects without loading the images, + use `dry_iter()`. + """ + + def __init__( + self, + subjects: Sequence[Subject], + transform: Callable[[Subject], Subject] | None = None, + load_getitem: bool = True, + ): + self._parse_subjects_list(subjects) + self._subjects = subjects + self._transform: Callable[[Subject], Subject] | None + self.set_transform(transform) + self.load_getitem = load_getitem + + def __len__(self): + return len(self._subjects) + + def __getitem__(self, index: int) -> Subject: + try: + index = int(index) + except (RuntimeError, TypeError) as err: + message = ( + f'Index "{index}" must be int or compatible dtype,' + f' but an object of type "{type(index)}" was passed' + ) + raise ValueError(message) from err + + subject = self._subjects[index] + subject = copy.deepcopy(subject) # cheap since images not loaded yet + if self.load_getitem: + subject.load() + + # Apply transform (this is usually the bottleneck) + if self._transform is not None: + subject = self._transform(subject) + return subject + + @classmethod + def from_batch(cls, batch: dict) -> SubjectsDataset: + """Instantiate a dataset from a batch generated by a data loader. + + Args: + batch: Dictionary generated by a data loader, containing data that + can be converted to instances of [Subject][.torchio.Subject]. + """ + subjects: list[Subject] = get_subjects_from_batch(batch) + return cls(subjects) + + def dry_iter(self) -> Sequence[Subject]: + """Return the internal list of subjects. + + This can be used to iterate over the subjects without loading the data + and applying any transforms. + + Examples: + >>> names = [subject.name for subject in dataset.dry_iter()] + """ + return self._subjects + + def set_transform( + self, + transform: Callable[[Subject], Subject] | None, + ) -> None: + """Set the `transform` attribute. + + Args: + transform: Callable object, typically an subclass of + [`torchio.transforms.Transform`][torchio.transforms.Transform]. + """ + if transform is not None and not callable(transform): + message = ( + 'The transform must be a callable object,' + f' but it has type {type(transform)}' + ) + raise ValueError(message) + self._transform = transform + + @staticmethod + def _parse_subjects_list(subjects_list: Iterable[Subject]) -> None: + # Check that it's an iterable + try: + iter(subjects_list) + except TypeError as e: + message = f'Subject list must be an iterable, not {type(subjects_list)}' + raise TypeError(message) from e + + # Check that it's not empty + if not subjects_list: + raise ValueError('Subjects list is empty') + + # Check each element + for subject in subjects_list: + if not isinstance(subject, Subject): + message = ( + 'Subjects list must contain instances of torchio.Subject,' + f' not "{type(subject)}"' + ) + raise TypeError(message) diff --git a/src/torchio/data/image.py b/src/torchio/data/image.py index 76837653c..fd7dabc80 100644 --- a/src/torchio/data/image.py +++ b/src/torchio/data/image.py @@ -1,1261 +1,1034 @@ -"""Image classes for TorchIO.""" - from __future__ import annotations -import copy -import io +import warnings +from collections import Counter from collections.abc import Callable +from collections.abc import Sequence from pathlib import Path +from typing import TYPE_CHECKING from typing import Any -from typing import TypeVar +from typing import TypeGuard +from typing import cast +from typing import overload +import humanize import nibabel as nib -import nibabel.spatialimages import numpy as np -import numpy.typing as npt import SimpleITK as sitk import torch -from einops import rearrange -from torch import Tensor -from typing_extensions import Self - -from ..external.imports import get_niizarr -from ..types import TypeImageData -from ..types import TypeSpatialShape -from ..types import TypeTensorShape -from .affine import AffineMatrix -from .backends import BackendRequest -from .backends import ImageDataBackend -from .backends import LazyReader -from .backends import NibabelBackend -from .backends import TensorBackend -from .backends import normalize_index -from .backends import resolve_backend -from .bboxes import BoundingBoxes -from .invertible import Invertible -from .io import ImageSource -from .io import default_reader -from .io import is_nifti_zarr -from .io import is_remote_nifti_zarr -from .io import resolve_source -from .points import Points - -_AnnotationType = TypeVar("_AnnotationType", Points, BoundingBoxes) - - -def _in_jupyter() -> bool: - """Check whether we are running inside a Jupyter notebook.""" - try: - from IPython import get_ipython - - shell = get_ipython() - return shell is not None and shell.__class__.__name__ == "ZMQInteractiveShell" - except ImportError: - return False - - -def _resolve_media_path( - output_path: str | Path | None, - *, - suffix: str, -) -> Path: - """Resolve an optional output path for media files. - - Args: - output_path: User-provided path, or `None`. - suffix: File extension (e.g., `".gif"`, `".mp4"`). - - Returns: - Resolved [`Path`][pathlib.Path]. - - Raises: - ValueError: If *output_path* is `None` outside Jupyter. - """ - if output_path is not None: - return Path(output_path) - if _in_jupyter(): - import tempfile - - with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: - return Path(f.name) - msg = ( - f"output_path is required outside Jupyter notebooks. " - f"Pass a path ending in {suffix!r}." - ) - raise ValueError(msg) - - -def _backend_label(backend: object | None) -> str: - """Short label for the backend type, used in `__repr__`.""" - if backend is None: - return "unknown" - name = type(backend).__name__ - match name: - case n if "Nibabel" in n: - return "NIfTI" - case n if "Zarr" in n: - return "NIfTI-Zarr" - case n if "Tensor" in n: - return "Tensor" - case _: - return name - - -class Image(Invertible): - r"""Base image class. - - TorchIO images are - [lazy loaders](https://en.wikipedia.org/wiki/Lazy_loading): - data is only read from disk when first accessed. - - Use [`ScalarImage`][torchio.ScalarImage] for intensity data and - [`LabelMap`][torchio.LabelMap] for segmentations. - Transforms use `isinstance` checks to decide behavior (e.g., nearest- - neighbor interpolation for [`LabelMap`][torchio.LabelMap]). - - The constructor accepts many source types and dispatches - automatically: - - | Source type | Behavior | - |---|---| - | `str`, `Path`, URL, `OpenFile`, file-like | Lazy load from file | - | `torch.Tensor`, `np.ndarray` | Eager, in-memory | - | `nib.Nifti1Image` | Lazy via `NibabelBackend` | - | `sitk.Image` | Eager, converted to tensor | - | `zarr.abc.store.Store` | Lazy via `zarr2nii` + `NibabelBackend` | - | `bytes`, `io.BytesIO` | Decoded via temp file | - | `None` (default) | Empty image, set data later | +from deprecated import deprecated +from nibabel.affines import apply_affine + +from ..constants import AFFINE +from ..constants import DATA +from ..constants import INTENSITY +from ..constants import LABEL +from ..constants import PATH +from ..constants import STEM +from ..constants import TYPE +from ..types import TypeAffineMatrix +from ..types import TypeData +from ..types import TypeDataAffine +from ..types import TypeDirection3D +from ..types import TypeImageDataAffine +from ..types import TypeImageTensor +from ..types import TypePath +from ..types import TypeQuartetInt +from ..types import TypeSlice +from ..types import TypeTripletFloat +from ..types import TypeTripletInt +from ..utils import get_stem +from ..utils import guess_external_viewer +from ..utils import is_iterable +from ..utils import to_tuple +from .io import check_uint_to_int +from .io import ensure_4d +from .io import get_rotation_and_spacing_from_affine +from .io import get_sitk_metadata_from_ras_affine +from .io import nib_to_sitk +from .io import read_affine +from .io import read_image +from .io import read_shape +from .io import sitk_to_nib +from .io import write_image + +if TYPE_CHECKING: + from matplotlib.figure import Figure + +PROTECTED_KEYS = DATA, AFFINE, TYPE, PATH, STEM +TypeBound = tuple[float, float] +TypeBounds = tuple[TypeBound, TypeBound, TypeBound] + +deprecation_message = ( + 'Setting the image data with the property setter is deprecated. Use the' + ' set_data() method instead' +) + + +class Image(dict[str, object]): + r"""TorchIO image. + + For information about medical image orientation, check out [NiBabel docs](https://nipy.org/nibabel/image_orientation.html), + the [3D Slicer wiki](https://www.slicer.org/wiki/Coordinate_systems), [Graham Wideman's website](http://www.grahamwideman.com/gw/brain/orientation/orientterms.htm), [FSL docs](https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/Orientation%20Explained) or + [SimpleITK docs](https://simpleitk.readthedocs.io/en/master/fundamentalConcepts.html). Args: - source: Image data or path. See the table above. - reader: Callable that takes a path and returns a tuple - `(tensor, affine_array)`. Overrides the default reader. - Only used for file-path sources. - reader_kwargs: Extra keyword arguments forwarded to the reader - function. For the default reader these are passed to - `nibabel.load()` or `SimpleITK.ReadImage()`. - affine: $4 \times 4$ affine matrix or - [`AffineMatrix`][torchio.AffineMatrix] instance. If given, overrides - the affine read from the file. - channels_last: If `True`, the tensor is assumed to have - shape $(I, J, K, C)$ and will be permuted to - $(C, I, J, K)$. Only used for tensor sources. - suffix: File suffix hint (e.g., `".nii.gz"`). Used for - file-like and bytes sources. - points: Named sets of [`Points`][torchio.Points] attached to - this image. - bounding_boxes: Named sets of - [`BoundingBoxes`][torchio.BoundingBoxes] attached to this - image. - **kwargs: Arbitrary metadata, accessible via attribute or - dict-style lookup (e.g., `protocol="MPRAGE"`). + path: Path to a file or sequence of paths to files that can be read by + [SimpleITK](https://simpleitk.org/doxygen/latest/html/namespaceitk_1_1simple.html) + or [nibabel][nibabel], or to a directory containing + DICOM files. If `tensor` is given, the data in + `path` will not be read. + If a sequence of paths is given, data + will be concatenated on the channel dimension so spatial + dimensions must match. + type: Type of image, such as `torchio.INTENSITY` or + `torchio.LABEL`. This will be used by the transforms to + decide whether to apply an operation, or which interpolation to use + when resampling. For example, [preprocessing](https://docs.torchio.org/transforms/preprocessing.html#intensity) and [augmentation](https://docs.torchio.org/transforms/augmentation.html#intensity) + intensity transforms will only be applied to images with type + `torchio.INTENSITY`. Spatial transforms will be applied to + all types, and nearest neighbor interpolation is always used to + resample images with type `torchio.LABEL`. + The type `torchio.SAMPLING_MAP` may be used with instances of + [`WeightedSampler`](../../patches/training/#torchio.data.WeightedSampler). + tensor: If `path` is not given, `tensor` must be a 4D + [`torch.Tensor`][torch.Tensor] or NumPy array with dimensions + $(C, W, H, D)$. + affine: $4 \times 4$ matrix to convert voxel coordinates to world + coordinates. If `None`, an identity matrix will be used. See the + [NiBabel docs on coordinates](https://nipy.org/nibabel/coordinate_systems.html#the-affine-matrix-as-a-transformation-between-spaces) for more information. + check_nans: If `True`, issues a warning if NaNs are found + in the image. If `False`, images will not be checked for the + presence of NaNs. + reader: Callable object that takes a path and returns a 4D tensor and a + 2D, $4 \times 4$ affine matrix. This can be used if your data + is saved in a custom format, such as `.npy` (see example below). + If the affine matrix is `None`, an identity matrix will be used. + **kwargs: Items that will be added to the image dictionary, e.g. + acquisition parameters or image ID. + verify_path: If `True`, the path will be checked to see if it exists. If + `False`, the path will not be verified. This is useful when it is + expensive to check the path, e.g., when reading a large dataset from a + mounted drive. + + TorchIO images are [lazy loaders](https://en.wikipedia.org/wiki/Lazy_loading), i.e. the data is only loaded from disk + when needed. Examples: >>> import torchio as tio - >>> image = tio.ScalarImage("t1.nii.gz") # from path (lazy) - >>> image = tio.ScalarImage(torch.randn(1, 256, 256, 176)) # from tensor - >>> image = tio.ScalarImage(nifti_image) # from nibabel (lazy) - """ + >>> import numpy as np + >>> image = tio.ScalarImage('t1.nii.gz') # subclass of Image + >>> image # not loaded yet + ScalarImage(path: t1.nii.gz; type: intensity) + >>> times_two = 2 * image.data # data is loaded and cached here + >>> image + ScalarImage(shape: (1, 256, 256, 176); spacing: (1.00, 1.00, 1.00); orientation: PIR+; memory: 44.0 MiB; type: intensity) + >>> image.save('doubled_image.nii.gz') + >>> def numpy_reader(path): + ... data = np.load(path).as_type(np.float32) + ... affine = np.eye(4) + ... return data, affine + >>> image = tio.ScalarImage('t1.npy', reader=numpy_reader) - #: Source types accepted by the constructor. - ImageInput = ( - ImageSource # str | Path | IOBase | OpenFile - | Tensor - | np.ndarray - | nib.Nifti1Image - | sitk.Image - | bytes - | io.BytesIO - # | zarr.abc.store.Store (optional, accepted at runtime) - | None - ) + """ def __init__( self, - source: ImageInput = None, - *, - reader: Callable[[Path], tuple[TypeImageData, np.ndarray]] | None = None, - reader_kwargs: dict[str, Any] | None = None, - affine: AffineMatrix | npt.ArrayLike | None = None, - channels_last: bool = False, - suffix: str | None = None, - points: dict[str, Points] | None = None, - bounding_boxes: dict[str, BoundingBoxes] | None = None, - **kwargs: Any, + path: TypePath | Sequence[TypePath] | None = None, + type: str | None = None, # noqa: A002 + tensor: TypeData | None = None, + affine: TypeData | None = None, + check_nans: bool = False, # removed by ITK by default + reader: Callable[[TypePath], TypeDataAffine] = read_image, + verify_path: bool = True, + **kwargs: object, ): - # Common state shared by all source types. - self._reader = reader or default_reader - self._reader_kwargs: dict[str, Any] = dict(reader_kwargs or {}) - self._channels_last = channels_last - self._metadata: dict[str, Any] = dict(kwargs) - self._data: Tensor | None = None - self._backend: ImageDataBackend | None = None - self._path: Path | None = None - self._remote_zarr_uri: str | None = None - self._zarr_store: Any = None - self._affine: AffineMatrix | None = ( - self._parse_affine(affine) if affine is not None else None - ) - self._points = self._parse_annotations(points, "Points") - self._bounding_boxes = self._parse_annotations( - bounding_boxes, - "BoundingBoxes", - ) - self.applied_transforms: list[Any] = [] + self.check_nans = check_nans + self.reader = reader + + if type is None: + warnings.warn( + 'Not specifying the image type is deprecated and will be' + ' mandatory in the future. You can probably use' + ' tio.ScalarImage or tio.LabelMap instead', + FutureWarning, + stacklevel=2, + ) + type = INTENSITY # noqa: A001 + + if path is None and tensor is None: + raise ValueError('A value for path or tensor must be given') + self._loaded = False + + tensor = self._parse_tensor(tensor) + affine = self._parse_affine(affine) + if tensor is not None: + self.set_data(tensor) + self.affine = affine + self._loaded = True + for key in PROTECTED_KEYS: + if key in kwargs: + message = f'Key "{key}" is reserved. Use a different one' + raise ValueError(message) + if 'channels_last' in kwargs: + message = ( + 'The "channels_last" keyword argument is deprecated after' + ' https://github.com/TorchIO-project/torchio/pull/685 and will be' + ' removed in the future' + ) + warnings.warn(message, FutureWarning, stacklevel=2) - # Dispatch based on source type. - self._dispatch_source( - source, affine=affine, channels_last=channels_last, suffix=suffix + super().__init__(**kwargs) + self.path: Path | list[Path] | None = self._parse_path( + path, + verify=verify_path, ) - def _dispatch_source( - self, - source: ImageInput, - *, - affine: AffineMatrix | npt.ArrayLike | None, - channels_last: bool, - suffix: str | None, - ) -> None: - """Route *source* to the appropriate init helper.""" - if source is None: - return - if isinstance(source, (Tensor, np.ndarray)): - self._init_from_tensor(source, affine=affine, channels_last=channels_last) - elif isinstance(source, nib.Nifti1Image): - self._init_from_nifti(source) - elif isinstance(source, sitk.Image): - self._init_from_sitk(source, affine=affine) - elif isinstance(source, (bytes, io.BytesIO)): - self._init_from_bytes(source, suffix=suffix or ".nii.gz") - elif isinstance(source, str) and is_remote_nifti_zarr(source): - self._remote_zarr_uri = source - elif self._is_zarr_store(source): - self._zarr_store = source + self[PATH] = '' if self.path is None else str(self.path) + self[STEM] = '' if self.path is None else get_stem(self.path) + self[TYPE] = type + + def __repr__(self): + properties = [] + properties.extend( + [ + f'shape: {self.shape}', + f'spacing: {self.get_spacing_string()}', + f'orientation: {self.orientation_str}+', + ] + ) + if self._loaded: + properties.append(f'dtype: {self.data.type()}') + natural = humanize.naturalsize(self.memory, binary=True) + properties.append(f'memory: {natural}') else: - # Path-like, URL, fsspec OpenFile, or file-like object. - self._path = resolve_source(source, suffix=suffix) - - # -- Private init helpers ------------------------------------------------- + properties.append(f'path: "{self.path}"') - def _init_from_tensor( - self, - tensor: TypeImageData | np.ndarray, - *, - affine: AffineMatrix | npt.ArrayLike | None, - channels_last: bool, - ) -> None: - parsed = self._parse_tensor(tensor) - if channels_last: - parsed = rearrange(parsed, "i j k c -> c i j k") - self._data = parsed - self._channels_last = False # already permuted - parsed_affine = self._parse_affine(affine) - self._affine = parsed_affine - self._backend = TensorBackend(self._data, affine=parsed_affine.data) - - def _init_from_nifti(self, nifti_image: nib.Nifti1Image) -> None: - affine_override = self._affine.data if self._affine is not None else None - self._backend = NibabelBackend(nifti_image, affine=affine_override) - - def _init_from_sitk( - self, - sitk_image: sitk.Image, - *, - affine: AffineMatrix | npt.ArrayLike | None, - ) -> None: - data = sitk.GetArrayFromImage(sitk_image) - n_components = sitk_image.GetNumberOfComponentsPerPixel() - data = data[np.newaxis] if n_components == 1 else np.moveaxis(data, -1, 0) - from .io import _numpy_to_tensor - - tensor = _numpy_to_tensor(data.copy()) - spacing = np.array(sitk_image.GetSpacing()) - origin = np.array(sitk_image.GetOrigin()) - direction = rearrange(np.array(sitk_image.GetDirection()), "(i j) -> i j", i=3) - affine_matrix = np.eye(4) - affine_matrix[:3, :3] = direction * spacing - affine_matrix[:3, 3] = origin - self._init_from_tensor( - tensor, - affine=affine if affine is not None else AffineMatrix(affine_matrix), - channels_last=False, - ) + properties = '; '.join(properties) + string = f'{self.__class__.__name__}({properties})' + return string - def _init_from_bytes( - self, - data: bytes | io.BytesIO, - *, - suffix: str, - ) -> None: - import tempfile - - if isinstance(data, io.BytesIO): - data = data.read() - with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp: - tmp.write(data) - tmp.flush() - tmp_path = Path(tmp.name) + def _repr_html_(self): try: - nii = nib.load(tmp_path) - if isinstance(nii, nib.Nifti1Image): - self._init_from_nifti(nii) - self.load() # materialize before temp file is deleted - return - # Non-NIfTI: fall back to SimpleITK - sitk_image = sitk.ReadImage(str(tmp_path)) - self._init_from_sitk(sitk_image, affine=self._affine) - finally: - tmp_path.unlink(missing_ok=True) + from matplotlib.figure import Figure + except ImportError: + return self.__repr__() - @staticmethod - def _is_zarr_store(obj: object) -> bool: - """Check if *obj* is a `zarr.abc.store.Store` without importing zarr.""" - try: - from zarr.abc.store import Store - except ImportError: # zarr not installed - return False - return isinstance(obj, Store) + fig = self.plot(return_fig=True, show=False) + assert isinstance(fig, Figure) - # -- Static helpers ------------------------------------------------------- + from ..visualization import _figure_to_html - @staticmethod - def _parse_tensor(tensor: Tensor | np.ndarray) -> Tensor: - if isinstance(tensor, np.ndarray): - from .io import _numpy_to_tensor + return _figure_to_html(fig) - tensor = _numpy_to_tensor(tensor.copy()) - if tensor.ndim != 4: - msg = f"Tensor must be 4D (C, I, J, K), got {tensor.ndim}D" - raise ValueError(msg) - return tensor + @overload + def __getitem__(self, item: str) -> object: ... - @staticmethod - def _parse_affine(affine: AffineMatrix | npt.ArrayLike | None) -> AffineMatrix: - if affine is None: - return AffineMatrix() - if isinstance(affine, AffineMatrix): - return affine - return AffineMatrix(affine) + @overload + def __getitem__(self, item: slice | int | tuple[object, ...]) -> Image: ... - @staticmethod - def _parse_annotations( - annotations: dict[str, _AnnotationType] | None, - type_name: str, - ) -> dict[str, _AnnotationType]: - """Validate and copy an annotation dict. + def __getitem__( + self, item: str | slice | int | tuple[object, ...] + ) -> object | Image: + if isinstance(item, (slice, int, tuple)): + return self._crop_from_slices(item) - Args: - annotations: Mapping of names to annotation objects, or `None`. - type_name: Expected class name (`"Points"` or - `"BoundingBoxes"`) used for validation and error messages. + if item in (DATA, AFFINE): + if item not in self: + self.load() + return super().__getitem__(item) + + def __array__(self): + return self.data.numpy() + + def __copy__(self): + extra_kwargs: dict[str, object] = {} + for key, value in self.items(): + if key in PROTECTED_KEYS: + continue + extra_kwargs[key] = value # should I copy? deepcopy? + new_image_class = type(self) + new_image = new_image_class( + path=self.path, + type=self.type, + tensor=self.data if self._loaded else None, + affine=self.affine if self._loaded else None, + check_nans=self.check_nans, + reader=self.reader, + **cast(dict[str, Any], extra_kwargs), + ) + return new_image - Returns: - A shallow copy of the dict, or an empty dict if *annotations* - is `None`. + @property + def data(self) -> TypeImageTensor: + """Tensor data (same as [Image.tensor][Image.tensor]).""" + value = self[DATA] + if not isinstance(value, torch.Tensor): + self.load() + value = self[DATA] + assert isinstance(value, torch.Tensor) + return value - Raises: - TypeError: If any value is not an instance of the expected class. - """ - if annotations is None: - return {} - expected_type = Points if type_name == "Points" else BoundingBoxes - for key, value in annotations.items(): - if not isinstance(value, expected_type): - msg = ( - f"Expected {type_name} for key {key!r}, got {type(value).__name__}" - ) - raise TypeError(msg) - return dict(annotations) + @data.setter + @deprecated(version='0.18.16', reason=deprecation_message) + def data(self, tensor: TypeData): + self.set_data(tensor) - def _deep_copy_annotations( - self, - ) -> tuple[dict[str, Points], dict[str, BoundingBoxes]]: - """Deep-copy both annotation dicts.""" - points_copy: dict[str, Points] = { - k: copy.deepcopy(v) for k, v in self._points.items() - } - bboxes_copy: dict[str, BoundingBoxes] = { - k: copy.deepcopy(v) for k, v in self._bounding_boxes.items() - } - return points_copy, bboxes_copy - - # --- Properties --- + def set_data(self, tensor: TypeData): + """Store a 4D tensor in the `data` key and attribute. - @property - def path(self) -> Path | None: - """Path to the image file, if any.""" - return self._path + Args: + tensor: 4D tensor with dimensions $(C, W, H, D)$. + """ + self[DATA] = self._parse_tensor(tensor, none_ok=False) + self._loaded = True @property - def is_loaded(self) -> bool: - """Whether the image data is loaded into memory.""" - return self._data is not None + def tensor(self) -> TypeImageTensor: + """Tensor data (same as [Image.data][Image.data]).""" + return self.data @property - def data(self) -> TypeImageData: - """Tensor data with shape (C, I, J, K). Triggers lazy load if needed.""" - if self._data is None: - self.load() - assert self._data is not None - return self._data + def affine(self) -> TypeAffineMatrix: + """Affine matrix to transform voxel indices into world coordinates.""" + # If path is a dir (probably DICOM), just load the data + # Same if it's a list of paths (used to create a 4D image) + # Finally, if we use a custom reader, SimpleITK probably won't be able + # to read the metadata, so we resort to loading everything into memory + is_custom_reader = self.reader is not read_image + if self._loaded or self._is_dir() or self._is_multipath() or is_custom_reader: + affine = self[AFFINE] + if not isinstance(affine, np.ndarray): + self.load() + affine = self[AFFINE] + else: + assert self.path is not None + assert isinstance(self.path, Path) + affine = read_affine(self.path) + assert isinstance(affine, np.ndarray) + return cast(TypeAffineMatrix, affine) + + @affine.setter + def affine(self, matrix: TypeData | None): + self[AFFINE] = self._parse_affine(matrix) @property - def affine(self) -> AffineMatrix: - """4x4 affine matrix mapping voxel indices to world coordinates.""" - if self._affine is None: - # Try existing backend first to avoid full data load - if self._backend is not None: - self._affine = AffineMatrix(self._backend.affine) - elif ( - self._remote_zarr_uri is not None - or self._zarr_store is not None - or (self._path is not None and self._reader_supports_lazy()) - ): - self._ensure_backend() - if self._backend is not None: - self._affine = AffineMatrix(self._backend.affine) - if self._affine is None: - self.load() - assert self._affine is not None - return self._affine + def type(self) -> str: # noqa: A003 + value = self[TYPE] + assert isinstance(value, str) + return value @property - def metadata(self) -> dict[str, Any]: - """Arbitrary metadata dict.""" - return self._metadata + def shape(self) -> TypeQuartetInt: + """Tensor shape as $(C, W, H, D)$.""" + custom_reader = self.reader is not read_image + multipath = self._is_multipath() + if isinstance(self.path, Path): + is_dir = self.path.is_dir() + shape: TypeQuartetInt + if self._loaded or custom_reader or multipath or is_dir: + channels, si, sj, sk = self.data.shape + shape = channels, si, sj, sk + else: + assert isinstance(self.path, (str, Path)) + shape = read_shape(self.path) + return shape @property - def dataobj(self) -> ImageDataBackend: - """Lazy data backend for advanced operations like slicing. + def spatial_shape(self) -> TypeTripletInt: + """Tensor spatial shape as $(W, H, D)$.""" + return self.shape[1:] - Returns the underlying backend without materializing the full tensor. - For NIfTI files this is a `NibabelBackend`; for NIfTI-Zarr files a - `ZarrBackend`; for in-memory images a `TensorBackend`. - """ - if self._backend is None: - self._ensure_backend() - assert self._backend is not None - return self._backend + def check_is_2d(self) -> None: + if not self.is_2d(): + message = f'Image is not 2D. Spatial shape: {self.spatial_shape}' + raise RuntimeError(message) @property - def shape(self) -> TypeTensorShape: - """Tensor shape as (C, I, J, K).""" - if self._data is not None: - c, si, sj, sk = self._data.shape - return (c, si, sj, sk) - if self._backend is not None: - s = self._backend.shape - return (int(s[0]), int(s[1]), int(s[2]), int(s[3])) - if self._remote_zarr_uri is not None or self._zarr_store is not None: - self._ensure_backend() - assert self._backend is not None - s = self._backend.shape - return (int(s[0]), int(s[1]), int(s[2]), int(s[3])) - if self._path is not None: - if not self._reader_supports_lazy(): - self.load() - return self.shape - # Try to create a lazy backend (NIfTI, Zarr, or custom reader) - self._ensure_backend() - if self._backend is not None: - s = self._backend.shape - return (int(s[0]), int(s[1]), int(s[2]), int(s[3])) - # Non-NIfTI: read shape from header via SimpleITK - return self._read_shape_sitk(self._path) - msg = "Cannot determine shape: no data or path" - raise RuntimeError(msg) + def height(self) -> int: + """Image height, if 2D.""" + self.check_is_2d() + return self.spatial_shape[1] @property - def spatial_shape(self) -> TypeSpatialShape: - """Spatial dimensions as (I, J, K).""" - return self.shape[1:] + def width(self) -> int: + """Image width, if 2D.""" + self.check_is_2d() + return self.spatial_shape[0] @property - def num_channels(self) -> int: - """Number of channels.""" - return self.shape[0] + def orientation(self) -> tuple[str, str, str]: + """Orientation codes.""" + return nib.orientations.aff2axcodes(self.affine) @property - def spacing(self) -> tuple[float, float, float]: - """Voxel spacing in mm, derived from the affine.""" - return self.affine.spacing + def orientation_str(self) -> str: + """Orientation as a string.""" + return ''.join(self.orientation) @property - def origin(self) -> tuple[float, float, float]: - """Center of the first voxel in world coordinates.""" - return self.affine.origin + def direction(self) -> TypeDirection3D: + _, _, direction = get_sitk_metadata_from_ras_affine( + self.affine, + lps=False, + ) + if len(direction) != 9: + raise RuntimeError(f'Expected a 3D direction, not {direction}') + return cast(TypeDirection3D, direction) @property - def memory(self) -> int: - """Number of bytes the tensor would occupy in RAM.""" - c, si, sj, sk = self.shape - return c * si * sj * sk * self.dtype.itemsize + def spacing(self) -> tuple[float, float, float]: + """Voxel spacing in mm.""" + _, spacing = get_rotation_and_spacing_from_affine(self.affine) + sx, sy, sz = spacing + return float(sx), float(sy), float(sz) @property - def dtype(self) -> torch.dtype | np.dtype: - """Data type of the image. - - Returns the PyTorch dtype if the image is loaded, otherwise - reads the on-disk dtype from the header without loading data. - """ - if self._data is not None: - return self._data.dtype - if self._backend is not None: - return self._backend.dtype - if self._remote_zarr_uri is not None: - self._ensure_backend() - if self._backend is not None: - return self._backend.dtype - if self._path is not None: - self._ensure_backend() - if self._backend is not None: - return self._backend.dtype - # Non-NIfTI fallback: read via SimpleITK header - return self._read_dtype_sitk(self._path) - msg = "Cannot determine dtype: no data or path" - raise RuntimeError(msg) + def origin(self) -> tuple[float, float, float]: + """Center of first voxel in array, in mm.""" + ox, oy, oz = self.affine[:3, 3] + return ox, oy, oz @property - def orientation(self) -> tuple[str, str, str]: - """Orientation codes from the affine.""" - return self.affine.orientation + def itemsize(self): + """Element size of the data type.""" + return self.data.element_size() @property - def points(self) -> dict[str, Points]: - """Named sets of points attached to this image.""" - return self._points + def memory(self) -> float: + """Number of Bytes that the tensor takes in the RAM.""" + return np.prod(self.shape) * self.itemsize @property - def bounding_boxes(self) -> dict[str, BoundingBoxes]: - """Named sets of bounding boxes attached to this image.""" - return self._bounding_boxes + def bounds(self) -> np.ndarray: + """Position of centers of voxels in smallest and largest indices.""" + ini = 0, 0, 0 + fin = np.array(self.spatial_shape) - 1 + point_ini = apply_affine(self.affine, ini) + point_fin = apply_affine(self.affine, fin) + return np.array((point_ini, point_fin)) - # --- Methods --- + @property + def num_channels(self) -> int: + """Get the number of channels in the associated 4D tensor.""" + return len(self.data) - def load(self) -> None: - """Load data from disk into memory.""" - if self._data is not None: - return - if self._try_load_via_backend(): - return - if self._path is None: - msg = "Cannot load: no path or backend set" - raise RuntimeError(msg) - tensor, affine_array = self._reader(self._path, **self._reader_kwargs) - self._data = tensor - if self._affine is None: - self._affine = AffineMatrix(affine_array) - self._apply_channels_last() - - def _reader_supports_lazy(self) -> bool: - """Whether the configured reader can provide a lazy backend. - - The default reader supports lazy NIfTI/NIfTI-Zarr access, and a custom - reader does too if it implements - [`LazyReader`][torchio.data.backends.LazyReader]. Simple `(tensor, - affine)` readers do not and trigger a full load. - """ - return self._reader is default_reader or isinstance(self._reader, LazyReader) - - def _try_load_via_backend(self) -> bool: - """Try to load data from an existing or newly-created backend.""" - if self._load_from_backend(): - return True - has_source = ( - self._remote_zarr_uri is not None - or self._zarr_store is not None - or self._path is not None - ) - if has_source and self._reader_supports_lazy(): - self._ensure_backend() - return self._load_from_backend() - return False - - def _load_from_backend(self) -> bool: - """Materialize data from the lazy backend if available.""" - if self._backend is None: - return False - self._data = self._backend.to_tensor() - if self._affine is None: - self._affine = AffineMatrix(self._backend.affine) - self._apply_channels_last() - return True - - def _apply_channels_last(self) -> None: - """Permute data from (I, J, K, C) to (C, I, J, K) if needed.""" - if self._channels_last and self._data is not None: - self._data = rearrange(self._data, "i j k c -> c i j k") - self._channels_last = False # only do it once - - def set_data(self, tensor: TypeImageData | np.ndarray) -> None: - """Replace the image data with a new tensor. - - The lazy backend is refreshed so that `dataobj`, `dtype`, and - `shape` stay coherent with the new tensor. The affine is preserved when - one is set or can be read from the source header; for an image created - without any source (e.g. `ScalarImage()` then `set_data`) it defaults to - the identity affine. + def axis_name_to_index(self, axis: str) -> int: + """Convert an axis name to an axis index. Args: - tensor: 4D tensor with shape (C, I, J, K). - """ - self._data = self._parse_tensor(tensor) - self._refresh_backend_from_data() - - def _refresh_backend_from_data(self) -> None: - """Rebuild the backend from the in-memory tensor. - - Once the data is held in memory it becomes the single source of - truth, so the backend is replaced with a `TensorBackend` wrapping it. - The affine is preserved: it is taken from the existing affine if set, - otherwise resolved (header-only) from a lazy source, and finally - defaults to identity for source-less images (e.g. created empty then - filled with `set_data`). This never triggers a full data load. + axis: Possible inputs are `'Left'`, `'Right'`, `'Anterior'`, + `'Posterior'`, `'Inferior'`, `'Superior'`. Lower-case + versions and first letters are also valid, as only the first + letter will be used. + + Note: + If you are working with animals, you should probably use + `'Superior'`, `'Inferior'`, `'Anterior'` and `'Posterior'` + for `'Dorsal'`, `'Ventral'`, `'Rostral'` and `'Caudal'`, + respectively. + + Note: + If your images are 2D, you can use `'Top'`, `'Bottom'`, + `'Left'` and `'Right'`. """ - assert self._data is not None - if self._affine is None: - self._affine = self._infer_affine_without_loading() - self._backend = TensorBackend(self._data, affine=self._affine.data) + # Top and bottom are used for the vertical 2D axis as the use of + # Height vs Horizontal might be ambiguous - def _infer_affine_without_loading(self) -> AffineMatrix: - """Best-effort affine that never materializes the full tensor. - - Uses an existing or header-only lazy backend when available, and falls - back to the identity affine when no affine source exists. - """ - has_source = ( - self._backend is not None - or self._path is not None - or self._remote_zarr_uri is not None - or self._zarr_store is not None - ) - if has_source: - if self._backend is None: - self._ensure_backend() - if self._backend is not None: - return AffineMatrix(self._backend.affine) - return AffineMatrix() - - @property - def device(self) -> torch.device: - """Device the image data resides on.""" - return self.data.device + if not isinstance(axis, str): + raise ValueError('Axis must be a string') + axis = axis[0].upper() - def to(self, *args: Any, **kwargs: Any) -> Self: - """Move image data and affine to a device and/or cast to a dtype. - - Accepts the same arguments as `torch.Tensor.to()`. - - Returns: - `self` (modified in-place). - """ - self._data = self.data.to(*args, **kwargs) - if self._affine is not None: - self._affine.to(*args, **kwargs) - self._refresh_backend_from_data() - return self - - def numpy(self) -> np.ndarray: - """Return the image data as a NumPy array. + # Generally, TorchIO tensors are (C, W, H, D) + if axis in 'TB': # Top, Bottom + return -2 + else: + try: + index = self.orientation.index(axis) + except ValueError: + index = self.orientation.index(self.flip_axis(axis)) + # Return negative indices so that it does not matter whether we + # refer to spatial dimensions or not + index = -3 + index + return index - If the data is not loaded, reads it from disk first. The returned - array shares memory with the tensor if possible (i.e., if the - tensor is on CPU and not a view). + @staticmethod + def flip_axis(axis: str) -> str: + """Return the opposite axis label. For example, `'L'` -> `'R'`. - Returns: - 4D array with shape (C, I, J, K). + Args: + axis: Axis label, such as `'L'` or `'left'`. """ - return self.data.cpu().numpy() - - def new_like( + labels = 'LRPAISTBDV' + first = labels[::2] + last = labels[1::2] + flip_dict = dict(zip(first + last, last + first, strict=True)) + axis = axis[0].upper() + flipped_axis = flip_dict.get(axis) + if flipped_axis is None: + values = ', '.join(labels) + message = f'Axis not understood. Please use one of: {values}' + raise ValueError(message) + return flipped_axis + + def get_spacing_string(self) -> str: + strings = [f'{n:.2f}' for n in self.spacing] + string = f'({", ".join(strings)})' + return string + + def get_bounds(self) -> TypeBounds: + """Get minimum and maximum world coordinates occupied by the image.""" + first_index = 3 * (-0.5,) + last_index = np.array(self.spatial_shape) - 0.5 + first_point = apply_affine(self.affine, first_index) + last_point = apply_affine(self.affine, last_index) + array = np.array((first_point, last_point)) + bounds_x, bounds_y, bounds_z = array.T.tolist() + x0, x1 = bounds_x + y0, y1 = bounds_y + z0, z1 = bounds_z + return (x0, x1), (y0, y1), (z0, z1) + + def _parse_single_path( self, + path: TypePath, *, - data: TypeImageData, - affine: AffineMatrix | npt.ArrayLike | None = None, - ) -> Self: - r"""Create a new image of the same class with new data. - - Preserves metadata, annotations, and affine. Uses the existing - affine unless a new one is provided. Works correctly with custom - subclasses. - - Args: - data: New 4D [`torch.Tensor`][torch.Tensor] with shape - $(C, I, J, K)$. - affine: New $4 \times 4$ affine. If `None`, uses `self.affine`. - """ - new_affine = ( - self._parse_affine(affine) if affine is not None else self.affine.clone() - ) - points_copy, bboxes_copy = self._deep_copy_annotations() - return type(self)( - data, - affine=new_affine, - points=points_copy, - bounding_boxes=bboxes_copy, - **dict(self._metadata), - ) - - def save( + verify: bool = True, + ) -> Path: + if isinstance(path, (torch.Tensor, np.ndarray)): + class_name = self.__class__.__name__ + message = ( + 'Expected type str or Path but found a tensor/array. Instead of' + f' {class_name}(your_tensor),' + f' use {class_name}(tensor=your_tensor).' + ) + raise TypeError(message) + try: + path = Path(path).expanduser() + except TypeError as err: + message = ( + f'Expected type str or Path but found an object with type' + f' {type(path)} instead' + ) + raise TypeError(message) from err + except RuntimeError as err: + message = f'Conversion to path not possible for variable: {path}' + raise RuntimeError(message) from err + if not verify: + return path + + if not (path.is_file() or path.is_dir()): # might be a dir with DICOM + raise FileNotFoundError(f'File not found: "{path}"') + return path + + def _parse_path( self, - path: str | Path, - **kwargs: Any, - ) -> None: - """Save the image to a file. - - NIfTI-Zarr (`.nii.zarr`) files are written via `niizarr` - (requires the `zarr` extra). All other formats are written - with [SimpleITK](https://simpleitk.org/). - - Args: - path: Output file path. The format is inferred from the - extension. - **kwargs: Extra keyword arguments forwarded to the - writer. For SimpleITK formats these are passed to - `SimpleITK.WriteImage()`. - """ - path = Path(path) - if is_nifti_zarr(path): - self._save_nii_zarr(path) - else: - self._save_sitk(path, **kwargs) - - def _save_sitk(self, path: Path, **kwargs: Any) -> None: - from .io import _RAS_TO_LPS + path: TypePath | Sequence[TypePath] | None, + *, + verify: bool = True, + ) -> Path | list[Path] | None: + if path is None: + return None + elif isinstance(path, dict): + # https://github.com/TorchIO-project/torchio/pull/838 + raise TypeError('The path argument cannot be a dictionary') + elif isinstance(path, (str, Path)): + return self._parse_single_path(path, verify=verify) + elif self._is_paths_sequence(path): + return [self._parse_single_path(p, verify=verify) for p in path] + message = f'Input path must be a path or sequence of paths, not {type(path)}' + raise TypeError(message) + + def _parse_tensor( + self, + tensor: TypeData | None, + none_ok: bool = True, + ) -> TypeImageTensor | None: + if tensor is None: + if none_ok: + return None + else: + raise RuntimeError('Input tensor cannot be None') + if isinstance(tensor, np.ndarray): + tensor = check_uint_to_int(tensor) + tensor = torch.as_tensor(tensor) + elif not isinstance(tensor, torch.Tensor): + message = ( + 'Input tensor must be a PyTorch tensor or NumPy array,' + f' but type "{type(tensor)}" was found' + ) + raise TypeError(message) + ndim = tensor.ndim + if ndim != 4: + raise ValueError(f'Input tensor must be 4D, but it is {ndim}D') + if tensor.dtype == torch.bool: + tensor = tensor.to(torch.uint8) + if self.check_nans and torch.isnan(tensor).any(): + warnings.warn('NaNs found in tensor', RuntimeWarning, stacklevel=2) + return tensor - data = self.numpy() - n_channels = data.shape[0] - if n_channels == 1: - array = rearrange(data, "1 i j k -> k j i") - sitk_image = sitk.GetImageFromArray(array) - else: - array = rearrange(data, "c i j k -> k j i c") - sitk_image = sitk.GetImageFromArray(array, isVector=True) - # Convert from RAS (TorchIO) to LPS (SimpleITK) before setting metadata. - lps_affine = _RAS_TO_LPS @ self.affine.numpy() - lps_spacing = np.sqrt(np.sum(lps_affine[:3, :3] ** 2, axis=0)) - lps_direction = lps_affine[:3, :3] / lps_spacing - lps_origin = lps_affine[:3, 3] - sitk_image.SetSpacing(lps_spacing.tolist()) - sitk_image.SetOrigin(lps_origin.tolist()) - sitk_image.SetDirection(lps_direction.ravel().tolist()) - sitk.WriteImage(sitk_image, str(path), **kwargs) - - def _save_nii_zarr(self, path: Path) -> None: - niizarr = get_niizarr() - data = self.data.numpy() - n_channels = data.shape[0] - if n_channels == 1: - array = rearrange(data, "1 i j k -> i j k") - else: - array = rearrange(data, "c i j k -> i j k c") - nii = nib.Nifti1Image(array, self.affine.numpy()) - niizarr.nii2zarr(nii, str(path)) - - def _ensure_backend(self) -> None: - """Create the lazy backend from path or zarr store, without loading data. - - Backend selection is delegated to - [`resolve_backend`][torchio.data.backends.resolve_backend], which - consults the backend registry. For NIfTI and NIfTI-Zarr sources this - yields a header-only lazy backend; for other formats (NRRD, MHA, etc.) - no lazy backend is available and `_backend` is left unset so callers - fall back to a full read. - """ - if self._backend is not None: - return - affine_override = self._affine.data if self._affine is not None else None - request = BackendRequest( - path=self._path, - remote_zarr_uri=self._remote_zarr_uri, - zarr_store=self._zarr_store, - affine=affine_override, - reader_kwargs=self._reader_kwargs, - reader=self._reader, - ) - backend = resolve_backend(request) - if backend is not None: - self._backend = backend - return - if ( - self._path is None - and self._remote_zarr_uri is None - and self._zarr_store is None - ): - msg = "Cannot create backend: no path or store set" - raise RuntimeError(msg) + @staticmethod + def _parse_tensor_shape(tensor: TypeData) -> TypeImageTensor: + return ensure_4d(tensor) @staticmethod - def _read_shape_sitk(path: Path) -> TypeTensorShape: - """Read shape from a SimpleITK-readable file without loading data.""" - reader = sitk.ImageFileReader() - reader.SetFileName(str(path)) - reader.ReadImageInformation() - size = reader.GetSize() - n_components = reader.GetNumberOfComponents() - ndim = reader.GetDimension() - if ndim == 3: - return (n_components, size[0], size[1], size[2]) - msg = f"Expected 3D image, got {ndim}D" - raise ValueError(msg) + def _parse_affine(affine: TypeData | None) -> TypeAffineMatrix: + if affine is None: + return cast(TypeAffineMatrix, np.eye(4)) + if isinstance(affine, torch.Tensor): + affine = affine.numpy() + if not isinstance(affine, np.ndarray): + bad_type = type(affine) + raise TypeError(f'Affine must be a NumPy array, not {bad_type}') + if affine.shape != (4, 4): + bad_shape = affine.shape + raise ValueError(f'Affine shape must be (4, 4), not {bad_shape}') + return cast(TypeAffineMatrix, affine.astype(np.float64)) @staticmethod - def _read_dtype_sitk(path: Path) -> np.dtype: - """Read dtype from a SimpleITK-readable file without loading data.""" - reader = sitk.ImageFileReader() - reader.SetFileName(str(path)) - reader.ReadImageInformation() - pixel_id = reader.GetPixelID() - # Map SimpleITK pixel IDs to numpy dtypes - sitk_to_numpy = { - sitk.sitkUInt8: np.dtype("uint8"), - sitk.sitkInt8: np.dtype("int8"), - sitk.sitkUInt16: np.dtype("uint16"), - sitk.sitkInt16: np.dtype("int16"), - sitk.sitkUInt32: np.dtype("uint32"), - sitk.sitkInt32: np.dtype("int32"), - sitk.sitkUInt64: np.dtype("uint64"), - sitk.sitkInt64: np.dtype("int64"), - sitk.sitkFloat32: np.dtype("float32"), - sitk.sitkFloat64: np.dtype("float64"), - sitk.sitkVectorUInt8: np.dtype("uint8"), - sitk.sitkVectorInt8: np.dtype("int8"), - sitk.sitkVectorUInt16: np.dtype("uint16"), - sitk.sitkVectorInt16: np.dtype("int16"), - sitk.sitkVectorUInt32: np.dtype("uint32"), - sitk.sitkVectorInt32: np.dtype("int32"), - sitk.sitkVectorFloat32: np.dtype("float32"), - sitk.sitkVectorFloat64: np.dtype("float64"), - } - return sitk_to_numpy.get(pixel_id, np.dtype("float32")) + def _is_paths_sequence( + path: TypePath | Sequence[TypePath] | None, + ) -> TypeGuard[Sequence[TypePath]]: + return ( + path is not None and not isinstance(path, (str, Path)) and is_iterable(path) + ) - def __getitem__( - self, - item: str | int | slice | tuple[int | slice, ...], - ) -> Any: - """Slice the image or look up metadata by key. - - When *item* is a `str`, the metadata value with that key is - returned. Otherwise, the image is sliced along channel and/or - spatial dimensions. - - Indexing follows the tensor layout `(C, I, J, K)`. Up to four - indices may be provided; unspecified trailing dimensions keep - their full extent. The affine origin is updated to reflect the - spatial crop. Ellipsis (`...`) expands to fill unspecified - dimensions with full slices. Negative indices and steps are - supported. - - When the image has not been loaded yet, slicing reads only the - requested region through the lazy backend. The full tensor is - never materialized. For uncompressed NIfTI (`.nii`) this uses - memory-mapping; for NIfTI-Zarr (`.nii.zarr`) chunked reads. - Even for compressed NIfTI (`.nii.gz`), nibabel's proxy avoids - materializing the full array, so slicing a small region from a - large volume is much faster than loading everything first. + def _is_multipath(self) -> bool: + return self._is_paths_sequence(self.path) + + def _as_path_list(self) -> list[Path]: + if self.path is None: + raise RuntimeError('Image path is not available') + if self._is_multipath(): + paths = self.path + assert isinstance(paths, list) + return paths + assert isinstance(self.path, Path) + return [self.path] + + def _is_dir(self) -> bool: + is_sequence = self._is_multipath() + if is_sequence: + return False + elif self.path is None: + return False + else: + assert isinstance(self.path, Path) + return self.path.is_dir() - Args: - item: Integer, slice, or tuple of slices/ints/ellipsis for - the C, I, J, K dimensions. + def load(self) -> None: + r"""Load the image from disk. Returns: - A new image of the same class containing the sliced data. - - Examples: - >>> image = tio.ScalarImage(torch.randn(3, 256, 256, 176)) - >>> image[0].shape # first channel - (1, 256, 256, 176) - >>> image[:, 100:200].shape # spatial range, all channels - (3, 100, 256, 176) - >>> image[..., 50:100].shape # last spatial dim - (3, 256, 256, 50) - >>> image[1:3, 10:20, 10:20, 10:20].shape - (2, 10, 10, 10) + Tuple containing a 4D tensor of size $(C, W, H, D)$ and a 2D + $4 \times 4$ affine matrix to convert voxel indices to world + coordinates. """ - # String key → metadata lookup - if isinstance(item, str): - if item in self._metadata: - return self._metadata[item] - msg = f"{type(self).__name__} has no metadata key {item!r}" - raise KeyError(msg) - - sc, si, sj, sk = normalize_index(item) - - full_shape = self.shape - cropped_data = self._slice_data(sc, si, sj, sk) - - # Update affine origin: shift by the spatial start offset - affine_matrix = self.affine.data.clone() - i_start, _, _ = si.indices(full_shape[1]) - j_start, _, _ = sj.indices(full_shape[2]) - k_start, _, _ = sk.indices(full_shape[3]) - start_voxel = torch.tensor( - [i_start, j_start, k_start], - dtype=torch.float64, - device=affine_matrix.device, - ) - affine_matrix[:3, 3] += affine_matrix[:3, :3] @ start_voxel + if self._loaded: + return - return self.new_like(data=cropped_data, affine=AffineMatrix(affine_matrix)) + paths = self._as_path_list() + tensor, affine = self.read_and_check(paths[0]) + tensors = [tensor] + for path in paths[1:]: + new_tensor, new_affine = self.read_and_check(path) + if not np.array_equal(affine, new_affine): + message = ( + 'Files have different affine matrices.' + f'\nMatrix of {paths[0]}:' + f'\n{affine}' + f'\nMatrix of {path}:' + f'\n{new_affine}' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + if not tensor.shape[1:] == new_tensor.shape[1:]: + message = ( + f'Files shape do not match, found {tensor.shape}' + f'and {new_tensor.shape}' + ) + raise RuntimeError(message) + tensors.append(new_tensor) + tensor = torch.cat(tensors) + self.set_data(tensor) + self.affine = affine + self._loaded = True - def _slice_data( - self, - sc: slice, - si: slice, - sj: slice, - sk: slice, - ) -> Tensor: - """Slice data, using the lazy backend if available.""" - if self._data is not None: - return self._data[sc, si, sj, sk] - self._ensure_backend() - if self._backend is not None: - return self._backend[sc, si, sj, sk] - return self.data[sc, si, sj, sk] - - def _repr_path_line(self) -> str: - """Build the `path:` line for `__repr__`.""" - if self._remote_zarr_uri is not None: - status = "loaded" if self.is_loaded else "lazy, NIfTI-Zarr" - return f" path: {self._remote_zarr_uri} ({status})" - if self._path is not None: - name = self._path.name - if self.is_loaded: - return f" path: {name} (loaded)" - fmt = _backend_label(self._backend) - return f" path: {name} (lazy, {fmt})" - return " path: (in memory)" - - def __repr__(self) -> str: - import humanize - - cls_name = type(self).__name__ - lines: list[str] = [] - try: - sp = ", ".join(f"{s:.2f}" for s in self.spacing) - ori = ", ".join(f"{o:.2f}" for o in self.origin) - angles = ", ".join(f"{a:.1f}°" for a in self.affine.euler_angles) - dt = str(self.dtype).replace("torch.", "") - mem = humanize.naturalsize(self.memory, binary=True) - - # Path / loading status (after header read so backend is set) - lines.append(self._repr_path_line()) - - lines.append(f" channels: {self.num_channels}") - lines.append(f" spatial: {self.spatial_shape}") - lines.append(f" spacing: ({sp}) mm") - lines.append(f" origin: ({ori}) mm") - lines.append(f" orientation: {''.join(self.orientation)}+") - lines.append(f" angles: ({angles})") - lines.append(f" dtype: {dt}") - if self.is_loaded: - lines.append(f" device: {self.device}") - lines.append(f" memory: {mem}") - except Exception: - if self._path is not None: - lines.append(f' path: "{self._path}"') - - if self._points: - names = ", ".join(self._points) - lines.append(f" points: {{{names}}}") - if self._bounding_boxes: - names = ", ".join(self._bounding_boxes) - lines.append(f" bboxes: {{{names}}}") - - body = "\n".join(lines) - return f"{cls_name}(\n{body}\n)" - - def _repr_html_(self) -> str: - """Rich HTML representation for Jupyter notebooks.""" - from ..repr_html import image_to_html - - return image_to_html(self) - - def plot(self, **kwargs: Any) -> Any: - """Plot 3 orthogonal slices of the image. - - Requires the `[plot]` extras (`pip install torchio[plot]`). - See [`plot_image`][torchio.visualization.plot_image] for the - full list of keyword arguments. + def unload(self) -> None: + """Unload the image from memory. + + Raises: + RuntimeError: If the images has not been loaded yet or if no path + is available. """ - from ..visualization import plot_image + if not self._loaded: + message = 'Image cannot be unloaded as it has not been loaded yet' + raise RuntimeError(message) + if self.path is None: + message = ( + 'Cannot unload image as no path is available' + ' from where the image could be loaded again' + ) + raise RuntimeError(message) + self[DATA] = None + self[AFFINE] = None + self._loaded = False + + def read_and_check(self, path: TypePath) -> TypeImageDataAffine: + tensor, affine = self.reader(path) + # Make sure the data type is compatible with PyTorch + if self.reader is not read_image and isinstance(tensor, np.ndarray): + tensor = check_uint_to_int(tensor) + tensor = self._parse_tensor_shape(tensor) + parsed_tensor = self._parse_tensor(tensor, none_ok=False) + assert parsed_tensor is not None + tensor = parsed_tensor + affine = self._parse_affine(affine) + if self.check_nans and torch.isnan(tensor).any(): + warnings.warn( + f'NaNs found in file "{path}"', + RuntimeWarning, + stacklevel=2, + ) + return tensor, affine - return plot_image(self, **kwargs) + def save(self, path: TypePath, squeeze: bool | None = None) -> None: + """Save image to disk. - def plot_interactive(self, *, height: int = 300) -> Any: - """Show an interactive NiiVue viewer in Jupyter. + Args: + path: String or instance of [pathlib.Path][pathlib.Path]. + squeeze: Whether to remove singleton dimensions before saving. + If `None`, the array will be squeezed if the output format is + JP(E)G, PNG, BMP or TIF(F). + """ + write_image( + self.data, + self.affine, + path, + squeeze=squeeze, + ) - Requires `ipyniivue` (`pip install torchio[niivue]`). + def is_2d(self) -> bool: + return self.shape[-1] == 1 - The widget supports scrolling through slices, zooming, and - crosshair navigation. Left hemisphere is displayed on the - right side of the screen (radiological convention). + def numpy(self) -> np.ndarray: + """Get a NumPy array containing the image data.""" + return np.asarray(self) - Args: - height: Height of the viewer in pixels. + def as_sitk(self, **kwargs) -> sitk.Image: + """Get the image as an instance of [sitk.Image][sitk.Image].""" + return nib_to_sitk(self.data, self.affine, **kwargs) - Returns: - An `ipyniivue.NiiVue` widget. + @classmethod + def from_sitk(cls, sitk_image): + """Instantiate a new TorchIO image from a [sitk.Image][sitk.Image]. - Raises: - ImportError: If `ipyniivue` is not installed. + Examples: + >>> import torchio as tio + >>> import SimpleITK as sitk + >>> sitk_image = sitk.Image(20, 30, 40, sitk.sitkUInt16) + >>> tio.LabelMap.from_sitk(sitk_image) + LabelMap(shape: (1, 20, 30, 40); spacing: (1.00, 1.00, 1.00); orientation: LPS+; memory: 93.8 KiB; dtype: torch.IntTensor) + >>> sitk_image = sitk.Image((224, 224), sitk.sitkVectorFloat32, 3) + >>> tio.ScalarImage.from_sitk(sitk_image) + ScalarImage(shape: (3, 224, 224, 1); spacing: (1.00, 1.00, 1.00); orientation: LPS+; memory: 588.0 KiB; dtype: torch.FloatTensor) """ - import tempfile - - import nibabel as nib - import numpy as np + tensor, affine = sitk_to_nib(sitk_image) + return cls(tensor=tensor, affine=affine) - from ..external.imports import get_ipyniivue + def as_pil(self, transpose=True): + """Get the image as an instance of [PIL.Image][PIL.Image]. - nv = get_ipyniivue() + Note: + Values will be clamped to 0-255 and cast to uint8. - data = self.data.cpu().float().numpy() - if data.ndim == 4 and data.shape[0] == 1: - data = data[0] - affine = np.asarray(self.affine.data, dtype=np.float64) - nii = nib.Nifti1Image(data, affine) - with tempfile.NamedTemporaryFile(suffix=".nii.gz", delete=False) as f: - nib.save(nii, f.name) - path = f.name - - widget = nv.NiiVue(height=height) - widget.opts.is_radiological_convention = True - widget.add_volume(nv.Volume(path=path)) - return widget + Note: + To use this method, Pillow needs to be installed: + `pip install Pillow`. + """ + try: + from PIL import Image as ImagePIL + except ModuleNotFoundError as e: + message = 'Please install Pillow to use Image.as_pil(): pip install Pillow' + raise RuntimeError(message) from e + + self.check_is_2d() + tensor = self.data + if len(tensor) not in (1, 3, 4): + raise NotImplementedError( + 'Only 1, 3 or 4 channels are supported for conversion to Pillow image' + ) + if len(tensor) == 1: + tensor = torch.cat(3 * [tensor]) + if transpose: + tensor = tensor.permute(3, 2, 1, 0) + else: + tensor = tensor.permute(3, 1, 2, 0) + array = tensor.clamp(0, 255).numpy()[0] + return ImagePIL.fromarray(array.astype(np.uint8)) def to_gif( self, - output_path: str | Path | None = None, - *, - seconds: float = 5.0, - direction: str = "I", + axis: int, + duration: float, # of full gif + output_path: TypePath, loop: int = 0, rescale: bool = True, optimize: bool = True, reverse: bool = False, - ) -> Any: - """Save an animated GIF sweeping through slices. - - Requires `Pillow` (`pip install torchio[plot]`). - - When *output_path* is `None` and the code is running inside - a Jupyter notebook, the GIF is written to a temporary file and - returned as an `IPython.display.Image` for inline display. - Outside Jupyter, *output_path* is required. + ) -> None: + """Save an animated GIF of the image. Args: - output_path: Path to the output `.gif` file. `None` - to auto-create a temporary file (Jupyter only). - seconds: Duration of the full animation in seconds. - direction: Anatomical sweep direction (`"I"`, `"S"`, - `"A"`, `"P"`, `"R"`, or `"L"`). - loop: Number of loops (0 = infinite). - rescale: Rescale intensities to `[0, 255]`. - optimize: Attempt to compress the GIF palette. + axis: Spatial axis (0, 1 or 2). + duration: Duration of the full animation in seconds. + output_path: Path to the output GIF file. + loop: Number of times the GIF should loop. + `0` means that it will loop forever. + rescale: Use [`RescaleIntensity`][torchio.transforms.preprocessing.intensity.rescale.RescaleIntensity] + to rescale the intensity values to $[0, 255]$. + optimize: If `True`, attempt to compress the palette by + eliminating unused colors. This is only useful if the palette + can be compressed to the next smaller power of 2 elements. reverse: Reverse the temporal order of frames. - - Returns: - `IPython.display.Image` when running in Jupyter, - `None` otherwise. - - Raises: - ValueError: If *output_path* is `None` and the code is - not running inside a Jupyter notebook. """ - output_path = _resolve_media_path(output_path, suffix=".gif") - from ..visualization import make_gif + from ..visualization import make_gif # avoid circular import make_gif( - self, + self.data, + axis, + duration, output_path, - seconds=seconds, - direction=direction, loop=loop, rescale=rescale, optimize=optimize, reverse=reverse, ) - if _in_jupyter(): - from IPython.display import Image as IPyImage - - return IPyImage(filename=str(output_path)) - return None - def to_video( - self, - output_path: str | Path | None = None, - *, - seconds: float = 5.0, - direction: str = "I", - verbosity: str = "error", - ) -> Any: - """Create an MP4 video sweeping through slices. + def to_ras(self) -> Image: + if self.orientation_str != 'RAS': + from ..transforms.preprocessing.spatial.to_canonical import ToCanonical - Requires `ffmpeg-python` (`pip install torchio[video]`). + return ToCanonical()(self) + return self - When *output_path* is `None` and the code is running inside - a Jupyter notebook, the video is written to a temporary file - and returned as an `IPython.display.Video` for inline - display. Outside Jupyter, *output_path* is required. + def get_center(self, lps: bool = False) -> TypeTripletFloat: + """Get image center in RAS+ or LPS+ coordinates. Args: - output_path: Path to the output `.mp4` file. `None` - to auto-create a temporary file (Jupyter only). - seconds: Duration of the full video in seconds. - direction: Anatomical sweep direction (`"I"`, `"S"`, - `"A"`, `"P"`, `"R"`, or `"L"`). - verbosity: ffmpeg log level. - - Returns: - `IPython.display.Video` when running in Jupyter, - `None` otherwise. - - Raises: - ValueError: If *output_path* is `None` and the code is - not running inside a Jupyter notebook. + lps: If `True`, the coordinates will be in LPS+ orientation, i.e. + the first dimension grows towards the left, etc. Otherwise, the + coordinates will be in RAS+ orientation. """ - output_path = _resolve_media_path(output_path, suffix=".mp4") - from ..visualization import make_video + size = np.array(self.spatial_shape) + center_index = (size - 1) / 2 + r, a, s = apply_affine(self.affine, center_index) + if lps: + return (-r, -a, s) + else: + return (r, a, s) - make_video( - self, - output_path, - seconds=seconds, - direction=direction, - verbosity=verbosity, - ) - if _in_jupyter(): - from IPython.display import Video + def set_check_nans(self, check_nans: bool) -> None: + self.check_nans = check_nans - return Video( - str(output_path), - embed=True, - html_attributes="controls autoplay loop muted", - ) + def plot(self, return_fig: bool = False, **kwargs) -> None | Figure: + """Plot image.""" + if self.is_2d(): + self.as_pil().show() + else: + from ..visualization import plot_volume # avoid circular import + + figure = plot_volume(self, **kwargs) + if return_fig: + assert figure is not None + return figure return None - def __getattr__(self, name: str) -> Any: - """Look up metadata by attribute name.""" - if name.startswith("_"): - raise AttributeError(name) - if name in self._metadata: - return self._metadata[name] - msg = f"{type(self).__name__} has no attribute {name!r}" - raise AttributeError(msg) - - def __copy__(self) -> Self: - return self.new_like(data=self.data.clone()) - - def __deepcopy__(self, memo: dict) -> Self: - affine_copy = copy.deepcopy(self._affine) if self._affine is not None else None - meta_copy = dict(self._metadata) - points_copy, bboxes_copy = self._deep_copy_annotations() - common_kwargs: dict[str, Any] = { - "affine": affine_copy, - "points": points_copy, - "bounding_boxes": bboxes_copy, - **meta_copy, - } - - if self._remote_zarr_uri is not None: - new = self._deepcopy_from_source( - self._remote_zarr_uri, - affine_copy, - common_kwargs, - with_reader=True, - ) - elif self._path is not None: - new = self._deepcopy_from_source( - self._path, - affine_copy, - common_kwargs, - with_reader=True, - ) - elif self._zarr_store is not None: - new = self._deepcopy_from_source( - self._zarr_store, - affine_copy, - common_kwargs, - with_reader=False, - ) - elif self._data is not None: - new = type(self)(self._data.clone(), **common_kwargs) - else: - new = type(self)(**common_kwargs) - if self._backend is not None: - # Backend-only lazy image (e.g. built from a nibabel image - # with no path or materialized data). Preserve the backend so - # the copy can still report its shape and load its data. - new._backend = copy.copy(self._backend) - memo[id(self)] = new - return new - - def _deepcopy_from_source( - self, - source: Any, - affine_copy: AffineMatrix | None, - common_kwargs: dict[str, Any], - *, - with_reader: bool, - ) -> Self: - """Build a deep copy from a path/store source, re-attaching payload. + def show(self, viewer_path: TypePath | None = None) -> None: + """Open the image using external software. Args: - source: The path, remote URI, or store to construct from. - affine_copy: Deep-copied affine to assign when materialized data - is carried over. - common_kwargs: Constructor keyword arguments shared by every - deep-copy branch (affine, annotations, metadata). - with_reader: Whether to forward the custom reader (paths and - remote URIs) or omit it (Zarr stores). + viewer_path: Path to the application used to view the image. If + `None`, the value of the environment variable + `SITK_SHOW_COMMAND` will be used. If this variable is also + not set, TorchIO will try to guess the location of + [ITK-SNAP ](http://www.itksnap.org/pmwiki/pmwiki.php) and + [3D Slicer ](https://www.slicer.org/). - Returns: - The reconstructed image. Materialized data is cloned when present; - otherwise a derived lazy backend (e.g. the cropped/padded view - installed by ``CropOrPad``) is shallow-copied so the crop/pad - survives instead of reverting to the full-resolution source. A - shallow copy is enough because the backend is a read-only access - adapter, avoiding deep copies of nibabel/zarr internals. + Raises: + RuntimeError: If the viewer is not found. """ - reader_kwargs = dict(self._reader_kwargs) - if with_reader: - new = type(self)( - source, - reader=self._reader, - reader_kwargs=reader_kwargs, - **common_kwargs, - ) - else: - new = type(self)(source, reader_kwargs=reader_kwargs, **common_kwargs) - if self._data is not None: - new._data = self._data.clone() - new._affine = affine_copy - elif self._backend is not None: - new._backend = copy.copy(self._backend) - return new + sitk_image = self.as_sitk() + image_viewer = sitk.ImageViewer() + # This is so that 3D Slicer creates segmentation nodes from label maps + if self.__class__.__name__ == 'LabelMap': + image_viewer.SetFileExtension('.seg.nrrd') + if viewer_path is not None: + image_viewer.SetApplication(str(viewer_path)) + try: + image_viewer.Execute(sitk_image) + except RuntimeError as e: + viewer_path = guess_external_viewer() + if viewer_path is None: + message = ( + 'No external viewer has been found. Please set the' + ' environment variable SITK_SHOW_COMMAND to a viewer of' + ' your choice' + ) + raise RuntimeError(message) from e + image_viewer.SetApplication(str(viewer_path)) + image_viewer.Execute(sitk_image) + def _crop_from_slices( + self, + slices: TypeSlice | tuple[object, ...], + ) -> Image: + from ..transforms import Crop + + slices_tuple = tuple(to_tuple(slices)) + cropping: list[int] = [] + for dim, slice_item in enumerate(slices_tuple): + if isinstance(slice_item, slice): + slice_object = slice_item + elif slice_item is Ellipsis: + message = 'Ellipsis slicing is not supported yet' + raise NotImplementedError(message) + elif isinstance(slice_item, int): + slice_object = slice(slice_item, slice_item + 1) + else: + message = f'Slice type not understood: "{type(slice_item)}"' + raise TypeError(message) + shape_dim = self.spatial_shape[dim] + start, stop, step = slice_object.indices(shape_dim) + if step != 1: + message = ( + 'Slicing with steps different from 1 is not supported yet.' + ' Use the Crop transform instead' + ) + raise ValueError(message) + crop_ini = start + crop_fin = shape_dim - stop + cropping.extend([crop_ini, crop_fin]) + while dim < 2: + cropping.extend([0, 0]) + dim += 1 + w_ini, w_fin, h_ini, h_fin, d_ini, d_fin = cropping + cropping_arg = w_ini, w_fin, h_ini, h_fin, d_ini, d_fin # making mypy happy + cropped = Crop(cropping_arg)(self) + assert isinstance(cropped, Image) + return cropped -class ScalarImage(Image): - """Image with scalar (intensity) data. - Transforms use `isinstance(image, ScalarImage)` to identify intensity - images for operations like normalization or augmentation. +class ScalarImage(Image): + """Image whose pixel values represent scalars. Examples: + >>> import torch >>> import torchio as tio - >>> image = tio.ScalarImage("t1.nii.gz") - >>> image = tio.ScalarImage(torch.randn(1, 256, 256, 176)) + >>> # Loading from a file + >>> t1_image = tio.ScalarImage('t1.nii.gz') + >>> dmri = tio.ScalarImage(tensor=torch.rand(32, 128, 128, 88)) + >>> image = tio.ScalarImage('safe_image.nrrd', check_nans=False) + >>> data, affine = image.data, image.affine + >>> affine.shape + (4, 4) + >>> image.data is image[tio.DATA] + True + >>> image.data is image.tensor + True + >>> type(image.data) + torch.Tensor + + See [`Image`][torchio.Image] for more information. """ + def __init__(self, *args, **kwargs): + if 'type' in kwargs and kwargs['type'] != INTENSITY: + raise ValueError('Type of ScalarImage is always torchio.INTENSITY') + kwargs.update({'type': INTENSITY}) + super().__init__(*args, **kwargs) + + def hist(self, **kwargs) -> None: + """Plot histogram.""" + from ..visualization import plot_histogram + + x = self.data.flatten().numpy() + plot_histogram(x, **kwargs) + + def to_video( + self, + output_path: TypePath, + frame_rate: float | None = 15, + seconds: float | None = None, + direction: str = 'I', + verbosity: str = 'error', + ) -> None: + """Create a video showing all image slices along a specified direction. + + Args: + output_path: Path to the output video file. + frame_rate: Number of frames per second (FPS). + seconds: Target duration of the full video. + direction: + verbosity: + + Note: + Only `frame_rate` or `seconds` may (and must) be specified. + """ + from ..visualization import make_video # avoid circular import + + ras_image = self.to_ras() + assert isinstance(ras_image, ScalarImage) + make_video( + ras_image, + output_path, + frame_rate=frame_rate, + seconds=seconds, + direction=direction, + verbosity=verbosity, + ) + class LabelMap(Image): - """Image with label (segmentation) data. + """Image whose pixel values represent segmentation labels. - Transforms use `isinstance(image, LabelMap)` to apply nearest-neighbor - interpolation during spatial transforms. + A sequence of paths to 3D images can be passed to create a 4D image. + This is useful to create a + `tissue probability map (TPM) `, + which contains the probability of each voxel belonging to a certain tissue type, + or to create a label map with overlapping labels. + + Intensity transforms are not applied to these images. + + Nearest neighbor interpolation is always used to resample label maps, + independently of the specified interpolation type in the transform + instantiation. Examples: + >>> import torch >>> import torchio as tio - >>> label = tio.LabelMap("seg.nii.gz") - >>> label = tio.LabelMap(torch.randint(0, 5, (1, 256, 256, 176))) + >>> binary_tensor = torch.rand(1, 128, 128, 68) > 0.5 + >>> label_map = tio.LabelMap(tensor=binary_tensor) # from a tensor + >>> label_map = tio.LabelMap('t1_seg.nii.gz') # from a file + >>> # Create a 4D tissue probability map from different 3D images + >>> tissues = 'gray_matter.nii.gz', 'white_matter.nii.gz', 'csf.nii.gz' + >>> tpm = tio.LabelMap(tissues) + + See [`Image`][torchio.Image] for more information. """ + + def __init__(self, *args, **kwargs): + if 'type' in kwargs and kwargs['type'] != LABEL: + raise ValueError('Type of LabelMap is always torchio.LABEL') + kwargs.update({'type': LABEL}) + super().__init__(*args, **kwargs) + + def count_nonzero(self) -> int: + """Get the number of voxels that are not 0.""" + return int(self.data.count_nonzero().item()) + + def count_labels(self) -> dict[int, int]: + """Get the number of voxels in each label.""" + values_list = self.data.flatten().tolist() + counter = Counter(values_list) + counts = {label: counter[label] for label in sorted(counter)} + return counts diff --git a/src/torchio/data/inference/__init__.py b/src/torchio/data/inference/__init__.py new file mode 100644 index 000000000..3817463bf --- /dev/null +++ b/src/torchio/data/inference/__init__.py @@ -0,0 +1,7 @@ +from ..sampler import GridSampler # for backward compatibility +from .aggregator import GridAggregator + +__all__ = [ + 'GridSampler', + 'GridAggregator', +] diff --git a/src/torchio/data/inference/aggregator.py b/src/torchio/data/inference/aggregator.py new file mode 100644 index 000000000..5c5665c43 --- /dev/null +++ b/src/torchio/data/inference/aggregator.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import warnings + +import numpy as np +import torch + +from ...constants import CHANNELS_DIMENSION +from ..sampler import GridSampler + + +class GridAggregator: + r"""Aggregate patches for dense inference. + + This class is typically used to build a volume made of patches after + inference of batches extracted by a [`GridSampler`](#torchio.data.GridSampler). + + Args: + sampler: Instance of [`GridSampler`](#torchio.data.GridSampler) used to + extract the patches. + overlap_mode: If `'crop'`, the overlapping predictions will be + cropped. If `'average'`, the predictions in the overlapping areas + will be averaged with equal weights. If `'hann'`, the predictions + in the overlapping areas will be weighted with a Hann window + function. See the [grid aggregator tests](https://github.com/TorchIO-project/torchio/blob/main/tests/data/inference/test_aggregator.py) for a raw visualization + of the three modes. + downsampling_factor: Factor by which the output volume is expected to + be smaller than the input volume in each spatial dimension. This is + useful when the model downsamples the input (e.g., with strided + convolutions or pooling layers). Currently, only a single integer + is supported, which applies the same downsampling factor to all + spatial dimensions. + + + Note: + Adapted from NiftyNet. See [this NiftyNet tutorial + ](https://niftynet.readthedocs.io/en/dev/window_sizes.html) for more + information about patch-based sampling. + """ + + def __init__( + self, + sampler: GridSampler, + overlap_mode: str = 'crop', + downsampling_factor: int = 1, # TODO: support one per dimension + ): + subject = sampler.subject + self.volume_padded = sampler.padding_mode is not None + self.spatial_shape = subject.spatial_shape + self._output_tensor: torch.Tensor | None = None + self.patch_overlap = sampler.patch_overlap + self.patch_size = sampler.patch_size + self._parse_overlap_mode(overlap_mode) + self.overlap_mode = overlap_mode + self._avgmask_tensor: torch.Tensor | None = None + self._hann_window: torch.Tensor | None = None + self._downsampling_factor = downsampling_factor + shape_array = np.array(subject.spatial_shape) // self._downsampling_factor + self.spatial_shape = tuple(shape_array.tolist()) + + @staticmethod + def _parse_overlap_mode(overlap_mode): + if overlap_mode not in ('crop', 'average', 'hann'): + message = ( + 'Overlap mode must be "crop", "average" or "hann" but ' + f' "{overlap_mode}" was passed' + ) + raise ValueError(message) + + def _crop_patch( + self, + patch: torch.Tensor, + location: np.ndarray, + overlap: np.ndarray, + ) -> tuple[torch.Tensor, np.ndarray]: + half_overlap = overlap // 2 # overlap is always even in grid sampler + index_ini, index_fin = location[:3], location[3:] + + # If the patch is not at the border, we crop half the overlap + crop_ini: np.ndarray = half_overlap.copy() + crop_fin: np.ndarray = half_overlap.copy() + + # If the volume has been padded, we don't need to worry about cropping + if self.volume_padded: + pass + else: + crop_ini *= index_ini > 0 + crop_fin *= index_fin != self.spatial_shape + + # Update the location of the patch in the volume + new_index_ini = index_ini + crop_ini + new_index_fin = index_fin - crop_fin + new_location = np.hstack((new_index_ini, new_index_fin)) + + patch_size = np.asarray(patch.shape[-3:], dtype=int) + crop_fin = crop_fin.astype(int) + i_ini, j_ini, k_ini = crop_ini + i_fin, j_fin, k_fin = patch_size - crop_fin + # Make type checkers happy + i_ini = int(i_ini) + j_ini = int(j_ini) + k_ini = int(k_ini) + i_fin = int(i_fin) + j_fin = int(j_fin) + k_fin = int(k_fin) + cropped_patch = patch[:, i_ini:i_fin, j_ini:j_fin, k_ini:k_fin] + return cropped_patch, new_location + + def _initialize_output_tensor(self, batch: torch.Tensor) -> None: + if self._output_tensor is not None: + return + num_channels = batch.shape[CHANNELS_DIMENSION] + self._output_tensor = torch.zeros( + num_channels, + *self.spatial_shape, + dtype=batch.dtype, + ) + + def _initialize_avgmask_tensor(self, batch: torch.Tensor) -> None: + if self._avgmask_tensor is not None: + return + num_channels = batch.shape[CHANNELS_DIMENSION] + self._avgmask_tensor = torch.zeros( + num_channels, + *self.spatial_shape, + dtype=batch.dtype, + ) + + @staticmethod + def _get_hann_window(patch_size) -> torch.Tensor: + hann_window_3d = torch.as_tensor([1]) + # create a n-dim hann window + for spatial_dim, size in enumerate(patch_size): + window_shape = np.ones_like(patch_size) + window_shape[spatial_dim] = size + hann_window_1d = torch.hann_window( + size + 2, + periodic=False, + ) + hann_window_1d = hann_window_1d[1:-1].view(*window_shape) + hann_window_3d = hann_window_3d * hann_window_1d + return hann_window_3d + + def _initialize_hann_window(self) -> None: + if self._hann_window is not None: + return + self._hann_window = self._get_hann_window(self.patch_size) + + def add_batch( + self, + batch_tensor: torch.Tensor, + locations: torch.Tensor, + ) -> None: + """Add batch processed by a network to the output prediction volume. + + Args: + batch_tensor: 5D tensor, typically the output of a convolutional + neural network, e.g. `batch['image'][torchio.DATA]`. + locations: 2D tensor with shape $(B, 6)$ representing the + patch indices in the original image. They are typically + extracted using `batch[torchio.LOCATION]`. + """ + batch = batch_tensor.cpu() + locations_array = locations.cpu().numpy() // self._downsampling_factor + target_shapes = locations_array[:, 3:] - locations_array[:, :3] + # There should be only one patch size + assert len(np.unique(target_shapes, axis=0)) == 1 + input_spatial_shape = tuple(batch.shape[-3:]) + target_spatial_shape_array = target_shapes[0] + target_spatial_shape = tuple(target_spatial_shape_array.tolist()) + if input_spatial_shape != target_spatial_shape: + message = ( + f'The shape of the input batch, {input_spatial_shape},' + ' does not match the shape of the target location,' + f' which is {target_spatial_shape}' + ) + raise RuntimeError(message) + self._initialize_output_tensor(batch) + assert isinstance(self._output_tensor, torch.Tensor) + if self.overlap_mode == 'crop': + for patch, location in zip(batch, locations_array, strict=True): + cropped_patch, new_location = self._crop_patch( + patch, + location, + self.patch_overlap, + ) + i_ini, j_ini, k_ini, i_fin, j_fin, k_fin = new_location + self._output_tensor[ + :, + i_ini:i_fin, + j_ini:j_fin, + k_ini:k_fin, + ] = cropped_patch + elif self.overlap_mode == 'average': + self._initialize_avgmask_tensor(batch) + assert isinstance(self._avgmask_tensor, torch.Tensor) + for patch, location in zip(batch, locations, strict=True): + i_ini, j_ini, k_ini, i_fin, j_fin, k_fin = location + self._output_tensor[ + :, + i_ini:i_fin, + j_ini:j_fin, + k_ini:k_fin, + ] += patch + self._avgmask_tensor[ + :, + i_ini:i_fin, + j_ini:j_fin, + k_ini:k_fin, + ] += 1 + elif self.overlap_mode == 'hann': + # To handle edge and corners avoid numerical problems, we save the + # hann window in a different tensor + # At the end, it will be filled with ones (or close values) where + # there is overlap and < 1 where there is not + # When we divide, the multiplication will be canceled in areas that + # do not overlap + self._initialize_avgmask_tensor(batch) + self._initialize_hann_window() + + if self._output_tensor.dtype != torch.float32: + self._output_tensor = self._output_tensor.float() + + assert isinstance(self._avgmask_tensor, torch.Tensor) # for mypy + if self._avgmask_tensor.dtype != torch.float32: + self._avgmask_tensor = self._avgmask_tensor.float() + + for patch, location in zip(batch, locations, strict=True): + i_ini, j_ini, k_ini, i_fin, j_fin, k_fin = location + + patch = patch * self._hann_window + self._output_tensor[ + :, + i_ini:i_fin, + j_ini:j_fin, + k_ini:k_fin, + ] += patch + assert self._hann_window is not None + self._avgmask_tensor[ + :, + i_ini:i_fin, + j_ini:j_fin, + k_ini:k_fin, + ] += self._hann_window + + def get_output_tensor(self) -> torch.Tensor: + """Get the aggregated volume after dense inference.""" + assert isinstance(self._output_tensor, torch.Tensor) + if self._output_tensor.dtype == torch.int64: + message = ( + 'Medical image frameworks such as ITK do not support int64.' + ' Casting to int32...' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + self._output_tensor = self._output_tensor.type(torch.int32) + if self.overlap_mode in ['average', 'hann']: + assert isinstance(self._avgmask_tensor, torch.Tensor) # for mypy + # true_divide is used instead of / in case the PyTorch version is + # old and one the operands is int: + # https://github.com/TorchIO-project/torchio/issues/526 + output = torch.true_divide( + self._output_tensor, + self._avgmask_tensor, + ) + else: + output = self._output_tensor + if self.volume_padded: + from ...transforms import Crop + + border = self.patch_overlap // 2 + cropping_values = [int(value) for value in border.repeat(2).tolist()] + cropping = ( + cropping_values[0], + cropping_values[1], + cropping_values[2], + cropping_values[3], + cropping_values[4], + cropping_values[5], + ) + crop = Crop(cropping) + cropped = crop(output) + assert isinstance(cropped, torch.Tensor) + return cropped + else: + return output diff --git a/src/torchio/data/invertible.py b/src/torchio/data/invertible.py deleted file mode 100644 index 65a890f6b..000000000 --- a/src/torchio/data/invertible.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Mixin for data classes that carry transform history.""" - -from __future__ import annotations - -from typing import Any - -from typing_extensions import Self - - -class Invertible: - """Mixin for objects that carry `applied_transforms` history. - - Provides `apply_inverse_transform()` and - `get_inverse_transform()` to undo recorded transforms. - - Classes that inherit from this mixin must initialise - `self.applied_transforms = []` in their constructor. - """ - - applied_transforms: list[Any] - - def get_inverse_transform( - self, - *, - warn: bool = True, - ignore_intensity: bool = False, - ) -> Any: - """Get a composed transform that inverts the applied history. - - Returns a [`Compose`][torchio.Compose] of the inverse of each - applied transform, in reverse order. Non-invertible transforms - are skipped (with a warning if `warn=True`). - - Args: - warn: Issue a warning for non-invertible transforms. - ignore_intensity: Skip all intensity transforms. - - Returns: - A `Compose` transform that undoes the history. - """ - from ..transforms.inverse import get_inverse_transform - - return get_inverse_transform( - self.applied_transforms, - warn=warn, - ignore_intensity=ignore_intensity, - ) - - def apply_inverse_transform(self, **kwargs: Any) -> Self: - """Apply the inverse of all applied transforms, in reverse order. - - Non-invertible transforms are skipped. Intensity transforms - can be ignored with `ignore_intensity=True`. - - Args: - **kwargs: Forwarded to - `get_inverse_transform()` (`warn`, - `ignore_intensity`). - - Returns: - Data with transforms undone. - - Examples: - >>> transformed = transform(subject) - >>> restored = transformed.apply_inverse_transform() - """ - inverse_transform = self.get_inverse_transform(**kwargs) - result = inverse_transform(self) - if hasattr(result, "applied_transforms"): - result.applied_transforms = [] - return result - - def clear_history(self) -> None: - """Remove all applied transform records.""" - self.applied_transforms = [] diff --git a/src/torchio/data/io.py b/src/torchio/data/io.py index d556c0126..a1f40402c 100644 --- a/src/torchio/data/io.py +++ b/src/torchio/data/io.py @@ -1,259 +1,469 @@ -"""Image I/O: readers, source resolution, format detection.""" - from __future__ import annotations -import tempfile -from io import IOBase +import warnings from pathlib import Path -from typing import Any from typing import cast -import fsspec import nibabel as nib import numpy as np +import numpy.typing as npt import SimpleITK as sitk import torch -from einops import rearrange - -from ..types import TypeImageData - -# NIfTI/TorchIO use RAS+; SimpleITK uses LPS+. Multiplying the first -# two rows of the 4x4 affine by -1 converts between the two conventions. -_RAS_TO_LPS = np.diag([-1.0, -1.0, 1.0, 1.0]) - -#: Input types accepted by the Image constructor. -ImageSource = str | Path | IOBase | fsspec.core.OpenFile +from nibabel.filebasedimages import ImageFileError +from nibabel.spatialimages import SpatialImage +from ..constants import REPO_URL +from ..types import TypeData +from ..types import TypeDataAffine +from ..types import TypeDirection +from ..types import TypeDoubletInt +from ..types import TypePath +from ..types import TypeQuartetInt +from ..types import TypeTripletFloat +from ..types import TypeTripletInt -# ── Source resolution ──────────────────────────────────────────────── +# Matrices used to switch between LPS and RAS +FLIPXY_33 = np.diag([-1, -1, 1]) +FLIPXY_44 = np.diag([-1, -1, 1, 1]) +# Image formats that are typically 2D +formats = ['.jpg', '.jpeg', '.bmp', '.png', '.tif', '.tiff'] +IMAGE_2D_FORMATS = formats + [s.upper() for s in formats] -def resolve_source( - source: ImageSource, - *, - suffix: str | None = None, -) -> Path: - """Resolve an ImageSource to a local Path. - - Local paths and `Path` objects are returned directly. - - Remote URIs (`http://`, `s3://`, `az://`, etc.) are - fetched via fsspec and cached to a temp file. - - `fsspec.core.OpenFile` objects are opened and written to a - temp file. - - File-like objects (`io.BytesIO`, open files) are written to a - temp file. A *suffix* is required so readers can detect the - format. - """ - if isinstance(source, Path): - return source - if isinstance(source, str): - if _is_remote(source): - return _fetch_remote(source) - return Path(source) - if isinstance(source, fsspec.core.OpenFile): - return _materialize_openfile(source) - if isinstance(source, IOBase): - if suffix is None: - msg = ( - "A 'suffix' (e.g. '.nii.gz') is required when passing" - " a file-like object so the reader can detect the format" +def read_image(path: TypePath) -> TypeDataAffine: + try: + result = _read_sitk(path) + except RuntimeError as e: # try with NiBabel + message = f'Error loading image with SimpleITK:\n{e}\n\nTrying NiBabel...' + warnings.warn(message, stacklevel=2) + try: + result = _read_nibabel(path) + except ImageFileError as e: + message = ( + f'File "{path}" not understood.' + ' Check supported formats by at' + ' https://simpleitk.readthedocs.io/en/master/IO.html#images' + ' and https://nipy.org/nibabel/api.html#file-formats' ) - raise ValueError(msg) - return _materialize_filelike(source, suffix=suffix) - msg = ( - "Expected path, URL, fsspec.OpenFile, or file-like," - f" got {type(source).__name__}" - ) - raise TypeError(msg) - - -# ── Format detection ───────────────────────────────────────────────── - - -def is_nifti(path: Path) -> bool: - """Check if a path looks like a NIfTI file.""" - name = path.name.lower() - return name.endswith(".nii") or name.endswith(".nii.gz") - - -def is_nifti_zarr(path: Path) -> bool: - """Check if a path looks like a NIfTI-Zarr file.""" - return str(path).endswith(".nii.zarr") - - -def is_remote_nifti_zarr(uri: str) -> bool: - """Check if a string is a remote NIfTI-Zarr URI.""" - clean = uri.split("?", maxsplit=1)[0].split("#", maxsplit=1)[0] - clean = clean.rstrip("/").lower() - return _is_remote(clean) and clean.endswith(".nii.zarr") - - -# ── Dtype helper ───────────────────────────────────────────────────── - - -# `torch.from_numpy` does not support these numpy dtypes directly, so we -# upcast to the next PyTorch-compatible signed/unsigned type while -# preserving integer semantics. -_NUMPY_DTYPE_PROMOTIONS: dict[np.dtype, np.dtype] = { - np.dtype("bool"): np.dtype("uint8"), - np.dtype("uint16"): np.dtype("int32"), - np.dtype("uint32"): np.dtype("int64"), - np.dtype("uint64"): np.dtype("int64"), -} - - -def _numpy_to_tensor(array: np.ndarray) -> torch.Tensor: - """Convert a numpy array to a torch tensor preserving dtype where possible. - - Unsigned integer dtypes (`uint16`, `uint32`, `uint64`) and - `bool` are not supported by `torch.from_numpy`; they are upcast - to the smallest signed/unsigned type that PyTorch supports while - preserving the integer range. All other dtypes keep their native - representation. - """ - promotion = _NUMPY_DTYPE_PROMOTIONS.get(array.dtype) - if promotion is not None: - array = array.astype(promotion, copy=False) - if not array.flags.writeable or not array.flags.c_contiguous: - array = np.ascontiguousarray(array) - return torch.from_numpy(array) - - -# ── Readers ────────────────────────────────────────────────────────── - - -def read_nibabel(path: Path, **kwargs: Any) -> tuple[TypeImageData, np.ndarray]: - """Read a NIfTI image using NiBabel. - - Args: - path: Path to the NIfTI file. - **kwargs: Forwarded to `nibabel.load()`. - """ - img = cast(nib.Nifti1Image, nib.load(path, **kwargs)) - data = np.asarray(img.dataobj) - affine = np.asarray(img.affine) - if data.ndim == 3: - data = rearrange(data, "i j k -> 1 i j k") - elif data.ndim == 4: - data = rearrange(data, "i j k c -> c i j k") - elif data.ndim == 5 and data.shape[3] == 1: - data = rearrange(data, "i j k 1 c -> c i j k") - else: - msg = f"Expected 3D or 4D data, got {data.ndim}D" - raise ValueError(msg) - tensor = _numpy_to_tensor(data.copy()) + raise RuntimeError(message) from e + return result + + +def _read_nibabel(path: TypePath) -> TypeDataAffine: + img = cast(SpatialImage, nib.load(str(path), mmap=False)) + data = img.get_fdata(dtype=np.float32) + if data.ndim == 5: + data = data[..., 0, :] + data = data.transpose(3, 0, 1, 2) + data = check_uint_to_int(data) + tensor = torch.as_tensor(data) + affine = img.affine + assert isinstance(affine, np.ndarray) return tensor, affine -def read_sitk(path: Path, **kwargs: Any) -> tuple[TypeImageData, np.ndarray]: - """Read an image using SimpleITK (for non-NIfTI formats). - - Args: - path: Path to the image file. - **kwargs: Forwarded to `SimpleITK.ReadImage()`. - """ - sitk_image = sitk.ReadImage(str(path), **kwargs) - data = sitk.GetArrayFromImage(sitk_image) - n_components = sitk_image.GetNumberOfComponentsPerPixel() - if data.ndim == 3 and n_components == 1: - data = rearrange(data, "k j i -> 1 i j k") - elif data.ndim == 4 and n_components > 1: - data = rearrange(data, "k j i c -> c i j k") +def _read_sitk(path: TypePath) -> TypeDataAffine: + if Path(path).is_dir(): # assume DICOM + image = _read_dicom(path) else: - msg = f"Expected 3D data, got {data.ndim}D with {n_components} components" - raise ValueError(msg) - spacing = np.array(sitk_image.GetSpacing()) - origin = np.array(sitk_image.GetOrigin()) - direction = rearrange(np.array(sitk_image.GetDirection()), "(i j) -> i j", i=3) - # SimpleITK returns LPS; convert to RAS for TorchIO's convention. - lps_affine = np.eye(4, dtype=np.float64) - lps_affine[:3, :3] = direction * spacing - lps_affine[:3, 3] = origin - affine = _RAS_TO_LPS @ lps_affine - tensor = _numpy_to_tensor(data.copy()) + image = sitk.ReadImage(str(path)) + data, affine = sitk_to_nib(image, keepdim=True) + data = check_uint_to_int(data) + tensor = torch.as_tensor(data) return tensor, affine -def default_reader(path: Path, **kwargs: Any) -> tuple[TypeImageData, np.ndarray]: - """Read an image, dispatching to NiBabel, NIfTI-Zarr, or SimpleITK. - - Args: - path: Path to the image file. - **kwargs: Forwarded to the underlying reader. - """ - if is_nifti_zarr(path): - return read_nifti_zarr(path, **kwargs) - if is_nifti(path): - return read_nibabel(path, **kwargs) - return read_sitk(path, **kwargs) - - -def read_nifti_zarr( - path: Path, - **kwargs: Any, -) -> tuple[TypeImageData, np.ndarray]: - """Read a NIfTI-Zarr image using `niizarr`. - - Requires the `zarr` extra: `pip install torchio[zarr]`. - - Args: - path: Path to a `.nii.zarr` directory. - **kwargs: Forwarded to `niizarr.zarr2nii()`. - """ - from ..external.imports import get_niizarr - - niizarr = get_niizarr() - nii = niizarr.zarr2nii(str(path), **kwargs) - data = np.asarray(nii.dataobj) - affine = np.asarray(nii.header.get_best_affine()) - if data.ndim == 3: - data = rearrange(data, "i j k -> 1 i j k") - elif data.ndim == 4: - data = rearrange(data, "i j k c -> c i j k") - else: - msg = f"Expected 3D or 4D NIfTI-Zarr data, got {data.ndim}D" - raise ValueError(msg) - tensor = _numpy_to_tensor(data.copy()) - return tensor, affine - - -# ── Internal helpers ───────────────────────────────────────────────── - - -def _is_remote(path_str: str) -> bool: - return "://" in path_str - - -def _write_to_tempfile(data: bytes, *, suffix: str) -> Path: - fd, name = tempfile.mkstemp(suffix=suffix) +def _read_dicom(directory: TypePath): + directory = Path(directory) + if not directory.is_dir(): # unreachable if called from _read_sitk + raise FileNotFoundError(f'Directory "{directory}" not found') + reader = sitk.ImageSeriesReader() + dicom_names = reader.GetGDCMSeriesFileNames(str(directory)) + if not dicom_names: + message = f'The directory "{directory}" does not seem to contain DICOM files' + raise FileNotFoundError(message) + reader.SetFileNames(dicom_names) + image = reader.Execute() + return image + + +def read_shape(path: TypePath) -> TypeQuartetInt: + reader = sitk.ImageFileReader() + reader.SetFileName(str(path)) + reader.ReadImageInformation() + num_channels = reader.GetNumberOfComponents() + num_dimensions = reader.GetDimension() + assert 2 <= num_dimensions <= 4 + if num_dimensions == 2: + spatial_shape_2d: TypeDoubletInt = reader.GetSize() + assert len(spatial_shape_2d) == 2 + si, sj = spatial_shape_2d + sk = 1 + elif num_dimensions == 4: + # We assume bad NIfTI file (channels encoded as spatial dimension) + spatial_shape_4d: TypeQuartetInt = reader.GetSize() + assert len(spatial_shape_4d) == 4 + si, sj, sk, num_channels = spatial_shape_4d + elif num_dimensions == 3: + spatial_shape_3d: TypeTripletInt = reader.GetSize() + assert len(spatial_shape_3d) == 3 + si, sj, sk = spatial_shape_3d + shape = num_channels, si, sj, sk + return shape + + +def read_affine(path: TypePath) -> np.ndarray: + reader = get_reader(path) + affine = get_ras_affine_from_sitk(reader) + return affine + + +def get_reader(path: TypePath, read: bool = True) -> sitk.ImageFileReader: + reader = sitk.ImageFileReader() + reader.SetFileName(str(path)) + if read: + reader.ReadImageInformation() + return reader + + +def write_image( + tensor: torch.Tensor, + affine: TypeData, + path: TypePath, + squeeze: bool | None = None, +) -> None: + args = tensor, affine, path try: - with open(fd, "wb") as f: - f.write(data) - except BaseException: - Path(name).unlink(missing_ok=True) - raise - return Path(name) - - -def _fetch_remote(uri: str) -> Path: - with fsspec.open(uri, "rb") as remote: - suffix = _guess_suffix(uri) - return _write_to_tempfile(remote.read(), suffix=suffix) + _write_sitk(*args, squeeze=squeeze) + except RuntimeError: # try with NiBabel + _write_nibabel(*args) -def _materialize_openfile(of: fsspec.core.OpenFile) -> Path: - with of as f: - suffix = _guess_suffix(of.path) - return _write_to_tempfile(f.read(), suffix=suffix) - - -def _materialize_filelike(f: IOBase, *, suffix: str) -> Path: - return _write_to_tempfile(f.read(), suffix=suffix) +def _write_nibabel( + tensor: torch.Tensor, + affine: TypeData, + path: TypePath, +) -> None: + """Write an image using NiBabel. + Expects a path with an extension that can be used by nibabel.save to + write a NIfTI-1 image, such as '.nii.gz' or '.img' + """ + assert tensor.ndim == 4 + num_components = tensor.shape[0] -def _guess_suffix(path_str: str) -> str: - clean = path_str.split("?")[0].split("#")[0] - if ".nii.gz" in clean: - return ".nii.gz" - p = Path(clean) - return p.suffix or ".nii.gz" + # NIfTI components must be at the end, in a 5D array + if num_components == 1: + tensor = tensor[0] + else: + tensor = tensor[np.newaxis].permute(2, 3, 4, 0, 1) + suffix = Path(str(path).replace('.gz', '')).suffix + img: nib.nifti1.Nifti1Image | nib.nifti1.Nifti1Pair + if '.nii' in suffix: + img = nib.nifti1.Nifti1Image(np.asarray(tensor), affine) + elif '.hdr' in suffix or '.img' in suffix: + img = nib.nifti1.Nifti1Pair(np.asarray(tensor), affine) + else: + raise ImageFileError + assert isinstance(img.header, nib.nifti1.Nifti1Header) + if num_components > 1: + img.header.set_intent('vector') + img.header['qform_code'] = 1 + img.header['sform_code'] = 0 + nib.loadsave.save(img, str(path)) + + +def _write_sitk( + tensor: torch.Tensor, + affine: TypeData, + path: TypePath, + use_compression: bool = True, + squeeze: bool | None = None, +) -> None: + assert tensor.ndim == 4 + path = Path(path) + array = tensor.numpy() + if path.suffix in ('.png', '.jpg', '.jpeg', '.bmp'): + warnings.warn( + f'Casting to uint 8 before saving to {path}', + RuntimeWarning, + stacklevel=2, + ) + array = array.astype(np.uint8) + if squeeze is None: + force_3d = path.suffix not in IMAGE_2D_FORMATS + else: + force_3d = not squeeze + image = nib_to_sitk(array, affine, force_3d=force_3d) + sitk.WriteImage(image, str(path), use_compression) + + +def read_matrix(path: TypePath): + """Read an affine transform and convert to tensor.""" + path = Path(path) + suffix = path.suffix + if suffix in ('.tfm', '.h5'): # ITK + tensor = _read_itk_matrix(path) + elif suffix in ('.txt', '.trsf'): # NiftyReg, blockmatching + tensor = _read_niftyreg_matrix(path) + else: + raise ValueError(f'Unknown suffix for transform file: "{suffix}"') + return tensor + + +def write_matrix(matrix: torch.Tensor, path: TypePath): + """Write an affine transform.""" + path = Path(path) + suffix = path.suffix + if suffix in ('.tfm', '.h5'): # ITK + _write_itk_matrix(matrix, path) + elif suffix in ('.txt', '.trsf'): # NiftyReg, blockmatching + _write_niftyreg_matrix(matrix, path) + + +def _to_itk_convention(matrix: TypeData) -> np.ndarray: + """RAS to LPS.""" + if isinstance(matrix, torch.Tensor): + matrix = matrix.numpy() + matrix = np.dot(FLIPXY_44, matrix) + matrix = np.dot(matrix, FLIPXY_44) + matrix = np.linalg.inv(matrix) + return matrix + + +def _from_itk_convention(matrix: TypeData) -> np.ndarray: + """LPS to RAS.""" + matrix = np.dot(matrix, FLIPXY_44) + matrix = np.dot(FLIPXY_44, matrix) + matrix = np.linalg.inv(matrix) + return matrix + + +def _read_itk_matrix(path: TypePath) -> torch.Tensor: + """Read an affine transform in ITK's .tfm format.""" + transform = sitk.ReadTransform(str(path)) + parameters = transform.GetParameters() + rotation_parameters = parameters[:9] + rotation_matrix = np.array(rotation_parameters).reshape(3, 3) + translation_parameters = parameters[9:] + translation_vector = np.array(translation_parameters).reshape(3, 1) + matrix = np.hstack([rotation_matrix, translation_vector]) + homogeneous_matrix_lps = np.vstack([matrix, [0, 0, 0, 1]]) + homogeneous_matrix_ras = _from_itk_convention(homogeneous_matrix_lps) + return torch.as_tensor(homogeneous_matrix_ras) + + +def _write_itk_matrix(matrix: TypeData, tfm_path: TypePath) -> None: + """The tfm file contains the matrix from floating to reference.""" + transform = _matrix_to_itk_transform(matrix) + transform.WriteTransform(str(tfm_path)) + + +def _matrix_to_itk_transform( + matrix: TypeData, + dimensions: int = 3, +) -> sitk.AffineTransform: + matrix = _to_itk_convention(matrix) + rotation = matrix[:dimensions, :dimensions].ravel().tolist() + translation = matrix[:dimensions, 3].tolist() + transform = sitk.AffineTransform(rotation, translation) + return transform + + +def _read_niftyreg_matrix(trsf_path: TypePath) -> torch.Tensor: + """Read a NiftyReg matrix and return it as a NumPy array.""" + read_matrix = np.loadtxt(trsf_path).astype(np.float64) + inverted = np.linalg.inv(read_matrix) + return torch.from_numpy(inverted) + + +def _write_niftyreg_matrix(matrix: TypeData, txt_path: TypePath) -> None: + """Write an affine transform in NiftyReg's .txt format (ref -> flo)""" + matrix = np.linalg.inv(matrix) + np.savetxt(txt_path, matrix, fmt='%.8f') + + +def get_rotation_and_spacing_from_affine( + affine: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + # From https://github.com/nipy/nibabel/blob/master/nibabel/orientations.py + rotation_zoom = affine[:3, :3] + spacing = np.sqrt(np.sum(rotation_zoom * rotation_zoom, axis=0)) + rotation = rotation_zoom / spacing + return rotation, spacing + + +def nib_to_sitk( + data: TypeData, + affine: TypeData, + force_3d: bool = False, + force_4d: bool = False, +) -> sitk.Image: + """Create a SimpleITK image from a tensor and a 4x4 affine matrix.""" + if data.ndim != 4: + shape = tuple(data.shape) + raise ValueError(f'Input must be 4D, but has shape {shape}') + # Possibilities + # (1, w, h, 1) + # (c, w, h, 1) + # (1, w, h, d) + # (c, w, h, d) + array = np.asarray(data) + affine = np.asarray(affine).astype(np.float64) + + is_multichannel = array.shape[0] > 1 and not force_4d + is_2d = array.shape[3] == 1 and not force_3d + if is_2d: + array = array[..., 0] + if not is_multichannel and not force_4d: + array = array[0] + array = array.transpose() # (W, H, D, C) or (W, H, D) + image = sitk.GetImageFromArray(array, isVector=is_multichannel) + + origin, spacing, direction = get_sitk_metadata_from_ras_affine( + affine, + is_2d=is_2d, + ) + image.SetOrigin(origin) # should I add a 4th value if force_4d? + image.SetSpacing(spacing) + image.SetDirection(direction) + + if data.ndim == 4: + assert image.GetNumberOfComponentsPerPixel() == data.shape[0] + num_spatial_dims = 2 if is_2d else 3 + assert image.GetSize() == data.shape[1 : 1 + num_spatial_dims] + + return image + + +def sitk_to_nib( + image: sitk.Image, + keepdim: bool = False, +) -> tuple[np.ndarray, np.ndarray]: + data = sitk.GetArrayFromImage(image).transpose() + data = check_uint_to_int(data) + num_components = image.GetNumberOfComponentsPerPixel() + if num_components == 1: + data = data[np.newaxis] # add channels dimension + input_spatial_dims = image.GetDimension() + if input_spatial_dims == 2: + data = data[..., np.newaxis] + elif input_spatial_dims == 4: # probably a bad NIfTI (1, sx, sy, sz, c) + # Try to fix it + num_components = data.shape[-1] + data = data[0] + data = data.transpose(3, 0, 1, 2) + input_spatial_dims = 3 + if not keepdim: + data = ensure_4d(data, num_spatial_dims=input_spatial_dims).numpy() + assert data.shape[0] == num_components + affine = get_ras_affine_from_sitk(image) + return data, affine + + +def get_ras_affine_from_sitk( + sitk_object: sitk.Image | sitk.ImageFileReader, +) -> np.ndarray: + spacing = np.array(sitk_object.GetSpacing(), dtype=np.float64) + direction_lps = np.array(sitk_object.GetDirection(), dtype=np.float64) + origin_lps = np.array(sitk_object.GetOrigin(), dtype=np.float64) + direction_length = len(direction_lps) + rotation_lps: npt.NDArray[np.float64] + if direction_length == 9: + rotation_lps = direction_lps.reshape(3, 3) + elif direction_length == 4: # ignore last dimension if 2D (1, W, H, 1) + rotation_lps_2d = direction_lps.reshape(2, 2) + rotation_lps = np.eye(3) + rotation_lps[:2, :2] = rotation_lps_2d + spacing = np.append(spacing, 1) + origin_lps = np.append(origin_lps, 0) + elif direction_length == 16: # probably a bad NIfTI. Let's try to fix it + rotation_lps = direction_lps.reshape(4, 4)[:3, :3] + spacing = spacing[:-1] + origin_lps = origin_lps[:-1] + rotation_ras = np.dot(FLIPXY_33, rotation_lps) + rotation_ras_zoom = rotation_ras * spacing + translation_ras = np.dot(FLIPXY_33, origin_lps) + affine = np.eye(4) + affine[:3, :3] = rotation_ras_zoom + affine[:3, 3] = translation_ras + return affine + + +def get_sitk_metadata_from_ras_affine( + affine: np.ndarray, + is_2d: bool = False, + lps: bool = True, +) -> tuple[TypeTripletFloat, TypeTripletFloat, TypeDirection]: + direction_ras, spacing_array = get_rotation_and_spacing_from_affine(affine) + origin_ras = affine[:3, 3] + origin_lps = np.dot(FLIPXY_33, origin_ras) + direction_lps = np.dot(FLIPXY_33, direction_ras) + if is_2d: # ignore orientation if 2D (1, W, H, 1) + direction_lps = np.diag((-1, -1)).astype(np.float64) + direction_ras = np.diag((1, 1)).astype(np.float64) + origin_array = origin_lps if lps else origin_ras + direction_array = direction_lps if lps else direction_ras + direction_array = direction_array.flatten() + # The following are to comply with mypy + # (although there must be prettier ways to do this) + ox, oy, oz = origin_array + sx, sy, sz = spacing_array + direction: TypeDirection + if is_2d: + d1, d2, d3, d4 = direction_array + direction = d1, d2, d3, d4 + else: + d1, d2, d3, d4, d5, d6, d7, d8, d9 = direction_array + direction = d1, d2, d3, d4, d5, d6, d7, d8, d9 + origin = ox, oy, oz + spacing = sx, sy, sz + return origin, spacing, direction + + +def ensure_4d(tensor: TypeData, num_spatial_dims=None) -> torch.Tensor: + # I wish named tensors were properly supported in PyTorch + tensor = torch.as_tensor(tensor) + num_dimensions = tensor.ndim + if num_dimensions == 4: + pass + elif num_dimensions == 5: # hope (W, H, D, 1, C) + if tensor.shape[-2] == 1: + tensor = tensor[..., 0, :] + tensor = tensor.permute(3, 0, 1, 2) + else: + raise ValueError('5D is not supported for shape[-2] > 1') + elif num_dimensions == 2: # assume 2D monochannel (W, H) + tensor = tensor[np.newaxis, ..., np.newaxis] # (1, W, H, 1) + elif num_dimensions == 3: # 2D multichannel or 3D monochannel? + if num_spatial_dims == 2: + tensor = tensor[..., np.newaxis] # (C, W, H, 1) + elif num_spatial_dims == 3: # (W, H, D) + tensor = tensor[np.newaxis] # (1, W, H, D) + else: # try to guess + shape = tensor.shape + maybe_rgb = 3 in (shape[0], shape[-1]) + if maybe_rgb: + if shape[-1] == 3: # (W, H, 3) + tensor = tensor.permute(2, 0, 1) # (3, W, H) + tensor = tensor[..., np.newaxis] # (3, W, H, 1) + else: # (W, H, D) + tensor = tensor[np.newaxis] # (1, W, H, D) + else: + message = ( + f'{num_dimensions}D images not supported yet. Please create an' + f' issue in {REPO_URL} if you would like support for them' + ) + raise ValueError(message) + assert tensor.ndim == 4 + return tensor + + +def check_uint_to_int(array: np.ndarray) -> np.ndarray: + # This is because PyTorch won't take uint16 nor uint32 + if array.dtype == np.uint16: + return array.astype(np.int32) + if array.dtype == np.uint32: + return array.astype(np.int64) + return array diff --git a/src/torchio/data/loader.py b/src/torchio/data/loader.py new file mode 100644 index 000000000..0d79f16a5 --- /dev/null +++ b/src/torchio/data/loader.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Any +from typing import TypeVar +from typing import cast + +import numpy as np +import torch +from torch.utils.data import DataLoader +from torch.utils.data import Dataset + +from .subject import Subject + +T = TypeVar('T') + + +class SubjectsLoader(DataLoader): + def __init__( + self, + dataset: Dataset, + collate_fn: Callable[[list[T]], Any] | None = None, + **kwargs, + ): + if collate_fn is None: + collate_fn = cast(Callable[[list[T]], Any], self._collate) + super().__init__( + dataset=dataset, + collate_fn=collate_fn, + **kwargs, + ) + + @staticmethod + def _collate(subjects: list[Subject]) -> dict[str, Any]: + first_subject = subjects[0] + batch_dict = {} + for key in first_subject.keys(): + collated_value = _stack([subject[key] for subject in subjects]) + batch_dict[key] = collated_value + return batch_dict + + +def _stack(x): + """Determine the type of the input and stack it accordingly. + + Args: + x: List of elements to stack. + Returns: + Stacked elements, as either a torch.Tensor, np.ndarray, dict or list. + """ + first_element = x[0] + if isinstance(first_element, torch.Tensor): + return torch.stack(x, dim=0) + elif isinstance(first_element, np.ndarray): + return np.stack(x, axis=0) + elif isinstance(first_element, dict): + # Assume that all elements have the same keys + collated_dict = {} + for key in first_element.keys(): + collated_dict[key] = _stack([element[key] for element in x]) + return collated_dict + else: + return x diff --git a/src/torchio/data/patch.py b/src/torchio/data/patch.py deleted file mode 100644 index 1fbe505cd..000000000 --- a/src/torchio/data/patch.py +++ /dev/null @@ -1,63 +0,0 @@ -"""Patch location metadata for patch-based pipelines.""" - -from __future__ import annotations - -from dataclasses import dataclass - -from ..types import TypeThreeInts - - -@dataclass(frozen=True) -class PatchLocation: - """Spatial location of an extracted patch within a volume. - - Attributes: - index: `(i, j, k)` voxel indices of the patch corner - (the corner closest to the origin). - size: `(si, sj, sk)` spatial shape of the patch. - subject_index: Optional identifier for multi-subject batches. - """ - - index: TypeThreeInts - size: TypeThreeInts - subject_index: int | None = None - - @property - def index_ini(self) -> TypeThreeInts: - """Starting voxel indices `(i, j, k)`.""" - return self.index - - @property - def index_fin(self) -> TypeThreeInts: - """One-past-the-end voxel indices.""" - return ( - self.index[0] + self.size[0], - self.index[1] + self.size[1], - self.index[2] + self.size[2], - ) - - def to_slices(self) -> tuple[slice, slice, slice]: - """Convert to spatial slices for tensor indexing.""" - ini = self.index_ini - fin = self.index_fin - return ( - slice(ini[0], fin[0]), - slice(ini[1], fin[1]), - slice(ini[2], fin[2]), - ) - - def scaled(self, factor: tuple[float, float, float]) -> PatchLocation: - """Return a new location with indices and size scaled by factor.""" - return PatchLocation( - index=( - round(self.index[0] * factor[0]), - round(self.index[1] * factor[1]), - round(self.index[2] * factor[2]), - ), - size=( - round(self.size[0] * factor[0]), - round(self.size[1] * factor[1]), - round(self.size[2] * factor[2]), - ), - subject_index=self.subject_index, - ) diff --git a/src/torchio/data/points.py b/src/torchio/data/points.py deleted file mode 100644 index 003f0ca0e..000000000 --- a/src/torchio/data/points.py +++ /dev/null @@ -1,263 +0,0 @@ -"""Points class for storing sets of 3D coordinates.""" - -from __future__ import annotations - -from typing import Any - -import numpy as np -import numpy.typing as npt -import torch -from torch import Tensor -from typing_extensions import Self - -from .affine import AffineMatrix -from .axes import AxesType -from .axes import axes_type -from .axes import get_axis_mapping -from .axes import validate_axes - - -class Points: - """A set of 3D points with a named axis convention. - - Stores an $(N, 3)$ tensor of coordinates alongside an affine matrix - and an axis string describing the coordinate system. - - The default axis convention is `"IJK"` (voxel indices). Points can - be converted to any other axis convention (including anatomical - systems like `"RAS"` or `"LPI"`) via - [`to_axes`][torchio.data.points.Points.to_axes]. - - Args: - data: $(N, 3)$ tensor or array of point coordinates. - axes: 3-character axis string (default `"IJK"`). - affine: $4 \\times 4$ affine matrix. Identity if not given. - metadata: Arbitrary metadata dict. - - Examples: - >>> import torch, torchio as tio - >>> pts = tio.Points(torch.tensor([[10.0, 20.0, 30.0]])) - >>> pts.axes - 'IJK' - >>> pts.num_points - 1 - """ - - def __init__( - self, - data: Tensor | npt.ArrayLike, - *, - axes: str = "IJK", - affine: AffineMatrix | npt.ArrayLike | None = None, - metadata: dict[str, Any] | None = None, - ) -> None: - self._data = self._parse_data(data) - self._axes = validate_axes(axes) - self._affine = self._parse_affine(affine) - self._metadata: dict[str, Any] = dict(metadata) if metadata else {} - - # --- Parsing --- - - @staticmethod - def _parse_data(data: Tensor | npt.ArrayLike) -> Tensor: - if not isinstance(data, Tensor): - data = torch.as_tensor(np.asarray(data), dtype=torch.float32) - if data.ndim != 2 or data.shape[1] != 3: - msg = f"Points must have shape (N, 3), got {tuple(data.shape)}" - raise ValueError(msg) - return data - - @staticmethod - def _parse_affine(affine: AffineMatrix | npt.ArrayLike | None) -> AffineMatrix: - if affine is None: - return AffineMatrix() - if isinstance(affine, AffineMatrix): - return affine - return AffineMatrix(affine) - - # --- Properties --- - - @property - def data(self) -> Tensor: - """$(N, 3)$ tensor of point coordinates.""" - return self._data - - @property - def axes(self) -> str: - """3-character axis string (e.g., `'IJK'`, `'RAS'`).""" - return self._axes - - @property - def affine(self) -> AffineMatrix: - """$4 \\times 4$ affine mapping voxel to world coordinates.""" - return self._affine - - @property - def metadata(self) -> dict[str, Any]: - """Arbitrary metadata dict.""" - return self._metadata - - @property - def num_points(self) -> int: - """Number of points.""" - return self._data.shape[0] - - @property - def device(self) -> torch.device: - """Device the point data resides on.""" - return self._data.device - - def to(self, *args: Any, **kwargs: Any) -> Self: - """Move point data to a device and/or cast to a dtype. - - Returns: - `self` (modified in-place). - """ - self._data = self._data.to(*args, **kwargs) - return self - - # --- Methods --- - - def to_world(self) -> Tensor: - """Transform points from voxel to world coordinates. - - Equivalent to `self.to_axes(orientation)` where *orientation* - is the anatomical orientation of the affine, but returns a raw - tensor instead of a new `Points` object. - - Returns: - $(N, 3)$ tensor in world (mm) coordinates. - """ - return self._affine.apply(self._data).to(torch.float32) - - def to_axes(self, target: str) -> Self: - """Convert points to a different axis convention. - - Handles permutations within the same type (voxel ↔ voxel, - anatomical ↔ anatomical) and cross-type conversions - (voxel ↔ anatomical) using the stored affine. - - Args: - target: Target axis string. - - Returns: - New `Points` in the target axis convention. - """ - target = validate_axes(target) - if target == self._axes: - return self._clone(axes=target) - - src_type = axes_type(self._axes) - tgt_type = axes_type(target) - - if src_type == tgt_type: - perm, flips = get_axis_mapping(self._axes, target) - converted = self._permute_and_flip(self._data, perm, flips) - else: - converted = self._cross_type(self._data, src_type, target, tgt_type) - - return self._clone(data=converted, axes=target) - - def new_like( - self, - *, - data: Tensor | npt.ArrayLike, - affine: AffineMatrix | npt.ArrayLike | None = None, - ) -> Self: - """Create a new Points with the same metadata and axes. - - Args: - data: New $(N, 3)$ coordinates. - affine: New affine. If `None`, uses `self.affine`. - """ - new_affine = ( - self._parse_affine(affine) if affine is not None else self._affine.clone() - ) - return type(self)( - data, - axes=self._axes, - affine=new_affine, - metadata=dict(self._metadata), - ) - - # --- Internal --- - - def _clone( - self, - *, - data: Tensor | None = None, - axes: str | None = None, - ) -> Self: - return type(self)( - data if data is not None else self._data.clone(), - axes=axes if axes is not None else self._axes, - affine=self._affine.clone(), - metadata=dict(self._metadata), - ) - - @staticmethod - def _permute_and_flip( - data: Tensor, - perm: tuple[int, int, int], - flips: tuple[bool, bool, bool], - ) -> Tensor: - result = data[:, list(perm)] - for col, flip in enumerate(flips): - if flip: - result[:, col] = -result[:, col] - return result - - def _cross_type( - self, - data: Tensor, - src_type: AxesType, - tgt_axes: str, - tgt_type: AxesType, - ) -> Tensor: - if src_type == AxesType.VOXEL: - # Voxel → anatomical. - # Normalise to IJK first. - if self._axes != "IJK": - perm, _ = get_axis_mapping(self._axes, "IJK") - data = data[:, list(perm)] - # Apply affine → world. - world = self._affine.apply(data).to(torch.float32) - # World system is the affine's orientation. - world_axes = "".join(self._affine.orientation) - if world_axes != tgt_axes: - perm, flips = get_axis_mapping(world_axes, tgt_axes) - world = self._permute_and_flip(world, perm, flips) - return world - else: - # Anatomical → voxel. - # Normalise to the affine's world system. - world_axes = "".join(self._affine.orientation) - if self._axes != world_axes: - perm, flips = get_axis_mapping(self._axes, world_axes) - data = self._permute_and_flip(data, perm, flips) - # Inverse affine → IJK. - inv = self._affine.inverse() - ijk = inv.apply(data).to(torch.float32) - # Reorder to target voxel axes. - if tgt_axes != "IJK": - perm, _ = get_axis_mapping("IJK", tgt_axes) - ijk = ijk[:, list(perm)] - return ijk - - # --- Dunder --- - - def __len__(self) -> int: - return self.num_points - - def __repr__(self) -> str: - return f"Points(num_points={self.num_points}, axes={self._axes!r})" - - def __deepcopy__(self, memo: dict) -> Self: - new = type(self)( - self._data.clone(), - axes=self._axes, - affine=self._affine.clone(), - metadata=dict(self._metadata), - ) - memo[id(self)] = new - return new diff --git a/src/torchio/data/queue.py b/src/torchio/data/queue.py index 21b8d6747..8d8530814 100644 --- a/src/torchio/data/queue.py +++ b/src/torchio/data/queue.py @@ -1,208 +1,389 @@ -"""Patch queue for efficient patch-based training.""" - from __future__ import annotations -import random as _random -from collections import deque from collections.abc import Iterator from collections.abc import Sequence from collections.abc import Sized -from concurrent.futures import Future -from concurrent.futures import ThreadPoolExecutor from itertools import islice -from typing import Any import humanize -from torch.utils.data import IterableDataset +import torch +from torch.utils.data import DataLoader +from torch.utils.data import Dataset from torch.utils.data import Sampler +from ..constants import NUM_SAMPLES +from .dataset import SubjectsDataset from .sampler import PatchSampler from .subject import Subject -class Queue(IterableDataset): - """Buffer of patches for stochastic patch-based training. +class Queue(Dataset): + r"""Queue used for stochastic patch-based training. + + A training iteration (i.e., forward and backward pass) performed on a + GPU is usually faster than loading, preprocessing, augmenting, and cropping + a volume on a CPU. + Most preprocessing operations could be performed using a GPU, + but these devices are typically reserved for training the CNN so that batch + size and input tensor size can be as large as possible. + Therefore, it is beneficial to prepare (i.e., load, preprocess and augment) + the volumes using multiprocessing CPU techniques in parallel with the + forward-backward passes of a training iteration. + Once a volume is appropriately prepared, it is computationally beneficial to + sample multiple patches from a volume rather than having to prepare the same + volume each time a patch needs to be extracted. + The sampled patches are then stored in a buffer or *queue* until + the next training iteration, at which point they are loaded onto the GPU + for inference. + For this, TorchIO provides the [`Queue`](#torchio.data.Queue) class, which + also inherits from the PyTorch [`Dataset`][torch.utils.data.Dataset]. + In this queueing system, + samplers behave as generators that yield patches from random locations + in volumes contained in the [`SubjectsDataset`](../../data/dataset/#torchio.data.SubjectsDataset). - Loads and preprocesses subjects in background threads, extracts - random patches via a sampler, and yields them one at a time. - Designed for use with `SubjectsLoader` or `DataLoader`. + The end of a training epoch is defined as the moment after which patches + from all subjects have been used for training. + At the beginning of each training epoch, + the subjects list in the [`SubjectsDataset`](../../data/dataset/#torchio.data.SubjectsDataset) is shuffled, + as is typically done in machine learning pipelines to increase variance + of training instances during model optimization. + A PyTorch loader queries the datasets copied in each process, + which load and process the volumes in parallel on the CPU. + A patches list is filled with patches extracted by the sampler, + and the queue is shuffled once it has reached a specified maximum length so + that batches are composed of patches from different subjects. + The internal data loader continues querying the + [`SubjectsDataset`](../../data/dataset/#torchio.data.SubjectsDataset) using multiprocessing. + The patches list, when emptied, is refilled with new patches. + A second data loader, external to the queue, + may be used to collate batches of patches stored in the queue, + which are passed to the neural network. Args: - subjects: Sequence of subjects to sample patches from. - patch_sampler: A sampler (e.g., - [`UniformSampler`][torchio.data.UniformSampler]) used to - extract patches from loaded subjects. The sampler must - accept a subject and `num_patches` in its `__call__`. - max_length: Maximum number of patches held in the buffer. - Larger values increase diversity but use more RAM. - patches_per_volume: Maximum patches to extract from each - subject. The sampler may yield fewer if valid positions - are exhausted. - num_workers: Number of background threads for loading and - preprocessing subjects. Set to 0 for synchronous loading. - shuffle_subjects: Shuffle the subject order at the start of - each epoch. - shuffle_patches: Shuffle the buffer after each refill. - transform: Optional transform applied to each subject after - loading and before patch extraction. - subject_sampler: A `torch.utils.data.Sampler` (e.g., - `DistributedSampler`) that yields subject indices. - When provided, `shuffle_subjects` must be `False`. + subjects_dataset: Instance of [`SubjectsDataset`](../../data/dataset/#torchio.data.SubjectsDataset). + max_length: Maximum number of patches that can be stored in the queue. + Using a large number means that the queue needs to be filled less + often, but more CPU memory is needed to store the patches. + samples_per_volume: Default number of patches to extract from each + volume. If a subject contains an attribute `num_samples`, it + will be used instead of `samples_per_volume`. + A small number of patches ensures a large variability in the queue, + but training will be slower. + sampler: A subclass of [`PatchSampler`](#torchio.data.PatchSampler) used + to extract patches from the volumes. + subject_sampler: Sampler to get subjects from the dataset. + It should be an instance of + [`DistributedSampler`][torch.utils.data.distributed.DistributedSampler] when + running [distributed training + ](https://pytorch.org/tutorials/beginner/dist_overview.html). + num_workers: Number of subprocesses to use for data loading + (as in [`torch.utils.data.DataLoader`][torch.utils.data.DataLoader]). + `0` means that the data will be loaded in the main process. + shuffle_subjects: If `True`, the subjects dataset is shuffled at the + beginning of each epoch, i.e. when all patches from all subjects + have been processed. + shuffle_patches: If `True`, patches are shuffled after filling the + queue. + start_background: If `True`, the loader will start working in the + background as soon as the queue is instantiated. + verbose: If `True`, some debugging messages will be printed. + + This diagram represents the connection between + a [`SubjectsDataset`](../../data/dataset/#torchio.data.SubjectsDataset), + a [`Queue`](#torchio.data.Queue) + and the [`DataLoader`][torch.utils.data.DataLoader] used to pop batches from the + queue. + + ![Training with patches](https://raw.githubusercontent.com/TorchIO-project/torchio/main/docs/images/diagram_patches.svg) + + This sketch can be used to experiment and understand how the queue works. + In this case, `shuffle_subjects` is `False` + and `shuffle_patches` is `True`. + + + + Note: + `num_workers` refers to the number of workers used to + load and transform the volumes. Multiprocessing is not needed to pop + patches from the queue, so you should always use `num_workers=0` for + the [`DataLoader`][torch.utils.data.DataLoader] you instantiate to generate + training batches. + + Examples: + >>> import torch + >>> import torchio as tio + >>> patch_size = 96 + >>> queue_length = 300 + >>> samples_per_volume = 10 + >>> sampler = tio.data.UniformSampler(patch_size) + >>> subject = tio.datasets.Colin27() + >>> subjects_dataset = tio.SubjectsDataset(10 * [subject]) + >>> patches_queue = tio.Queue( + ... subjects_dataset, + ... queue_length, + ... samples_per_volume, + ... sampler, + ... num_workers=4, + ... ) + >>> patches_loader = tio.SubjectsLoader( + ... patches_queue, + ... batch_size=16, + ... num_workers=0, # this must be 0 + ... ) + >>> num_epochs = 2 + >>> model = torch.nn.Identity() + >>> for epoch_index in range(num_epochs): + ... for patches_batch in patches_loader: + ... inputs = patches_batch['t1'][tio.DATA] + ... targets = patches_batch['brain'][tio.DATA] + ... logits = model(inputs) Examples: - >>> queue = tio.Queue( - ... subjects, - ... patch_sampler=tio.UniformSampler(subject, patch_size=64), - ... max_length=300, - ... patches_per_volume=10, + >>> # Usage with distributed training + >>> import torch.distributed as dist + >>> from torch.utils.data.distributed import DistributedSampler + >>> # Assume a process running on distributed node 3 + >>> rank = 3 + >>> patch_sampler = tio.data.UniformSampler(patch_size) + >>> subject = tio.datasets.Colin27() + >>> subjects_dataset = tio.SubjectsDataset(10 * [subject]) + >>> subject_sampler = dist.DistributedSampler( + ... subjects_dataset, + ... rank=local_rank, + ... shuffle=True, + ... drop_last=True, + ... ) + >>> # Each process is assigned (len(subjects_dataset) // num_processes) subjects + >>> patches_queue = tio.Queue( + ... subjects_dataset, + ... queue_length, + ... samples_per_volume, + ... patch_sampler, ... num_workers=4, + ... subject_sampler=subject_sampler, + ... ) + >>> patches_loader = tio.SubjectsLoader( + ... patches_queue, + ... batch_size=16, + ... num_workers=0, # this must be 0 ... ) - >>> loader = SubjectsLoader(queue, batch_size=16) - >>> for batch in loader: - ... outputs = model(batch.t1.data) + >>> num_epochs = 2 + >>> model = torch.nn.Identity() + >>> for epoch_index in range(num_epochs): + ... subject_sampler.set_epoch(epoch_index) + ... for patches_batch in patches_loader: + ... inputs = patches_batch['t1'][tio.DATA] + ... targets = patches_batch['brain'][tio.DATA] + ... logits = model(inputs) """ def __init__( self, - subjects: Sequence[Subject], - patch_sampler: PatchSampler, - max_length: int = 300, - patches_per_volume: int = 10, + subjects_dataset: SubjectsDataset, + max_length: int, + samples_per_volume: int, + sampler: PatchSampler, + subject_sampler: Sampler | None = None, num_workers: int = 0, shuffle_subjects: bool = True, shuffle_patches: bool = True, - transform: Any | None = None, - subject_sampler: Sampler | None = None, - ) -> None: - if subject_sampler is not None and shuffle_subjects: - msg = ( - "shuffle_subjects must be False when subject_sampler" - " is provided (the sampler controls the order)" - ) - raise ValueError(msg) - self.subjects = subjects - self.patch_sampler = patch_sampler + start_background: bool = True, + verbose: bool = False, + ): + self.subjects_dataset = subjects_dataset self.max_length = max_length - self.patches_per_volume = patches_per_volume - self.num_workers = num_workers self.shuffle_subjects = shuffle_subjects self.shuffle_patches = shuffle_patches - self.transform = transform + self.samples_per_volume = samples_per_volume + self.sampler = sampler self.subject_sampler = subject_sampler + self.num_workers = num_workers + self.verbose = verbose + self._subjects_iterable = None + self._incomplete_subject: Subject | None = None + self._num_patches_incomplete = 0 + self._num_sampled_subjects = 0 + if start_background: + self._initialize_subjects_iterable() + self.patches_list: list[Subject] = [] + + if self.shuffle_subjects and self.subject_sampler is not None: + raise ValueError( + 'The flag shuffle_subjects cannot be set' + ' when a subject sampler is passed', + ) + + def __len__(self): + return self.iterations_per_epoch + + def __getitem__(self, index): + # There are probably more elegant ways of doing this + if not self.patches_list: + self._print('Patches list is empty.') + self._fill() + self.patches_list.reverse() + sample_patch = self.patches_list.pop() + return sample_patch + + def __repr__(self): + attributes = [ + f'max_length={self.max_length}', + f'num_subjects={self.num_subjects}', + f'num_patches={self.num_patches}', + f'samples_per_volume={self.samples_per_volume}', + f'iterations_per_epoch={self.iterations_per_epoch}', + ] + attributes_string = ', '.join(attributes) + return f'Queue({attributes_string})' - def __iter__(self) -> Iterator[Subject]: - """Yield patches, loading subjects in the background.""" - buffer: list[Subject] = [] - subject_iter = self._make_subject_iter() + def _print(self, *args): + if self.verbose: + print(*args) # noqa: T201 - if self.num_workers > 0: - yield from self._iter_threaded(subject_iter, buffer) + def _initialize_subjects_iterable(self): + self._subjects_iterable = self._get_subjects_iterable() + + @property + def subjects_iterable(self): + if self._subjects_iterable is None: + self._initialize_subjects_iterable() + return self._subjects_iterable + + @property + def num_subjects(self) -> int: + if self.subject_sampler is not None: + if not isinstance(self.subject_sampler, Sized): + raise ValueError( + 'The subject sampler passed to the queue must have a' + ' __len__ method', + ) + num_subjects = len(self.subject_sampler) else: - yield from self._iter_sync(subject_iter, buffer) + num_subjects = len(self.subjects_dataset) + return num_subjects - def _iter_sync( - self, - subject_iter: Iterator[Subject], - buffer: list[Subject], - ) -> Iterator[Subject]: - for raw in subject_iter: - prepared = self._prepare(raw) - buffer.extend(self._sample_patches(prepared)) - yield from self._drain_if_full(buffer) - yield from self._flush(buffer) - - def _iter_threaded( - self, - subject_iter: Iterator[Subject], - buffer: list[Subject], - ) -> Iterator[Subject]: - with ThreadPoolExecutor(max_workers=self.num_workers) as pool: - futures: deque[Future[Subject]] = deque() + @property + def num_patches(self) -> int: + return len(self.patches_list) + + @property + def iterations_per_epoch(self) -> int: + all_subjects_list = self.subjects_dataset.dry_iter() + subjects_list: Sequence[Subject] + if self.subject_sampler is not None: + subjects_list = [] + for subject_index in self.subject_sampler: + subject = all_subjects_list[subject_index] + subjects_list.append(subject) + else: + subjects_list = all_subjects_list - for raw in subject_iter: - futures.append(pool.submit(self._prepare, raw)) - yield from self._collect_ready(futures, buffer) - yield from self._drain_if_full(buffer) + total_num_patches = sum( + self._get_subject_num_samples(subject) for subject in subjects_list + ) + return total_num_patches - # Drain remaining futures - for future in futures: - prepared = future.result() - buffer.extend(self._sample_patches(prepared)) + def _get_subject_num_samples(self, subject): + num_samples = getattr( + subject, + NUM_SAMPLES, + self.samples_per_volume, + ) + return num_samples - yield from self._flush(buffer) + def _fill(self) -> None: + assert self.sampler is not None + + if self._incomplete_subject is not None: + subject = self._incomplete_subject + iterable = self.sampler(subject) + patches = list(islice(iterable, self._num_patches_incomplete)) + self.patches_list.extend(patches) + self._incomplete_subject = None + + while True: + subject = self._get_next_subject() + iterable = self.sampler(subject) + num_samples = self._get_subject_num_samples(subject) + num_free_slots = self.max_length - len(self.patches_list) + if num_free_slots < num_samples: + self._incomplete_subject = subject + self._num_patches_incomplete = num_samples - num_free_slots + num_samples = min(num_samples, num_free_slots) + patches = list(islice(iterable, num_samples)) + self.patches_list.extend(patches) + self._num_sampled_subjects += 1 + list_full = len(self.patches_list) >= self.max_length + all_sampled = self._num_sampled_subjects >= self.num_subjects + if list_full or all_sampled: + break - def _collect_ready( - self, - futures: deque[Future[Subject]], - buffer: list[Subject], - ) -> Iterator[Subject]: - """Move patches from completed futures into the buffer.""" - while futures and futures[0].done(): - prepared = futures.popleft().result() - buffer.extend(self._sample_patches(prepared)) - return iter(()) # nothing to yield yet - - def _drain_if_full(self, buffer: list[Subject]) -> Iterator[Subject]: - """Yield all patches from buffer if it reached max_length.""" - if len(buffer) >= self.max_length: - yield from self._flush(buffer) - - def _flush(self, buffer: list[Subject]) -> Iterator[Subject]: - """Shuffle (if enabled) and yield all patches from buffer.""" if self.shuffle_patches: - _random.shuffle(buffer) - while buffer: - yield buffer.pop() - - def _prepare(self, subject: Subject) -> Subject: - """Load images and apply transform (may run in a thread).""" - subject.load() - if self.transform is not None: - subject = self.transform(subject) + self._shuffle_patches_list() + + def _shuffle_patches_list(self): + indices = torch.randperm(self.num_patches) + self.patches_list = [self.patches_list[i] for i in indices] + + def _get_next_subject(self) -> Subject: + # A StopIteration exception is expected when the queue is empty + try: + subject = next(self.subjects_iterable) + except StopIteration as exception: + self._print('Queue is empty:', exception) + self._initialize_subjects_iterable() + subject = next(self.subjects_iterable) + except AssertionError as exception: + if 'can only test a child process' in str(exception): + message = ( + 'The number of workers for the data loader used to pop' + ' patches from the queue should be 0. Is it?' + ) + raise RuntimeError(message) from exception + raise exception return subject - def _sample_patches(self, subject: Subject) -> list[Subject]: - """Extract up to patches_per_volume patches.""" - gen = iter(self.patch_sampler(subject)) - return list(islice(gen, self.patches_per_volume)) + @staticmethod + def _get_first_item(batch): + return batch[0] - def _make_subject_iter(self) -> Iterator[Subject]: - """Build the subject iterator for one epoch.""" - if self.subject_sampler is not None: - indices = list(self.subject_sampler) - return (self.subjects[i] for i in indices) - subjects = list(self.subjects) - if self.shuffle_subjects: - _random.shuffle(subjects) - return iter(subjects) + def _get_subjects_iterable(self) -> Iterator: + # I need a DataLoader to handle parallelism + # But this loader is always expected to yield single subject samples + self._print( + f'\nCreating subjects loader with {self.num_workers} workers', + ) + subjects_loader = DataLoader( + self.subjects_dataset, + num_workers=self.num_workers, + batch_size=1, + collate_fn=self._get_first_item, + sampler=self.subject_sampler, + shuffle=self.shuffle_subjects, + ) + self._num_sampled_subjects = 0 + return iter(subjects_loader) - @property - def num_subjects(self) -> int: - """Number of subjects per epoch.""" - sampler = self.subject_sampler - if sampler is not None: - if not isinstance(sampler, Sized): - msg = "subject_sampler must have a __len__ method" - raise TypeError(msg) - return len(sampler) - return len(self.subjects) + def get_max_memory(self, subject: Subject | None = None) -> int: + """Get the maximum RAM occupied by the patches queue in bytes. - @property - def patches_per_epoch(self) -> int: - """Total patches yielded per epoch (upper bound).""" - return self.num_subjects * self.patches_per_volume + Args: + subject: Sample subject to compute the size of a patch. + """ + images_channels = 0 + if subject is None: + subject = self.subjects_dataset[0] + for image in subject.get_images(intensity_only=False): + images_channels += len(image.data) + voxels_in_patch = int(self.sampler.patch_size.prod() * images_channels) + bytes_per_patch = 4 * voxels_in_patch # assume float32 + return int(bytes_per_patch * self.max_length) - @property - def max_memory(self) -> int: - """Estimated max RAM for the patch buffer in bytes.""" - sample = self.subjects[0] - channels = sum(img.num_channels for img in sample.images.values()) - voxels = 1 - for s in self.patch_sampler.patch_size: - voxels *= s - return 4 * channels * voxels * self.max_length + def get_max_memory_pretty(self, subject: Subject | None = None) -> str: + """Get human-readable maximum RAM occupied by the patches queue. - @property - def max_memory_pretty(self) -> str: - """Human-readable max memory estimate.""" - return humanize.naturalsize(self.max_memory, binary=True) + Args: + subject: Sample subject to compute the size of a patch. + """ + memory = self.get_max_memory(subject=subject) + return humanize.naturalsize(memory, binary=True) diff --git a/src/torchio/data/sampler.py b/src/torchio/data/sampler.py deleted file mode 100644 index b1778dcfd..000000000 --- a/src/torchio/data/sampler.py +++ /dev/null @@ -1,375 +0,0 @@ -"""Patch samplers for training and inference.""" - -from __future__ import annotations - -from collections.abc import Iterator -from typing import TYPE_CHECKING -from typing import Any - -import numpy as np -import torch -from torch import Tensor -from torch.utils.data import Dataset -from torch.utils.data import IterableDataset - -from ..types import TypeThreeInts -from .patch import PatchLocation -from .subject import Subject - -if TYPE_CHECKING: - from ..transforms.spatial._padding import PaddingMode - - -class PatchSampler: - """Base class for patch samplers. - - Args: - patch_size: Spatial size of each patch. A single `int` - is broadcast to all three axes. - """ - - def __init__(self, patch_size: int | TypeThreeInts) -> None: - if isinstance(patch_size, int): - patch_size = (patch_size, patch_size, patch_size) - self.patch_size: TypeThreeInts = patch_size - - def __call__( - self, - subject: Subject, - num_patches: int | None = None, - ) -> Iterator[Subject]: - """Sample patches from a subject. - - Override in subclasses for custom sampling strategies. - - Args: - subject: Subject to extract patches from. - num_patches: Number of patches to yield. If `None`, - yields indefinitely (for random samplers) or yields - all grid positions (for GridSampler). - """ - msg = f"{type(self).__name__} must implement __call__" - raise NotImplementedError(msg) - - def _extract_patch( - self, - subject: Subject, - location: PatchLocation, - ) -> Subject: - """Extract a patch from a subject at the given location.""" - si, sj, sk = location.to_slices() - kwargs: dict[str, Any] = {} - for name, image in subject.images.items(): - kwargs[name] = image[:, si, sj, sk] - for name in subject.metadata: - kwargs[name] = subject.metadata[name] - kwargs["patch_location"] = location - return Subject(**kwargs) - - -class GridSampler(PatchSampler, Dataset): - """Extract patches on a regular grid for dense inference. - - A map-style `Dataset` with known length and random access. - Pass directly to a `DataLoader` for batched inference. - Typically used with - [`PatchAggregator`][torchio.data.PatchAggregator]. - - Args: - subject: Subject to extract patches from. - patch_size: Spatial size of each patch. - patch_overlap: Overlap between adjacent patches. Must be even. - A single `int` is broadcast to all axes. - padding_mode: If not `None`, pad the volume by - `overlap // 2` on each side before sampling. - fill: Fill value when `padding_mode='constant'`. - - Examples: - >>> sampler = tio.GridSampler(subject, patch_size=64, patch_overlap=8) - >>> loader = DataLoader(sampler, batch_size=4) - >>> aggregator = tio.PatchAggregator(subject.spatial_shape, overlap_mode="hann") - >>> for batch in loader: - ... outputs = model(batch.t1.data) - ... aggregator.add_batch(outputs, batch.patch_location) - >>> volume = aggregator.get_output() - """ - - def __init__( - self, - subject: Subject, - patch_size: int | TypeThreeInts, - patch_overlap: int | TypeThreeInts = 0, - padding_mode: PaddingMode | None = None, - fill: float = 0, - ) -> None: - super().__init__(patch_size) - if isinstance(patch_overlap, int): - patch_overlap = (patch_overlap, patch_overlap, patch_overlap) - self.patch_overlap: TypeThreeInts = patch_overlap - self.padding_mode = padding_mode - self.fill = fill - self.subject = self._maybe_pad(subject) - self.locations = self._compute_locations(self.subject.spatial_shape) - - def __len__(self) -> int: - return len(self.locations) - - def __getitem__(self, index: int) -> Subject: - return self._extract_patch(self.subject, self.locations[index]) - - def _maybe_pad(self, subject: Subject) -> Subject: - if self.padding_mode is None: - return subject - from ..transforms.spatial.pad import Pad - - border = tuple(v // 2 for v in self.patch_overlap) - padding = ( - border[0], - border[0], - border[1], - border[1], - border[2], - border[2], - ) - pad = Pad( - padding=padding, - padding_mode=self.padding_mode, - fill=self.fill, - copy=False, - ) - return pad(subject) - - def _compute_locations( - self, - spatial_shape: TypeThreeInts, - ) -> list[PatchLocation]: - """Compute grid locations covering the volume.""" - locations: list[PatchLocation] = [] - indices_per_axis: list[list[int]] = [] - for dim in range(3): - size = spatial_shape[dim] - patch = self.patch_size[dim] - overlap = self.patch_overlap[dim] - step = max(patch - overlap, 1) - indices = list(range(0, size - patch + 1, step)) - if not indices or indices[-1] != size - patch: - indices.append(max(size - patch, 0)) - indices_per_axis.append(indices) - - for i in indices_per_axis[0]: - for j in indices_per_axis[1]: - for k in indices_per_axis[2]: - locations.append( - PatchLocation( - index=(i, j, k), - size=self.patch_size, - ), - ) - return locations - - -class UniformSampler(PatchSampler, IterableDataset): - """Random patches with uniform spatial probability. - - An `IterableDataset` for training. Also callable for use with - [`Queue`][torchio.data.Queue]. - - Args: - subject: Subject to sample patches from (for Dataset use). - patch_size: Spatial size of each patch. - num_patches: Number of patches per epoch. If `None`, - yields indefinitely. - - Examples: - >>> sampler = tio.UniformSampler(subject, patch_size=64, num_patches=100) - >>> loader = DataLoader(sampler, batch_size=8) - """ - - def __init__( - self, - subject: Subject, - patch_size: int | TypeThreeInts, - num_patches: int | None = None, - ) -> None: - super().__init__(patch_size) - self.subject = subject - self.num_patches = num_patches - - def __call__( - self, - subject: Subject, - num_patches: int | None = None, - ) -> Iterator[Subject]: - """Sample random patches from a given subject.""" - limit = num_patches or self.num_patches - count = 0 - while limit is None or count < limit: - index = self._random_index(subject.spatial_shape) - loc = PatchLocation(index=index, size=self.patch_size) - yield self._extract_patch(subject, loc) - count += 1 - - def __iter__(self) -> Iterator[Subject]: - return self(self.subject, self.num_patches) - - def _random_index( - self, - spatial_shape: TypeThreeInts, - ) -> TypeThreeInts: - def _rand(d: int) -> int: - hi = max(spatial_shape[d] - self.patch_size[d], 0) + 1 - return int(torch.randint(0, hi, (1,)).item()) - - return (_rand(0), _rand(1), _rand(2)) - - -class WeightedSampler(PatchSampler, IterableDataset): - """Random patches weighted by a probability map. - - An `IterableDataset` for training with spatial priors. - - Args: - subject: Subject to sample patches from. - patch_size: Spatial size of each patch. - probability_map: Name of the image in the subject to use - as sampling weights. - num_patches: Number of patches per epoch. If `None`, - yields indefinitely. - """ - - def __init__( - self, - subject: Subject, - patch_size: int | TypeThreeInts, - probability_map: str, - num_patches: int | None = None, - ) -> None: - super().__init__(patch_size) - self.subject = subject - self.probability_map = probability_map - self.num_patches = num_patches - - def __call__( - self, - subject: Subject, - num_patches: int | None = None, - ) -> Iterator[Subject]: - """Sample weighted patches from a given subject.""" - prob_data = self._build_probability_map_for(subject) - flat = prob_data.flatten() - if flat.sum() == 0: - msg = f"Probability map '{self.probability_map}' is all zeros" - raise RuntimeError(msg) - - limit = num_patches or self.num_patches - count = 0 - while limit is None or count < limit: - idx_flat = torch.multinomial(flat, 1).item() - center = tuple( - int(x) for x in np.unravel_index(int(idx_flat), prob_data.shape) - ) - index = _center_to_corner(center, subject.spatial_shape, self.patch_size) - loc = PatchLocation(index=index, size=self.patch_size) - yield self._extract_patch(subject, loc) - count += 1 - - def __iter__(self) -> Iterator[Subject]: - return self(self.subject, self.num_patches) - - def _build_probability_map_for(self, subject: Subject) -> Tensor: - prob_image = subject.images[self.probability_map] - prob_data = prob_image.data[0].float() - return _mask_borders(prob_data, subject.spatial_shape, self.patch_size) - - def _build_probability_map(self) -> Tensor: - return self._build_probability_map_for(self.subject) - - -class LabelSampler(WeightedSampler): - """Random patches centered on labeled voxels. - - An `IterableDataset` for training with class imbalance. - - Args: - subject: Subject to sample patches from. - patch_size: Spatial size of each patch. - label_name: Name of the label image in the subject. - label_probabilities: Dict mapping label values to sampling - weights. If `None`, all non-zero labels have equal - weight. - num_patches: Number of patches per epoch. - """ - - def __init__( - self, - subject: Subject, - patch_size: int | TypeThreeInts, - label_name: str, - label_probabilities: dict[int, float] | None = None, - num_patches: int | None = None, - ) -> None: - super().__init__( - subject, - patch_size, - probability_map=label_name, - num_patches=num_patches, - ) - self.label_name = label_name - self.label_probabilities = label_probabilities - - def _build_probability_map_for(self, subject: Subject) -> Tensor: - label_image = subject.images[self.label_name] - label_data = label_image.data[0] - - if self.label_probabilities is not None: - prob = torch.zeros_like(label_data, dtype=torch.float32) - for label, weight in self.label_probabilities.items(): - prob[label_data == label] = weight - else: - prob = (label_data > 0).float() - - return _mask_borders(prob, subject.spatial_shape, self.patch_size) - - def _build_probability_map(self) -> Tensor: - return self._build_probability_map_for(self.subject) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _mask_borders( - prob: Tensor, - spatial_shape: TypeThreeInts, - patch_size: TypeThreeInts, -) -> Tensor: - """Zero probability near borders where a patch center can't be placed.""" - prob = prob.clone() - for d in range(3): - half = patch_size[d] // 2 - if half > 0: - slices_lo: list[slice] = [slice(None)] * 3 - slices_lo[d] = slice(0, half) - prob[tuple(slices_lo)] = 0 - tail = spatial_shape[d] - half - if tail < spatial_shape[d]: - slices_hi: list[slice] = [slice(None)] * 3 - slices_hi[d] = slice(tail, None) - prob[tuple(slices_hi)] = 0 - return prob - - -def _center_to_corner( - center: tuple[int, ...], - spatial_shape: TypeThreeInts, - patch_size: TypeThreeInts, -) -> TypeThreeInts: - """Convert a center voxel to the patch corner index.""" - result: list[int] = [] - for d in range(3): - half = patch_size[d] // 2 - corner = max(0, center[d] - half) - corner = min(corner, spatial_shape[d] - patch_size[d]) - result.append(corner) - return (result[0], result[1], result[2]) diff --git a/src/torchio/data/sampler/__init__.py b/src/torchio/data/sampler/__init__.py new file mode 100644 index 000000000..849904277 --- /dev/null +++ b/src/torchio/data/sampler/__init__.py @@ -0,0 +1,15 @@ +from .grid import GridSampler +from .label import LabelSampler +from .sampler import PatchSampler +from .sampler import RandomSampler +from .uniform import UniformSampler +from .weighted import WeightedSampler + +__all__ = [ + 'GridSampler', + 'LabelSampler', + 'UniformSampler', + 'WeightedSampler', + 'PatchSampler', + 'RandomSampler', +] diff --git a/src/torchio/data/sampler/grid.py b/src/torchio/data/sampler/grid.py new file mode 100644 index 000000000..cb165ec5c --- /dev/null +++ b/src/torchio/data/sampler/grid.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +from collections.abc import Generator + +import numpy as np + +from ...data.subject import Subject +from ...types import TypeSpatialShape +from ...types import TypeTripletInt +from ...utils import to_tuple +from .sampler import PatchSampler + + +class GridSampler(PatchSampler): + r"""Extract patches across a whole volume. + + Grid samplers are useful to perform inference using all patches from a + volume. It is often used with a [`GridAggregator`](../inference/#torchio.data.GridAggregator). + + Args: + subject: Instance of [`Subject`](../../data/subject/#torchio.Subject) + from which patches will be extracted. + patch_size: Tuple of integers $(w, h, d)$ to generate patches + of size $w \times h \times d$. + If a single number $n$ is provided, + $w = h = d = n$. + patch_overlap: Tuple of even integers $(w_o, h_o, d_o)$ + specifying the overlap between patches for dense inference. If a + single number $n$ is provided, $w_o = h_o = d_o = n$. + padding_mode: Same as `padding_mode` in + [`Pad`][torchio.transforms.Pad]. If `None`, the volume will not + be padded before sampling and patches at the border will not be + cropped by the aggregator. + Otherwise, the volume will be padded with + $\left(\frac{w_o}{2}, \frac{h_o}{2}, \frac{d_o}{2} \right)$ + on each side before sampling. If the sampler is passed to a + [`GridAggregator`](../inference/#torchio.data.GridAggregator), it will crop the output + to its original size. + + Examples: + >>> import torchio as tio + >>> colin = tio.datasets.Colin27() + >>> sampler = tio.GridSampler(colin, patch_size=88) + >>> for i, patch in enumerate(sampler()): + ... patch.t1.save(f'patch_{i}.nii.gz') + ... + >>> # To figure out the number of patches beforehand: + >>> sampler = tio.GridSampler(colin, patch_size=88) + >>> len(sampler) + 8 + + Note: + Adapted from NiftyNet. See [this NiftyNet tutorial + ](https://niftynet.readthedocs.io/en/dev/window_sizes.html) for more + information about patch based sampling. Note that + `patch_overlap` is twice `border` in NiftyNet + tutorial. + """ + + def __init__( + self, + subject: Subject, + patch_size: TypeSpatialShape, + patch_overlap: TypeSpatialShape = (0, 0, 0), + padding_mode: str | float | None = None, + ): + super().__init__(patch_size) + self.patch_overlap = np.array(to_tuple(patch_overlap, length=3)) + self.padding_mode = padding_mode + self.subject = self._pad(subject) + self.locations = self._compute_locations(self.subject) + + def __len__(self): + return len(self.locations) + + def __getitem__(self, index): + # Assume 3D + location = self.locations[index] + index_ini = ( + int(location[0]), + int(location[1]), + int(location[2]), + ) + si, sj, sk = (int(value) for value in self.patch_size.tolist()) + patch_size = si, sj, sk + cropped_subject = self.crop(self.subject, index_ini, patch_size) + return cropped_subject + + def __call__( + self, + subject: Subject | None = None, + num_patches: int | None = None, + ) -> Generator[Subject]: + subject = self.subject if subject is None else subject + return super().__call__(subject, num_patches=num_patches) + + def _pad(self, subject: Subject) -> Subject: + if self.padding_mode is not None: + from ...transforms import Pad + + border = self.patch_overlap // 2 + padding_values = [int(value) for value in border.repeat(2).tolist()] + padding = ( + padding_values[0], + padding_values[1], + padding_values[2], + padding_values[3], + padding_values[4], + padding_values[5], + ) + pad = Pad(padding, padding_mode=self.padding_mode) + transformed = pad(subject) + assert isinstance(transformed, Subject) + subject = transformed + return subject + + def _compute_locations(self, subject: Subject): + patch_size_values = [int(value) for value in self.patch_size.tolist()] + patch_overlap_values = [int(value) for value in self.patch_overlap.tolist()] + patch_size = ( + patch_size_values[0], + patch_size_values[1], + patch_size_values[2], + ) + patch_overlap = ( + patch_overlap_values[0], + patch_overlap_values[1], + patch_overlap_values[2], + ) + self._parse_sizes(subject.spatial_shape, patch_size, patch_overlap) + return self._get_patches_locations( + subject.spatial_shape, patch_size, patch_overlap + ) + + def _generate_patches( + self, + subject: Subject, + num_patches: int | None = None, + ) -> Generator[Subject]: + if num_patches is not None: + message = 'GridSampler does not support limiting the number of patches' + raise ValueError(message) + subject = self._pad(subject) + locations = self._compute_locations(subject) + for location in locations: + index_ini = ( + int(location[0]), + int(location[1]), + int(location[2]), + ) + yield self.extract_patch(subject, index_ini) + + @staticmethod + def _parse_sizes( + image_size: TypeTripletInt, + patch_size: TypeTripletInt, + patch_overlap: TypeTripletInt, + ) -> None: + image_size_array = np.array(image_size) + patch_size_array = np.array(patch_size) + patch_overlap_array = np.array(patch_overlap) + if np.any(patch_size_array > image_size_array): + message = ( + f'Patch size {tuple(patch_size_array)} cannot be' + f' larger than image size {tuple(image_size_array)}' + ) + raise ValueError(message) + if np.any(patch_overlap_array >= patch_size_array): + message = ( + f'Patch overlap {tuple(patch_overlap_array)} must be smaller' + f' than patch size {tuple(patch_size_array)}' + ) + raise ValueError(message) + if np.any(patch_overlap_array % 2): + message = ( + 'Patch overlap must be a tuple of even integers,' + f' not {tuple(patch_overlap_array)}' + ) + raise ValueError(message) + + @staticmethod + def _get_patches_locations( + image_size: TypeTripletInt, + patch_size: TypeTripletInt, + patch_overlap: TypeTripletInt, + ) -> np.ndarray: + # Example with image_size 10, patch_size 5, overlap 2: + # [0 1 2 3 4 5 6 7 8 9] + # [0 0 0 0 0] + # [1 1 1 1 1] + # [2 2 2 2 2] + # Locations: + # [[0, 5], + # [3, 8], + # [5, 10]] + indices = [] + zipped = zip(image_size, patch_size, patch_overlap, strict=True) + for im_size_dim, patch_size_dim, patch_overlap_dim in zipped: + end = im_size_dim + 1 - patch_size_dim + step = patch_size_dim - patch_overlap_dim + indices_dim = list(range(0, end, step)) + if indices_dim[-1] != im_size_dim - patch_size_dim: + indices_dim.append(im_size_dim - patch_size_dim) + indices.append(indices_dim) + indices_ini = np.array(np.meshgrid(*indices)).reshape(3, -1).T + indices_ini = np.unique(indices_ini, axis=0) + indices_fin = indices_ini + np.array(patch_size) + locations = np.hstack((indices_ini, indices_fin)) + return np.array(sorted(locations.tolist())) diff --git a/src/torchio/data/sampler/label.py b/src/torchio/data/sampler/label.py new file mode 100644 index 000000000..6fa4d14c2 --- /dev/null +++ b/src/torchio/data/sampler/label.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import numpy as np +import torch + +from ...constants import LABEL +from ...constants import TYPE +from ...data.image import Image +from ...data.subject import Subject +from ...types import TypeSpatialShape +from .weighted import WeightedSampler + + +class LabelSampler(WeightedSampler): + r"""Extract random patches with labeled voxels at their center. + + This sampler yields patches whose center value is greater than 0 + in the `label_name`. + + Args: + patch_size: See [`PatchSampler`](#torchio.data.PatchSampler). + label_name: Name of the label image in the subject that will be used to + generate the sampling probability map. If `None`, the first image + of type `torchio.LABEL` found in the subject subject will be + used. + label_probabilities: Dictionary containing the probability that each + class will be sampled. Probabilities do not need to be normalized. + For example, a value of `{0: 0, 1: 2, 2: 1, 3: 1}` will create a + sampler whose patches centers will have 50% probability of being + labeled as `1`, 25% of being `2` and 25% of being `3`. + If `None`, the label map is binarized and the value is set to + `{0: 0, 1: 1}`. + If the input has multiple channels, a value of + `{0: 0, 1: 2, 2: 1, 3: 1}` will create a + sampler whose patches centers will have 50% probability of being + taken from a non zero value of channel `1`, 25% from channel + `2` and 25% from channel `3`. + + Examples: + >>> import torchio as tio + >>> subject = tio.datasets.Colin27() + >>> subject + Colin27(Keys: ('t1', 'head', 'brain'); images: 3) + >>> probabilities = {0: 0.5, 1: 0.5} + >>> sampler = tio.data.LabelSampler( + ... patch_size=64, + ... label_name='brain', + ... label_probabilities=probabilities, + ... ) + >>> generator = sampler(subject) + >>> for patch in generator: + ... print(patch.shape) + + If you want a specific number of patches from a volume, e.g. 10: + + >>> generator = sampler(subject, num_patches=10) + >>> for patch in iterator: + ... print(patch.shape) + """ + + def __init__( + self, + patch_size: TypeSpatialShape, + label_name: str | None = None, + label_probabilities: dict[int, float] | None = None, + ): + super().__init__(patch_size, probability_map=label_name) + self.label_probabilities_dict = label_probabilities + + def get_probability_map_image(self, subject: Subject) -> Image: + if self.probability_map_name is None: + for image in subject.get_images(intensity_only=False): + if image[TYPE] == LABEL: + label_map = image + break + else: + images = subject.get_images(intensity_only=False) + message = ( + f'No label maps found in subject {subject} with image' + f' paths {[image.path for image in images]}' + ) + raise RuntimeError(message) + elif self.probability_map_name in subject: + label_map = subject.get_image(self.probability_map_name) + else: + message = ( + f'Image "{self.probability_map_name}"' + f' not found in subject subject: {subject}' + ) + raise KeyError(message) + return label_map + + def get_probability_map(self, subject: Subject) -> torch.Tensor: + label_map_tensor = self.get_probability_map_image(subject).data.float() + + if self.label_probabilities_dict is None: + return label_map_tensor > 0 + probability_map = self.get_probabilities_from_label_map( + label_map_tensor, + self.label_probabilities_dict, + self.patch_size, + ) + return probability_map + + @staticmethod + def get_probabilities_from_label_map( + label_map: torch.Tensor, + label_probabilities_dict: dict[int, float], + patch_size: np.ndarray, + ) -> torch.Tensor: + """Create probability map according to label map probabilities.""" + patch_size = patch_size.astype(int) + ini_i, ini_j, ini_k = patch_size // 2 + spatial_shape = np.array(label_map.shape[1:]) + if np.any(patch_size > spatial_shape): + message = f'Patch size {patch_size}larger than label map {spatial_shape}' + raise RuntimeError(message) + crop_fin_i, crop_fin_j, crop_fin_k = crop_fin = (patch_size - 1) // 2 + fin_i, fin_j, fin_k = spatial_shape - crop_fin + # See https://github.com/TorchIO-project/torchio/issues/458 + label_map = label_map[:, ini_i:fin_i, ini_j:fin_j, ini_k:fin_k] + + multichannel = label_map.shape[0] > 1 + probability_map = torch.zeros_like(label_map) + label_probs = torch.Tensor(list(label_probabilities_dict.values())) + normalized_probs = label_probs / label_probs.sum() + iterable = zip(label_probabilities_dict, normalized_probs, strict=True) + for label, label_probability in iterable: + if multichannel: + mask = label_map[label] + else: + mask = label_map == label + label_size = mask.sum() + if not label_size: + continue + prob_voxels = label_probability / label_size + if multichannel: + probability_map[label] = prob_voxels * mask + else: + probability_map[mask] = prob_voxels + if multichannel: + probability_map = probability_map.sum(dim=0, keepdim=True) + + # See https://github.com/TorchIO-project/torchio/issues/458 + padding = ini_k, crop_fin_k, ini_j, crop_fin_j, ini_i, crop_fin_i + probability_map = torch.nn.functional.pad( + probability_map, + padding, + ) + return probability_map diff --git a/src/torchio/data/sampler/sampler.py b/src/torchio/data/sampler/sampler.py new file mode 100644 index 000000000..c1c2eca08 --- /dev/null +++ b/src/torchio/data/sampler/sampler.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +from collections.abc import Generator + +import numpy as np +import torch + +from ...constants import LOCATION +from ...data.subject import Subject +from ...types import TypeSpatialShape +from ...types import TypeTripletInt +from ...utils import to_tuple + + +class PatchSampler: + r"""Base class for TorchIO samplers. + + Args: + patch_size: Tuple of integers $(w, h, d)$ to generate patches + of size $w \times h \times d$. + If a single number $n$ is provided, $w = h = d = n$. + + Warning: + This is an abstract class that should only be instantiated + using child classes such as [`UniformSampler`](#torchio.data.UniformSampler) and + [`WeightedSampler`](#torchio.data.WeightedSampler). + """ + + def __init__(self, patch_size: TypeSpatialShape): + patch_size_array = np.array(to_tuple(patch_size, length=3)) + for n in patch_size_array: + if n < 1 or not isinstance(n, (int, np.integer)): + message = ( + 'Patch dimensions must be positive integers,' + f' not {patch_size_array}' + ) + raise ValueError(message) + self.patch_size = patch_size_array.astype(np.uint16) + + def extract_patch( + self, + subject: Subject, + index_ini: TypeTripletInt, + ) -> Subject: + si, sj, sk = (int(value) for value in self.patch_size.tolist()) + patch_size = si, sj, sk + cropped_subject = self.crop(subject, index_ini, patch_size) + return cropped_subject + + def crop( + self, + subject: Subject, + index_ini: TypeTripletInt, + patch_size: TypeTripletInt, + ) -> Subject: + transform = self._get_crop_transform(subject, index_ini, patch_size) + cropped_subject = transform(subject) + index_ini_array = np.asarray(index_ini) + patch_size_array = np.asarray(patch_size) + index_fin = index_ini_array + patch_size_array + location = index_ini_array.tolist() + index_fin.tolist() + cropped_subject[LOCATION] = torch.as_tensor(location) + cropped_subject.update_attributes() + return cropped_subject + + @staticmethod + def _get_crop_transform( + subject, + index_ini: TypeTripletInt, + patch_size: TypeSpatialShape, + ): + from ...transforms.preprocessing.spatial.crop import Crop + + shape = np.array(subject.spatial_shape, dtype=np.uint16) + index_ini_array = np.array(index_ini, dtype=np.uint16) + patch_size_array = np.array(patch_size, dtype=np.uint16) + assert len(index_ini_array) == 3 + assert len(patch_size_array) == 3 + index_fin = index_ini_array + patch_size_array + crop_ini = index_ini_array.tolist() + crop_fin = (shape - index_fin).tolist() + cropping_values = [ + int(value) + for pair in zip(crop_ini, crop_fin, strict=True) + for value in pair + ] + cropping = ( + cropping_values[0], + cropping_values[1], + cropping_values[2], + cropping_values[3], + cropping_values[4], + cropping_values[5], + ) + return Crop(cropping) + + def __call__( + self, + subject: Subject, + num_patches: int | None = None, + ) -> Generator[Subject]: + subject.check_consistent_space() + if np.any(self.patch_size > subject.spatial_shape): + message = ( + f'Patch size {tuple(self.patch_size)} cannot be' + f' larger than image size {tuple(subject.spatial_shape)}' + ) + raise RuntimeError(message) + kwargs = {} if num_patches is None else {'num_patches': num_patches} + return self._generate_patches(subject, **kwargs) + + def _generate_patches( + self, + subject: Subject, + num_patches: int | None = None, + ) -> Generator[Subject]: + raise NotImplementedError + + +class RandomSampler(PatchSampler): + r"""Base class for random samplers. + + Args: + patch_size: Tuple of integers $(w, h, d)$ to generate patches + of size $w \times h \times d$. + If a single number $n$ is provided, $w = h = d = n$. + """ + + def get_probability_map(self, subject: Subject): + raise NotImplementedError diff --git a/src/torchio/data/sampler/uniform.py b/src/torchio/data/sampler/uniform.py new file mode 100644 index 000000000..43bff7cca --- /dev/null +++ b/src/torchio/data/sampler/uniform.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from collections.abc import Generator + +import torch + +from ...data.subject import Subject +from .sampler import RandomSampler + + +class UniformSampler(RandomSampler): + """Randomly extract patches from a volume with uniform probability. + + Args: + patch_size: See [`PatchSampler`](#torchio.data.PatchSampler). + """ + + def get_probability_map(self, subject: Subject) -> torch.Tensor: + return torch.ones(1, *subject.spatial_shape) + + def _generate_patches( + self, + subject: Subject, + num_patches: int | None = None, + ) -> Generator[Subject]: + valid_range = subject.spatial_shape - self.patch_size + patches_left = num_patches if num_patches is not None else True + while patches_left: + i, j, k = tuple(int(torch.randint(x + 1, (1,)).item()) for x in valid_range) + index_ini = i, j, k + yield self.extract_patch(subject, index_ini) + if num_patches is not None: + patches_left -= 1 diff --git a/src/torchio/data/sampler/weighted.py b/src/torchio/data/sampler/weighted.py new file mode 100644 index 000000000..3e05b8b9d --- /dev/null +++ b/src/torchio/data/sampler/weighted.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +from collections.abc import Generator +from typing import overload + +import numpy as np +import torch + +from ...constants import MIN_FLOAT_32 +from ...types import TypeSpatialShape +from ...types import TypeTripletInt +from ..image import Image +from ..subject import Subject +from .sampler import RandomSampler + + +class WeightedSampler(RandomSampler): + r"""Randomly extract patches from a volume given a probability map. + + The probability of sampling a patch centered on a specific voxel is the + value of that voxel in the probability map. The probabilities need not be + normalized. For example, voxels can have values 0, 1 and 5. Voxels with + value 0 will never be at the center of a patch. Voxels with value 5 will + have 5 times more chance of being at the center of a patch that voxels + with a value of 1. + + Args: + patch_size: See [`PatchSampler`](#torchio.data.PatchSampler). + probability_map: Name of the image in the input subject that will be + used as a sampling probability map. + + Raises: + RuntimeError: If the probability map is empty. + + Examples: + >>> import torchio as tio + >>> subject = tio.Subject( + ... t1=tio.ScalarImage('t1_mri.nii.gz'), + ... sampling_map=tio.Image('sampling.nii.gz', type=tio.SAMPLING_MAP), + ... ) + >>> patch_size = 64 + >>> sampler = tio.data.WeightedSampler(patch_size, 'sampling_map') + >>> for patch in sampler(subject): + ... print(patch[tio.LOCATION]) + + Note: + The index of the center of a patch with even size $s$ is + arbitrarily set to $s/2$. This is an implementation detail that + will typically not make any difference in practice. + + Note: + Values of the probability map near the border will be set to 0 as + the center of the patch cannot be at the border (unless the patch has + size 1 or 2 along that axis). + """ + + def __init__( + self, + patch_size: TypeSpatialShape, + probability_map: str | None, + ): + super().__init__(patch_size) + self.probability_map_name = probability_map + self.cdf = None + + def _generate_patches( + self, + subject: Subject, + num_patches: int | None = None, + ) -> Generator[Subject]: + probability_map = self.get_probability_map(subject) + probability_map_array = self.process_probability_map( + probability_map, + subject, + ) + cdf = self.get_cumulative_distribution_function(probability_map_array) + + patches_left = num_patches if num_patches is not None else True + while patches_left: + yield self.extract_patch(subject, probability_map_array, cdf) + if num_patches is not None: + patches_left -= 1 + + def get_probability_map_image(self, subject: Subject) -> Image: + assert self.probability_map_name is not None + if self.probability_map_name in subject: + return subject.get_image(self.probability_map_name) + else: + message = ( + f'Image "{self.probability_map_name}" not found in subject: {subject}' + ) + raise KeyError(message) + + def get_probability_map(self, subject: Subject) -> torch.Tensor: + data = self.get_probability_map_image(subject).data + if torch.any(data < 0): + message = ( + 'Negative values found' + f' in probability map "{self.probability_map_name}"' + ) + raise ValueError(message) + return data + + def process_probability_map( + self, + probability_map: torch.Tensor, + subject: Subject, + ) -> np.ndarray: + # Using float32 can create cdf with maximum very far from 1, e.g. 0.92! + data = probability_map[0].numpy().astype(np.float64) + assert data.ndim == 3 + self.clear_probability_borders(data, self.patch_size) + total = data.sum() + if total == 0: + half_patch_size = tuple(n // 2 for n in self.patch_size) + message = ( + 'Empty probability map found:' + f' {self.get_probability_map_image(subject).path}' + '\nVoxels with positive probability might be near the image' + ' border.\nIf you suspect that this is the case, try adding a' + ' padding transform\nwith half the patch size:' + f' torchio.Pad({half_patch_size})' + ) + raise RuntimeError(message) + data /= total # normalize probabilities + return data + + @staticmethod + def clear_probability_borders( + probability_map: np.ndarray, + patch_size: np.ndarray, + ) -> None: + # Set probability to 0 on voxels that wouldn't possibly be sampled + # given the current patch size + # We will arbitrarily define the center of an array with even length + # using the // Python operator + # For example, the center of an array (3, 4) will be on (1, 2) + # + # Patch center + # . . . . . . . . + # . . . . -> . . x . + # . . . . . . . . + # + # + # Prob. map After preprocessing + # + # x x x x x x x . . . . . . . + # x x x x x x x . . x x x x . + # x x x x x x x --> . . x x x x . + # x x x x x x x --> . . x x x x . + # x x x x x x x . . x x x x . + # x x x x x x x . . . . . . . + # + # The dots represent removed probabilities, x mark possible locations + crop_ini = patch_size // 2 + crop_fin = (patch_size - 1) // 2 + crop_i, crop_j, crop_k = crop_ini + probability_map[:crop_i, :, :] = 0 + probability_map[:, :crop_j, :] = 0 + probability_map[:, :, :crop_k] = 0 + + # The call tolist() is very important. Using np.uint16 as negative + # index will not work because e.g. -np.uint16(2) == 65534 + crop_i, crop_j, crop_k = (int(value) for value in crop_fin.tolist()) + if crop_i: + probability_map[-crop_i:, :, :] = 0 + if crop_j: + probability_map[:, -crop_j:, :] = 0 + if crop_k: + probability_map[:, :, -crop_k:] = 0 + + @staticmethod + def get_cumulative_distribution_function( + probability_map: np.ndarray, + ) -> np.ndarray: + """Return the cumulative distribution function of a probability map.""" + flat_map = probability_map.flatten() + flat_map_normalized = flat_map / flat_map.sum() + cdf = np.cumsum(flat_map_normalized) + return cdf + + @overload + def extract_patch( + self, + subject: Subject, + index_ini: TypeTripletInt, + cdf: None = None, + ) -> Subject: ... + + @overload + def extract_patch( + self, + subject: Subject, + probability_map: np.ndarray, + cdf: np.ndarray, + ) -> Subject: ... + + def extract_patch( + self, + subject: Subject, + index_ini_or_probability_map: TypeTripletInt | np.ndarray, + cdf: np.ndarray | None = None, + ) -> Subject: + if cdf is None: + if isinstance(index_ini_or_probability_map, np.ndarray): + index_ini = ( + int(index_ini_or_probability_map[0]), + int(index_ini_or_probability_map[1]), + int(index_ini_or_probability_map[2]), + ) + else: + index_ini = index_ini_or_probability_map + return super().extract_patch(subject, index_ini) + + if not isinstance(index_ini_or_probability_map, np.ndarray): + message = 'Probability map must be a NumPy array when using a CDF' + raise TypeError(message) + + i, j, k = self.get_random_index_ini(index_ini_or_probability_map, cdf) + index_ini = i, j, k + return super().extract_patch(subject, index_ini) + + def get_random_index_ini( + self, + probability_map: np.ndarray, + cdf: np.ndarray, + ) -> np.ndarray: + center = self.sample_probability_map(probability_map, cdf) + assert np.all(center >= 0) + # See self.clear_probability_borders + index_ini = center - self.patch_size // 2 + assert np.all(index_ini >= 0) + return index_ini + + @classmethod + def sample_probability_map( + cls, + probability_map: np.ndarray, + cdf: np.ndarray, + ) -> np.ndarray: + """Inverse transform sampling. + + Examples: + >>> probability_map = np.array( + ... ((0,0,1,1,5,2,1,1,0), + ... (2,2,2,2,2,2,2,2,2))) + >>> probability_map + array([[0, 0, 1, 1, 5, 2, 1, 1, 0], + [2, 2, 2, 2, 2, 2, 2, 2, 2]]) + >>> histogram = np.zeros_like(probability_map) + >>> for _ in range(100000): + ... histogram[WeightedSampler.sample_probability_map(probability_map, cdf)] += 1 # doctest:+SKIP + ... + >>> histogram # doctest:+SKIP + array([[ 0, 0, 3479, 3478, 17121, 7023, 3355, 3378, 0], + [ 6808, 6804, 6942, 6809, 6946, 6988, 7002, 6826, 7041]]) + """ + # Get first value larger than random number ensuring the random number + # is not exactly 0 (see https://github.com/TorchIO-project/torchio/issues/510) + random_number = max(MIN_FLOAT_32, torch.rand(1).item()) * cdf[-1] + + random_location_index = np.searchsorted(cdf, random_number) + + center = np.unravel_index( + random_location_index, + probability_map.shape, + ) + + probability = probability_map[center] + if probability <= 0: + message = ( + 'Error retrieving probability in weighted sampler.' + ' Please report this issue at' + ' https://github.com/TorchIO-project/torchio/issues/new?labels=bug&template=bug_report.md' + ) + raise RuntimeError(message) + + return np.array(center) diff --git a/src/torchio/data/subject.py b/src/torchio/data/subject.py index fe0d5e543..4f83ba3a5 100644 --- a/src/torchio/data/subject.py +++ b/src/torchio/data/subject.py @@ -1,367 +1,539 @@ -"""Subject class.""" - from __future__ import annotations -import types -from collections.abc import Iterator +import copy +import pprint +from collections.abc import Mapping +from collections.abc import Sequence +from typing import TYPE_CHECKING from typing import Any +from typing import Literal +from typing import TypeAlias +from typing import overload import numpy as np -import torch -from typing_extensions import Self -from ..types import TypeSpacing -from ..types import TypeSpatialShape -from ..types import TypeTensorShape -from .bboxes import BoundingBoxes +from ..utils import get_subclasses from .image import Image -from .invertible import Invertible -from .points import Points +from .image import LabelMap +from .image import ScalarImage -# Union of all spatial data types stored by Subject -_SpatialData = Image | Points | BoundingBoxes +if TYPE_CHECKING: + from matplotlib.figure import Figure + from ..transforms import Compose + from ..transforms import Transform -class Subject(Invertible): - """Container for images, points, bounding boxes, and metadata. - A `Subject` holds one or more named data entries and optional - metadata. Data entries are classified automatically: +AppliedTransformParameters: TypeAlias = dict[str, object] +AppliedTransform: TypeAlias = tuple[str, AppliedTransformParameters] - - [`Image`][torchio.Image] (including `ScalarImage`, `LabelMap`) - - [`Points`][torchio.Points] - - [`BoundingBoxes`][torchio.BoundingBoxes] - - Everything else is stored as metadata. - At least one `Image` must be provided. +class Subject(dict[str, object]): + """Class to store information about the images corresponding to a subject. Args: - **kwargs: Named data entries and/or metadata values. + *args: If provided, a dictionary of items. + **kwargs: Items that will be added to the subject sample. Examples: - >>> import torch, torchio as tio + >>> import torchio as tio + >>> # One way: >>> subject = tio.Subject( - ... t1=tio.ScalarImage("t1.nii.gz"), - ... seg=tio.LabelMap("seg.nii.gz"), - ... landmarks=tio.Points(torch.randn(10, 3)), - ... tumors=tio.BoundingBoxes( - ... torch.tensor([[10, 20, 30, 50, 60, 70]]), - ... format=tio.BoundingBoxFormat.IJKIJK, - ... ), + ... one_image=tio.ScalarImage('path_to_image.nii.gz'), + ... a_segmentation=tio.LabelMap('path_to_seg.nii.gz'), ... age=45, + ... name='John Doe', + ... hospital='Hospital Juan Negrín', ... ) - >>> subject.t1 # Image access - >>> subject.landmarks # Points access - >>> subject.tumors # BoundingBoxes access - >>> subject.age # metadata access (returns 45) + >>> # If you want to create the mapping before, or have spaces in the keys: + >>> subject_dict = { + ... 'one image': tio.ScalarImage('path_to_image.nii.gz'), + ... 'a segmentation': tio.LabelMap('path_to_seg.nii.gz'), + ... 'age': 45, + ... 'name': 'John Doe', + ... 'hospital': 'Hospital Juan Negrín', + ... } + >>> subject = tio.Subject(subject_dict) """ - def __init__(self, **kwargs: Any) -> None: - images: dict[str, Image] = {} - points: dict[str, Points] = {} - bounding_boxes: dict[str, BoundingBoxes] = {} - metadata: dict[str, Any] = {} - - for k, v in kwargs.items(): - if isinstance(v, Image): - images[k] = v - elif isinstance(v, Points): - points[k] = v - elif isinstance(v, BoundingBoxes): - bounding_boxes[k] = v + def __init__(self, *args: Mapping[str, object], **kwargs: object): + if args: + if len(args) == 1 and isinstance(args[0], Mapping): + kwargs.update(args[0]) else: - metadata[k] = v - - if not images and not points and not bounding_boxes and not metadata: - msg = "A Subject must contain at least one entry" - raise ValueError(msg) - - self._images: dict[str, Image] = images - self._points: dict[str, Points] = points - self._bounding_boxes: dict[str, BoundingBoxes] = bounding_boxes - self._metadata: dict[str, Any] = metadata - self.applied_transforms: list[Any] = [] - - # --- Access --- - - def __getattr__(self, name: str) -> Any: - if name.startswith("_"): - raise AttributeError(name) - for store in (self._images, self._points, self._bounding_boxes): - if name in store: - return store[name] - if name in self._metadata: - return self._metadata[name] - msg = f"{type(self).__name__} has no attribute {name!r}" - raise AttributeError(msg) + message = 'Only one dictionary as positional argument is allowed' + raise ValueError(message) + super().__init__(**kwargs) + self._parse_images(self.get_images(intensity_only=False)) + self.update_attributes() # this allows me to do e.g. subject.t1 + self.applied_transforms: list[AppliedTransform] = [] + + def __repr__(self): + num_images = len(self.get_images(intensity_only=False)) + string = ( + f'{self.__class__.__name__}' + f'(Keys: {tuple(self.keys())}; images: {num_images})' + ) + return string - def __getitem__( - self, - item: str | int | slice | tuple[int | slice, ...], - ) -> _SpatialData | Subject: - """Look up a named entry, or spatially slice all images. + def _repr_html_(self): + try: + from matplotlib.figure import Figure + except ImportError: + return self.__repr__() - When *item* is a `str`, the corresponding data entry is - returned (image, points, or bounding boxes). + fig = self.plot(return_fig=True, show=False) + assert isinstance(fig, Figure) - When *item* is an `int`, `slice`, or `tuple` of - slices/ints, a **new** [`Subject`][torchio.Subject] is returned with every - image sliced identically. All images must be spatially - consistent (same `spatial_shape`). Only the **spatial** - dimensions `(I, J, K)` are sliced. The channel dimension of - each image is preserved. + from ..visualization import _figure_to_html - Args: - item: A string key, or an int/slice/tuple for spatial - indexing. + return _figure_to_html(fig) - Returns: - A single data entry (when *item* is `str`), or a new - [`Subject`][torchio.Subject] with sliced images. - - Examples: - >>> subject["t1"] # lookup by name - >>> subject[10:20] # slice dim I - >>> subject[10:20, 30:60] # slice I and J - >>> subject[..., 50:100] # slice dim K - >>> subject[10:20, 10:20, 10:20] # all three spatial dims - """ - if isinstance(item, str): - for store in (self._images, self._points, self._bounding_boxes): - if item in store: - return store[item] - raise KeyError(item) - - return self._spatial_slice(item) - - def __contains__(self, name: object) -> bool: - return any( - name in store - for store in ( - self._images, - self._points, - self._bounding_boxes, - ) - ) + def __len__(self): + return len(self.get_images(intensity_only=False)) - def __iter__(self) -> Iterator[str]: - yield from self._images - yield from self._points - yield from self._bounding_boxes + @overload + def __getitem__(self, item: str) -> object: ... - def __len__(self) -> int: - return len(self._images) + len(self._points) + len(self._bounding_boxes) + @overload + def __getitem__(self, item: slice | int | tuple[object, ...]) -> Subject: ... - # --- Properties --- + def __getitem__( + self, item: str | slice | int | tuple[object, ...] + ) -> object | Subject: + if isinstance(item, (slice, int, tuple)): + try: + self.check_consistent_spatial_shape() + except RuntimeError as e: + message = ( + 'To use indexing, all images in the subject must have the' + ' same spatial shape' + ) + raise RuntimeError(message) from e + copied = copy.deepcopy(self) + for image_name, image in copied.items(): + if isinstance(image, Image): + copied[image_name] = image[item] + return copied + else: + return super().__getitem__(item) + + def __getattr__(self, item: str) -> Any: + try: + return self[item] + except KeyError as error: + raise AttributeError( + f'{self.__class__.__name__!s} has no attribute {item!r}', + ) from error + + @staticmethod + def _parse_images(images: list[Image]) -> None: + # Check that it's not empty + if not images: + raise TypeError('A subject without images cannot be created') @property - def metadata(self) -> dict[str, Any]: - """Non-spatial metadata.""" - return self._metadata + def shape(self): + """Return shape of first image in subject. - @property - def spatial_shape(self) -> TypeSpatialShape: - """Spatial shape, checked for consistency across all images.""" - self._check_consistent_attribute("spatial_shape") - return self._first_image().spatial_shape + Consistency of shapes across images in the subject is checked first. - @property - def shape(self) -> TypeTensorShape: - """Shape of the first image, checked for consistency.""" - self._check_consistent_attribute("shape") - return self._first_image().shape + Examples: + >>> import torchio as tio + >>> colin = tio.datasets.Colin27() + >>> colin.shape + (1, 181, 217, 181) + """ + self.check_consistent_attribute('shape') + return self.get_first_image().shape @property - def spacing(self) -> TypeSpacing: - """Spacing from the first image, checked for consistency.""" - self._check_consistent_attribute("spacing") - return self._first_image().spacing + def spatial_shape(self): + """Return spatial shape of first image in subject. - @property - def device(self) -> torch.device: - """Device of the data, checked for consistency across all entries.""" - devices: list[torch.device] = [] - for image in self._images.values(): - devices.append(image.device) - for pts in self._points.values(): - devices.append(pts.device) - for boxes in self._bounding_boxes.values(): - devices.append(boxes.device) - if not devices: - return torch.device("cpu") - ref = devices[0] - if not all(d == ref for d in devices): - msg = f"Inconsistent devices: {devices}" - raise RuntimeError(msg) - return ref - - # --- Methods --- + Consistency of spatial shapes across images in the subject is checked + first. - @property - def images(self) -> dict[str, Image]: - """Dict of all `Image` entries.""" - return dict(self._images) + Examples: + >>> import torchio as tio + >>> colin = tio.datasets.Colin27() + >>> colin.spatial_shape + (181, 217, 181) + """ + self.check_consistent_spatial_shape() + return self.get_first_image().spatial_shape @property - def points(self) -> dict[str, Points]: - """Dict of all `Points` entries.""" - return dict(self._points) + def spacing(self): + """Return spacing of first image in subject. + + Consistency of spacings across images in the subject is checked first. + + Examples: + >>> import torchio as tio + >>> colin = tio.datasets.Slicer() + >>> colin.spacing + (1.0, 1.0, 1.2999954223632812) + """ + self.check_consistent_attribute('spacing') + return self.get_first_image().spacing @property - def bounding_boxes(self) -> dict[str, BoundingBoxes]: - """Dict of all `BoundingBoxes` entries.""" - return dict(self._bounding_boxes) + def history(self): + # Kept for backwards compatibility + return self.get_applied_transforms() - def all_points(self) -> dict[str | tuple[str, str], Points]: - """Collect points from both subject-level and image-level. + def is_2d(self): + return all(i.is_2d() for i in self.get_images(intensity_only=False)) - Subject-level points are keyed by their name (`str`). - Image-level points are keyed by a `(image_name, points_name)` - tuple. + def get_applied_transforms( + self, + ignore_intensity: bool = False, + image_interpolation: str | None = None, + ) -> list[Transform]: + from ..transforms.intensity_transform import IntensityTransform + from ..transforms.transform import Transform + + name_to_transform = {cls.__name__: cls for cls in get_subclasses(Transform)} + transforms_list = [] + for transform_name, arguments in self.applied_transforms: + transform = name_to_transform[transform_name](**arguments) + if ignore_intensity and isinstance(transform, IntensityTransform): + continue + resamples = hasattr(transform, 'image_interpolation') + if resamples and image_interpolation is not None: + parsed = transform.parse_interpolation(image_interpolation) + transform.image_interpolation = parsed + transforms_list.append(transform) + return transforms_list + + def get_composed_history( + self, + ignore_intensity: bool = False, + image_interpolation: str | None = None, + ) -> Compose: + from ..transforms.augmentation.composition import Compose + + transforms = self.get_applied_transforms( + ignore_intensity=ignore_intensity, + image_interpolation=image_interpolation, + ) + return Compose(transforms) - Returns: - Merged dict of all points across both levels. - """ - result: dict[str | tuple[str, str], Points] = {} - result.update(self._points) - for image_name, image in self._images.items(): - for points_name, pts in image.points.items(): - result[(image_name, points_name)] = pts - return result - - def all_bounding_boxes( + def get_inverse_transform( self, - ) -> dict[str | tuple[str, str], BoundingBoxes]: - """Collect bounding boxes from both subject-level and image-level. + warn: bool = True, + ignore_intensity: bool = False, + image_interpolation: str | None = None, + ) -> Compose: + """Get a reversed list of the inverses of the applied transforms. + + Args: + warn: Issue a warning if some transforms are not invertible. + ignore_intensity: If `True`, all instances of + `IntensityTransform` + will be ignored. + image_interpolation: Modify interpolation for scalar images inside + transforms that perform resampling. + """ + history_transform = self.get_composed_history( + ignore_intensity=ignore_intensity, + image_interpolation=image_interpolation, + ) + inverse_transform = history_transform.inverse(warn=warn) + return inverse_transform - Subject-level boxes are keyed by their name (`str`). - Image-level boxes are keyed by a `(image_name, boxes_name)` - tuple. + def apply_inverse_transform(self, **kwargs) -> Subject: + """Apply the inverse of all applied transforms, in reverse order. - Returns: - Merged dict of all bounding boxes across both levels. + Args: + **kwargs: Keyword arguments passed on to + [`get_inverse_transform()`][torchio.data.subject.Subject.get_inverse_transform]. """ - result: dict[str | tuple[str, str], BoundingBoxes] = {} - result.update(self._bounding_boxes) - for image_name, image in self._images.items(): - for box_name, boxes in image.bounding_boxes.items(): - result[(image_name, box_name)] = boxes - return result + inverse_transform = self.get_inverse_transform(**kwargs) + transformed = inverse_transform(self) + transformed.clear_history() + return transformed - def load(self) -> None: - """Load all images from disk.""" - for image in self._images.values(): - image.load() + def clear_history(self) -> None: + self.applied_transforms = [] - def to(self, *args: Any, **kwargs: Any) -> Self: - """Move all data to a device and/or cast to a dtype. + def check_consistent_attribute( + self, + attribute: str, + relative_tolerance: float = 1e-6, + absolute_tolerance: float = 1e-6, + message: str | None = None, + ) -> None: + r"""Check for consistency of an attribute across all images. - Calls `.to()` on every Image, Points, and BoundingBoxes. + Args: + attribute: Name of the image attribute to check + relative_tolerance: Relative tolerance for `numpy.allclose()` + absolute_tolerance: Absolute tolerance for `numpy.allclose()` - Returns: - `self` (modified in-place). + Examples: + >>> import numpy as np + >>> import torch + >>> import torchio as tio + >>> scalars = torch.randn(1, 512, 512, 100) + >>> mask = torch.tensor(scalars > 0).type(torch.int16) + >>> af1 = np.eye([0.8, 0.8, 2.50000000000001, 1]) + >>> af2 = np.eye([0.8, 0.8, 2.49999999999999, 1]) # small difference here (e.g. due to different reader) + >>> subject = tio.Subject( + ... image = tio.ScalarImage(tensor=scalars, affine=af1), + ... mask = tio.LabelMap(tensor=mask, affine=af2) + ... ) + >>> subject.check_consistent_attribute('spacing') # no error as tolerances are > 0 + + Note: + To check that all values for a specific attribute are close + between all images in the subject, `numpy.allclose()` is used. + This function returns `True` if + $|a_i - b_i| \leq t_{abs} + t_{rel} * |b_i|$, where + $a_i$ and $b_i$ are the $i$-th element of the same + attribute of two images being compared, + $t_{abs}$ is the `absolute_tolerance` and + $t_{rel}$ is the `relative_tolerance`. """ - for image in self._images.values(): - image.to(*args, **kwargs) - for pts in self._points.values(): - pts.to(*args, **kwargs) - for boxes in self._bounding_boxes.values(): - boxes.to(*args, **kwargs) - return self + message = ( + f'More than one value for "{attribute}" found in subject images:\n{{}}' + ) - # --- Internal --- + names_images = self.get_images_dict(intensity_only=False).items() + try: + first_attribute = None + first_image = None + + for image_name, image in names_images: + if first_attribute is None: + first_attribute = getattr(image, attribute) + first_image = image_name + continue + current_attribute = getattr(image, attribute) + all_close = np.allclose( + current_attribute, + first_attribute, + rtol=relative_tolerance, + atol=absolute_tolerance, + ) + if not all_close: + message = message.format( + pprint.pformat( + { + first_image: first_attribute, + image_name: current_attribute, + } + ), + ) + raise RuntimeError(message) + except TypeError: + # fallback for non-numeric values + values_dict = {} + for image_name, image in names_images: + values_dict[image_name] = getattr(image, attribute) + num_unique_values = len(set(values_dict.values())) + if num_unique_values > 1: + message = message.format(pprint.pformat(values_dict)) + raise RuntimeError(message) from None + + def check_consistent_spatial_shape(self) -> None: + self.check_consistent_attribute('spatial_shape') + + def check_consistent_orientation(self) -> None: + self.check_consistent_attribute('orientation') + + def check_consistent_affine(self) -> None: + self.check_consistent_attribute('affine') + + def check_consistent_space(self) -> None: + try: + self.check_consistent_attribute('spacing') + self.check_consistent_attribute('direction') + self.check_consistent_attribute('origin') + self.check_consistent_spatial_shape() + except RuntimeError as e: + message = ( + 'As described above, some images in the subject are not in the' + ' same space. You probably can use the transforms ToCanonical' + ' and Resample to fix this, as explained at' + ' https://github.com/TorchIO-project/torchio/issues/647#issuecomment-913025695' + ) + raise RuntimeError(message) from e - def _first_image(self) -> Image: - return next(iter(self._images.values())) + def get_images_names(self) -> list[str]: + return list(self.get_images_dict(intensity_only=False).keys()) - def _check_consistent_attribute( + @overload + def get_images_dict( self, - attribute: str, - rtol: float = 1e-6, - atol: float = 1e-6, - ) -> None: - """Check that an attribute is consistent across all images.""" - values = [] - names = [] - for name, image in self._images.items(): - values.append(getattr(image, attribute)) - names.append(name) - - if len(values) < 2: - return - - ref = values[0] - for i, val in enumerate(values[1:], 1): - if not np.allclose(ref, val, rtol=rtol, atol=atol): - msg = f"Inconsistent {attribute}: {names[0]}={ref}, {names[i]}={val}" - raise RuntimeError(msg) - - def _spatial_slice( + intensity_only: Literal[True] = True, + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, + ) -> dict[str, ScalarImage]: ... + + @overload + def get_images_dict( self, - item: int | slice | tuple[int | slice, ...], - ) -> Subject: - """Slice all images along spatial dimensions (I, J, K). + intensity_only: Literal[False] = False, + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, + ) -> dict[str, Image]: ... - The channel dimension of each image is preserved. All images - must have the same `spatial_shape`. - """ - if not self._images: - msg = "Cannot spatially slice a Subject with no images" - raise RuntimeError(msg) + def get_images_dict( + self, + intensity_only: bool = True, + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, + ) -> dict[str, ScalarImage] | dict[str, Image]: + if intensity_only: + scalar_images: dict[str, ScalarImage] = {} + for image_name, image in self.items(): + if not isinstance(image, ScalarImage): + continue + if include is not None and image_name not in include: + continue + if exclude is not None and image_name in exclude: + continue + scalar_images[image_name] = image + return scalar_images + + all_images: dict[str, Image] = {} + for image_name, image in self.items(): + if not isinstance(image, Image): + continue + if include is not None and image_name not in include: + continue + if exclude is not None and image_name in exclude: + continue + all_images[image_name] = image + return all_images + + @overload + def get_images( + self, + intensity_only: Literal[True] = True, + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, + ) -> list[ScalarImage]: ... - self._check_consistent_attribute("spatial_shape") + @overload + def get_images( + self, + intensity_only: Literal[False] = False, + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, + ) -> list[Image]: ... - # Normalise to tuple - if isinstance(item, (int, slice)) or item is Ellipsis: - items: tuple[int | slice | types.EllipsisType, ...] = (item,) - elif isinstance(item, tuple): - items = item - else: - msg = f"Index type {type(item).__name__} not understood" - raise TypeError(msg) - - # Slice each image, prepending slice(None) for channels - sliced_images: dict[str, Image] = {} - for name, image in self._images.items(): - sliced_images[name] = image[(slice(None), *items)] - - kwargs: dict[str, Any] = dict(sliced_images) - kwargs.update(self._points) - kwargs.update(self._bounding_boxes) - kwargs.update(self._metadata) - new = type(self)(**kwargs) - new.applied_transforms = list(self.applied_transforms) - return new - - def __repr__(self) -> str: - parts = [] - if self._images: - parts.append(f"images: {tuple(self._images.keys())}") - if self._points: - parts.append(f"points: {tuple(self._points.keys())}") - if self._bounding_boxes: - parts.append(f"bboxes: {tuple(self._bounding_boxes.keys())}") - return f"{type(self).__name__}({'; '.join(parts)})" - - def _repr_html_(self) -> str: - """Rich HTML representation for Jupyter notebooks.""" - from ..repr_html import subject_to_html - - return subject_to_html(self) - - def plot(self, **kwargs: Any) -> Any: - """Plot all images as a grid of orthogonal slices. - - Requires the `[plot]` extras (`pip install torchio[plot]`). - See [`plot_subject`][torchio.visualization.plot_subject] for the - full list of keyword arguments. - """ - from ..visualization import plot_subject + @overload + def get_images( + self, + intensity_only: bool, + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, + ) -> list[ScalarImage] | list[Image]: ... - return plot_subject(self, **kwargs) + def get_images( + self, + intensity_only: bool = True, + include: Sequence[str] | None = None, + exclude: Sequence[str] | None = None, + ) -> list[ScalarImage] | list[Image]: + if intensity_only: + scalar_images = self.get_images_dict( + intensity_only=True, + include=include, + exclude=exclude, + ) + return list(scalar_images.values()) + + all_images = self.get_images_dict( + intensity_only=False, + include=include, + exclude=exclude, + ) + return list(all_images.values()) + + def get_image(self, image_name: str) -> Image: + """Get a single image by its name.""" + return self.get_images_dict(intensity_only=False)[image_name] + + def get_scalar_image(self, image_name: str) -> ScalarImage: + image = self.get_image(image_name) + if not isinstance(image, ScalarImage): + message = f'Image "{image_name}" is not a scalar image' + raise TypeError(message) + return image + + def get_label_map(self, image_name: str) -> LabelMap: + image = self.get_image(image_name) + if not isinstance(image, LabelMap): + message = f'Image "{image_name}" is not a label map' + raise TypeError(message) + return image + + def get_first_image(self) -> Image: + return self.get_images(intensity_only=False)[0] + + def add_transform( + self, + transform: Transform, + parameters_dict: AppliedTransformParameters, + ) -> None: + self.applied_transforms.append((transform.name, parameters_dict)) + + def load(self) -> None: + """Load images in subject on RAM.""" + for image in self.get_images(intensity_only=False): + image.load() + + def unload(self) -> None: + """Unload images in subject.""" + for image in self.get_images(intensity_only=False): + image.unload() + def update_attributes(self) -> None: + # This allows to get images using attribute notation, e.g. subject.t1 + self.__dict__.update(self) + + @staticmethod + def _check_image_name(image_name: object) -> str: + if not isinstance(image_name, str): + message = ( + f'The image name must be a string, but it has type "{type(image_name)}"' + ) + raise ValueError(message) + return image_name + + def add_image(self, image: Image, image_name: str) -> None: + """Add an image to the subject instance.""" + if not isinstance(image, Image): + message = ( + 'Image must be an instance of torchio.Image,' + f' but its type is "{type(image)}"' + ) + raise ValueError(message) + self._check_image_name(image_name) + self[image_name] = image + self.update_attributes() + + def remove_image(self, image_name: str) -> None: + """Remove an image from the subject instance.""" + self._check_image_name(image_name) + del self[image_name] + delattr(self, image_name) + + def plot(self, return_fig: bool = False, **kwargs) -> None | Figure: + """Plot images using matplotlib. + + Args: + return_fig: If ``True``, return the figure instead of showing it. + **kwargs: Keyword arguments that will be passed on to + [`plot()`][torchio.Image.plot]. + """ + from ..visualization import plot_subject # avoid circular import -# In DICOM terminology, a "study" contains a set of "series" (volumes). -# This maps directly to Subject (a container of named images + metadata). -# Both names are provided so users can pick whichever fits their application. -Study = Subject + figure = plot_subject(self, **kwargs) + if return_fig: + return figure + return None diff --git a/src/torchio/datasets/__init__.py b/src/torchio/datasets/__init__.py index 73ddfc7b7..e5ef17ed8 100644 --- a/src/torchio/datasets/__init__.py +++ b/src/torchio/datasets/__init__.py @@ -1,41 +1,49 @@ -"""Built-in demo subjects and datasets.""" - +from .bite import BITE3 +from .ct_rate import CtRate +from .episurg import EPISURG from .fpg import FPG from .itk_snap import T1T2 from .itk_snap import AorticValve from .itk_snap import BrainTumor -from .ixi import ixi -from .ixi import ixi_tiny -from .medmnist import adrenal_mnist_3d -from .medmnist import fracture_mnist_3d -from .medmnist import nodule_mnist_3d -from .medmnist import organ_mnist_3d -from .medmnist import synapse_mnist_3d -from .medmnist import vessel_mnist_3d +from .ixi import IXI +from .ixi import IXITiny +from .medmnist import AdrenalMNIST3D +from .medmnist import FractureMNIST3D +from .medmnist import NoduleMNIST3D +from .medmnist import OrganMNIST3D +from .medmnist import SynapseMNIST3D +from .medmnist import VesselMNIST3D from .mni import Colin27 from .mni import ICBM2009CNonlinearSymmetric from .mni import Pediatric from .mni import Sheep +from .rsna_miccai import RSNAMICCAI +from .rsna_spine_fracture import RSNACervicalSpineFracture from .slicer import Slicer from .zone_plate import ZonePlate __all__ = [ - "FPG", - "T1T2", - "AorticValve", - "BrainTumor", - "Colin27", - "ICBM2009CNonlinearSymmetric", - "Pediatric", - "Sheep", - "Slicer", - "ZonePlate", - "adrenal_mnist_3d", - "fracture_mnist_3d", - "ixi", - "ixi_tiny", - "nodule_mnist_3d", - "organ_mnist_3d", - "synapse_mnist_3d", - "vessel_mnist_3d", + 'FPG', + 'Slicer', + 'BITE3', + 'CtRate', + 'IXI', + 'IXITiny', + 'RSNAMICCAI', + 'RSNACervicalSpineFracture', + 'EPISURG', + 'BrainTumor', + 'T1T2', + 'AorticValve', + 'Colin27', + 'Sheep', + 'Pediatric', + 'ICBM2009CNonlinearSymmetric', + 'OrganMNIST3D', + 'NoduleMNIST3D', + 'AdrenalMNIST3D', + 'FractureMNIST3D', + 'VesselMNIST3D', + 'SynapseMNIST3D', + 'ZonePlate', ] diff --git a/src/torchio/datasets/bite.py b/src/torchio/datasets/bite.py new file mode 100644 index 000000000..3b12ae7d7 --- /dev/null +++ b/src/torchio/datasets/bite.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import abc +from pathlib import Path + +from ..data import Image +from ..data import LabelMap +from ..data import ScalarImage +from ..data import Subject +from ..data import SubjectsDataset +from ..download import download_and_extract_archive +from ..transforms import Transform +from ..types import TypePath + + +class BITE(SubjectsDataset, abc.ABC): + base_url = 'http://www.bic.mni.mcgill.ca/uploads/Services/' + + def __init__( + self, + root: TypePath, + transform: Transform | None = None, + download: bool = False, + **kwargs, + ): + root = Path(root).expanduser().absolute() + if download: + self._download(root) + subjects_list = self._get_subjects_list(root) + self.kwargs = kwargs + super().__init__(subjects_list, transform=transform, **kwargs) + + @abc.abstractmethod + def _download(self, root: Path): + raise NotImplementedError + + @abc.abstractmethod + def _get_subjects_list(self, root: Path): + raise NotImplementedError + + +class BITE3(BITE): + """Pre- and post-resection MR images in BITE. + + *The goal of BITE is to share in vivo medical images of patients wtith + brain tumors to facilitate the development and validation of new image + processing algorithms.* + + Please check the [BITE website](https://nist.mni.mcgill.ca/bite-brain-images-of-tumors-for-evaluation-database/) for more information and + acknowledgments instructions. + + + Args: + root: Root directory to which the dataset will be downloaded. + transform: An instance of + [`Transform`][torchio.transforms.transform.Transform]. + download: If set to `True`, will download the data into `root`. + """ + + dirname = 'group3' + + def _download(self, root: Path): + if (root / self.dirname).is_dir(): + return + root.mkdir(exist_ok=True, parents=True) + filename = f'{self.dirname}.tar.gz' + url = self.base_url + filename + download_and_extract_archive( + url, + download_root=root, + md5='e415b63887c40b727c45552614b44634', + ) + (root / filename).unlink() # cleanup + + def _get_subjects_list(self, root: Path): + subjects_dir = root / self.dirname + subjects = [] + for i in range(1, 15): + if i == 13: + continue # no MRI for this subject + subject_id = f'{i:02d}' + subject_dir = subjects_dir / subject_id + preop_path = subject_dir / f'{subject_id}_preop_mri.mnc' + postop_path = subject_dir / f'{subject_id}_postop_mri.mnc' + images_dict: dict[str, Image] = {} + images_dict['preop'] = ScalarImage(preop_path) + images_dict['postop'] = ScalarImage(postop_path) + for fp in subject_dir.glob('*tumor*'): + images_dict[fp.stem[3:]] = LabelMap(fp) + subject = Subject(images_dict) + subjects.append(subject) + return subjects diff --git a/src/torchio/datasets/ct_rate.py b/src/torchio/datasets/ct_rate.py new file mode 100644 index 000000000..b3f2a4ae8 --- /dev/null +++ b/src/torchio/datasets/ct_rate.py @@ -0,0 +1,389 @@ +from __future__ import annotations + +import enum +import multiprocessing +from pathlib import Path +from typing import TYPE_CHECKING +from typing import Literal +from typing import Union + +from tqdm.contrib.concurrent import thread_map + +from ..data.dataset import SubjectsDataset +from ..data.image import ScalarImage +from ..data.subject import Subject +from ..external.imports import get_pandas +from ..types import TypePath + +if TYPE_CHECKING: + from collections.abc import Hashable + + import pandas as pd + + +TypeSplit = Union[ + Literal['train'], + Literal['valid'], + Literal['validation'], +] + +TypeParallelism = Literal['thread', 'process', None] + + +class MetadataIndexColumn(str, enum.Enum): + SUBJECT_ID = 'subject_id' + SCAN_ID = 'scan_id' + RECONSTRUCTION_ID = 'reconstruction_id' + + +class CtRate(SubjectsDataset): + """CT-RATE dataset. + + This class helps loading the [CT-RATE dataset + ](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE), + which contains chest CT scans with associated radiology reports and + abnormality labels. + + The dataset must have been downloaded previously. + + Args: + root: Root directory where the dataset has been downloaded. + split: Dataset split to use, either `'train'` or `'validation'`. + num_subjects: Optional limit on the number of subjects to load (useful for + debugging). If `None`, all subjects in the split are loaded. + report_key: Key to use for storing radiology reports in the Subject metadata. + sizes: List of image sizes (in-plane, in voxels) to include. + load_fixed: If `True`, load the files with fixed spatial metadata + added in [this pull request + ](https://huggingface.co/datasets/ibrahimhamamci/CT-RATE/discussions/85). + Otherwise, load the original files with incorrect spatial metadata. + verify_paths: If `True`, verify that the paths to the images exist + during instantiation of the dataset. This might be slow for large that are + not stored locally. + **kwargs: Additional arguments for SubjectsDataset. + + Examples: + >>> from torchio.datasets import CtRate + >>> dataset = CtRate('/path/to/CT-RATE', sizes=[512]) + """ + + _REPO_ID = 'ibrahimhamamci/CT-RATE' + _FILENAME_KEY = 'VolumeName' + _SIZES = [512, 768, 1024] + ABNORMALITIES = [ + 'Medical material', + 'Arterial wall calcification', + 'Cardiomegaly', + 'Pericardial effusion', + 'Coronary artery wall calcification', + 'Hiatal hernia', + 'Lymphadenopathy', + 'Emphysema', + 'Atelectasis', + 'Lung nodule', + 'Lung opacity', + 'Pulmonary fibrotic sequela', + 'Pleural effusion', + 'Mosaic attenuation pattern', + 'Peribronchial thickening', + 'Consolidation', + 'Bronchiectasis', + 'Interlobular septal thickening', + ] + REPORT_KEYS = [ + 'ClinicalInformation_EN', + 'Findings_EN', + 'Impressions_EN', + 'Technique_EN', + ] + + def __init__( + self, + root: TypePath, + split: TypeSplit = 'train', + *, + num_subjects: int | None = None, + report_key: str = 'report', + sizes: list[int] | None = None, + load_fixed: bool = True, + verify_paths: bool = False, + **kwargs, + ): + self._root_dir = Path(root) + self._num_subjects = num_subjects + self._report_key = report_key + self._sizes = self._SIZES if sizes is None else sizes + + self._split = self._parse_split(split) + self.metadata = self._get_metadata() + self._load_fixed = load_fixed + self._verify_paths = verify_paths + subjects_list = self._get_subjects_list(self.metadata) + super().__init__(subjects_list, **kwargs) + + @staticmethod + def _parse_split(split: str) -> str: + """Normalize the split name. + + Converts 'validation' to 'valid' and validates that the split name + is one of the allowed values. + + Args: + split: The split name to parse ('train', 'valid', or 'validation'). + + Returns: + str: Normalized split name ('train' or 'valid'). + + Raises: + ValueError: If the split name is not one of the allowed values. + """ + if split in ['valid', 'validation']: + return 'valid' + if split not in ['train', 'valid']: + raise ValueError(f"Invalid split '{split}'. Use 'train' or 'valid'") + return split + + def _get_csv( + self, + dirname: str, + filename: str, + ) -> pd.DataFrame: + """Load a CSV file from the specified directory within the dataset. + + Args: + dirname: Directory name within 'dataset/' where the CSV is located. + filename: Name of the CSV file to load. + """ + subfolder = Path(f'dataset/{dirname}') + path = Path(self._root_dir, subfolder, filename) + pd = get_pandas() + table = pd.read_csv(path) + return table + + def _get_csv_prefix(self, expand_validation: bool = True) -> str: + """Get the prefix for CSV filenames based on the current split. + + Returns the appropriate prefix for CSV filenames based on the current split. + For the validation split, can either return 'valid' or 'validation' depending + on the expand_validation parameter. + + Args: + expand_validation: If `True` and split is `'valid'`, return + `'validation'`. Otherwise, return the split name as is. + """ + if expand_validation and self._split == 'valid': + prefix = 'validation' + else: + prefix = self._split + return prefix + + def _get_metadata(self) -> pd.DataFrame: + """Load and process the dataset metadata. + + Loads metadata from the appropriate CSV file, filters images by size, + extracts subject, scan, and reconstruction IDs from filenames, and + merges in reports and abnormality labels. + """ + dirname = 'metadata' + prefix = self._get_csv_prefix() + filename = f'{prefix}_metadata.csv' + metadata = self._get_csv(dirname, filename) + + # Exclude images with size not in self._sizes + rows_int = metadata['Rows'].astype(int) + metadata = metadata[rows_int.isin(self._sizes)] + + index_columns = [ + MetadataIndexColumn.SUBJECT_ID.value, + MetadataIndexColumn.SCAN_ID.value, + MetadataIndexColumn.RECONSTRUCTION_ID.value, + ] + pattern = r'\w+_(\d+)_(\w+)_(\d+)\.nii\.gz' + metadata[index_columns] = metadata[self._FILENAME_KEY].str.extract(pattern) + + if self._num_subjects is not None: + metadata = self._keep_n_subjects(metadata, self._num_subjects) + + # Add reports and abnormality labels to metadata, keeping only the rows for the + # images in the metadata table + metadata = self._merge(metadata, self._get_reports()) + metadata = self._merge(metadata, self._get_labels()) + + metadata.set_index(index_columns, inplace=True) + return metadata + + def _merge(self, base_df: pd.DataFrame, new_df: pd.DataFrame) -> pd.DataFrame: + """Merge a new dataframe into the base dataframe using the filename as the key. + + This method performs a left join between `base_df` and `new_df` using the + volume filename as the join key, ensuring that all records from `base_df` are + preserved while matching data from `new_df` is added. + + Args: + base_df: The primary dataframe to merge into. + new_df: The dataframe containing additional data to be merged. + + Returns: + pd.DataFrame: The merged dataframe with all rows from base_df and + matching columns from new_df. + """ + pd = get_pandas() + return pd.merge( + base_df, + new_df, + on=self._FILENAME_KEY, + how='left', + ) + + def _keep_n_subjects(self, metadata: pd.DataFrame, n: int) -> pd.DataFrame: + """Limit the metadata to the first `n` subjects. + + Args: + metadata: The complete metadata dataframe. + n: Maximum number of subjects to keep. + """ + unique_subjects = metadata['subject_id'].unique() + selected_subjects = unique_subjects[:n] + return metadata[metadata['subject_id'].isin(selected_subjects)] + + def _get_reports(self) -> pd.DataFrame: + """Load the radiology reports associated with the CT scans. + + Retrieves the CSV file containing radiology reports for the current split + (train or validation). + """ + dirname = 'radiology_text_reports' + prefix = self._get_csv_prefix() + filename = f'{prefix}_reports.csv' + return self._get_csv(dirname, filename) + + def _get_labels(self) -> pd.DataFrame: + """Load the abnormality labels for the CT scans. + + Retrieves the CSV file containing predicted abnormality labels for the + current split. + """ + dirname = 'multi_abnormality_labels' + prefix = self._get_csv_prefix(expand_validation=False) + filename = f'{prefix}_predicted_labels.csv' + return self._get_csv(dirname, filename) + + def _get_subjects_list(self, metadata: pd.DataFrame) -> list[Subject]: + """Create a list of Subject instances from the metadata. + + Processes the metadata to create Subject objects, each containing one or more + CT images. Processing is performed in parallel. + + Note: + This method uses parallelization to improve performance when creating + multiple Subject instances. + """ + df_no_index = metadata.reset_index() + num_subjects = df_no_index['subject_id'].nunique() + iterable = df_no_index.groupby('subject_id') + subjects = thread_map( + self._get_subject, + iterable, + max_workers=multiprocessing.cpu_count(), + total=num_subjects, + ) + return subjects + + def _get_subject( + self, + subject_id_and_metadata: tuple[Hashable, pd.DataFrame], + ) -> Subject: + """Create a Subject instance for a specific subject. + + Processes all images belonging to a single subject and creates a Subject + object containing those images. + + Args: + subject_id_and_metadata: A tuple containing the subject ID (string) and a + DataFrame containing metadata for all images associated to that subject. + """ + subject_id, subject_df = subject_id_and_metadata + subject_dict: dict[str, object] = {'subject_id': str(subject_id)} + for _, image_row in subject_df.iterrows(): + image = self._instantiate_image(image_row) + scan_id = image_row['scan_id'] + reconstruction_id = image_row['reconstruction_id'] + image_key = f'scan_{scan_id}_reconstruction_{reconstruction_id}' + subject_dict[image_key] = image + return Subject(subject_dict) + + def _instantiate_image(self, image_row: pd.Series) -> ScalarImage: + """Create a ScalarImage object for a specific image. + + Processes a row from the metadata DataFrame to create a ScalarImage object, + + Args: + image_row: A pandas Series representing a row from the metadata DataFrame, + containing information about a single image. + """ + image_dict = {str(key): value for key, value in image_row.to_dict().items()} + filename = image_dict[self._FILENAME_KEY] + if not isinstance(filename, str): + message = ( + f'Expected {self._FILENAME_KEY} to be a string, not {type(filename)!r}' + ) + raise TypeError(message) + relative_image_path = self._get_image_path( + filename, + load_fixed=self._load_fixed, + ) + image_path = self._root_dir / relative_image_path + report_dict = self._extract_report_dict(image_dict) + image_dict[self._report_key] = report_dict + image = ScalarImage(image_path, verify_path=self._verify_paths, **image_dict) + return image + + def _extract_report_dict(self, subject_dict: dict[str, object]) -> dict[str, str]: + """Extract radiology report information from the subject dictionary. + + Extracts the English radiology report components (clinical information, + findings, impressions, and technique) from the subject dictionary and + removes these keys from the original dictionary. + + Args: + subject_dict: Image metadata including report fields. + + Note: + This method modifies the input subject_dict by removing the report keys. + """ + report_dict = {} + for key in self.REPORT_KEYS: + value = subject_dict.pop(key) + if not isinstance(value, str): + message = ( + f'Expected report field {key!r} to be a string, not {type(value)!r}' + ) + raise TypeError(message) + report_dict[key] = value + return report_dict + + @staticmethod + def _get_image_path(filename: str, load_fixed: bool) -> Path: + """Construct the relative path to an image file within the dataset structure. + + Parses the filename to determine the hierarchical directory structure + where the image is stored in the CT-RATE dataset. + + Args: + filename: The name of the image file (e.g., 'train_2_a_1.nii.gz'). + + Returns: + Path: The relative path to the image file within the dataset directory. + + Examples: + >>> path = CtRate._get_image_path('train_2_a_1.nii.gz') + # Returns Path('dataset/train/train_2/train_2_a/train_2_a_1.nii.gz') + """ + parts = filename.split('_') + base_dir = 'dataset' + split_dir = parts[0] + if load_fixed: + split_dir = f'{split_dir}_fixed' + level1 = f'{parts[0]}_{parts[1]}' + level2 = f'{level1}_{parts[2]}' + return Path(base_dir, split_dir, level1, level2, filename) diff --git a/src/torchio/datasets/episurg.py b/src/torchio/datasets/episurg.py new file mode 100644 index 000000000..1938111fa --- /dev/null +++ b/src/torchio/datasets/episurg.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +import csv +from pathlib import Path + +from ..data import LabelMap +from ..data import ScalarImage +from ..data import Subject +from ..data import SubjectsDataset +from ..download import download_and_extract_archive +from ..transforms import Transform +from ..types import TypePath + + +class EPISURG(SubjectsDataset): + """ + [EPISURG ](https://doi.org/10.5522/04/9996158.v1) is a clinical dataset of + $T_1$-weighted MRI from 430 epileptic patients who underwent + resective brain surgery at the National Hospital of Neurology and + Neurosurgery (Queen Square, London, United Kingdom) between 1990 and 2018. + + The dataset comprises 430 postoperative MRI. The corresponding preoperative + MRI is present for 268 subjects. + + Three human raters segmented the resection cavity on partially overlapping + subsets of EPISURG. + + If you use this dataset for your research, you agree with the *Data use + agreement* presented at the EPISURG entry on the [UCL Research Data + Repository ](https://doi.org/10.5522/04/9996158.v1) and you must cite the + corresponding publications. + + Args: + root: Root directory to which the dataset will be downloaded. + transform: An instance of + [`Transform`][torchio.transforms.transform.Transform]. + download: If set to `True`, will download the data into `root`. + + Warning: + The size of this dataset is multiple GB. + If you set `download` to `True`, it will take some time + to be downloaded if it is not already present. + """ + + data_url = 'https://s3-eu-west-1.amazonaws.com/pstorage-ucl-2748466690/26153588/EPISURG.zip' + md5 = '5ec5831a2c6fbfdc8489ba2910a6504b' + + def __init__( + self, + root: TypePath, + transform: Transform | None = None, + download: bool = False, + **kwargs, + ): + root = Path(root).expanduser().absolute() + if download: + self._download(root) + subjects_list = self._get_subjects_list(root) + self.kwargs = kwargs + super().__init__(subjects_list, transform=transform, **kwargs) + + @staticmethod + def _check_exists(root, modalities): + for modality in modalities: + modality_dir = root / modality + if not modality_dir.is_dir(): + exists = False + break + else: + exists = True + return exists + + @staticmethod + def _get_subjects_list(root): + episurg_dir = root / 'EPISURG' + subjects_dir = episurg_dir / 'subjects' + csv_path = episurg_dir / 'subjects.csv' + with open(csv_path) as csvfile: + reader = csv.DictReader(csvfile) + subjects = [] + for row in reader: + subject_id = row['Subject'] + subject_dir = subjects_dir / subject_id + subject_dict = { + 'subject_id': subject_id, + 'hemisphere': row['Hemisphere'], + 'surgery_type': row['Type'], + } + preop_dir = subject_dir / 'preop' + preop_paths = list(preop_dir.glob('*preop*')) + assert len(preop_paths) <= 1 + if preop_paths: + subject_dict['preop_mri'] = ScalarImage(preop_paths[0]) + postop_dir = subject_dir / 'postop' + postop_path = list(postop_dir.glob('*postop-t1mri*'))[0] + subject_dict['postop_mri'] = ScalarImage(postop_path) + for seg_path in postop_dir.glob('*seg*'): + seg_id = seg_path.name[-8] + subject_dict[f'seg_{seg_id}'] = LabelMap(seg_path) + subjects.append(Subject(**subject_dict)) + return subjects + + def _download(self, root: Path): + """Download the EPISURG data if it does not exist already.""" + if (root / 'EPISURG').is_dir(): + return + root.mkdir(exist_ok=True, parents=True) + download_and_extract_archive( + self.data_url, + download_root=root, + md5=self.md5, + ) + (root / 'EPISURG.zip').unlink() # cleanup + + def _glob_subjects(self, string): + subjects = [] + for subject in self._subjects: + for image_name in subject: + if string in image_name: + subjects.append(subject) + break + return subjects + + def _get_labeled_subjects(self): + return self._glob_subjects('seg') + + def _get_paired_subjects(self): + return self._glob_subjects('preop') + + def _get_subset(self, subjects): + dataset = SubjectsDataset( + subjects, + transform=self._transform, + **(self.kwargs), + ) + return dataset + + def get_labeled(self) -> SubjectsDataset: + """Get dataset from subjects with manual annotations.""" + return self._get_subset(self._get_labeled_subjects()) + + def get_unlabeled(self) -> SubjectsDataset: + """Get dataset from subjects without manual annotations.""" + subjects = [s for s in self._subjects if s not in self._get_labeled_subjects()] + return self._get_subset(subjects) + + def get_paired(self) -> SubjectsDataset: + """Get dataset from subjects with pre- and post-op MRI.""" + return self._get_subset(self._get_paired_subjects()) diff --git a/src/torchio/datasets/fpg.py b/src/torchio/datasets/fpg.py index 5aaf6a122..4c6e68060 100644 --- a/src/torchio/datasets/fpg.py +++ b/src/torchio/datasets/fpg.py @@ -1,84 +1,76 @@ -"""FPG brain MRI dataset.""" - -from __future__ import annotations - import urllib.parse -from typing import Any -from typing import ClassVar +from ..constants import DATA_REPO from ..data import LabelMap from ..data import ScalarImage +from ..data.io import read_matrix from ..data.subject import Subject from ..download import download_url -from ..download import get_torchio_cache_dir -from ..io import read_matrix - -DATA_REPO = "https://github.com/TorchIO-project/torchio-data/raw/main/data/" +from ..utils import get_torchio_cache_dir class FPG(Subject): - r"""3T :math:`T_1`-weighted brain MRI and corresponding parcellation. + """3T $T_1$-weighted brain MRI and corresponding parcellation. Args: load_all: If `True`, three more images will be loaded: a - :math:`T_2`-weighted MRI, a diffusion MRI and a functional MRI. + $T_2$-weighted MRI, a diffusion MRI and a functional MRI. """ - def __init__(self, load_all: bool = False) -> None: - repo_dir = urllib.parse.urljoin(DATA_REPO, "fernando/") + def __init__(self, load_all: bool = False): + repo_dir = urllib.parse.urljoin(DATA_REPO, 'fernando/') - filenames: dict[str, str] = { - "t1": "t1.nii.gz", - "seg": "t1_seg_gif.nii.gz", - "rigid": "t1_to_mni.tfm", - "affine": "t1_to_mni_affine.h5", + self.filenames = { + 't1': 't1.nii.gz', + 'seg': 't1_seg_gif.nii.gz', + 'rigid': 't1_to_mni.tfm', + 'affine': 't1_to_mni_affine.h5', } if load_all: - filenames["t2"] = "t2.nii.gz" - filenames["fmri"] = "fmri.nrrd" - filenames["dmri"] = "dmri.nrrd" + self.filenames['t2'] = 't2.nii.gz' + self.filenames['fmri'] = 'fmri.nrrd' + self.filenames['dmri'] = 'dmri.nrrd' - download_root = get_torchio_cache_dir() / "fpg" + download_root = get_torchio_cache_dir() / 'fpg' - for filename in filenames.values(): + for filename in self.filenames.values(): download_url( urllib.parse.urljoin(repo_dir, filename), download_root, filename=filename, ) - rigid = read_matrix(download_root / filenames["rigid"]) - affine_matrix = read_matrix(download_root / filenames["affine"]) - - subject_dict: dict[str, ScalarImage | LabelMap] = { - "t1": ScalarImage( - download_root / filenames["t1"], + rigid = read_matrix(download_root / self.filenames['rigid']) + affine = read_matrix(download_root / self.filenames['affine']) + subject_dict = { + 't1': ScalarImage( + download_root / self.filenames['t1'], rigid_matrix=rigid, - affine_matrix=affine_matrix, + affine_matrix=affine, ), - "seg": LabelMap( - download_root / filenames["seg"], + 'seg': LabelMap( + download_root / self.filenames['seg'], rigid_matrix=rigid, - affine_matrix=affine_matrix, - color_map=FPG.GIF_COLORS, + affine_matrix=affine, ), } if load_all: - subject_dict["t2"] = ScalarImage(download_root / filenames["t2"]) - subject_dict["fmri"] = ScalarImage( - download_root / filenames["fmri"], + subject_dict['t2'] = ScalarImage( + download_root / self.filenames['t2'], ) - subject_dict["dmri"] = ScalarImage( - download_root / filenames["dmri"], + subject_dict['fmri'] = ScalarImage( + download_root / self.filenames['fmri'], ) - super().__init__(**subject_dict) + subject_dict['dmri'] = ScalarImage( + download_root / self.filenames['dmri'], + ) + super().__init__(subject_dict) + self.gif_colors = self.GIF_COLORS - def plot(self, **kwargs: Any) -> Any: - """Plot with GIF parcellation colors for the seg image.""" - kwargs.setdefault("cmap_dict", {"seg": self.GIF_COLORS}) - return super().plot(**kwargs) + def plot(self, *args, **kwargs): + super().plot(*args, **kwargs, cmap_dict={'seg': self.gif_colors}) - GIF_COLORS: ClassVar[dict[int, tuple[int, int, int]]] = { + GIF_COLORS = { 0: (0, 0, 0), 1: (0, 0, 0), 5: (127, 255, 212), @@ -241,3 +233,7 @@ def plot(self, **kwargs: Any) -> Any: 207: (0, 0, 128), 208: (0, 0, 128), } + + +# For backward compatibility +GIF_COLORS = FPG.GIF_COLORS diff --git a/src/torchio/datasets/itk_snap/__init__.py b/src/torchio/datasets/itk_snap/__init__.py index b464b7278..130a7a068 100644 --- a/src/torchio/datasets/itk_snap/__init__.py +++ b/src/torchio/datasets/itk_snap/__init__.py @@ -1,13 +1,9 @@ -"""ITK-SNAP sample datasets.""" - from .itk_snap import T1T2 from .itk_snap import AorticValve from .itk_snap import BrainTumor -from .itk_snap import SubjectITKSNAP __all__ = [ - "T1T2", - "AorticValve", - "BrainTumor", - "SubjectITKSNAP", + 'BrainTumor', + 'T1T2', + 'AorticValve', ] diff --git a/src/torchio/datasets/itk_snap/itk_snap.py b/src/torchio/datasets/itk_snap/itk_snap.py index bde125f39..b6ab2d730 100644 --- a/src/torchio/datasets/itk_snap/itk_snap.py +++ b/src/torchio/datasets/itk_snap/itk_snap.py @@ -1,30 +1,25 @@ -"""ITK-SNAP Image Data Downloads.""" - -from __future__ import annotations - import urllib.parse from ...data import LabelMap from ...data import ScalarImage from ...data.subject import Subject from ...download import download_and_extract_archive -from ...download import get_torchio_cache_dir +from ...utils import get_torchio_cache_dir class SubjectITKSNAP(Subject): """ITK-SNAP Image Data Downloads. - See `the ITK-SNAP website - `_ - for more information. + See [the ITK-SNAP website](http://www.itksnap.org/pmwiki/pmwiki.php?n=Downloads.Data) for more information. + """ - url_base = "https://www.nitrc.org/frs/download.php/" + url_base = 'https://www.nitrc.org/frs/download.php/' - def __init__(self, name: str, code: str) -> None: + def __init__(self, name, code): self.name = name - self.url_dir = urllib.parse.urljoin(self.url_base, f"{code}/") - self.filename = f"{self.name}.zip" + self.url_dir = urllib.parse.urljoin(self.url_base, f'{code}/') + self.filename = f'{self.name}.zip' self.url = urllib.parse.urljoin(self.url_dir, self.filename) self.download_root = get_torchio_cache_dir() / self.name if not self.download_root.is_dir(): @@ -33,61 +28,55 @@ def __init__(self, name: str, code: str) -> None: download_root=self.download_root, filename=self.filename, ) - super().__init__(**self._get_kwargs()) + super().__init__(**self.get_kwargs()) - def _get_kwargs(self) -> dict: + def get_kwargs(self): raise NotImplementedError class BrainTumor(SubjectITKSNAP): - """BRATS brain tumor sample data.""" + def __init__(self): + super().__init__('braintumor', '6161') - def __init__(self) -> None: - super().__init__("braintumor", "6161") - - def _get_kwargs(self) -> dict: + def get_kwargs(self): t1, t1c, t2, flair, seg = ( - self.download_root / self.name / f"BRATS_HG0015_{name}.mha" - for name in ("T1", "T1C", "T2", "FLAIR", "truth") + self.download_root / self.name / f'BRATS_HG0015_{name}.mha' + for name in ('T1', 'T1C', 'T2', 'FLAIR', 'truth') ) return { - "t1": ScalarImage(t1), - "t1c": ScalarImage(t1c), - "t2": ScalarImage(t2), - "flair": ScalarImage(flair), - "seg": LabelMap(seg), + 't1': ScalarImage(t1), + 't1c': ScalarImage(t1c), + 't2': ScalarImage(t2), + 'flair': ScalarImage(flair), + 'seg': LabelMap(seg), } class T1T2(SubjectITKSNAP): - """Multi-site T1 and T2 brain MRI.""" - - def __init__(self) -> None: - super().__init__("ashs_test", "10983") + def __init__(self): + super().__init__('ashs_test', '10983') - def _get_kwargs(self) -> dict: - mprage = self.download_root / self.name / "mprage_3T_bet_dr.nii" - tse = self.download_root / self.name / "tse_3t_dr.nii" + def get_kwargs(self): + mprage = self.download_root / self.name / 'mprage_3T_bet_dr.nii' + tse = self.download_root / self.name / 'tse_3t_dr.nii' return { - "mprage": ScalarImage(mprage), - "tse": ScalarImage(tse), + 'mprage': ScalarImage(mprage), + 'tse': ScalarImage(tse), } class AorticValve(SubjectITKSNAP): - """Cardiac aortic valve CT frames with segmentation.""" - - def __init__(self) -> None: - super().__init__("bav_example", "11021") + def __init__(self): + super().__init__('bav_example', '11021') - def _get_kwargs(self) -> dict: + def get_kwargs(self): b14, b14_seg, b25, b25_seg = ( - self.download_root / self.name / f"bav_frame_{name}.nii.gz" - for name in ("14", "14_manseg", "25", "25_manseg") + self.download_root / self.name / f'bav_frame_{name}.nii.gz' + for name in ('14', '14_manseg', '25', '25_manseg') ) return { - "b14": ScalarImage(b14), - "b14_seg": LabelMap(b14_seg), - "b25": ScalarImage(b25), - "b25_seg": LabelMap(b25_seg), + 'b14': ScalarImage(b14), + 'b14_seg': LabelMap(b14_seg), + 'b25': ScalarImage(b25), + 'b25_seg': LabelMap(b25_seg), } diff --git a/src/torchio/datasets/ixi.py b/src/torchio/datasets/ixi.py index 8af1bee20..d5a2cf754 100644 --- a/src/torchio/datasets/ixi.py +++ b/src/torchio/datasets/ixi.py @@ -1,184 +1,238 @@ -"""IXI dataset: ~600 brain MRIs from healthy subjects. - -The `Information eXtraction from Images (IXI) -`_ dataset contains -nearly 600 MR images from normal, healthy subjects. - -This data is made available under the Creative Commons CC BY-SA 3.0 -license. If you use it, please acknowledge the source. +"""The [Information eXtraction from Images (IXI)](https://brain-development.org/ixi-dataset/) +dataset contains "nearly 600 MR images from normal, healthy subjects", +including "T1, T2 and PD-weighted images, MRA images and Diffusion-weighted +images (15 directions)". + +Note: + This data is made available under the + Creative Commons CC BY-SA 3.0 license. + If you use it please acknowledge the source of the IXI data, e.g. + [the IXI website](https://brain-development.org/ixi-dataset/). """ +# Adapted from +# https://pytorch.org/docs/stable/_modules/torchvision/datasets/mnist.html#MNIST from __future__ import annotations import shutil from collections.abc import Sequence from pathlib import Path from tempfile import NamedTemporaryFile -from typing import Any -from ..data.image import LabelMap -from ..data.image import ScalarImage -from ..data.subject import Subject +from ..data import LabelMap +from ..data import ScalarImage +from ..data import Subject +from ..data import SubjectsDataset from ..download import download_and_extract_archive +from ..transforms import Transform from ..types import TypePath -def ixi( - root: TypePath, - *, - download: bool = False, - modalities: Sequence[str] = ("T1", "T2"), -) -> list[Subject]: - """Download and load the full IXI dataset. +class IXI(SubjectsDataset): + """Full IXI dataset. Args: - root: Root directory for the dataset. - download: If `True`, download the data into `root`. - modalities: Modalities to include. Must be a subset of + root: Root directory to which the dataset will be downloaded. + transform: An instance of + [`Transform`][torchio.transforms.transform.Transform]. + download: If set to `True`, will download the data into `root`. + modalities: List of modalities to be downloaded. They must be in `('T1', 'T2', 'PD', 'MRA', 'DTI')`. - Returns: - List of subjects, one per scan with all requested modalities. - Warning: - The dataset is several GB. Downloading may take a while. + The size of this dataset is multiple GB. + If you set `download` to `True`, it will take some time + to be downloaded if it is not already present. + + Examples: + >>> import torchio as tio + >>> transforms = [ + ... tio.ToCanonical(), # to RAS + ... tio.Resample((1, 1, 1)), # to 1 mm iso + ... ] + >>> ixi_dataset = tio.datasets.IXI( + ... 'path/to/ixi_root/', + ... modalities=('T1', 'T2'), + ... transform=tio.Compose(transforms), + ... download=True, + ... ) + >>> print('Number of subjects in dataset:', len(ixi_dataset)) # 577 + >>> sample_subject = ixi_dataset[0] + >>> print('Keys in subject:', tuple(sample_subject.keys())) # ('T1', 'T2') + >>> print('Shape of T1 data:', sample_subject['T1'].shape) # [1, 180, 268, 268] + >>> print('Shape of T2 data:', sample_subject['T2'].shape) # [1, 241, 257, 188] """ - root = Path(root) - md5s = _IXI_MD5 - for m in modalities: - if m not in md5s: - msg = f'Modality "{m}" must be one of {tuple(md5s.keys())}' - raise ValueError(msg) - if download: - _download_ixi(root, modalities, md5s) - if not all((root / m).is_dir() for m in modalities): - msg = "Dataset not found. Use download=True to download it" - raise RuntimeError(msg) - return _load_ixi_subjects(root, modalities) - - -def ixi_tiny( - root: TypePath, - *, - download: bool = False, -) -> list[Subject]: - r"""Download and load IXITiny (566 $T_1$ images + segmentations). - - All images have shape $83 \times 44 \times 55$. Useful as a - medical image MNIST for quick experiments. + + base_url = 'http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-{modality}.tar' + md5_dict = { + 'T1': '34901a0593b41dd19c1a1f746eac2d58', + 'T2': 'e3140d78730ecdd32ba92da48c0a9aaa', + 'PD': '88ecd9d1fa33cb4a2278183b42ffd749', + 'MRA': '29be7d2fee3998f978a55a9bdaf3407e', + 'DTI': '636573825b1c8b9e8c78f1877df3ee66', + } + + def __init__( + self, + root: TypePath, + download: bool = False, + modalities: Sequence[str] = ('T1', 'T2'), + **kwargs, + ): + root = Path(root) + for modality in modalities: + if modality not in self.md5_dict: + message = ( + f'Modality "{modality}" must be' + f' one of {tuple(self.md5_dict.keys())}' + ) + raise ValueError(message) + if download: + self._download(root, modalities) + if not self._check_exists(root, modalities): + message = 'Dataset not found. You can use download=True to download it' + raise RuntimeError(message) + subjects_list = self._get_subjects_list(root, modalities) + super().__init__(subjects_list, **kwargs) + + @staticmethod + def _check_exists(root, modalities): + for modality in modalities: + modality_dir = root / modality + if not modality_dir.is_dir(): + exists = False + break + else: + exists = True + return exists + + @staticmethod + def _get_subjects_list(root: Path, modalities: Sequence[str]) -> list[Subject]: + # The number of files for each modality is not the same + # E.g. 581 for T1, 578 for T2 + # Let's just use the first modality as reference for now + # I.e. only subjects with all modalities will be included + one_modality = modalities[0] + paths = sglob(root / one_modality, '*.nii.gz') + subjects = [] + for filepath in paths: + subject_id = get_subject_id(filepath) + images_dict: dict[str, str | ScalarImage] = {'subject_id': subject_id} + images_dict[one_modality] = ScalarImage(filepath) + for modality in modalities[1:]: + globbed = sglob( + root / modality, + f'{subject_id}-{modality}.nii.gz', + ) + if globbed: + assert len(globbed) == 1 + images_dict[modality] = ScalarImage(globbed[0]) + else: + skip_subject = True + break + else: + skip_subject = False + if skip_subject: + continue + subjects.append(Subject(images_dict)) + return subjects + + def _download(self, root, modalities): + """Download the IXI data if it does not exist already.""" + for modality in modalities: + modality_dir = root / modality + if modality_dir.is_dir(): + continue + modality_dir.mkdir(exist_ok=True, parents=True) + + # download files + url = self.base_url.format(modality=modality) + md5 = self.md5_dict[modality] + + with NamedTemporaryFile(suffix='.tar', delete=False) as f: + download_and_extract_archive( + url, + download_root=modality_dir, + filename=f.name, + md5=md5, + ) + + +class IXITiny(SubjectsDataset): + r"""This is the dataset used in the main [notebook](https://github.com/TorchIO-project/torchio/blob/main/tutorials/README.md). It is a tiny version + of IXI, containing 566 $T_1$-weighted brain MR images and their + corresponding brain segmentations, all with size $83 \times 44 \times + 55$. + + It can be used as a medical image MNIST. Args: - root: Root directory for the dataset. - download: If `True`, download the data into `root`. + root: Root directory to which the dataset will be downloaded. + transform: An instance of + [`Transform`][torchio.transforms.transform.Transform]. + download: If set to `True`, will download the data into `root`. - Returns: - List of subjects with `image` and `label` keys. """ - root = Path(root) - if download: - _download_ixi_tiny(root) - if not root.is_dir(): - msg = "Dataset not found. Use download=True to download it" - raise RuntimeError(msg) - return _load_ixi_tiny_subjects(root) - - -# --------------------------------------------------------------------------- -# Internal helpers -# --------------------------------------------------------------------------- - -_IXI_BASE_URL = ( - "http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-{modality}.tar" -) -_IXI_MD5: dict[str, str] = { - "T1": "34901a0593b41dd19c1a1f746eac2d58", - "T2": "e3140d78730ecdd32ba92da48c0a9aaa", - "PD": "88ecd9d1fa33cb4a2278183b42ffd749", - "MRA": "29be7d2fee3998f978a55a9bdaf3407e", - "DTI": "636573825b1c8b9e8c78f1877df3ee66", -} - -_IXI_TINY_URL = "https://www.dropbox.com/s/ogxjwjxdv5mieah/ixi_tiny.zip?dl=1" -_IXI_TINY_MD5 = "bfb60f4074283d78622760230bfa1f98" - - -def _download_ixi( - root: Path, - modalities: Sequence[str], - md5s: dict[str, str], -) -> None: - for modality in modalities: - modality_dir = root / modality - if modality_dir.is_dir(): - continue - modality_dir.mkdir(exist_ok=True, parents=True) - url = _IXI_BASE_URL.format(modality=modality) - with NamedTemporaryFile(suffix=".tar", delete=False) as f: + + url = 'https://www.dropbox.com/s/ogxjwjxdv5mieah/ixi_tiny.zip?dl=1' + md5 = 'bfb60f4074283d78622760230bfa1f98' + + def __init__( + self, + root: TypePath, + transform: Transform | None = None, + download: bool = False, + **kwargs, + ): + root = Path(root) + if download: + self._download(root) + if not root.is_dir(): + message = 'Dataset not found. You can use download=True to download it' + raise RuntimeError(message) + subjects_list = self._get_subjects_list(root) + super().__init__(subjects_list, transform=transform, **kwargs) + + @staticmethod + def _get_subjects_list(root): + image_paths = sglob(root / 'image', '*.nii.gz') + label_paths = sglob(root / 'label', '*.nii.gz') + if not (image_paths and label_paths): + message = ( + f'Images not found. Remove the root directory ({root}) and try again' + ) + raise FileNotFoundError(message) + + subjects = [] + for image_path, label_path in zip(image_paths, label_paths, strict=True): + subject_id = get_subject_id(image_path) + subject_dict = {} + subject_dict['image'] = ScalarImage(image_path) + subject_dict['label'] = LabelMap(label_path) + subject_dict['subject_id'] = subject_id + subjects.append(Subject(**subject_dict)) + return subjects + + def _download(self, root): + """Download the tiny IXI data if it doesn't exist already.""" + if root.is_dir(): # assume it's been downloaded + return + with NamedTemporaryFile(suffix='.zip', delete=False) as f: download_and_extract_archive( - url, - download_root=modality_dir, + self.url, + download_root=root, filename=f.name, - md5=md5s[modality], + md5=self.md5, ) + ixi_tiny_dir = root / 'ixi_tiny' + (ixi_tiny_dir / 'image').rename(root / 'image') + (ixi_tiny_dir / 'label').rename(root / 'label') + shutil.rmtree(ixi_tiny_dir) -def _load_ixi_subjects( - root: Path, - modalities: Sequence[str], -) -> list[Subject]: - first = modalities[0] - paths = sorted((root / first).glob("*.nii.gz")) - subjects: list[Subject] = [] - for filepath in paths: - sid = _subject_id(filepath) - images: dict[str, Any] = {"subject_id": sid} - images[first] = ScalarImage(filepath) - skip = False - for m in modalities[1:]: - matches = sorted((root / m).glob(f"{sid}-{m}.nii.gz")) - if matches: - images[m] = ScalarImage(matches[0]) - else: - skip = True - break - if not skip: - subjects.append(Subject(**images)) - return subjects - - -def _download_ixi_tiny(root: Path) -> None: - if root.is_dir(): - return - with NamedTemporaryFile(suffix=".zip", delete=False) as f: - download_and_extract_archive( - _IXI_TINY_URL, - download_root=root, - filename=f.name, - md5=_IXI_TINY_MD5, - ) - ixi_tiny_dir = root / "ixi_tiny" - (ixi_tiny_dir / "image").rename(root / "image") - (ixi_tiny_dir / "label").rename(root / "label") - shutil.rmtree(ixi_tiny_dir) - - -def _load_ixi_tiny_subjects(root: Path) -> list[Subject]: - image_paths = sorted((root / "image").glob("*.nii.gz")) - label_paths = sorted((root / "label").glob("*.nii.gz")) - if not (image_paths and label_paths): - msg = f"Images not found. Remove {root} and try again" - raise FileNotFoundError(msg) - subjects: list[Subject] = [] - for img_path, lbl_path in zip(image_paths, label_paths, strict=True): - subjects.append( - Subject( - image=ScalarImage(img_path), - label=LabelMap(lbl_path), - subject_id=_subject_id(img_path), - ), - ) - return subjects - - -def _subject_id(path: Path) -> str: - return "-".join(path.name.split("-")[:-1]) +def sglob(directory, pattern): + return sorted(Path(directory).glob(pattern)) + + +def get_subject_id(path): + return '-'.join(path.name.split('-')[:-1]) diff --git a/src/torchio/datasets/medmnist.py b/src/torchio/datasets/medmnist.py index fc9cb417c..8e726e0ba 100644 --- a/src/torchio/datasets/medmnist.py +++ b/src/torchio/datasets/medmnist.py @@ -1,113 +1,77 @@ -"""3D MedMNIST v2 datasets. - -Datasets from `MedMNIST v2: A Large-Scale Lightweight Benchmark for -2D and 3D Biomedical Image Classification -`_. - -Check the `MedMNIST website `_ for details -and licensing. -""" - -from __future__ import annotations - import numpy as np import torch -from ..data.image import ScalarImage -from ..data.subject import Subject +from ..data import ScalarImage +from ..data import Subject +from ..data import SubjectsDataset from ..download import download_url -from ..download import get_torchio_cache_dir - - -def _load_medmnist( - class_name: str, - split: str, -) -> list[Subject]: - """Shared loader for all MedMNIST 3D datasets.""" - valid = ("train", "training", "val", "validation", "test", "testing") - if split not in valid: - msg = f"Split must be one of {valid}, got {split!r}" - raise ValueError(msg) - match split: - case "training": - split = "train" - case "validation": - split = "val" - case "testing": - split = "test" - - filename = f"{class_name}.npz" - url = f"https://zenodo.org/record/5208230/files/{filename}?download=1" - download_root = get_torchio_cache_dir() / "MedMNIST" - download_url(url, download_root, filename=filename) - path = download_root / filename - npz = np.load(path) - images = npz[f"{split}_images"] - labels = npz[f"{split}_labels"] - subjects: list[Subject] = [] - for image, label in zip(images, labels, strict=True): - tensor = torch.from_numpy(image[np.newaxis].copy()).float() - scalar = ScalarImage.from_tensor( # ty: ignore[unresolved-attribute] - tensor, - ) - subjects.append( - Subject( - image=scalar, - labels=torch.from_numpy(label.copy()), - ), - ) - return subjects +from ..utils import get_torchio_cache_dir -def organ_mnist_3d(split: str = "train") -> list[Subject]: - """3D organ segmentation dataset. +class MedMNIST(SubjectsDataset): + """3D MedMNIST v2 datasets. - Args: - split: `'train'`, `'val'`, or `'test'`. - """ - return _load_medmnist("organmnist3d", split) + Datasets from [MedMNIST v2: A Large-Scale Lightweight Benchmark for 2D and + 3D Biomedical Image Classification ](https://arxiv.org/abs/2110.14795). - -def nodule_mnist_3d(split: str = "train") -> list[Subject]: - """3D lung nodule dataset. + Please check the [MedMNIST website ](https://medmnist.com/) for more + information, inclusing the license. Args: - split: `'train'`, `'val'`, or `'test'`. + split: Dataset split. Should be `'train'`, `'val'` or `'test'`. """ - return _load_medmnist("nodulemnist3d", split) + BASE_URL = 'https://zenodo.org/record/5208230/files' + SPLITS = 'train', 'training', 'val', 'validation', 'test', 'testing' + + def __init__(self, split, **kwargs): + if split not in self.SPLITS: + raise ValueError(f'The split must be one of {self.SPLITS}') + split = 'train' if split == 'training' else split + split = 'val' if split == 'validation' else split + split = 'test' if split == 'testing' else split + url = f'{self.BASE_URL}/{self.filename}?download=1' + download_root = get_torchio_cache_dir() / 'MedMNIST' + download_url( + url, + download_root, + filename=self.filename, + ) + path = download_root / self.filename + npz_file = np.load(path) + images = npz_file[f'{split}_images'] + labels = npz_file[f'{split}_labels'] + subjects = [] + for image, label in zip(images, labels, strict=True): + image = ScalarImage(tensor=image[np.newaxis]) + subject = Subject(image=image, labels=torch.from_numpy(label)) + subjects.append(subject) + super().__init__(subjects, **kwargs) + + @property + def filename(self): + return f'{self.__class__.__name__.lower()}.npz' -def adrenal_mnist_3d(split: str = "train") -> list[Subject]: - """3D adrenal gland dataset. - Args: - split: `'train'`, `'val'`, or `'test'`. - """ - return _load_medmnist("adrenalmnist3d", split) +class OrganMNIST3D(MedMNIST): + __doc__ = MedMNIST.__doc__ -def fracture_mnist_3d(split: str = "train") -> list[Subject]: - """3D bone fracture dataset. +class NoduleMNIST3D(MedMNIST): + __doc__ = MedMNIST.__doc__ - Args: - split: `'train'`, `'val'`, or `'test'`. - """ - return _load_medmnist("fracturemnist3d", split) +class AdrenalMNIST3D(MedMNIST): + __doc__ = MedMNIST.__doc__ -def vessel_mnist_3d(split: str = "train") -> list[Subject]: - """3D vessel dataset. - Args: - split: `'train'`, `'val'`, or `'test'`. - """ - return _load_medmnist("vesselmnist3d", split) +class FractureMNIST3D(MedMNIST): + __doc__ = MedMNIST.__doc__ -def synapse_mnist_3d(split: str = "train") -> list[Subject]: - """3D synapse dataset. +class VesselMNIST3D(MedMNIST): + __doc__ = MedMNIST.__doc__ - Args: - split: `'train'`, `'val'`, or `'test'`. - """ - return _load_medmnist("synapsemnist3d", split) + +class SynapseMNIST3D(MedMNIST): + __doc__ = MedMNIST.__doc__ diff --git a/src/torchio/datasets/mni/__init__.py b/src/torchio/datasets/mni/__init__.py index 4aa34f193..62db01d5d 100644 --- a/src/torchio/datasets/mni/__init__.py +++ b/src/torchio/datasets/mni/__init__.py @@ -1,15 +1,11 @@ -"""MNI atlases.""" - from .colin import Colin27 from .icbm import ICBM2009CNonlinearSymmetric -from .mni import SubjectMNI from .pediatric import Pediatric from .sheep import Sheep __all__ = [ - "Colin27", - "ICBM2009CNonlinearSymmetric", - "Pediatric", - "Sheep", - "SubjectMNI", + 'Sheep', + 'Colin27', + 'Pediatric', + 'ICBM2009CNonlinearSymmetric', ] diff --git a/src/torchio/datasets/mni/colin.py b/src/torchio/datasets/mni/colin.py index 00fee7bf7..880c23bc8 100644 --- a/src/torchio/datasets/mni/colin.py +++ b/src/torchio/datasets/mni/colin.py @@ -1,61 +1,70 @@ -"""Colin27 MNI template.""" - -from __future__ import annotations - import urllib.parse -from typing import ClassVar from ...data import LabelMap from ...data import ScalarImage -from ...download import compress from ...download import download_and_extract_archive +from ...utils import compress from .mni import SubjectMNI TISSUES_2008 = { - 1: "Cerebro-spinal fluid", - 2: "Gray Matter", - 3: "White Matter", - 4: "Fat", - 5: "Muscles", - 6: "Skin and Muscles", - 7: "Skull", - 9: "Fat 2", - 10: "Dura", - 11: "Marrow", - 12: "Vessels", + 1: 'Cerebro-spinal fluid', + 2: 'Gray Matter', + 3: 'White Matter', + 4: 'Fat', + 5: 'Muscles', + 6: 'Skin and Muscles', + 7: 'Skull', + 9: 'Fat 2', + 10: 'Dura', + 11: 'Marrow', + 12: 'Vessels', } class Colin27(SubjectMNI): - """Colin27 MNI template. + NAME_TO_LABEL = {name: label for label, name in TISSUES_2008.items()} + r"""Colin27 MNI template. More information can be found in the website of the - `1998 `_ and - `2008 `_ + [1998 ](https://nist.mni.mcgill.ca/colin-27-average-brain/) and + [2008 ](http://www.bic.mni.mcgill.ca/ServicesAtlases/Colin27Highres) versions. - Args: + ![MNI Colin 27 2008 version](http://www.bic.mni.mcgill.ca/uploads/ServicesAtlases/mni_colin27_2008.jpg) + + Arguments: version: Template year. It can be `1998` or `2008`. Warning: The resolution of the `2008` version is quite high. The subject instance will contain four images of size - 362 x 434 x 362, therefore applying a transform to + $362 \times 434 \times 362$, therefore applying a transform to it might take longer than expected. - """ - NAME_TO_LABEL: ClassVar[dict[str, int]] = { - name: label for label, name in TISSUES_2008.items() - } + Examples: + >>> import torchio as tio + >>> colin_1998 = tio.datasets.Colin27(version=1998) + >>> colin_1998 + Colin27(Keys: ('t1', 'head', 'brain'); images: 3) + >>> colin_1998.load() + >>> colin_1998.t1 + ScalarImage(shape: (1, 181, 217, 181); spacing: (1.00, 1.00, 1.00); orientation: RAS+; memory: 27.1 MiB; type: intensity) + >>> + >>> colin_2008 = tio.datasets.Colin27(version=2008) + >>> colin_2008 + Colin27(Keys: ('t1', 't2', 'pd', 'cls'); images: 4) + >>> colin_2008.load() + >>> colin_2008.t1 + ScalarImage(shape: (1, 362, 434, 362); spacing: (0.50, 0.50, 0.50); orientation: RAS+; memory: 217.0 MiB; type: intensity) + """ - def __init__(self, version: int = 1998) -> None: + def __init__(self, version=1998): if version not in (1998, 2008): - msg = f'Version must be 1998 or 2008, not "{version}"' - raise ValueError(msg) + raise ValueError(f'Version must be 1998 or 2008, not "{version}"') self.version = version - self.name = f"mni_colin27_{version}_nifti" - self.url_dir = urllib.parse.urljoin(self.url_base, "colin27/") - self.filename = f"{self.name}.zip" + self.name = f'mni_colin27_{version}_nifti' + self.url_dir = urllib.parse.urljoin(self.url_base, 'colin27/') + self.filename = f'{self.name}.zip' self.url = urllib.parse.urljoin(self.url_dir, self.filename) if not self.download_root.is_dir(): download_and_extract_archive( @@ -63,51 +72,65 @@ def __init__(self, version: int = 1998) -> None: download_root=self.download_root, filename=self.filename, ) - # Fix label map - # https://github.com/TorchIO-project/torchio/issues/220 + + # Fix label map (https://github.com/TorchIO-project/torchio/issues/220) if version == 2008: - path = self.download_root / "colin27_cls_tal_hires.nii" + path = self.download_root / 'colin27_cls_tal_hires.nii' cls_image = LabelMap(path) cls_image.set_data(cls_image.data.round().byte()) cls_image.save(path) (self.download_root / self.filename).unlink() - for path in self.download_root.glob("*.nii"): + for path in self.download_root.glob('*.nii'): compress(path) path.unlink() - subject_kwargs = self._get_subject_kwargs( - self.download_root, - extension=".nii.gz", - ) - super().__init__(**subject_kwargs) + try: + subject_dict = self.get_subject_dict( + self.download_root, + extension='.nii.gz', + ) + except FileNotFoundError: # for backward compatibility + subject_dict = self.get_subject_dict( + self.download_root, + extension='.nii', + ) + super().__init__(subject_dict) - def _get_subject_kwargs(self, download_root, extension): + def get_subject_dict(self, download_root, extension): if self.version == 1998: - return _get_colin1998_kwargs(download_root, extension) - return _get_colin2008_kwargs(download_root, extension) - - -def _get_colin1998_kwargs(download_root, extension): - t1, head, mask = ( - download_root / f"colin27_t1_tal_lin{suffix}{extension}" - for suffix in ("", "_headmask", "_mask") - ) - return { - "t1": ScalarImage(t1), - "head": LabelMap(head), - "brain": LabelMap(mask), - } - - -def _get_colin2008_kwargs(download_root, extension): - t1, t2, pd, label = ( - download_root / f"colin27_{name}_tal_hires{extension}" - for name in ("t1", "t2", "pd", "cls") - ) - return { - "t1": ScalarImage(t1), - "t2": ScalarImage(t2), - "pd": ScalarImage(pd), - "cls": LabelMap(label), - } + subject_dict = Colin1998.get_subject_dict(download_root, extension) + elif self.version == 2008: + subject_dict = Colin2008.get_subject_dict(download_root, extension) + return subject_dict + + +class Colin1998: + @staticmethod + def get_subject_dict(download_root, extension): + t1, head, mask = ( + download_root / f'colin27_t1_tal_lin{suffix}{extension}' + for suffix in ('', '_headmask', '_mask') + ) + subject_dict = { + 't1': ScalarImage(t1), + 'head': LabelMap(head), + 'brain': LabelMap(mask), + } + return subject_dict + + +class Colin2008: + @staticmethod + def get_subject_dict(download_root, extension): + t1, t2, pd, label = ( + download_root / f'colin27_{name}_tal_hires{extension}' + for name in ('t1', 't2', 'pd', 'cls') + ) + subject_dict = { + 't1': ScalarImage(t1), + 't2': ScalarImage(t2), + 'pd': ScalarImage(pd), + 'cls': LabelMap(label, labels=TISSUES_2008), + } + return subject_dict diff --git a/src/torchio/datasets/mni/icbm.py b/src/torchio/datasets/mni/icbm.py index 196669655..8559ad966 100644 --- a/src/torchio/datasets/mni/icbm.py +++ b/src/torchio/datasets/mni/icbm.py @@ -1,36 +1,42 @@ -"""ICBM 2009c Nonlinear Symmetric template.""" - -from __future__ import annotations - import urllib.parse -from pathlib import Path import torch from ...data import LabelMap from ...data import ScalarImage -from ...download import compress from ...download import download_and_extract_archive -from ...download import get_torchio_cache_dir +from ...utils import compress +from ...utils import get_torchio_cache_dir from .mni import SubjectMNI class ICBM2009CNonlinearSymmetric(SubjectMNI): r"""ICBM template. - More information can be found in the - `website `_. + More information can be found in the [website + ](http://www.bic.mni.mcgill.ca/ServicesAtlases/ICBM152NLin2009). + + ![ICBM 2009c Nonlinear Symmetric](http://www.bic.mni.mcgill.ca/uploads/ServicesAtlases/mni_icbm152_sym_09c_small.jpg) Args: - load_4d_tissues: If `True`, the tissue probability maps will be - loaded together into a 4D image. Otherwise, they will be loaded - into independent images. + load_4d_tissues: If `True`, the tissue probability maps will be loaded + together into a 4D image. Otherwise, they will be loaded into + independent images. + + Examples: + >>> import torchio as tio + >>> icbm = tio.datasets.ICBM2009CNonlinearSymmetric() + >>> icbm + ICBM2009CNonlinearSymmetric(Keys: ('t1', 'eyes', 'face', 'brain', 't2', 'pd', 'tissues'); images: 7) + >>> icbm = tio.datasets.ICBM2009CNonlinearSymmetric(load_4d_tissues=False) + >>> icbm + ICBM2009CNonlinearSymmetric(Keys: ('t1', 'eyes', 'face', 'brain', 't2', 'pd', 'gm', 'wm', 'csf'); images: 9) """ - def __init__(self, load_4d_tissues: bool = False) -> None: - self.name = "mni_icbm152_nlin_sym_09c_nifti" - self.url_base = "http://www.bic.mni.mcgill.ca/~vfonov/icbm/2009/" - self.filename = f"{self.name}.zip" + def __init__(self, load_4d_tissues: bool = True): + self.name = 'mni_icbm152_nlin_sym_09c_nifti' + self.url_base = 'http://www.bic.mni.mcgill.ca/~vfonov/icbm/2009/' + self.filename = f'{self.name}.zip' self.url = urllib.parse.urljoin(self.url_base, self.filename) download_root = get_torchio_cache_dir() / self.name if not download_root.is_dir(): @@ -41,43 +47,40 @@ def __init__(self, load_4d_tissues: bool = False) -> None: remove_finished=True, ) - files_dir = download_root / "mni_icbm152_nlin_sym_09c" + files_dir = download_root / 'mni_icbm152_nlin_sym_09c' - p = str(files_dir / "mni_icbm152") - m = "tal_nlin_sym_09c" - s = ".nii.gz" + p = files_dir / 'mni_icbm152' + m = 'tal_nlin_sym_09c' + s = '.nii.gz' - tissues_path = f"{p}_tissues_{m}.nii.gz" - if not Path(tissues_path).is_file(): - gm = LabelMap(f"{p}_gm_{m}.nii") - wm = LabelMap(f"{p}_wm_{m}.nii") - csf = LabelMap(f"{p}_csf_{m}.nii") - gm.load() - wm.load() - csf.load() + tissues_path = files_dir / f'{p}_tissues_{m}.nii.gz' + if not tissues_path.is_file(): + gm = LabelMap(f'{p}_gm_{m}.nii') + wm = LabelMap(f'{p}_wm_{m}.nii') + csf = LabelMap(f'{p}_csf_{m}.nii') gm.set_data(torch.cat((gm.data, wm.data, csf.data))) gm.save(tissues_path) - for fp in files_dir.glob("*.nii"): - compress(fp, fp.with_suffix(".nii.gz")) + for fp in files_dir.glob('*.nii'): + compress(fp, fp.with_suffix('.nii.gz')) fp.unlink() - subject_kwargs: dict[str, ScalarImage | LabelMap] = { - "t1": ScalarImage(f"{p}_t1_{m}{s}"), - "eyes": LabelMap(f"{p}_t1_{m}_eye_mask{s}"), - "face": LabelMap(f"{p}_t1_{m}_face_mask{s}"), - "brain": LabelMap(f"{p}_t1_{m}_mask{s}"), - "t2": ScalarImage(f"{p}_t2_{m}{s}"), - "pd": ScalarImage(f"{p}_csf_{m}{s}"), + subject_dict = { + 't1': ScalarImage(f'{p}_t1_{m}{s}'), + 'eyes': LabelMap(f'{p}_t1_{m}_eye_mask{s}'), + 'face': LabelMap(f'{p}_t1_{m}_face_mask{s}'), + 'brain': LabelMap(f'{p}_t1_{m}_mask{s}'), + 't2': ScalarImage(f'{p}_t2_{m}{s}'), + 'pd': ScalarImage(f'{p}_csf_{m}{s}'), } if load_4d_tissues: - subject_kwargs["tissues"] = LabelMap( + subject_dict['tissues'] = LabelMap( tissues_path, channels_last=True, ) else: - subject_kwargs["gm"] = LabelMap(f"{p}_gm_{m}{s}") - subject_kwargs["wm"] = LabelMap(f"{p}_wm_{m}{s}") - subject_kwargs["csf"] = LabelMap(f"{p}_csf_{m}{s}") + subject_dict['gm'] = LabelMap(f'{p}_gm_{m}{s}') + subject_dict['wm'] = LabelMap(f'{p}_wm_{m}{s}') + subject_dict['csf'] = LabelMap(f'{p}_csf_{m}{s}') - super().__init__(**subject_kwargs) + super().__init__(subject_dict) diff --git a/src/torchio/datasets/mni/mni.py b/src/torchio/datasets/mni/mni.py index 72f7df3c5..1ffc414ff 100644 --- a/src/torchio/datasets/mni/mni.py +++ b/src/torchio/datasets/mni/mni.py @@ -1,22 +1,16 @@ -"""Base class for MNI atlases.""" - -from __future__ import annotations - from ...data.subject import Subject -from ...download import get_torchio_cache_dir +from ...utils import get_torchio_cache_dir class SubjectMNI(Subject): """Atlases from the Montreal Neurological Institute (MNI). - See `the website `_ for more + See [the website ](https://nist.mni.mcgill.ca/?page_id=714) for more information. """ - url_base = "http://packages.bic.mni.mcgill.ca/mni-models/" - name: str + url_base = 'http://packages.bic.mni.mcgill.ca/mni-models/' @property def download_root(self): - """Return the download root directory for this atlas.""" return get_torchio_cache_dir() / self.name diff --git a/src/torchio/datasets/mni/pediatric.py b/src/torchio/datasets/mni/pediatric.py index d163344bb..297764831 100644 --- a/src/torchio/datasets/mni/pediatric.py +++ b/src/torchio/datasets/mni/pediatric.py @@ -1,13 +1,9 @@ -"""MNI Pediatric atlases.""" - -from __future__ import annotations - import urllib.parse from ...data import LabelMap from ...data import ScalarImage -from ...download import compress from ...download import download_and_extract_archive +from ...utils import compress from .mni import SubjectMNI SUPPORTED_YEARS = ( @@ -20,41 +16,41 @@ ) -def _format_age(n: float) -> str: +def format_age(n): integer = int(n) decimal = int(10 * (n - integer)) - return f"{integer:02d}.{decimal}" + return f'{integer:02d}.{decimal}' class Pediatric(SubjectMNI): """MNI pediatric atlases. - See `the MNI website - `_ + See [the MNI website ](https://nist.mni.mcgill.ca/pediatric-atlases-4-5-18-5y/) for more information. - Args: + ![Pediatric MNI template](https://nist.mni.mcgill.ca/wp-content/uploads/2016/04/nihpd_asym_all_sm.jpg) + + Arguments: years: Tuple of 2 ages. Possible values are: `(4.5, 18.5)`, - `(4.5, 8.5)`, `(7, 11)`, `(7.5, 13.5)`, - `(10, 14)` and `(13, 18.5)`. + `(4.5, 8.5)`, + `(7, 11)`, + `(7.5, 13.5)`, + `(10, 14)` and + `(13, 18.5)`. symmetric: If `True`, the left-right symmetric templates will be - used. Otherwise, the asymmetric (natural) templates will be used. + used. Else, the asymmetric (natural) templates will be used. """ - def __init__( - self, - years: tuple[float, float], - symmetric: bool = False, - ) -> None: - self.url_dir = "http://www.bic.mni.mcgill.ca/~vfonov/nihpd/obj1/" - sym_string = "sym" if symmetric else "asym" + def __init__(self, years, symmetric=False): + self.url_dir = 'http://www.bic.mni.mcgill.ca/~vfonov/nihpd/obj1/' + sym_string = 'sym' if symmetric else 'asym' if not isinstance(years, tuple) or years not in SUPPORTED_YEARS: - message = f"Years must be a tuple in {SUPPORTED_YEARS}" + message = f'Years must be a tuple in {SUPPORTED_YEARS}' raise ValueError(message) a, b = years - self.file_id = f"{sym_string}_{_format_age(a)}-{_format_age(b)}" - self.name = f"nihpd_{self.file_id}_nifti" - self.filename = f"{self.name}.zip" + self.file_id = f'{sym_string}_{format_age(a)}-{format_age(b)}' + self.name = f'nihpd_{self.file_id}_nifti' + self.filename = f'{self.name}.zip' self.url = urllib.parse.urljoin(self.url_dir, self.filename) if not self.download_root.is_dir(): download_and_extract_archive( @@ -63,18 +59,22 @@ def __init__( filename=self.filename, ) (self.download_root / self.filename).unlink() - for path in self.download_root.glob("*.nii"): + for path in self.download_root.glob('*.nii'): compress(path) path.unlink() - subject_kwargs = self._get_subject_kwargs(".nii.gz") - super().__init__(**subject_kwargs) + try: + subject_dict = self.get_subject_dict('.nii.gz') + except FileNotFoundError: # for backward compatibility + subject_dict = self.get_subject_dict('.nii') + super().__init__(subject_dict) - def _get_subject_kwargs(self, extension: str) -> dict: + def get_subject_dict(self, extension): root = self.download_root - return { - "t1": ScalarImage(root / f"nihpd_{self.file_id}_t1w{extension}"), - "t2": ScalarImage(root / f"nihpd_{self.file_id}_t2w{extension}"), - "pd": ScalarImage(root / f"nihpd_{self.file_id}_pdw{extension}"), - "mask": LabelMap(root / f"nihpd_{self.file_id}_mask{extension}"), + subject_dict = { + 't1': ScalarImage(root / f'nihpd_{self.file_id}_t1w{extension}'), + 't2': ScalarImage(root / f'nihpd_{self.file_id}_t2w{extension}'), + 'pd': ScalarImage(root / f'nihpd_{self.file_id}_pdw{extension}'), + 'mask': LabelMap(root / f'nihpd_{self.file_id}_mask{extension}'), } + return subject_dict diff --git a/src/torchio/datasets/mni/sheep.py b/src/torchio/datasets/mni/sheep.py index 50fc4b61f..d150bc2ea 100644 --- a/src/torchio/datasets/mni/sheep.py +++ b/src/torchio/datasets/mni/sheep.py @@ -1,39 +1,33 @@ -"""MNI Sheep atlas.""" - -from __future__ import annotations - import shutil import urllib.parse from ...data import ScalarImage -from ...download import compress from ...download import download_and_extract_archive +from ...utils import compress from .mni import SubjectMNI class Sheep(SubjectMNI): - """Ovine brain atlas at 0.5 mm resolution. - - See `the MNI website - `_ for more information. - """ - - def __init__(self) -> None: - self.name = "NIFTI_ovine_05mm" - self.url_dir = urllib.parse.urljoin(self.url_base, "sheep/") - self.filename = f"{self.name}.zip" + def __init__(self): + self.name = 'NIFTI_ovine_05mm' + self.url_dir = urllib.parse.urljoin(self.url_base, 'sheep/') + self.filename = f'{self.name}.zip' self.url = urllib.parse.urljoin(self.url_dir, self.filename) - t1_nii_path = self.download_root / "ovine_model_05.nii" - t1_niigz_path = self.download_root / "ovine_model_05.nii.gz" + t1_nii_path = self.download_root / 'ovine_model_05.nii' + t1_niigz_path = self.download_root / 'ovine_model_05.nii.gz' if not self.download_root.is_dir(): download_and_extract_archive( self.url, download_root=self.download_root, filename=self.filename, ) - shutil.rmtree(self.download_root / "masks") + shutil.rmtree(self.download_root / 'masks') for path in self.download_root.iterdir(): if path == t1_nii_path: compress(t1_nii_path, t1_niigz_path) path.unlink() - super().__init__(t1=ScalarImage(t1_niigz_path)) + try: + subject_dict = {'t1': ScalarImage(t1_niigz_path)} + except FileNotFoundError: # for backward compatibility + subject_dict = {'t1': ScalarImage(t1_nii_path)} + super().__init__(subject_dict) diff --git a/src/torchio/datasets/rsna_miccai.py b/src/torchio/datasets/rsna_miccai.py new file mode 100644 index 000000000..4816ddf2e --- /dev/null +++ b/src/torchio/datasets/rsna_miccai.py @@ -0,0 +1,115 @@ +import csv +import warnings +from collections.abc import Sequence +from pathlib import Path + +from ..data import ScalarImage +from ..data import Subject +from ..data import SubjectsDataset +from ..types import TypePath + + +class RSNAMICCAI(SubjectsDataset): + """RSNA-MICCAI Brain Tumor Radiogenomic Classification challenge dataset. + + This is a helper class for the dataset used in the + [RSNA-MICCAI Brain Tumor Radiogenomic Classification challenge](https://www.kaggle.com/c/rsna-miccai-brain-tumor-radiogenomic-classification) hosted on + [kaggle ](https://www.kaggle.com/). The dataset must be downloaded before + instantiating this class (as opposed to, e.g., [`torchio.datasets.IXI`][torchio.datasets.IXI]). + + This [kaggle kernel ](https://www.kaggle.com/fepegar/preprocessing-mri-with-torchio/) + includes a usage example including preprocessing of all the scans. + + If you reference or use the dataset in any form, include the following + citation: + + U.Baid, et al., "The RSNA-ASNR-MICCAI BraTS 2021 Benchmark on Brain Tumor + Segmentation and Radiogenomic Classification", arXiv:2107.02314, 2021. + + Args: + root_dir: Directory containing the dataset (`train` directory, + `test` directory, etc.). + train: If `True`, the `train` set will be used. Otherwise the + `test` set will be used. + ignore_empty: If `True`, the three subjects flagged as "presenting + issues" (empty images) by the challenge organizers will be ignored. + The subject IDs are `00109`, `00123` and `00709`. + + Examples: + >>> import torchio as tio + >>> from subprocess import call + >>> call('kaggle competitions download -c rsna-miccai-brain-tumor-radiogenomic-classification'.split()) + >>> root_dir = 'rsna-miccai-brain-tumor-radiogenomic-classification' + >>> train_set = tio.datasets.RSNAMICCAI(root_dir, train=True) + >>> test_set = tio.datasets.RSNAMICCAI(root_dir, train=False) + >>> len(train_set), len(test_set) + (582, 87) + + + """ + + id_key = 'BraTS21ID' + label_key = 'MGMT_value' + bad_subjects = '00109', '00123', '00709' + + def __init__( + self, + root_dir: TypePath, + train: bool = True, + ignore_empty: bool = True, + modalities: Sequence[str] = ('T1w', 'T1wCE', 'T2w', 'FLAIR'), + **kwargs, + ): + self.root_dir = Path(root_dir).expanduser().resolve() + if isinstance(modalities, str): + modalities = [modalities] + self.modalities = modalities + subjects = self._get_subjects(self.root_dir, train, ignore_empty) + super().__init__(subjects, **kwargs) + self.train = train + + def _get_subjects( + self, + root_dir: Path, + train: bool, + ignore_empty: bool, + ) -> list[Subject]: + subjects = [] + if train: + csv_path = root_dir / 'train_labels.csv' + try: + with open(csv_path) as csvfile: + reader = csv.DictReader(csvfile) + labels_dict = { + row[self.id_key]: int(row[self.label_key]) for row in reader + } + except FileNotFoundError: + warnings.warn( + 'Labels CSV not found. Ignoring MGMT labels', + stacklevel=2, + ) + labels_dict = {} + subjects_dir = root_dir / 'train' + else: + subjects_dir = root_dir / 'test' + + for subject_dir in sorted(subjects_dir.iterdir()): + subject_id = subject_dir.name + if ignore_empty and subject_id in self.bad_subjects: + continue + try: + int(subject_id) + except ValueError: + continue + images_dict: dict[str, object] = {self.id_key: subject_dir.name} + if train and labels_dict: + images_dict[self.label_key] = labels_dict[subject_id] + for modality in self.modalities: + image_dir = subject_dir / modality + filepaths = list(image_dir.iterdir()) + num_files = len(filepaths) + path = filepaths[0] if num_files == 1 else image_dir + images_dict[modality] = ScalarImage(path) + subject = Subject(images_dict) + subjects.append(subject) + return subjects diff --git a/src/torchio/datasets/rsna_spine_fracture.py b/src/torchio/datasets/rsna_spine_fracture.py new file mode 100644 index 000000000..291521fa4 --- /dev/null +++ b/src/torchio/datasets/rsna_spine_fracture.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any +from typing import Union + +from ..data import LabelMap +from ..data import ScalarImage +from ..data import Subject +from ..data import SubjectsDataset +from ..external.imports import get_pandas +from ..types import TypePath +from ..utils import normalize_path + +TypeBoxes = list[dict[str, Union[str, float, int]]] + + +class RSNACervicalSpineFracture(SubjectsDataset): + """RSNA 2022 Cervical Spine Fracture Detection dataset. + + This is a helper class for the dataset used in the + [RSNA 2022 Cervical Spine Fracture Detection](https://www.kaggle.com/competitions/rsna-2022-cervical-spine-fracture-detection/overview/evaluation) hosted on + [kaggle ](https://www.kaggle.com/). The dataset must be downloaded before + instantiating this class. + + """ + + UID = 'StudyInstanceUID' + + def __init__( + self, + root_dir: TypePath, + add_segmentations: bool = False, + add_bounding_boxes: bool = False, + **kwargs, + ): + self.root_dir = normalize_path(root_dir) + subjects = self._get_subjects( + add_segmentations, + add_bounding_boxes, + ) + super().__init__(subjects, **kwargs) + + @staticmethod + def _get_image_dirs_dict(images_dir: Path) -> dict[str, Path]: + dirs_dict = {} + for dicom_dir in sorted(images_dir.iterdir()): + dirs_dict[dicom_dir.name] = dicom_dir + return dirs_dict + + @staticmethod + def _get_segs_paths_dict(segs_dir: Path) -> dict[str, Path]: + paths_dict = {} + for image_path in sorted(segs_dir.iterdir()): + key = image_path.name.replace('.gz', '').replace('.nii', '') + paths_dict[key] = image_path + return paths_dict + + def _get_subjects( + self, + add_segmentations: bool, + add_bounding_boxes: bool, + ) -> list[Subject]: + subjects = [] + pd = get_pandas() + from tqdm.auto import tqdm + + split_name = 'train' + images_dirname = f'{split_name}_images' + images_dir = self.root_dir / images_dirname + image_dirs_dict = self._get_image_dirs_dict(images_dir) + + segmentations_dir = self.root_dir / 'segmentations' + seg_paths_dict = self._get_segs_paths_dict(segmentations_dir) + + bboxes_path = self.root_dir / 'train_bounding_boxes.csv' + bounding_boxes_df = pd.read_csv(bboxes_path) + grouped_boxes = bounding_boxes_df.groupby(self.UID) + + df = pd.read_csv(self.root_dir / f'{split_name}.csv') + + for _, row in tqdm(list(df.iterrows())): + uid = row[self.UID] + image_dir = image_dirs_dict[uid] + seg_path = None + if add_segmentations: + seg_path = seg_paths_dict.get(uid, None) + boxes = [] + if add_bounding_boxes: + try: + boxes_df = grouped_boxes.get_group(uid) + boxes = [dict(row) for _, row in boxes_df.iterrows()] + except KeyError: + pass + subject = self._get_subject( + dict(row), + image_dir, + seg_path, + boxes, + ) + subjects.append(subject) + return subjects + + @staticmethod + def _filter_list(iterable: list[Path], target: str): + def _filter(path: Path): + if path.is_dir(): + return target == path.name + else: + name = path.name.replace('.gz', '').replace('.nii', '') + return target == name + + found = list(filter(_filter, iterable)) + if found: + assert len(found) == 1 + result = found[0] + else: + result = None + return result + + def _get_subject( + self, + csv_row_dict: dict[str, str | int], + image_dir: Path, + seg_path: Path | None, + boxes: TypeBoxes, + ) -> Subject: + subject_dict: dict[str, Any] = {} + subject_dict.update(csv_row_dict) + subject_dict['ct'] = ScalarImage(image_dir) + if seg_path is not None: + subject_dict['seg'] = LabelMap(seg_path) + if boxes: + subject_dict['boxes'] = boxes + return Subject(**subject_dict) diff --git a/src/torchio/datasets/slicer.py b/src/torchio/datasets/slicer.py index d1d04accb..48d7b0343 100644 --- a/src/torchio/datasets/slicer.py +++ b/src/torchio/datasets/slicer.py @@ -1,69 +1,72 @@ -"""3D Slicer sample data.""" - -from __future__ import annotations - import urllib.parse from ..data import ScalarImage from ..data.subject import Subject from ..download import download_url -from ..download import get_torchio_cache_dir +from ..utils import get_torchio_cache_dir -SLICER_URL = "https://github.com/Slicer/SlicerTestingData/releases/download/" +SLICER_URL = 'https://github.com/Slicer/SlicerTestingData/releases/download/' URLS_DICT = { - "MRHead": ( - ("MRHead.nrrd",), - ("SHA256/cc211f0dfd9a05ca3841ce1141b292898b2dd2d3f08286affadf823a7e58df93",), + 'MRHead': ( + ('MRHead.nrrd',), + ('SHA256/cc211f0dfd9a05ca3841ce1141b292898b2dd2d3f08286affadf823a7e58df93',), ), - "DTIBrain": ( - ("DTI-Brain.nrrd",), - ("SHA256/5c78d00c86ae8d968caa7a49b870ef8e1c04525b1abc53845751d8bce1f0b91a",), + 'DTIBrain': ( + ('DTI-Brain.nrrd',), + ('SHA256/5c78d00c86ae8d968caa7a49b870ef8e1c04525b1abc53845751d8bce1f0b91a',), ), - "DTIVolume": ( + 'DTIVolume': ( ( - "DTIVolume.raw.gz", - "DTIVolume.nhdr", + 'DTIVolume.raw.gz', + 'DTIVolume.nhdr', ), ( - "SHA256/d785837276758ddd9d21d76a3694e7fd866505a05bc305793517774c117cb38d", - "SHA256/67564aa42c7e2eec5c3fd68afb5a910e9eab837b61da780933716a3b922e50fe", + 'SHA256/d785837276758ddd9d21d76a3694e7fd866505a05bc305793517774c117cb38d', + 'SHA256/67564aa42c7e2eec5c3fd68afb5a910e9eab837b61da780933716a3b922e50fe', ), ), - "CTChest": ( - ("CT-chest.nrrd",), - ("SHA256/4507b664690840abb6cb9af2d919377ffc4ef75b167cb6fd0f747befdb12e38e",), + 'CTChest': ( + ('CT-chest.nrrd',), + ('SHA256/4507b664690840abb6cb9af2d919377ffc4ef75b167cb6fd0f747befdb12e38e',), ), - "CTACardio": ( - ("CTA-cardio.nrrd",), - ("SHA256/3b0d4eb1a7d8ebb0c5a89cc0504640f76a030b4e869e33ff34c564c3d3b88ad2",), + 'CTACardio': ( + ('CTA-cardio.nrrd',), + ('SHA256/3b0d4eb1a7d8ebb0c5a89cc0504640f76a030b4e869e33ff34c564c3d3b88ad2',), ), } class Slicer(Subject): - """Sample data provided by `3D Slicer `_. + """Sample data provided by [3D Slicer ](https://www.slicer.org/). - See `the Slicer wiki `_ + See [the Slicer wiki ](https://www.slicer.org/wiki/SampleData) for more information. - For information about licensing and permissions, check the `Sample Data - module `_. + For information about licensing and permissions, check the [Sample Data + module ](https://github.com/Slicer/Slicer/blob/31c89f230919a953e56f6722718281ce6da49e06/Modules/Scripted/SampleData/SampleData.py#L75-L81). Args: - name: One of the keys in - `torchio.datasets.slicer.URLS_DICT`. + name: One of the keys in `torchio.datasets.slicer.URLS_DICT`. """ - def __init__(self, name: str = "MRHead") -> None: + def __init__(self, name='MRHead'): try: filenames, url_files = URLS_DICT[name] except KeyError as e: message = f'Invalid name "{name}". Valid names are: {", ".join(URLS_DICT)}' raise ValueError(message) from e - download_root = get_torchio_cache_dir() / "slicer" for filename, url_file in zip(filenames, url_files, strict=True): - filename = filename.replace("-", "_") + filename = filename.replace('-', '_') url = urllib.parse.urljoin(SLICER_URL, url_file) - download_url(url, download_root, filename=filename) - stem = filename.split(".")[0] - super().__init__(**{stem: ScalarImage(download_root / filename)}) + download_root = get_torchio_cache_dir() / 'slicer' + stem = filename.split('.')[0] + download_url( + url, + download_root, + filename=filename, + ) + super().__init__( + { + stem: ScalarImage(download_root / filename), # use last filename + } + ) diff --git a/src/torchio/datasets/zone_plate.py b/src/torchio/datasets/zone_plate.py index e5c8dcd8b..8d007f92e 100644 --- a/src/torchio/datasets/zone_plate.py +++ b/src/torchio/datasets/zone_plate.py @@ -1,7 +1,3 @@ -"""Synthetic data generated from a zone plate.""" - -from __future__ import annotations - import numpy as np from ..data import ScalarImage @@ -16,19 +12,17 @@ class ZonePlate(Subject): processing algorithms, particularly those related to frequency analysis and interpolation. - See equation 10.63 in `Practical Handbook on Image Processing for - Scientific Applications - `_ - by Bernd Jähne. + See equation 10.63 in [Practical Handbook on Image Processing for + Scientific Applications ](https://www.routledge.com/Practical-Handbook-on-Image-Processing-for-Scientific-and-Technical-Applications/Jahne/p/book/9780849319006?srsltid=AfmBOoptrtzILIlMx9FYqvx6UrGbevfD66x2k242iprFdn_CfyOWXjjH) + by Bernd Jähne. Args: size: The size of the generated image along all dimensions. """ - def __init__(self, size: int = 501) -> None: + def __init__(self, size: int = 501): if size < 3: - msg = "Size must be at least 3." - raise ValueError(msg) + raise ValueError('Size must be at least 3.') self.size = size image = self._generate_image(size) super().__init__(image=image) @@ -41,11 +35,11 @@ def _generate_image(size: int) -> ScalarImage: else: fin = size // 2 ini = -fin + 1 - x = np.arange(ini, fin + 1) - y = np.arange(ini, fin + 1) - z = np.arange(ini, fin + 1) - xx, yy, zz = np.meshgrid(x, y, z) - r = np.sqrt(xx**2 + yy**2 + zz**2) + x = np.arange(ini, fin) + y = np.arange(ini, fin) + z = np.arange(ini, fin) + X, Y, Z = np.meshgrid(x, y, z) + r = np.sqrt(X**2 + Y**2 + Z**2) km = 0.8 * np.pi rm = ini w = rm / 10 @@ -55,4 +49,4 @@ def _generate_image(size: int) -> ScalarImage: affine = np.eye(4) origin = np.array([ini, ini, ini]) affine[:3, 3] = origin - return ScalarImage(g[np.newaxis], affine=affine) + return ScalarImage(tensor=g[np.newaxis], affine=affine) diff --git a/src/torchio/download.py b/src/torchio/download.py index e6978d4ee..db60b5c6d 100644 --- a/src/torchio/download.py +++ b/src/torchio/download.py @@ -1,200 +1,175 @@ -"""Download utilities for built-in datasets.""" +"""Most of this code is from torchvision. + +I will remove all this once verbosity is reduced. More info: +https://github.com/pytorch/vision/issues/2830 +""" from __future__ import annotations import gzip import hashlib import os -import shutil import tarfile import zipfile -from pathlib import Path from urllib import error from urllib import request -from loguru import logger -from platformdirs import user_cache_dir -from rich.progress import BarColumn -from rich.progress import DownloadColumn -from rich.progress import Progress -from rich.progress import TransferSpeedColumn +from torch.hub import tqdm from .types import TypePath -def get_torchio_cache_dir() -> Path: - """Return the default cache directory for TorchIO data.""" - return Path(user_cache_dir("torchio")) - - -def calculate_md5(fpath: TypePath, chunk_size: int = 1024 * 1024) -> str: - """Calculate the MD5 checksum of a file.""" +def calculate_md5(fpath, chunk_size=1024 * 1024): md5 = hashlib.md5() - with open(fpath, "rb") as f: - for chunk in iter(lambda: f.read(chunk_size), b""): + with open(fpath, 'rb') as f: + for chunk in iter(lambda: f.read(chunk_size), b''): md5.update(chunk) return md5.hexdigest() -def check_integrity(fpath: TypePath, md5: str | None = None) -> bool: - """Check whether a file exists and optionally matches a checksum.""" +def check_md5(fpath, md5, **kwargs): + return md5 == calculate_md5(fpath, **kwargs) + + +def check_integrity(fpath, md5=None): if not os.path.isfile(fpath): return False if md5 is None: return True - return md5 == calculate_md5(fpath) + return check_md5(fpath, md5) -def download_url( +def gen_bar_updater(): + pbar = tqdm(total=None) + + def bar_update(count, block_size, total_size): + if pbar.total is None and total_size: + pbar.total = total_size + progress_bytes = count * block_size + pbar.update(progress_bytes - pbar.n) + + return bar_update + + +# Adapted from torchvision, removing print statements +def download_and_extract_archive( url: str, - root: TypePath, - *, - filename: str | None = None, + download_root: TypePath, + extract_root: TypePath | None = None, + filename: TypePath | None = None, md5: str | None = None, -) -> Path: - """Download a file from a URL and place it in *root*. - - Args: - url: URL to download file from. - root: Directory to place downloaded file in. - filename: Name to save the file under. If `None`, use the - basename of the URL. - md5: MD5 checksum of the download. If `None`, skip check. - - Returns: - Path to the downloaded file. - """ - root = os.path.expanduser(root) + remove_finished: bool = False, +) -> None: + download_root = os.path.expanduser(download_root) + if extract_root is None: + extract_root = download_root if not filename: filename = os.path.basename(url) - fpath = os.path.join(root, filename) - os.makedirs(root, exist_ok=True) - if not check_integrity(fpath, md5): - progress = Progress( - "[progress.description]{task.description}", - BarColumn(), - DownloadColumn(), - TransferSpeedColumn(), - ) - - def _try_download(download_url: str) -> None: - with progress: - task = progress.add_task( - f"Downloading [cyan]{filename}", - total=None, - ) + download_url(url, download_root, filename, md5) + archive = os.path.join(download_root, filename) + extract_archive(archive, extract_root, remove_finished) - def _reporthook( - count: int, - block_size: int, - total_size: int, - ) -> None: - if total_size > 0: - progress.update(task, total=total_size) - progress.update(task, completed=count * block_size) - request.urlretrieve(download_url, fpath, reporthook=_reporthook) +def _is_tarxz(filename): + return filename.endswith('.tar.xz') - try: - _try_download(url) - except (error.URLError, OSError): - if url.startswith("https"): - http_url = url.replace("https:", "http:") - logger.info("Retrying with http: {}", http_url) - _try_download(http_url) - else: - raise - if not check_integrity(fpath, md5): - msg = f"File not found or corrupted: {fpath}" - raise RuntimeError(msg) - return Path(fpath) +def _is_tar(filename): + return filename.endswith('.tar') -def extract_archive( - from_path: TypePath, - to_path: TypePath | None = None, - *, - remove_finished: bool = False, -) -> None: - """Extract an archive file. - Supports `.zip`, `.tar`, `.tar.gz`, `.tgz`, `.tar.xz`, - and `.gz` (single-file gzip). - """ - from_path = str(from_path) +def _is_targz(filename): + return filename.endswith('.tar.gz') + + +def _is_tgz(filename): + return filename.endswith('.tgz') + + +def _is_gzip(filename): + return filename.endswith('.gz') and not filename.endswith('.tar.gz') + + +def _is_zip(filename): + return filename.endswith('.zip') + + +def extract_archive(from_path, to_path=None, remove_finished=False): if to_path is None: to_path = os.path.dirname(from_path) - if from_path.endswith(".tar"): - with tarfile.open(from_path, "r") as tar: + if _is_tar(from_path): + with tarfile.open(from_path, 'r') as tar: tar.extractall(path=to_path) - elif from_path.endswith((".tar.gz", ".tgz")): - with tarfile.open(from_path, "r:gz") as tar: + elif _is_targz(from_path) or _is_tgz(from_path): + with tarfile.open(from_path, 'r:gz') as tar: tar.extractall(path=to_path) - elif from_path.endswith(".tar.xz"): - with tarfile.open(from_path, "r:xz") as tar: + elif _is_tarxz(from_path): + with tarfile.open(from_path, 'r:xz') as tar: tar.extractall(path=to_path) - elif from_path.endswith(".gz"): + elif _is_gzip(from_path): stem = os.path.splitext(os.path.basename(from_path))[0] - out = os.path.join(str(to_path), stem) - with open(out, "wb") as out_f, gzip.GzipFile(from_path) as zip_f: + to_path = os.path.join(to_path, stem) + with open(to_path, 'wb') as out_f, gzip.GzipFile(from_path) as zip_f: out_f.write(zip_f.read()) - elif from_path.endswith(".zip"): - with zipfile.ZipFile(from_path, "r") as z: + elif _is_zip(from_path): + with zipfile.ZipFile(from_path, 'r') as z: z.extractall(to_path) else: - msg = f"Extraction of {from_path} not supported" - raise ValueError(msg) + raise ValueError(f'Extraction of {from_path} not supported') if remove_finished: os.remove(from_path) -def download_and_extract_archive( +# Adapted from torchvision, removing print statements +def download_url( url: str, - download_root: TypePath, - *, - extract_root: TypePath | None = None, - filename: str | None = None, + root: TypePath, + filename: TypePath | None = None, md5: str | None = None, - remove_finished: bool = False, ) -> None: - """Download an archive and extract it. + """Download a file from a url and place it in root. Args: - url: URL to download. - download_root: Directory to place the downloaded archive. - extract_root: Directory to extract to. Defaults to - *download_root*. - filename: Archive filename. Defaults to the URL basename. - md5: MD5 checksum of the archive. - remove_finished: Delete the archive after extraction. + url: URL to download file from + root: Directory to place downloaded file in + filename: Name to save the file under. + If `None`, use the basename of the URL + md5: MD5 checksum of the download. If None, do not check """ - if extract_root is None: - extract_root = download_root + + root = os.path.expanduser(root) if not filename: filename = os.path.basename(url) - download_url(url, download_root, filename=filename, md5=md5) - archive = os.path.join(os.path.expanduser(download_root), filename) - extract_archive(archive, extract_root, remove_finished=remove_finished) - - -def compress( - input_path: TypePath, - output_path: TypePath | None = None, -) -> Path: - """Compress a NIfTI file with gzip. - - Args: - input_path: Path to the `.nii` file. - output_path: Path for the compressed file. Defaults to - replacing the suffix with `.nii.gz`. - - Returns: - Path to the compressed file. - """ - if output_path is None: - output_path = Path(input_path).with_suffix(".nii.gz") - with open(input_path, "rb") as f_in, gzip.open(output_path, "wb") as f_out: - shutil.copyfileobj(f_in, f_out) - return Path(output_path) + fpath = os.path.join(root, filename) + os.makedirs(root, exist_ok=True) + # check if file is already present locally + if not check_integrity(fpath, md5): + try: + print('Downloading ' + url + ' to ' + fpath) # noqa: T201 + request.urlretrieve( + url, + fpath, + reporthook=gen_bar_updater(), + ) + except (error.URLError, OSError) as e: + if url[:5] == 'https': + url = url.replace('https:', 'http:') + message = ( + 'Failed download. Trying https -> http instead. Downloading ' + + url + + ' to ' + + fpath + ) + print(message) # noqa: T201 + request.urlretrieve( + url, + fpath, + reporthook=gen_bar_updater(), + ) + else: + raise e + # check integrity of downloaded file + if not check_integrity(fpath, md5): + raise RuntimeError('File not found or corrupted.') diff --git a/src/torchio/external/__init__.py b/src/torchio/external/__init__.py index 57b731e1b..e69de29bb 100644 --- a/src/torchio/external/__init__.py +++ b/src/torchio/external/__init__.py @@ -1 +0,0 @@ -"""External utilities: optional dependency imports and duecredit stubs.""" diff --git a/src/torchio/external/due.py b/src/torchio/external/due.py index d800e1ee5..ad5b7956e 100644 --- a/src/torchio/external/due.py +++ b/src/torchio/external/due.py @@ -1,29 +1,36 @@ -"""Stub for a guaranteed safe import of duecredit constructs. +"""Stub file for a guaranteed safe import of duecredit constructs: if +duecredit is not available. -If duecredit is not available, a no-op collector is used instead. +To use it, place it into your project codebase to be imported, e.g., copy as:: + + cp stub.py /path/tomodule/module/due.py + +Note that it might be better to avoid naming it duecredit.py to avoid +shadowing installed duecredit. + +Then use in your code as:: + + from .due import due, Doi, BibTeX, Text + +See https://github.com/duecredit/duecredit/blob/master/README.md for examples. Origin: Originally a part of the duecredit Copyright: 2015-2019 DueCredit developers License: BSD-2 - -See https://github.com/duecredit/duecredit/blob/master/README.md for examples. """ -from __future__ import annotations - -import importlib -import logging -from typing import Any +__version__ = '0.0.8' class InactiveDueCreditCollector: - """Stub Collector that does nothing.""" + """Just a stub at the Collector which would not do anything.""" - def _donothing(self, *args: Any, **kwargs: Any) -> None: + def _donothing(self, *args, **kwargs): + """Perform no good and no bad.""" pass - def dcite(self, *args: Any, **kwargs: Any) -> Any: - def nondecorating_decorator(func: Any) -> Any: + def dcite(self, *args, **kwargs): + def nondecorating_decorator(func): return func return nondecorating_decorator @@ -31,30 +38,36 @@ def nondecorating_decorator(func: Any) -> Any: active = False activate = add = cite = dump = load = _donothing - def __repr__(self) -> str: - return f"{self.__class__.__name__}()" + def __repr__(self): + return self.__class__.__name__ + '()' -def _donothing_func(*args: Any, **kwargs: Any) -> None: +def _donothing_func(*args, **kwargs): + """Perform no good and no bad.""" pass try: - _duecredit = importlib.import_module("duecredit") - BibTeX = _duecredit.BibTeX - Doi = _duecredit.Doi - Text = _duecredit.Text - Url = _duecredit.Url - due = _duecredit.due - - if "due" in locals() and not hasattr(due, "cite"): - msg = "Imported due lacks .cite. DueCredit is now disabled" - raise RuntimeError(msg) -except Exception as _exc: - if not isinstance(_exc, ImportError): - logging.getLogger("duecredit").error( - "Failed to import duecredit due to %s", - _exc, + import importlib + + duecredit = importlib.import_module('duecredit') + BibTeX = duecredit.BibTeX + Doi = duecredit.Doi + Text = duecredit.Text + Url = duecredit.Url + due = duecredit.due + + if 'due' in locals() and not hasattr(due, 'cite'): + raise RuntimeError( + 'Imported due lacks .cite. DueCredit is now disabled', + ) +except Exception as e: + if not isinstance(e, ImportError): + import logging + + logging.getLogger('duecredit').error( + f'Failed to import duecredit due to {e}', ) + # Initiate due stub due = InactiveDueCreditCollector() BibTeX = Doi = Url = Text = _donothing_func diff --git a/src/torchio/external/imports.py b/src/torchio/external/imports.py index 648900baa..b933e1b17 100644 --- a/src/torchio/external/imports.py +++ b/src/torchio/external/imports.py @@ -1,68 +1,53 @@ -"""Helpers for optional dependency imports.""" - from __future__ import annotations from importlib import import_module from importlib.util import find_spec +from shutil import which from types import ModuleType -from typing import Any def _check_module(*, module: str, extra: str, package: str | None = None) -> None: if find_spec(module) is None: name = module if package is None else package - msg = ( - f"The `{name}` package is required for this." - f" Install TorchIO with the `{extra}` extra:" - f" `pip install torchio[{extra}]`." + message = ( + f'The `{name}` package is required for this.' + f' Install TorchIO with the `{extra}` extra:' + f' `pip install torchio[{extra}]`.' ) - raise ImportError(msg) + raise ImportError(message) -def _check_and_import(module: str, extra: str, **kwargs: Any) -> ModuleType: +def _check_and_import(module: str, extra: str, **kwargs) -> ModuleType: _check_module(module=module, extra=extra, **kwargs) return import_module(module) -def get_niizarr() -> ModuleType: - return _check_and_import( - module="niizarr", - extra="zarr", - package="nifti-zarr", - ) - - -def get_matplotlib() -> ModuleType: - return _check_and_import(module="matplotlib", extra="plot") - +def get_pandas() -> ModuleType: + return _check_and_import(module='pandas', extra='csv') -def get_matplotlib_pyplot() -> ModuleType: - _check_module(module="matplotlib", extra="plot") - return import_module("matplotlib.pyplot") +def get_colorcet() -> ModuleType: + return _check_and_import(module='colorcet', extra='plot') -def get_colorcet() -> ModuleType | None: - """Return colorcet if installed, else None (fallback to tab10).""" - if find_spec("colorcet") is None: - return None - return import_module("colorcet") - -def get_pillow() -> ModuleType: - return _check_and_import(module="PIL", extra="plot", package="Pillow") +def get_ffmpeg() -> ModuleType: + ffmpeg = _check_and_import(module='ffmpeg', extra='video', package='ffmpeg-python') + _check_executable('ffmpeg') + return ffmpeg -def get_ffmpeg() -> ModuleType: - return _check_and_import(module="ffmpeg", extra="video", package="ffmpeg-python") +def get_sklearn() -> ModuleType: + return _check_and_import(module='sklearn', extra='sklearn', package='scikit-learn') def get_monai() -> ModuleType: - return _check_and_import(module="monai", extra="monai") + return _check_and_import(module='monai', extra='monai') -def get_ipyniivue() -> ModuleType: - return _check_and_import( - module="ipyniivue", - extra="niivue", - package="ipyniivue", - ) +def _check_executable(executable: str) -> None: + if which(executable) is None: + message = ( + f'The `{executable}` executable is required for this. Install it from your' + ' package manager or download it from the official website.' + ) + raise FileNotFoundError(message) diff --git a/src/torchio/io.py b/src/torchio/io.py deleted file mode 100644 index ac6a55d36..000000000 --- a/src/torchio/io.py +++ /dev/null @@ -1,119 +0,0 @@ -"""I/O helpers for spatial transforms and matrices.""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import SimpleITK as sitk -import torch -from einops import rearrange -from torch import Tensor - -from .types import TypePath - -# Matrices used to switch between LPS and RAS -FLIPXY_33 = np.diag([-1.0, -1.0, 1.0]) -FLIPXY_44 = np.diag([-1.0, -1.0, 1.0, 1.0]) - - -def read_matrix(path: TypePath) -> Tensor: - """Read an affine transform from a file and return a 4x4 tensor in RAS. - - Supported formats: - - - `.tfm`, `.h5`: ITK transforms (read via SimpleITK) - - `.txt`, `.trsf`: NiftyReg / blockmatching matrices - - Args: - path: Path to the transform file. - - Returns: - A `(4, 4)` float64 tensor representing the affine in RAS - convention. - """ - path = Path(path) - suffix = path.suffix - if suffix in (".tfm", ".h5"): - return _read_itk_matrix(path) - if suffix in (".txt", ".trsf"): - return _read_niftyreg_matrix(path) - msg = f'Unknown suffix for transform file: "{suffix}"' - raise ValueError(msg) - - -def write_matrix(matrix: Tensor, path: TypePath) -> None: - """Write a 4x4 affine matrix to a file. - - Args: - matrix: A `(4, 4)` tensor in RAS convention. - path: Destination path. Suffix determines format. - """ - path = Path(path) - suffix = path.suffix - if suffix in (".tfm", ".h5"): - _write_itk_matrix(matrix, path) - elif suffix in (".txt", ".trsf"): - _write_niftyreg_matrix(matrix, path) - else: - msg = f'Unknown suffix for transform file: "{suffix}"' - raise ValueError(msg) - - -# --- ITK (LPS convention) --------------------------------------------------- - - -def _to_itk_convention(matrix: Tensor | np.ndarray) -> np.ndarray: - """Convert a RAS affine to ITK's LPS convention.""" - if isinstance(matrix, Tensor): - matrix = matrix.numpy() - matrix = FLIPXY_44 @ matrix @ FLIPXY_44 - return np.linalg.inv(matrix) - - -def _from_itk_convention(matrix: np.ndarray) -> np.ndarray: - """Convert an ITK LPS affine to RAS convention.""" - matrix = matrix @ FLIPXY_44 - matrix = FLIPXY_44 @ matrix - return np.linalg.inv(matrix) - - -def _read_itk_matrix(path: TypePath) -> Tensor: - """Read an affine transform in ITK's `.tfm` or `.h5` format.""" - transform = sitk.ReadTransform(str(path)) - parameters = transform.GetParameters() - rotation_parameters = parameters[:9] - rotation_matrix = rearrange(np.array(rotation_parameters), "(i j) -> i j", i=3) - translation_parameters = parameters[9:] - translation_vector = rearrange(np.array(translation_parameters), "i -> i 1") - matrix = np.hstack([rotation_matrix, translation_vector]) - homogeneous_matrix_lps = np.vstack([matrix, [0, 0, 0, 1]]) - homogeneous_matrix_ras = _from_itk_convention(homogeneous_matrix_lps) - return torch.as_tensor(homogeneous_matrix_ras) - - -def _write_itk_matrix(matrix: Tensor, path: TypePath) -> None: - """Write a RAS affine as an ITK `.tfm` file.""" - itk_matrix = _to_itk_convention(matrix) - rotation = itk_matrix[:3, :3].ravel().tolist() - translation = itk_matrix[:3, 3].tolist() - transform = sitk.AffineTransform(rotation, translation) - transform.WriteTransform(str(path)) - - -# --- NiftyReg / blockmatching ------------------------------------------------ - - -def _read_niftyreg_matrix(path: TypePath) -> Tensor: - """Read a NiftyReg matrix (reference → floating, inverted to RAS).""" - matrix = np.loadtxt(path).astype(np.float64) - inverted = np.linalg.inv(matrix) - return torch.from_numpy(inverted) - - -def _write_niftyreg_matrix(matrix: Tensor, path: TypePath) -> None: - """Write a RAS affine as a NiftyReg `.txt` file.""" - if isinstance(matrix, Tensor): - matrix = matrix.numpy() - inverted = np.linalg.inv(matrix) - np.savetxt(path, inverted, fmt="%.8f") diff --git a/src/torchio/loader.py b/src/torchio/loader.py deleted file mode 100644 index d26f53e63..000000000 --- a/src/torchio/loader.py +++ /dev/null @@ -1,95 +0,0 @@ -"""DataLoader wrappers for Subject and Image collation.""" - -from __future__ import annotations - -from collections.abc import Sequence -from typing import Any - -from torch.utils.data import DataLoader -from torch.utils.data import Dataset - -from .data.batch import ImagesBatch -from .data.batch import SubjectsBatch - - -def collate_subjects(batch: Sequence[Any]) -> SubjectsBatch: - """Collate a list of Subjects into a SubjectsBatch. - - Args: - batch: Sequence of `Subject` instances. - - Returns: - A `SubjectsBatch` with stacked 5D tensors. - """ - return SubjectsBatch.from_subjects(list(batch)) - - -def collate_images(batch: Sequence[Any]) -> ImagesBatch: - """Collate a list of Images into an ImagesBatch. - - Args: - batch: Sequence of `Image` instances. - - Returns: - An `ImagesBatch` with a stacked 5D tensor. - """ - return ImagesBatch.from_images(list(batch)) - - -class SubjectsLoader(DataLoader): - """DataLoader that returns `SubjectsBatch` instances. - - A thin wrapper around `torch.utils.data.DataLoader` that - collates `Subject` instances into `SubjectsBatch`. - - Args: - dataset: A dataset that returns `Subject` instances. - **kwargs: Passed to `DataLoader.__init__`. - - Examples: - >>> loader = tio.SubjectsLoader(dataset, batch_size=4) - >>> batch = next(iter(loader)) - >>> batch.t1.data.shape - torch.Size([4, 1, 256, 256, 176]) - """ - - def __init__(self, dataset: Dataset, **kwargs: Any) -> None: - if "collate_fn" in kwargs: - msg = ( - "SubjectsLoader sets collate_fn automatically; " - "pass a plain DataLoader if you need a custom collate_fn" - ) - raise ValueError(msg) - super().__init__(dataset, collate_fn=collate_subjects, **kwargs) - - -class ImagesLoader(DataLoader): - """DataLoader that returns `ImagesBatch` instances. - - A thin wrapper around `torch.utils.data.DataLoader` that - collates `Image` instances into `ImagesBatch`. - - Args: - dataset: A dataset that returns `Image` instances. - **kwargs: Passed to `DataLoader.__init__`. - - Examples: - >>> loader = tio.ImagesLoader(dataset, batch_size=4) - >>> batch = next(iter(loader)) - >>> batch.data.shape - torch.Size([4, 1, 256, 256, 176]) - """ - - def __init__(self, dataset: Dataset, **kwargs: Any) -> None: - if "collate_fn" in kwargs: - msg = ( - "ImagesLoader sets collate_fn automatically; " - "pass a plain DataLoader if you need a custom collate_fn" - ) - raise ValueError(msg) - super().__init__(dataset, collate_fn=collate_images, **kwargs) - - -# Aliases for radiology users (see Subject/Study note in subject.py). -StudiesLoader = SubjectsLoader -collate_studies = collate_subjects diff --git a/src/torchio/logging.py b/src/torchio/logging.py deleted file mode 100644 index d2b4afa52..000000000 --- a/src/torchio/logging.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Logging configuration for TorchIO. - -Logging is **disabled by default**. Call -[`enable_logging`][torchio.enable_logging] to opt in. - -Example:: - - import torchio as tio - tio.enable_logging("DEBUG") # rich-formatted debug output -""" - -from __future__ import annotations - -import sys - -from loguru import logger -from rich.logging import RichHandler - -# Libraries must not emit logs unless the user opts in -logger.disable("torchio") - - -def enable_logging(level: str = "INFO", *, rich: bool = True) -> None: - """Enable TorchIO logging. - - Args: - level: Minimum log level (`DEBUG`, `INFO`, `WARNING`, - `ERROR`). - rich: If `True` (default), use `rich.logging.RichHandler` - for colourful, markup-enabled output with pretty tracebacks. - Set to `False` for plain stderr output. - """ - logger.enable("torchio") - logger.remove() - if rich: - logger.add( - RichHandler(markup=True, rich_tracebacks=True), - format="{message}", - level=level, - ) - else: - logger.add(sys.stderr, level=level) diff --git a/src/torchio/reference.py b/src/torchio/reference.py new file mode 100644 index 000000000..b5de0d1a9 --- /dev/null +++ b/src/torchio/reference.py @@ -0,0 +1,40 @@ +# Use duecredit to provide a citation to relevant work to +# be cited. This does nothing unless the user has duecredit installed +# and calls this with duecredit (as in `python -m duecredit script.py`): +from .external.due import BibTeX +from .external.due import Doi +from .external.due import due + +BIBTEX = r"""@article{perez-garcia_torchio_2021, + title = {{TorchIO}: a {Python} library for efficient loading, preprocessing, augmentation and patch-based sampling of medical images in deep learning}, + journal = {Computer Methods and Programs in Biomedicine}, + pages = {106236}, + year = {2021}, + issn = {0169-2607}, + doi = {https://doi.org/10.1016/j.cmpb.2021.106236}, + url = {https://www.sciencedirect.com/science/article/pii/S0169260721003102}, + author = {P{\'e}rez-Garc{\'i}a, Fernando and Sparks, Rachel and Ourselin, S{\'e}bastien}, + keywords = {Medical image computing, Deep learning, Data augmentation, Preprocessing}, +} """ + +TITLE = ( + 'TorchIO: a Python library for efficient loading, preprocessing,' + ' augmentation and patch-based sampling of medical images in deep learning' +) + +DESCRIPTION = 'Tools for loading, augmenting and writing 3D medical images on PyTorch' + +due.cite( + BibTeX(BIBTEX), + description=TITLE, + path='torchio', + cite_module=True, +) + +due.cite( + Doi('10.5281/zenodo.3739230'), + description=DESCRIPTION, + path='torchio', + tags=['implementation'], + cite_module=True, +) diff --git a/src/torchio/repr_html.py b/src/torchio/repr_html.py deleted file mode 100644 index c7966d574..000000000 --- a/src/torchio/repr_html.py +++ /dev/null @@ -1,232 +0,0 @@ -"""HTML table representations for Jupyter notebooks.""" - -from __future__ import annotations - -from html import escape -from typing import TYPE_CHECKING -from typing import Any - -import humanize - -if TYPE_CHECKING: - from .data.image import Image - from .data.subject import Subject - -# Minimal inline CSS so the tables look decent in any Jupyter theme. -_STYLE = """\ -""" - - -def _pluralize(word: str, n: int) -> str: - if n == 1: - return f"1 {word}" - return f"{n} {word}s" if not word.endswith("x") else f"{n} {word}es" - - -def _row(key: str, value: str) -> str: - return f"{escape(key)}{escape(value)}" - - -def image_to_html(image: Image) -> str: - """Build an HTML representation for an Image with an embedded plot.""" - cls_name = type(image).__name__ - rows: list[str] = [] - - rows.append(_row("Type", cls_name)) - - try: - sp = "({})".format(", ".join(f"{s:.2f}" for s in image.spacing)) - ori = "({})".format(", ".join(f"{o:.2f}" for o in image.origin)) - angles = "({})".format( - ", ".join(f"{a:.1f}°" for a in image.affine.euler_angles) - ) - dt = str(image.dtype).replace("torch.", "") - rows.append(_row("Channels", str(image.num_channels))) - rows.append(_row("Spatial shape", str(image.spatial_shape))) - rows.append(_row("Spacing", f"{sp} mm")) - rows.append(_row("Origin", f"{ori} mm")) - rows.append(_row("Orientation", "".join(image.orientation) + "+")) - rows.append(_row("Euler angles", angles)) - rows.append(_row("dtype", dt)) - rows.append(_row("Memory", humanize.naturalsize(image.memory, binary=True))) - except Exception: - if image.path is not None: - rows.append(_row("Path", str(image.path))) - - for name, pts in image.points.items(): - rows.append(_row(f"Points '{name}'", _pluralize("point", pts.num_points))) - for name, boxes in image.bounding_boxes.items(): - rows.append(_row(f"BBoxes '{name}'", _pluralize("box", boxes.num_boxes))) - - table = f'{_STYLE}\n\n' + "\n".join(rows) + "\n
" - - if plot_html := _try_plot_base64(image): - return f"{table}\n{plot_html}" - return table - - -def _render_fig_base64(render_fn: Any) -> str | None: - """Call *render_fn* and return a base64 tag. - - Renders via the Agg canvas without changing the global matplotlib - backend, so interactive plotting still works afterwards. - """ - import base64 - import io - - try: - import matplotlib.pyplot as plt - from matplotlib.backends.backend_agg import FigureCanvasAgg - - fig = render_fn() - if fig is None: - return None - FigureCanvasAgg(fig) - buf = io.BytesIO() - fig.savefig(buf, format="png", bbox_inches="tight", dpi=100) - plt.close(fig) - buf.seek(0) - b64 = base64.b64encode(buf.read()).decode("ascii") - return f'' - except Exception: - return None - - -def _try_plot_base64(image: Image) -> str | None: - """Render a 3-slice plot as an inline base64 `` tag.""" - try: - from .visualization import plot_image - except ImportError: - return None - - return _render_fig_base64( - lambda: plot_image(image, show=False), - ) - - -def subject_to_html(subject: Subject) -> str: - """Build an HTML table for a Subject.""" - parts: list[str] = [_STYLE] - - if subject.images: - parts.append(_images_table_html(subject)) - if subject.points: - parts.append(_points_table_html(subject)) - if subject.bounding_boxes: - parts.append(_bboxes_table_html(subject)) - if subject.metadata: - parts.append(_metadata_table_html(subject)) - - if plot_html := _try_subject_plot_base64(subject): - parts.append(plot_html) - - return "\n".join(parts) - - -def _images_table_html(subject: Subject) -> str: - header = ( - "NameTypeShape" - "SpacingOrientation" - ) - rows: list[str] = [header] - for name, image in subject.images.items(): - img_type = type(image).__name__ - try: - shape = str(image.shape) - sp = "({})".format(", ".join(f"{s:.2f}" for s in image.spacing)) - orient = "".join(image.orientation) + "+" - except Exception: - shape = sp = orient = "?" - rows.append( - f"{escape(name)}{escape(img_type)}" - f"{escape(shape)}{escape(sp)}" - f"{escape(orient)}" - ) - return ( - '
Images
\n' - '\n' + "\n".join(rows) + "\n
" - ) - - -def _points_table_html(subject: Subject) -> str: - rows = ["NameCountAxes"] - for name, pts in subject.points.items(): - rows.append( - f"{escape(name)}" - f"{_pluralize('point', pts.num_points)}" - f"{escape(pts.axes)}" - ) - return ( - '
Points
\n' - '\n' + "\n".join(rows) + "\n
" - ) - - -def _bboxes_table_html(subject: Subject) -> str: - rows = ["NameCountFormat"] - for name, boxes in subject.bounding_boxes.items(): - fmt = f"{boxes.format.axes} ({boxes.format.representation.value})" - rows.append( - f"{escape(name)}" - f"{_pluralize('box', boxes.num_boxes)}" - f"{escape(fmt)}" - ) - return ( - '
Bounding Boxes
\n' - '\n' + "\n".join(rows) + "\n
" - ) - - -def _metadata_table_html(subject: Subject) -> str: - rows = ["KeyValue"] - for key, value in subject.metadata.items(): - rows.append(_row(key, str(value))) - return ( - '
Metadata
\n' - '\n' + "\n".join(rows) + "\n
" - ) - - -def _try_subject_plot_base64(subject: Subject) -> str | None: - """Render a subject grid plot as an inline base64 `` tag.""" - try: - from .visualization import plot_subject - except ImportError: - return None - - num_images = len(subject.images) - return _render_fig_base64( - lambda: plot_subject( - subject, - show=False, - figsize=(12, 3 * num_images), - ), - ) diff --git a/src/torchio/transforms/__init__.py b/src/torchio/transforms/__init__.py index 46b5c9eb4..8c62e88a8 100644 --- a/src/torchio/transforms/__init__.py +++ b/src/torchio/transforms/__init__.py @@ -1,105 +1,125 @@ -"""TorchIO transforms.""" - -from .compose import Compose -from .compose import OneOf -from .compose import SomeOf -from .cornucopia_adapter import CornucopiaAdapter -from .intensity.bias_field import BiasField -from .intensity.blur import Blur -from .intensity.clamp import Clamp -from .intensity.gamma import Gamma -from .intensity.ghosting import Ghosting -from .intensity.histogram_standardization import HistogramStandardization -from .intensity.labels_to_image import LabelsToImage -from .intensity.mask import Mask -from .intensity.motion import Motion -from .intensity.noise import Noise -from .intensity.normalize import Normalize -from .intensity.normalize import RescaleIntensity -from .intensity.pca import PCA -from .intensity.spike import Spike -from .intensity.standardize import Standardize -from .intensity.standardize import ZNormalization -from .intensity.swap import Swap -from .label.contour import Contour -from .label.keep_largest import KeepLargestComponent -from .label.one_hot import OneHot -from .label.remap_labels import RemapLabels -from .label.remove_labels import RemoveLabels -from .label.sequential_labels import SequentialLabels +from .augmentation.composition import Compose +from .augmentation.composition import OneOf +from .augmentation.intensity import BiasField +from .augmentation.intensity import Blur +from .augmentation.intensity import Gamma +from .augmentation.intensity import Ghosting +from .augmentation.intensity import LabelsToImage +from .augmentation.intensity import Motion +from .augmentation.intensity import Noise +from .augmentation.intensity import RandomBiasField +from .augmentation.intensity import RandomBlur +from .augmentation.intensity import RandomGamma +from .augmentation.intensity import RandomGhosting +from .augmentation.intensity import RandomLabelsToImage +from .augmentation.intensity import RandomMotion +from .augmentation.intensity import RandomNoise +from .augmentation.intensity import RandomSpike +from .augmentation.intensity import RandomSwap +from .augmentation.intensity import Spike +from .augmentation.intensity import Swap +from .augmentation.spatial import Affine +from .augmentation.spatial import AffineElasticDeformation +from .augmentation.spatial import ElasticDeformation +from .augmentation.spatial import Flip +from .augmentation.spatial import RandomAffine +from .augmentation.spatial import RandomAffineElasticDeformation +from .augmentation.spatial import RandomAnisotropy +from .augmentation.spatial import RandomElasticDeformation +from .augmentation.spatial import RandomFlip +from .fourier import FourierTransform +from .intensity_transform import IntensityTransform from .lambda_transform import Lambda from .monai_adapter import MonaiAdapter -from .parameter_range import Choice -from .spatial.anisotropy import Anisotropy -from .spatial.copy_affine import CopyAffine -from .spatial.crop import Crop -from .spatial.crop_or_pad import CropOrPad -from .spatial.ensure_shape_multiple import EnsureShapeMultiple -from .spatial.flip import Flip -from .spatial.pad import Pad -from .spatial.reorient import Reorient -from .spatial.resize import Resize -from .spatial.spatial import Affine -from .spatial.spatial import ElasticDeformation -from .spatial.spatial import Resample -from .spatial.spatial import Spatial -from .spatial.to_reference_space import ToReferenceSpace -from .spatial.transpose import Transpose -from .to import To -from .transform import AppliedTransform -from .transform import IntensityTransform -from .transform import SpatialTransform +from .preprocessing import PCA +from .preprocessing import Clamp +from .preprocessing import Contour +from .preprocessing import CopyAffine +from .preprocessing import Crop +from .preprocessing import CropOrPad +from .preprocessing import EnsureShapeMultiple +from .preprocessing import HistogramStandardization +from .preprocessing import KeepLargestComponent +from .preprocessing import Mask +from .preprocessing import OneHot +from .preprocessing import Pad +from .preprocessing import RemapLabels +from .preprocessing import RemoveLabels +from .preprocessing import Resample +from .preprocessing import RescaleIntensity +from .preprocessing import Resize +from .preprocessing import SequentialLabels +from .preprocessing import To +from .preprocessing import ToCanonical +from .preprocessing import ToOrientation +from .preprocessing import ToReferenceSpace +from .preprocessing import Transpose +from .preprocessing import ZNormalization +from .preprocessing.intensity.histogram_standardization import train_histogram +from .preprocessing.label.label_transform import LabelTransform +from .spatial_transform import SpatialTransform from .transform import Transform __all__ = [ - "PCA", - "Affine", - "Anisotropy", - "AppliedTransform", - "BiasField", - "Blur", - "Choice", - "Clamp", - "Compose", - "Contour", - "CopyAffine", - "CornucopiaAdapter", - "Crop", - "CropOrPad", - "ElasticDeformation", - "EnsureShapeMultiple", - "Flip", - "Gamma", - "Ghosting", - "HistogramStandardization", - "IntensityTransform", - "KeepLargestComponent", - "LabelsToImage", - "Lambda", - "Mask", - "MonaiAdapter", - "Motion", - "Noise", - "Normalize", - "OneHot", - "OneOf", - "Pad", - "RemapLabels", - "RemoveLabels", - "Reorient", - "Resample", - "RescaleIntensity", - "Resize", - "SequentialLabels", - "SomeOf", - "Spatial", - "SpatialTransform", - "Spike", - "Standardize", - "Swap", - "To", - "ToReferenceSpace", - "Transform", - "Transpose", - "ZNormalization", + 'Transform', + 'FourierTransform', + 'SpatialTransform', + 'IntensityTransform', + 'LabelTransform', + 'Lambda', + 'MonaiAdapter', + 'OneOf', + 'Compose', + 'RandomFlip', + 'Flip', + 'RandomAffine', + 'Affine', + 'RandomAnisotropy', + 'RandomElasticDeformation', + 'ElasticDeformation', + 'RandomAffineElasticDeformation', + 'AffineElasticDeformation', + 'RandomSwap', + 'Swap', + 'RandomBlur', + 'Blur', + 'RandomNoise', + 'Noise', + 'RandomSpike', + 'Spike', + 'RandomGamma', + 'Gamma', + 'RandomMotion', + 'Motion', + 'RandomGhosting', + 'Ghosting', + 'RandomBiasField', + 'BiasField', + 'RandomLabelsToImage', + 'LabelsToImage', + 'Pad', + 'Crop', + 'Resize', + 'Resample', + 'To', + 'ToCanonical', + 'ToOrientation', + 'ToReferenceSpace', + 'Transpose', + 'ZNormalization', + 'HistogramStandardization', + 'RescaleIntensity', + 'PCA', + 'Clamp', + 'Mask', + 'CropOrPad', + 'CopyAffine', + 'EnsureShapeMultiple', + 'train_histogram', + 'OneHot', + 'Contour', + 'RemapLabels', + 'RemoveLabels', + 'SequentialLabels', + 'KeepLargestComponent', ] diff --git a/src/torchio/transforms/_statistics.py b/src/torchio/transforms/_statistics.py deleted file mode 100644 index 0ea04498b..000000000 --- a/src/torchio/transforms/_statistics.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Shared statistical helpers for transforms.""" - -from __future__ import annotations - -import math - -import torch -from torch import Tensor - - -def compute_quantile(values: Tensor, q: float) -> Tensor: - """Compute a single quantile of a one-dimensional tensor. - - `torch.quantile` raises `RuntimeError: quantile() input tensor is too - large` for inputs with more than `2**24` elements, which is easily - reached by high-resolution volumes. `torch.kthvalue` has no such size - limit and is much faster on large tensors, so it is used here with linear - interpolation to reproduce the default `torch.quantile` behavior. - - This is adapted from the solution by Elie Goudout (`@ego-thales`): - https://github.com/pytorch/pytorch/issues/157431#issuecomment-3026856373 - - Args: - values: One-dimensional tensor of values. - q: Quantile to compute, in the `[0, 1]` range. - - Returns: - Zero-dimensional tensor with the computed quantile. - - Raises: - ValueError: If `q` is outside the `[0, 1]` range. - """ - if not 0 <= q <= 1: - msg = f"Only values 0 <= q <= 1 are supported, but got {q!r}" - raise ValueError(msg) - index = q * (values.numel() - 1) - lower = math.floor(index) - lower_value = torch.kthvalue(values, lower + 1).values - if index == lower: - return lower_value - upper_value = torch.kthvalue(values, lower + 2).values - weight = index - lower - return lower_value.lerp(upper_value, weight) diff --git a/src/torchio/transforms/augmentation/__init__.py b/src/torchio/transforms/augmentation/__init__.py new file mode 100644 index 000000000..68bf56270 --- /dev/null +++ b/src/torchio/transforms/augmentation/__init__.py @@ -0,0 +1,5 @@ +from .random_transform import RandomTransform + +__all__ = [ + 'RandomTransform', +] diff --git a/src/torchio/transforms/augmentation/composition.py b/src/torchio/transforms/augmentation/composition.py new file mode 100644 index 000000000..c1d865685 --- /dev/null +++ b/src/torchio/transforms/augmentation/composition.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import warnings +from collections.abc import Sequence +from typing import TypeAlias + +import numpy as np +import torch + +from ...data.subject import Subject +from ..transform import Transform +from . import RandomTransform + +TypeTransformsDict: TypeAlias = dict[Transform, float] | Sequence[Transform] +HydraConfig: TypeAlias = dict[str, object] +HydraConfigDict: TypeAlias = dict[str, HydraConfig] + + +class Compose(Transform): + """Compose several transforms together. + + Args: + transforms: Sequence of instances of + [`Transform`][torchio.transforms.Transform]. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__(self, transforms: Sequence[Transform], **kwargs): + super().__init__(parse_input=False, **kwargs) + for transform in transforms: + if not callable(transform): + message = ( + 'One or more of the objects passed to the Compose' + f' transform are not callable: "{transform}"' + ) + raise TypeError(message) + self.transforms = list(transforms) + + def __len__(self): + return len(self.transforms) + + def __getitem__(self, index) -> Transform: + return self.transforms[index] + + def __repr__(self) -> str: + return f'{self.name}({self.transforms})' + + def _get_base_args(self) -> dict[str, object]: + init_args = super()._get_base_args() + if 'parse_input' in init_args: + init_args.pop('parse_input') + return init_args + + def apply_transform(self, subject: Subject) -> Subject: + for transform in self.transforms: + subject = transform(subject) + return subject + + def is_invertible(self) -> bool: + return all(t.is_invertible() for t in self.transforms) + + def inverse(self, warn: bool = True) -> Compose: + """Return a composed transform with inverted order and transforms. + + Args: + warn: Issue a warning if some transforms are not invertible. + """ + transforms = [] + for transform in self.transforms: + if transform.is_invertible(): + transforms.append(transform.inverse()) + elif warn: + message = f'Skipping {transform.name} as it is not invertible' + warnings.warn(message, RuntimeWarning, stacklevel=2) + transforms.reverse() + result = Compose(transforms, **self._get_base_args()) + if not transforms and warn: + warnings.warn( + 'No invertible transforms found', + RuntimeWarning, + stacklevel=2, + ) + return result + + def to_hydra_config(self) -> HydraConfig: + """Return a dictionary representation of the transform for Hydra instantiation.""" + transform_dict: HydraConfig = {'_target_': self._get_name_with_module()} + transform_dict.update(self._get_reproducing_arguments()) + transforms_config: list[HydraConfig] = [] + for transform in self.transforms: + transforms_config.append(transform.to_hydra_config()) + transform_dict['transforms'] = transforms_config + return self._tuples_to_lists(transform_dict) + + +class OneOf(RandomTransform): + """Apply only one of the given transforms. + + Args: + transforms: Dictionary with instances of + [`Transform`][torchio.transforms.Transform] as keys and + probabilities as values. Probabilities are normalized so they sum + to one. If a sequence is given, the same probability will be + assigned to each transform. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> colin = tio.datasets.Colin27() + >>> transforms_dict = { + ... tio.RandomAffine(): 0.75, + ... tio.RandomElasticDeformation(): 0.25, + ... } # Using 3 and 1 as probabilities would have the same effect + >>> transform = tio.OneOf(transforms_dict) + >>> transformed = transform(colin) + """ + + def __init__(self, transforms: TypeTransformsDict, **kwargs): + super().__init__(parse_input=False, **kwargs) + self.transforms_dict = self._get_transforms_dict(transforms) + + def _get_base_args(self) -> dict: + init_args = super()._get_base_args() + if 'parse_input' in init_args: + init_args.pop('parse_input') + return init_args + + def apply_transform(self, subject: Subject) -> Subject: + weights = torch.Tensor(list(self.transforms_dict.values())) + index = torch.multinomial(weights, 1) + transforms = list(self.transforms_dict.keys()) + transform = transforms[index] + transformed = transform(subject) + return transformed + + def _get_transforms_dict( + self, + transforms: TypeTransformsDict, + ) -> dict[Transform, float]: + if isinstance(transforms, dict): + transforms_dict = dict(transforms) + self._normalize_probabilities(transforms_dict) + else: + try: + p = 1 / len(transforms) + except TypeError as e: + message = ( + 'Transforms argument must be a dictionary or a sequence,' + f' not {type(transforms)}' + ) + raise ValueError(message) from e + transforms_dict = {transform: p for transform in transforms} + for transform in transforms_dict: + if not isinstance(transform, Transform): + message = ( + 'All keys in transform_dict must be instances of' + f'torchio.Transform, not "{type(transform)}"' + ) + raise ValueError(message) + return transforms_dict + + @staticmethod + def _normalize_probabilities( + transforms_dict: dict[Transform, float], + ) -> None: + probabilities = np.array(list(transforms_dict.values()), dtype=float) + if np.any(probabilities < 0): + message = ( + f'Probabilities must be greater or equal to zero, not "{probabilities}"' + ) + raise ValueError(message) + if np.all(probabilities == 0): + message = ( + 'At least one probability must be greater than zero,' + f' but they are "{probabilities}"' + ) + raise ValueError(message) + for transform, probability in transforms_dict.items(): + transforms_dict[transform] = probability / probabilities.sum() diff --git a/src/torchio/transforms/augmentation/intensity/__init__.py b/src/torchio/transforms/augmentation/intensity/__init__.py new file mode 100644 index 000000000..343a83591 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/__init__.py @@ -0,0 +1,39 @@ +from .random_bias_field import BiasField +from .random_bias_field import RandomBiasField +from .random_blur import Blur +from .random_blur import RandomBlur +from .random_gamma import Gamma +from .random_gamma import RandomGamma +from .random_ghosting import Ghosting +from .random_ghosting import RandomGhosting +from .random_labels_to_image import LabelsToImage +from .random_labels_to_image import RandomLabelsToImage +from .random_motion import Motion +from .random_motion import RandomMotion +from .random_noise import Noise +from .random_noise import RandomNoise +from .random_spike import RandomSpike +from .random_spike import Spike +from .random_swap import RandomSwap +from .random_swap import Swap + +__all__ = [ + 'RandomSwap', + 'Swap', + 'RandomBlur', + 'Blur', + 'RandomNoise', + 'Noise', + 'RandomSpike', + 'Spike', + 'RandomGamma', + 'Gamma', + 'RandomMotion', + 'Motion', + 'RandomGhosting', + 'Ghosting', + 'RandomBiasField', + 'BiasField', + 'RandomLabelsToImage', + 'LabelsToImage', +] diff --git a/src/torchio/transforms/augmentation/intensity/random_bias_field.py b/src/torchio/transforms/augmentation/intensity/random_bias_field.py new file mode 100644 index 000000000..3340f9b59 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_bias_field.py @@ -0,0 +1,172 @@ +from collections.abc import Sequence + +import numpy as np +import torch + +from ....data.subject import Subject +from ....types import TypeData +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + + +class RandomBiasField(RandomTransform, IntensityTransform): + r"""Add random MRI bias field artifact. + + MRI magnetic field inhomogeneity creates intensity + variations of very low frequency across the whole image. + + The bias field is modeled as a linear combination of + polynomial basis functions, as in K. Van Leemput et al., 1999, + *Automated model-based tissue classification of MR images of the brain*. + + It was implemented in NiftyNet by Carole Sudre and used in + [Sudre et al., 2017, Longitudinal segmentation of age-related + white matter hyperintensities + ](https://www.sciencedirect.com/science/article/pii/S1361841517300257?via%3Dihub). + + Args: + coefficients: Maximum magnitude $n$ of polynomial coefficients. + If a tuple $(a, b)$ is specified, then + $n \sim \mathcal{U}(a, b)$. + order: Order of the basis polynomial functions. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + coefficients: float | tuple[float, float] = 0.5, + order: int = 3, + **kwargs, + ): + super().__init__(**kwargs) + self.coefficients_range = self._parse_range( + coefficients, + 'coefficients_range', + ) + self.order = _parse_order(order) + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + coefficients_by_name: dict[str, list[float]] = {} + orders_by_name: dict[str, int] = {} + for image_name in images_dict: + coefficients = self.get_params(self.order, self.coefficients_range) + coefficients_by_name[image_name] = coefficients + orders_by_name[image_name] = self.order + transform = BiasField( + coefficients=coefficients_by_name, + order=orders_by_name, + **self._get_base_args(), + ) + transformed = transform(subject) + return transformed + + def get_params( + self, + order: int, + coefficients_range: tuple[float, float], + ) -> list[float]: + # Sampling of the appropriate number of coefficients for the creation + # of the bias field map + random_coefficients = [] + for x_order in range(0, order + 1): + for y_order in range(0, order + 1 - x_order): + for _ in range(0, order + 1 - (x_order + y_order)): + sample = self.sample_uniform(*coefficients_range) + random_coefficients.append(sample) + return random_coefficients + + +class BiasField(IntensityTransform): + r"""Add MRI bias field artifact. + + Args: + coefficients: Magnitudes of the polinomial coefficients. + order: Order of the basis polynomial functions. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + coefficients: list[float] | dict[str, list[float]], + order: int | dict[str, int], + **kwargs, + ): + super().__init__(**kwargs) + self.coefficients = coefficients + self.order = order + self.invert_transform = False + self.args_names = ['coefficients', 'order'] + + def arguments_are_dict(self): + coefficients_dict = isinstance(self.coefficients, dict) + order_dict = isinstance(self.order, dict) + if coefficients_dict != order_dict: + message = 'If one of the arguments is a dict, all must be' + raise ValueError(message) + return coefficients_dict and order_dict + + def apply_transform(self, subject: Subject) -> Subject: + for name, image in self.get_images_dict(subject).items(): + coefficients = self.get_parameter(self.coefficients, name) + order = self.get_parameter(self.order, name) + bias_field = self.generate_bias_field( + image.data, + order, + coefficients, + ) + if self.invert_transform: + np.divide(1, bias_field, out=bias_field) + image.set_data(image.data * torch.as_tensor(bias_field)) + return subject + + @staticmethod + def generate_bias_field( + data: TypeData, + order: int, + coefficients: Sequence[float], + ) -> np.ndarray: + # Create the bias field map using a linear combination of polynomial + # functions and the coefficients previously sampled + shape = np.array(data.shape[1:]) # first axis is channels + half_shape = shape / 2 + + ranges = [np.arange(-n, n) + 0.5 for n in half_shape] + + bias_field = np.zeros(shape) + meshes = np.asarray(np.meshgrid(*ranges)) + + for mesh in meshes: + mesh_max = mesh.max() + if mesh_max > 0: + mesh /= mesh_max + x_mesh, y_mesh, z_mesh = meshes + + i = 0 + for x_order in range(order + 1): + for y_order in range(order + 1 - x_order): + for z_order in range(order + 1 - (x_order + y_order)): + coefficient = coefficients[i] + new_map = ( + coefficient + * x_mesh**x_order + * y_mesh**y_order + * z_mesh**z_order + ) + bias_field += np.transpose(new_map, (1, 0, 2)) # why? + i += 1 + bias_field = np.exp(bias_field).astype(np.float32) + return bias_field + + +def _parse_order(order): + if not isinstance(order, int): + raise TypeError(f'Order must be an int, not {type(order)}') + if order < 0: + raise ValueError(f'Order must be a positive int, not {order}') + return order diff --git a/src/torchio/transforms/augmentation/intensity/random_blur.py b/src/torchio/transforms/augmentation/intensity/random_blur.py new file mode 100644 index 000000000..d2353e889 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_blur.py @@ -0,0 +1,103 @@ +import numpy as np +import scipy.ndimage as ndi +import torch + +from ....data.subject import Subject +from ....types import TypeData +from ....types import TypeSextetFloat +from ....types import TypeTripletFloat +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + + +class RandomBlur(RandomTransform, IntensityTransform): + r"""Blur an image using a random-sized Gaussian filter. + + Args: + std: Tuple $(a_1, b_1, a_2, b_2, a_3, b_3)$ representing the + ranges (in mm) of the standard deviations + $(\sigma_1, \sigma_2, \sigma_3)$ of the Gaussian kernels used + to blur the image along each axis, where + $\sigma_i \sim \mathcal{U}(a_i, b_i)$. + If two values $(a, b)$ are provided, + then $\sigma_i \sim \mathcal{U}(a, b)$. + If only one value $x$ is provided, + then $\sigma_i \sim \mathcal{U}(0, x)$. + If three values $(x_1, x_2, x_3)$ are provided, + then $\sigma_i \sim \mathcal{U}(0, x_i)$. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__(self, std: float | tuple[float, float] = (0, 2), **kwargs): + super().__init__(**kwargs) + self.std_ranges = self.parse_params(std, None, 'std', min_constraint=0) + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + std_by_name: dict[str, TypeTripletFloat] = {} + for name in images_dict: + std_by_name[name] = self.get_params(self.std_ranges) + transform = Blur(std=std_by_name, **self._get_base_args()) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + def get_params(self, std_ranges: TypeSextetFloat) -> TypeTripletFloat: + sx, sy, sz = self.sample_uniform_sextet(std_ranges) + return sx, sy, sz + + +class Blur(IntensityTransform): + r"""Blur an image using a Gaussian filter. + + Args: + std: Tuple $(\sigma_1, \sigma_2, \sigma_3)$ representing the + the standard deviations (in mm) of the Gaussian kernels used to + blur the image along each axis. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + std: TypeTripletFloat | dict[str, TypeTripletFloat], + **kwargs, + ): + super().__init__(**kwargs) + self.std = std + self.args_names = ['std'] + + def apply_transform(self, subject: Subject) -> Subject: + for name, image in self.get_images_dict(subject).items(): + stds = self.get_parameter(self.std, name) + repets = image.num_channels, 1 + stds_channels: np.ndarray = np.tile(stds, repets) + transformed_tensors = [] + for std, channel in zip(stds_channels, image.data, strict=True): + transformed_tensor = blur( + channel, + image.spacing, + std, + ) + transformed_tensors.append(transformed_tensor) + image.set_data(torch.stack(transformed_tensors)) + return subject + + +def blur( + data: TypeData, + spacing: TypeTripletFloat, + std_physical: TypeTripletFloat, +) -> torch.Tensor: + assert data.ndim == 3 + # For example, if the standard deviation of the kernel is 2 mm and the + # image spacing is 0.5 mm/voxel, the kernel should be + # (2 mm / 0.5 mm/voxel) = 4 voxels wide + std_voxel = np.array(std_physical) / np.array(spacing) + blurred = ndi.gaussian_filter(data, std_voxel) + tensor = torch.as_tensor(blurred) + return tensor diff --git a/src/torchio/transforms/augmentation/intensity/random_gamma.py b/src/torchio/transforms/augmentation/intensity/random_gamma.py new file mode 100644 index 000000000..ba8488db8 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_gamma.py @@ -0,0 +1,140 @@ +from collections.abc import Sequence + +import numpy as np +import torch + +from ....data.subject import Subject +from ....types import TypeRangeFloat +from ....utils import to_tuple +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + + +class RandomGamma(RandomTransform, IntensityTransform): + r"""Randomly change contrast of an image by raising its values to the power + $\gamma$. + + Args: + log_gamma: Tuple $(a, b)$ to compute the exponent + $\gamma = e ^ \beta$, + where $\beta \sim \mathcal{U}(a, b)$. + If a single value $d$ is provided, then + $\beta \sim \mathcal{U}(-d, d)$. + Negative and positive values for this argument perform gamma + compression and expansion, respectively. + See the [Gamma correction](https://en.wikipedia.org/wiki/Gamma_correction) Wikipedia entry for more information. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + + Note: + Fractional exponentiation of negative values is generally not + well-defined for non-complex numbers. + If negative values are found in the input image $I$, + the applied transform is $\text{sign}(I) |I|^\gamma$, + instead of the usual $I^\gamma$. The + [`RescaleIntensity`][torchio.transforms.RescaleIntensity] + transform may be used to ensure that all values are positive. This is + generally not problematic, but it is recommended to visualize results + on images with negative values. More information can be found on + [this StackExchange question](https://math.stackexchange.com/questions/317528/how-do-you-compute-negative-numbers-to-fractional-powers). + + + + Examples: + >>> import torchio as tio + >>> subject = tio.datasets.FPG() + >>> transform = tio.RandomGamma(log_gamma=(-0.3, 0.3)) # gamma between 0.74 and 1.34 + >>> transformed = transform(subject) + + """ + + def __init__(self, log_gamma: TypeRangeFloat = (-0.3, 0.3), **kwargs): + super().__init__(**kwargs) + self.log_gamma_range = self._parse_range(log_gamma, 'log_gamma') + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + gamma_by_name: dict[str, float | Sequence[float]] = {} + for name, image in images_dict.items(): + gammas = [self.get_params(self.log_gamma_range) for _ in image.data] + gamma_by_name[name] = gammas + transform = Gamma(gamma=gamma_by_name, **self._get_base_args()) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + def get_params(self, log_gamma_range: tuple[float, float]) -> float: + gamma = np.exp(self.sample_uniform(*log_gamma_range)) + return gamma + + +class Gamma(IntensityTransform): + r"""Change contrast of an image by raising its values to the power + $\gamma$. + + Args: + gamma: Exponent to which values in the image will be raised. + Negative and positive values for this argument perform gamma + compression and expansion, respectively. + See the [Gamma correction](https://en.wikipedia.org/wiki/Gamma_correction) Wikipedia entry for more information. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + + Note: + Fractional exponentiation of negative values is generally not + well-defined for non-complex numbers. + If negative values are found in the input image $I$, + the applied transform is $\text{sign}(I) |I|^\gamma$, + instead of the usual $I^\gamma$. The + [`RescaleIntensity`][torchio.transforms.preprocessing.intensity.rescale.RescaleIntensity] + transform may be used to ensure that all values are positive. This is + generally not problematic, but it is recommended to visualize results + on image with negative values. More information can be found on + [this StackExchange question](https://math.stackexchange.com/questions/317528/how-do-you-compute-negative-numbers-to-fractional-powers). + + + Examples: + >>> import torchio as tio + >>> subject = tio.datasets.FPG() + >>> transform = tio.Gamma(0.8) + >>> transformed = transform(subject) + """ + + def __init__( + self, + gamma: float | Sequence[float] | dict[str, float | Sequence[float]], + **kwargs, + ): + super().__init__(**kwargs) + self.gamma = gamma + self.args_names = ['gamma'] + self.invert_transform = False + + def apply_transform(self, subject: Subject) -> Subject: + for name, image in self.get_images_dict(subject).items(): + gamma = self.get_parameter(self.gamma, name) + gammas = to_tuple(gamma, length=len(image.data)) + transformed_tensors = [] + image.set_data(image.data.float()) + for gamma, tensor in zip(gammas, image.data, strict=True): + if self.invert_transform: + correction = power(tensor, 1 - gamma) + transformed_tensor = tensor * correction + else: + transformed_tensor = power(tensor, gamma) + transformed_tensors.append(transformed_tensor) + image.set_data(torch.stack(transformed_tensors)) + return subject + + +def power(tensor, gamma): + if tensor.min() < 0: + output = tensor.sign() * tensor.abs() ** gamma + else: + output = tensor**gamma + return output diff --git a/src/torchio/transforms/augmentation/intensity/random_ghosting.py b/src/torchio/transforms/augmentation/intensity/random_ghosting.py new file mode 100644 index 000000000..f239a38d5 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_ghosting.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +from collections.abc import Iterable +from collections.abc import Sequence + +import numpy as np +import torch + +from ....data.subject import Subject +from ...fourier import FourierTransform +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + + +class RandomGhosting(RandomTransform, IntensityTransform): + r"""Add random MRI ghosting artifact. + + Discrete "ghost" artifacts may occur along the phase-encode direction + whenever the position or signal intensity of imaged structures within the + field-of-view vary or move in a regular (periodic) fashion. Pulsatile flow + of blood or CSF, cardiac motion, and respiratory motion are the most + important patient-related causes of ghost artifacts in clinical MR imaging + (from [mriquestions.com](http://mriquestions.com/why-discrete-ghosts.html)). + + + Args: + num_ghosts: Number of 'ghosts' $n$ in the image. + If `num_ghosts` is a tuple $(a, b)$, then + $n \sim \mathcal{U}(a, b) \cap \mathbb{N}$. + If only one value $d$ is provided, + $n \sim \mathcal{U}(0, d) \cap \mathbb{N}$. + axes: Axis along which the ghosts will be created. If + `axes` is a tuple, the axis will be randomly chosen + from the passed values. Anatomical labels may also be used (see + [`RandomFlip`](../RandomFlip/#torchio.transforms.RandomFlip)). + intensity: Positive number representing the artifact strength + $s$ with respect to the maximum of the $k$-space. + If `0`, the ghosts will not be visible. If a tuple + $(a, b)$ is provided then $s \sim \mathcal{U}(a, b)$. + If only one value $d$ is provided, + $s \sim \mathcal{U}(0, d)$. + restore: Number between `0` and `1` indicating how much of the + $k$-space center should be restored after removing the planes + that generate the artifact. If `None`, only the central slice + will be restored. If a tuple $(a, b)$ is provided then + $r \sim \mathcal{U}(a, b)$. If only one value $d$ is + provided, $r \sim \mathcal{U}(0, d)$. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Note: + The execution time of this transform does not depend on the + number of ghosts. + """ + + def __init__( + self, + num_ghosts: int | tuple[int, int] = (4, 10), + axes: int | str | Sequence[int | str] = (0, 1, 2), + intensity: float | tuple[float, float] = (0.5, 1), + restore: float | None = None, + **kwargs, + ): + super().__init__(**kwargs) + if axes is None: + raise ValueError('Axes cannot be None') + if isinstance(axes, (int, str)): + axes = (axes,) + else: + axes = tuple(axes) + assert isinstance(axes, Iterable) + for axis in axes: + if not isinstance(axis, str) and axis not in (0, 1, 2): + raise ValueError(f'Axes must be in (0, 1, 2), not "{axes}"') + self.axes = axes + self.num_ghosts_range = self._parse_range( + num_ghosts, + 'num_ghosts', + min_constraint=0, + type_constraint=int, + ) + self.intensity_range = self._parse_range( + intensity, + 'intensity_range', + min_constraint=0, + ) + if restore is None: + self.restore = None + else: + self.restore = self._parse_range( + restore, + 'restore', + min_constraint=0, + max_constraint=1, + ) + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + if any(isinstance(axis, str) for axis in self.axes): + subject.check_consistent_orientation() + + num_ghosts_by_name: dict[str, int] = {} + axis_by_name: dict[str, int] = {} + intensity_by_name: dict[str, float] = {} + restore_by_name: dict[str, float | None] = {} + for name, image in images_dict.items(): + is_2d = image.is_2d() + axes = tuple(a for a in self.axes if a != 2) if is_2d else self.axes + min_ghosts, max_ghosts = self.num_ghosts_range + params = self.get_params( + (int(min_ghosts), int(max_ghosts)), + tuple(int(axis) for axis in axes if not isinstance(axis, str)), + self.intensity_range, + self.restore, + ) + num_ghosts_param, axis_param, intensity_param, restore_param = params + num_ghosts_by_name[name] = num_ghosts_param + axis_by_name[name] = axis_param + intensity_by_name[name] = intensity_param + restore_by_name[name] = restore_param + transform = Ghosting( + num_ghosts=num_ghosts_by_name, + axis=axis_by_name, + intensity=intensity_by_name, + restore=restore_by_name, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + def get_params( + self, + num_ghosts_range: tuple[int, int], + axes: tuple[int, ...], + intensity_range: tuple[float, float], + restore_range: tuple[float, float] | None, + ) -> tuple[int, int, float, float | None]: + ng_min, ng_max = num_ghosts_range + num_ghosts = int(torch.randint(ng_min, ng_max + 1, (1,)).item()) + axis = axes[torch.randint(0, len(axes), (1,))] + intensity = self.sample_uniform(*intensity_range) + if restore_range is None: + restore = None + else: + restore = self.sample_uniform(*restore_range) + return num_ghosts, axis, intensity, restore + + +class Ghosting(IntensityTransform, FourierTransform): + r"""Add MRI ghosting artifact. + + Discrete "ghost" artifacts may occur along the phase-encode direction + whenever the position or signal intensity of imaged structures within the + field-of-view vary or move in a regular (periodic) fashion. Pulsatile flow + of blood or CSF, cardiac motion, and respiratory motion are the most + important patient-related causes of ghost artifacts in clinical MR imaging + (from [mriquestions.com](http://mriquestions.com/why-discrete-ghosts.html)). + + + Args: + num_ghosts: Number of 'ghosts' $n$ in the image. + axes: Axis along which the ghosts will be created. + intensity: Positive number representing the artifact strength + $s$ with respect to the maximum of the $k$-space. + If `0`, the ghosts will not be visible. + restore: Number between `0` and `1` indicating how much of the + $k$-space center should be restored after removing the planes + that generate the artifact. If `None`, only the central slice + will be restored. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Note: + The execution time of this transform does not depend on the + number of ghosts. + """ + + def __init__( + self, + num_ghosts: int | dict[str, int], + axis: int | dict[str, int], + intensity: float | dict[str, float], + restore: float | None | dict[str, float | None], + **kwargs, + ): + super().__init__(**kwargs) + self.axis = axis + self.num_ghosts = num_ghosts + self.intensity = intensity + self.restore = restore + self.args_names = ['num_ghosts', 'axis', 'intensity', 'restore'] + + def apply_transform(self, subject: Subject) -> Subject: + for name, image in self.get_images_dict(subject).items(): + axis = self.get_parameter(self.axis, name) + num_ghosts = self.get_parameter(self.num_ghosts, name) + intensity = self.get_parameter(self.intensity, name) + restore = self.get_parameter(self.restore, name) + transformed_tensors = [] + for tensor in image.data: + transformed_tensor = self.add_artifact( + tensor, + num_ghosts, + axis, + float(intensity), + restore, + ) + transformed_tensors.append(transformed_tensor) + image.set_data(torch.stack(transformed_tensors)) + return subject + + def add_artifact( + self, + tensor: torch.Tensor, + num_ghosts: int, + axis: int, + intensity: float, + restore_center: float | None, + ): + if not num_ghosts or not intensity: + return tensor + + spectrum = self.fourier_transform(tensor) + + # Variable "planes" is the part of the spectrum that will be modified + # Variable "restore" is the part of the spectrum that will be restored + planes = self._get_planes_to_modify(spectrum, axis, num_ghosts) + tensor_restore, slices = self._get_slices_to_restore( + spectrum, axis, restore_center + ) + tensor_restore = tensor_restore.clone() + + # Multiply by 0 if intensity is 1 + planes *= 1 - intensity + + # Restore the center of k-space to avoid extreme artifacts + spectrum[slices] = tensor_restore + + tensor_ghosts = self.inv_fourier_transform(spectrum) + return tensor_ghosts.real.float() + + @staticmethod + def _get_planes_to_modify( + spectrum: torch.Tensor, + axis: int, + num_ghosts: int, + ) -> torch.Tensor: + slices = [slice(None)] * spectrum.ndim + slices[axis] = slice(None, None, num_ghosts) + slices_tuple = tuple(slices) + return spectrum[slices_tuple] + + @staticmethod + def _get_slices_to_restore( + spectrum: torch.Tensor, + axis: int, + restore_center: float | None, + ) -> tuple[torch.Tensor, tuple[slice, ...]]: + dim_shape = spectrum.shape[axis] + mid_idx = dim_shape // 2 + slices = [slice(None)] * spectrum.ndim + if restore_center is None: + slice_ = slice(mid_idx, mid_idx + 1) + else: + size_restore = int(np.round(restore_center * dim_shape)) + slice_ = slice(mid_idx - size_restore // 2, mid_idx + size_restore // 2) + slices[axis] = slice_ + slices_tuple = tuple(slices) + restore_tensor = spectrum[slices_tuple] + return restore_tensor, slices_tuple diff --git a/src/torchio/transforms/augmentation/intensity/random_labels_to_image.py b/src/torchio/transforms/augmentation/intensity/random_labels_to_image.py new file mode 100644 index 000000000..658e28458 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_labels_to_image.py @@ -0,0 +1,494 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import TypeVar + +import torch + +from ....data.image import LabelMap +from ....data.image import ScalarImage +from ....data.subject import Subject +from ....types import TypeRangeFloat +from ....utils import check_sequence +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + +GaussianParameterT = TypeVar('GaussianParameterT') + + +class RandomLabelsToImage(RandomTransform, IntensityTransform): + r"""Randomly generate an image from a segmentation. + + Based on the work by Billot et al.: [A Learning Strategy for Contrast-agnostic MRI Segmentation](http://proceedings.mlr.press/v121/billot20a.html) + and [Partial Volume Segmentation of Brain MRI Scans of any Resolution and Contrast](https://link.springer.com/chapter/10.1007/978-3-030-59728-3_18). + + + + + Args: + label_key: String designating the label map in the subject + that will be used to generate the new image. + used_labels: Sequence of integers designating the labels used + to generate the new image. If categorical encoding is used, + `label_channels` refers to the values of the + categorical encoding. If one hot encoding or partial-volume + label maps are used, `label_channels` refers to the + channels of the label maps. + Default uses all labels. Missing voxels will be filled with zero + or with voxels from an already existing volume, + see `image_key`. + image_key: String designating the key to which the new volume will be + saved. If this key corresponds to an already existing volume, + missing voxels will be filled with the corresponding values + in the original volume. + mean: Sequence of means for each label. + For each value $v$, if a tuple $(a, b)$ is + provided then $v \sim \mathcal{U}(a, b)$. + If `None`, `default_mean` range will be used for every + label. + If not `None` and `label_channels` is not `None`, + `mean` and `label_channels` must have the + same length. + std: Sequence of standard deviations for each label. + For each value $v$, if a tuple $(a, b)$ is + provided then $v \sim \mathcal{U}(a, b)$. + If `None`, `default_std` range will be used for every + label. + If not `None` and `label_channels` is not `None`, + `std` and `label_channels` must have the + same length. + default_mean: Default mean range. + default_std: Default standard deviation range. + discretize: If `True`, partial-volume label maps will be discretized. + Does not have any effects if not using partial-volume label maps. + Discretization is done taking the class of the highest value per + voxel in the different partial-volume label maps using + `torch.argmax()` on the channel dimension (i.e. 0). + ignore_background: If `True`, input voxels labeled as `0` will not + be modified. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Tip: + It is recommended to blur the new images in order to simulate + partial volume effects at the borders of the synthetic structures. See + [`RandomBlur`][torchio.transforms.augmentation.intensity.random_blur.RandomBlur]. + + Examples: + >>> import torchio as tio + >>> subject = tio.datasets.ICBM2009CNonlinearSymmetric() + >>> # Using the default parameters + >>> transform = tio.RandomLabelsToImage(label_key='tissues') + >>> # Using custom mean and std + >>> transform = tio.RandomLabelsToImage( + ... label_key='tissues', mean=[0.33, 0.66, 1.], std=[0, 0, 0] + ... ) + >>> # Discretizing the partial volume maps and blurring the result + >>> simulation_transform = tio.RandomLabelsToImage( + ... label_key='tissues', mean=[0.33, 0.66, 1.], std=[0, 0, 0], discretize=True + ... ) + >>> blurring_transform = tio.RandomBlur(std=0.3) + >>> transform = tio.Compose([simulation_transform, blurring_transform]) + >>> transformed = transform(subject) # subject has a new key 'image_from_labels' with the simulated image + >>> # Filling holes of the simulated image with the original T1 image + >>> rescale_transform = tio.RescaleIntensity( + ... out_min_max=(0, 1), percentiles=(1, 99)) # Rescale intensity before filling holes + >>> simulation_transform = tio.RandomLabelsToImage( + ... label_key='tissues', + ... image_key='t1', + ... used_labels=[0, 1] + ... ) + >>> transform = tio.Compose([rescale_transform, simulation_transform]) + >>> transformed = transform(subject) # subject's key 't1' has been replaced with the simulated image + + !!! note "See also" + [`RemapLabels`][torchio.transforms.preprocessing.label.remap_labels.RemapLabels]. + + """ + + def __init__( + self, + label_key: str | None = None, + used_labels: Sequence[int] | None = None, + image_key: str = 'image_from_labels', + mean: Sequence[TypeRangeFloat] | None = None, + std: Sequence[TypeRangeFloat] | None = None, + default_mean: TypeRangeFloat = (0.1, 0.9), + default_std: TypeRangeFloat = (0.01, 0.1), + discretize: bool = False, + ignore_background: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.label_key = _parse_label_key(label_key) + self.used_labels = _parse_used_labels(used_labels) + self.mean_ranges: list[tuple[float, float]] | None + self.std_ranges: list[tuple[float, float]] | None + self.mean_ranges, self.std_ranges = self.parse_mean_and_std(mean, std) + self.default_mean = self.parse_gaussian_parameter( + default_mean, + 'default_mean', + ) + self.default_std = self.parse_gaussian_parameter( + default_std, + 'default_std', + ) + self.image_key = image_key + self.discretize = discretize + self.ignore_background = ignore_background + + def parse_mean_and_std( + self, + mean: Sequence[TypeRangeFloat] | None, + std: Sequence[TypeRangeFloat] | None, + ) -> tuple[list[tuple[float, float]] | None, list[tuple[float, float]] | None]: + if mean is not None: + mean = self.parse_gaussian_parameters(mean, 'mean') + if std is not None: + std = self.parse_gaussian_parameters(std, 'std') + if mean is not None and std is not None: + message = ( + 'If both "mean" and "std" are defined they must have the samelength' + ) + assert len(mean) == len(std), message + return mean, std + + def parse_gaussian_parameters( + self, + params: Sequence[TypeRangeFloat], + name: str, + ) -> list[tuple[float, float]]: + check_sequence(params, name) + parsed_params: list[tuple[float, float]] = [ + self.parse_gaussian_parameter(p, f'{name}[{i}]') + for i, p in enumerate(params) + ] + if self.used_labels is not None: + message = ( + f'If both "{name}" and "used_labels" are defined, ' + 'they must have the same length' + ) + assert len(parsed_params) == len(self.used_labels), message + return parsed_params + + @staticmethod + def parse_gaussian_parameter( + nums_range: TypeRangeFloat, + name: str, + ) -> tuple[float, float]: + if isinstance(nums_range, (int, float)): + return nums_range, nums_range + + if len(nums_range) != 2: + raise ValueError( + f'If {name} is a sequence, it must be of len 2, not {nums_range}', + ) + min_value, max_value = nums_range + if min_value > max_value: + raise ValueError( + f'If {name} is a sequence, the second value must be' + f' equal or greater than the first, not {nums_range}', + ) + return min_value, max_value + + def _guess_label_key(self, subject: Subject) -> None: + if self.label_key is None: + iterable = subject.get_images_dict(intensity_only=False).items() + for name, image in iterable: + if isinstance(image, LabelMap): + self.label_key = name + break + else: + message = f'No label maps found in subject: {subject}' + raise RuntimeError(message) + + def apply_transform(self, subject: Subject) -> Subject: + self._guess_label_key(subject) + assert self.label_key is not None + + means: list[float] = [] + stds: list[float] = [] + label_map = subject.get_label_map(self.label_key).data + + # Find out if we face a partial-volume image or a label map. + # One-hot-encoded label map is considered as a partial-volume image + all_discrete = label_map.eq(label_map.float().round()).all() + same_num_dims = label_map.squeeze().dim() < label_map.dim() + is_discretized = all_discrete and same_num_dims + + if not is_discretized and self.discretize: + # Take label with highest value in voxel + max_label, label_map = label_map.max(dim=0, keepdim=True) + # Remove values where all labels are 0 (i.e. missing labels) + label_map[max_label == 0] = -1 + is_discretized = True + + if is_discretized: + labels = label_map.unique().long().tolist() + if -1 in labels: + labels.remove(-1) + else: + labels = range(label_map.shape[0]) + + # Raise error if mean and std are not defined for every label + _check_mean_and_std_length(labels, self.mean_ranges, self.std_ranges) + + for label in labels: + mean, std = self.get_params(label) + means.append(mean) + stds.append(std) + + transform = LabelsToImage( + label_key=self.label_key, + mean=means, + std=stds, + image_key=self.image_key, + used_labels=self.used_labels, + ignore_background=self.ignore_background, + discretize=self.discretize, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + def get_params(self, label: int) -> tuple[float, float]: + if self.mean_ranges is None: + mean_range = self.default_mean + else: + mean_range = self.mean_ranges[label] + if self.std_ranges is None: + std_range = self.default_std + else: + std_range = self.std_ranges[label] + mean = self.sample_uniform(*mean_range) + std = self.sample_uniform(*std_range) + return mean, std + + def _get_named_arguments(self) -> dict[str, object]: + return { + 'label_key': self.label_key, + 'used_labels': self.used_labels, + 'image_key': self.image_key, + 'mean': self.mean_ranges, + 'std': self.std_ranges, + 'default_mean': self.default_mean, + 'default_std': self.default_std, + 'discretize': self.discretize, + 'ignore_background': self.ignore_background, + } + + +class LabelsToImage(IntensityTransform): + r"""Generate an image from a segmentation. + + Args: + label_key: String designating the label map in the subject + that will be used to generate the new image. + used_labels: Sequence of integers designating the labels used + to generate the new image. If categorical encoding is used, + `label_channels` refers to the values of the + categorical encoding. If one hot encoding or partial-volume + label maps are used, `label_channels` refers to the + channels of the label maps. + Default uses all labels. Missing voxels will be filled with zero + or with voxels from an already existing volume, + see `image_key`. + image_key: String designating the key to which the new volume will be + saved. If this key corresponds to an already existing volume, + missing voxels will be filled with the corresponding values + in the original volume. + mean: Sequence of means for each label. + If not `None` and `label_channels` is not `None`, + `mean` and `label_channels` must have the + same length. + std: Sequence of standard deviations for each label. + If not `None` and `label_channels` is not `None`, + `std` and `label_channels` must have the + same length. + discretize: If `True`, partial-volume label maps will be discretized. + Does not have any effects if not using partial-volume label maps. + Discretization is done taking the class of the highest value per + voxel in the different partial-volume label maps using + `torch.argmax()` on the channel dimension (i.e. 0). + ignore_background: If `True`, input voxels labeled as `0` will not + be modified. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Note: + It is recommended to blur the new images to make the result more + realistic. See + [`RandomBlur`][torchio.transforms.augmentation.RandomBlur]. + """ + + def __init__( + self, + label_key: str, + mean: Sequence[float] | None, + std: Sequence[float] | None, + image_key: str = 'image_from_labels', + used_labels: Sequence[int] | None = None, + ignore_background: bool = False, + discretize: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + parsed_label_key = _parse_label_key(label_key) + assert parsed_label_key is not None + self.label_key: str = parsed_label_key + self.used_labels = _parse_used_labels(used_labels) + self.means: Sequence[float] | None = mean + self.stds: Sequence[float] | None = std + self.image_key = image_key + self.ignore_background = ignore_background + self.discretize = discretize + self.args_names = [ + 'label_key', + 'mean', + 'std', + 'image_key', + 'used_labels', + 'ignore_background', + 'discretize', + ] + + def apply_transform(self, subject: Subject) -> Subject: + original_image = ( + subject.get_scalar_image(self.image_key) + if self.image_key in subject + else None + ) + + label_map_image = subject.get_label_map(self.label_key) + label_map = label_map_image.data + affine = label_map_image.affine + + # Find out if we face a partial-volume image or a label map. + # One-hot-encoded label map is considered as a partial-volume image + all_discrete = label_map.eq(label_map.float().round()).all() + same_num_dims = label_map.squeeze().dim() < label_map.dim() + is_discretized = all_discrete and same_num_dims + + if not is_discretized and self.discretize: + # Take label with highest value in voxel + max_label, label_map = label_map.max(dim=0, keepdim=True) + # Remove values where all labels are 0 (i.e. missing labels) + label_map[max_label == 0] = -1 + is_discretized = True + + tissues = torch.zeros(1, *label_map_image.spatial_shape).float() + if is_discretized: + labels_in_image = label_map.unique().long().tolist() + if -1 in labels_in_image: + labels_in_image.remove(-1) + else: + labels_in_image = range(label_map.shape[0]) + + # Raise error if mean and std are not defined for every label + _check_mean_and_std_length( + labels_in_image, + self.means, + self.stds, + ) + + for i, label in enumerate(labels_in_image): + if label == 0 and self.ignore_background: + continue + if self.used_labels is None or label in self.used_labels: + assert self.means is not None + assert self.stds is not None + mean = self.means[i] + std = self.stds[i] + if is_discretized: + mask = label_map == label + else: + mask = label_map[label] + tissues += self.generate_tissue(mask, mean, std) + + else: + # Modify label map to easily compute background mask + if is_discretized: + label_map[label_map == label] = -1 + else: + label_map[label] = 0 + + final_image = ScalarImage(affine=affine, tensor=tissues) + + if original_image is not None: + if is_discretized: + bg_mask = label_map == -1 + else: + bg_mask = label_map.sum(dim=0, keepdim=True) < 0.5 + final_image.data[bg_mask] = original_image.data[bg_mask].float() + + subject.add_image(final_image, self.image_key) + return subject + + def _get_named_arguments(self) -> dict[str, object]: + return { + 'label_key': self.label_key, + 'mean': self.means, + 'std': self.stds, + 'image_key': self.image_key, + 'used_labels': self.used_labels, + 'ignore_background': self.ignore_background, + 'discretize': self.discretize, + } + + @staticmethod + def generate_tissue( + data: torch.Tensor, + mean: float, + std: float, + ) -> torch.Tensor: + # Create the simulated tissue using a gaussian random variable + gaussian = torch.randn(data.shape) * std + mean + return gaussian * data + + +def _parse_label_key(label_key: str | None) -> str | None: + if label_key is not None and not isinstance(label_key, str): + message = f'"label_key" must be a string or None, not {type(label_key)}' + raise TypeError(message) + return label_key + + +def _parse_used_labels( + used_labels: Sequence[int] | None, +) -> Sequence[int] | None: + if used_labels is None: + return None + check_sequence(used_labels, 'used_labels') + for e in used_labels: + if not isinstance(e, int): + message = ( + 'Items in "used_labels" must be integers,' + f' but some are not: {used_labels}' + ) + raise ValueError(message) + return used_labels + + +def _check_mean_and_std_length( + labels: Sequence[int], + means: Sequence[GaussianParameterT] | None, + stds: Sequence[GaussianParameterT] | None, +) -> None: + num_labels = len(labels) + if means is not None: + num_means = len(means) + message = ( + '"mean" must define a value for each label but length of "mean"' + f' is {num_means} while {num_labels} labels were found' + ) + if num_means != num_labels: + raise RuntimeError(message) + if stds is not None: + num_stds = len(stds) + message = ( + '"std" must define a value for each label but length of "std"' + f' is {num_stds} while {num_labels} labels were found' + ) + if num_stds != num_labels: + raise RuntimeError(message) diff --git a/src/torchio/transforms/augmentation/intensity/random_motion.py b/src/torchio/transforms/augmentation/intensity/random_motion.py new file mode 100644 index 000000000..90eda0f4a --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_motion.py @@ -0,0 +1,298 @@ +from collections.abc import Sequence + +import numpy as np +import SimpleITK as sitk +import torch + +from ....data.io import nib_to_sitk +from ....data.subject import Subject +from ...fourier import FourierTransform +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + + +class RandomMotion(RandomTransform, IntensityTransform, FourierTransform): + r"""Add random MRI motion artifact. + + Magnetic resonance images suffer from motion artifacts when the subject + moves during image acquisition. This transform follows + [Shaw et al., 2019 ](http://proceedings.mlr.press/v102/shaw19a.html) to + simulate motion artifacts for data augmentation. + + Args: + degrees: Tuple $(a, b)$ defining the rotation range in degrees of + the simulated movements. The rotation angles around each axis are + $(\theta_1, \theta_2, \theta_3)$, + where $\theta_i \sim \mathcal{U}(a, b)$. + If only one value $d$ is provided, + $\theta_i \sim \mathcal{U}(-d, d)$. + Larger values generate more distorted images. + translation: Tuple $(a, b)$ defining the translation in mm of + the simulated movements. The translations along each axis are + $(t_1, t_2, t_3)$, + where $t_i \sim \mathcal{U}(a, b)$. + If only one value $t$ is provided, + $t_i \sim \mathcal{U}(-t, t)$. + Larger values generate more distorted images. + num_transforms: Number of simulated movements. + Larger values generate more distorted images. + image_interpolation: See Interpolation. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Warning: + Large numbers of movements lead to longer execution times for + 3D images. + """ + + def __init__( + self, + degrees: float | tuple[float, float] = 10, + translation: float | tuple[float, float] = 10, # in mm + num_transforms: int = 2, + image_interpolation: str = 'linear', + **kwargs, + ): + super().__init__(**kwargs) + self.degrees_range = self.parse_degrees(degrees) + self.translation_range = self.parse_translation(translation) + if num_transforms < 1 or not isinstance(num_transforms, int): + message = ( + 'Number of transforms must be a strictly positive natural' + f'number, not {num_transforms}' + ) + raise ValueError(message) + self.num_transforms = num_transforms + self.image_interpolation = self.parse_interpolation( + image_interpolation, + ) + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + times_by_name: dict[str, np.ndarray] = {} + degrees_by_name: dict[str, np.ndarray] = {} + translation_by_name: dict[str, np.ndarray] = {} + interpolation_by_name: dict[str, str] = {} + for name, image in images_dict.items(): + params = self.get_params( + self.degrees_range, + self.translation_range, + self.num_transforms, + is_2d=image.is_2d(), + ) + times_params, degrees_params, translation_params = params + times_by_name[name] = times_params + degrees_by_name[name] = degrees_params + translation_by_name[name] = translation_params + interpolation_by_name[name] = self.image_interpolation + transform = Motion( + degrees=degrees_by_name, + translation=translation_by_name, + times=times_by_name, + image_interpolation=interpolation_by_name, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + def get_params( + self, + degrees_range: tuple[float, float], + translation_range: tuple[float, float], + num_transforms: int, + perturbation: float = 0.3, + is_2d: bool = False, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + # If perturbation is 0, time intervals between movements are constant + degrees_params = self.get_params_array( + degrees_range, + num_transforms, + ) + translation_params = self.get_params_array( + translation_range, + num_transforms, + ) + if is_2d: # imagine sagittal (1, A, S) + degrees_params[:, :-1] = 0 # rotate around Z axis only + translation_params[:, 2] = 0 # translate in XY plane only + step = 1 / (num_transforms + 1) + times = torch.arange(0, 1, step)[1:] + noise = torch.FloatTensor(num_transforms) + noise.uniform_(-step * perturbation, step * perturbation) + times += noise + times_params = times.numpy() + return times_params, degrees_params, translation_params + + @staticmethod + def get_params_array(nums_range: tuple[float, float], num_transforms: int): + tensor = torch.FloatTensor(num_transforms, 3).uniform_(*nums_range) + return tensor.numpy() + + +class Motion(IntensityTransform, FourierTransform): + r"""Add MRI motion artifact. + + Magnetic resonance images suffer from motion artifacts when the subject + moves during image acquisition. This transform follows + [Shaw et al., 2019 ](http://proceedings.mlr.press/v102/shaw19a.html) to + simulate motion artifacts for data augmentation. + + Args: + degrees: Sequence of rotations $(\theta_1, \theta_2, \theta_3)$. + translation: Sequence of translations $(t_1, t_2, t_3)$ in mm. + times: Sequence of times from 0 to 1 at which the motions happen. + image_interpolation: See Interpolation. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + degrees: np.ndarray | dict[str, np.ndarray], + translation: np.ndarray | dict[str, np.ndarray], + times: np.ndarray | dict[str, np.ndarray], + image_interpolation: str | dict[str, str], + **kwargs, + ): + super().__init__(**kwargs) + self.degrees = degrees + self.translation = translation + self.times = times + self.image_interpolation = image_interpolation + self.args_names = [ + 'degrees', + 'translation', + 'times', + 'image_interpolation', + ] + + def apply_transform(self, subject: Subject) -> Subject: + for image_name, image in self.get_images_dict(subject).items(): + degrees = self.get_parameter(self.degrees, image_name) + translation = self.get_parameter(self.translation, image_name) + times = self.get_parameter(self.times, image_name) + image_interpolation = self.get_parameter( + self.image_interpolation, image_name + ) + result_arrays = [] + for channel in image.data: + sitk_image = nib_to_sitk( + channel[np.newaxis], + image.affine, + force_3d=True, + ) + transforms = self.get_rigid_transforms( + degrees, + translation, + sitk_image, + ) + transformed_channel = self.add_artifact( + sitk_image, + transforms, + times, + image_interpolation, + ) + result_arrays.append(transformed_channel) + result = np.stack(result_arrays) + image.set_data(torch.as_tensor(result)) + return subject + + def get_rigid_transforms( + self, + degrees_params: np.ndarray, + translation_params: np.ndarray, + image: sitk.Image, + ) -> list[sitk.Euler3DTransform]: + center_ijk = np.array(image.GetSize()) / 2 + center_lps = image.TransformContinuousIndexToPhysicalPoint(center_ijk) + identity = np.eye(4) + matrices = [identity] + zipped = zip(degrees_params, translation_params, strict=True) + for degrees, translation in zipped: + radians = np.radians(degrees).tolist() + motion = sitk.Euler3DTransform() + motion.SetCenter(center_lps) + motion.SetRotation(*radians) + motion.SetTranslation(translation.tolist()) + motion_matrix = self.transform_to_matrix(motion) + matrices.append(motion_matrix) + transforms = [self.matrix_to_transform(m) for m in matrices] + return transforms + + @staticmethod + def transform_to_matrix(transform: sitk.Euler3DTransform) -> np.ndarray: + matrix = np.eye(4) + rotation = np.array(transform.GetMatrix()).reshape(3, 3) + matrix[:3, :3] = rotation + matrix[:3, 3] = transform.GetTranslation() + return matrix + + @staticmethod + def matrix_to_transform(matrix: np.ndarray) -> sitk.Euler3DTransform: + transform = sitk.Euler3DTransform() + rotation = matrix[:3, :3].flatten().tolist() + transform.SetMatrix(rotation) + transform.SetTranslation(matrix[:3, 3]) + return transform + + def resample_images( + self, + image: sitk.Image, + transforms: Sequence[sitk.Euler3DTransform], + interpolation: str, + ) -> list[sitk.Image]: + floating = reference = image + default_value = np.float64(sitk.GetArrayViewFromImage(image).min()) + transforms = transforms[1:] # first is identity + images = [image] # first is identity + for transform in transforms: + interpolator = self.get_sitk_interpolator(interpolation) + resampler = sitk.ResampleImageFilter() + resampler.SetInterpolator(interpolator) + resampler.SetReferenceImage(reference) + resampler.SetOutputPixelType(sitk.sitkFloat32) + resampler.SetDefaultPixelValue(default_value) + resampler.SetTransform(transform) + resampled = resampler.Execute(floating) + images.append(resampled) + return images + + @staticmethod + def sort_spectra(spectra: list[torch.Tensor], times: np.ndarray): + """Use original spectrum to fill the center of k-space.""" + num_spectra = len(spectra) + if np.any(times > 0.5): + index = np.where(times > 0.5)[0].min() + else: + index = num_spectra - 1 + spectra[0], spectra[index] = spectra[index], spectra[0] + + def add_artifact( + self, + image: sitk.Image, + transforms: Sequence[sitk.Euler3DTransform], + times: np.ndarray, + interpolation: str, + ): + images = self.resample_images(image, transforms, interpolation) + spectra = [] + for image in images: + array = sitk.GetArrayFromImage(image).transpose() # sitk to np + spectrum = self.fourier_transform(torch.from_numpy(array)) + spectra.append(spectrum) + self.sort_spectra(spectra, times) + result_spectrum = torch.empty_like(spectra[0]) + last_index = result_spectrum.shape[2] + indices_array = (last_index * times).astype(int) + indices = [int(value) for value in indices_array.tolist()] + indices.append(last_index) + ini = 0 + for spectrum, fin in zip(spectra, indices, strict=True): + result_spectrum[..., ini:fin] = spectrum[..., ini:fin] + ini = fin + result_image = self.inv_fourier_transform(result_spectrum).real.float() + return result_image diff --git a/src/torchio/transforms/augmentation/intensity/random_noise.py b/src/torchio/transforms/augmentation/intensity/random_noise.py new file mode 100644 index 000000000..ebd8c9f74 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_noise.py @@ -0,0 +1,124 @@ +import torch + +from ....data.subject import Subject +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + + +class RandomNoise(RandomTransform, IntensityTransform): + r"""Add Gaussian noise with random parameters. + + Add noise sampled from a normal distribution with random parameters. + + Args: + mean: Mean $\mu$ of the Gaussian distribution + from which the noise is sampled. + If two values $(a, b)$ are provided, + then $\mu \sim \mathcal{U}(a, b)$. + If only one value $d$ is provided, + $\mu \sim \mathcal{U}(-d, d)$. + std: Standard deviation $\sigma$ of the Gaussian distribution + from which the noise is sampled. + If two values $(a, b)$ are provided, + then $\sigma \sim \mathcal{U}(a, b)$. + If only one value $d$ is provided, + $\sigma \sim \mathcal{U}(0, d)$. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + mean: float | tuple[float, float] = 0, + std: float | tuple[float, float] = (0, 0.25), + **kwargs, + ): + super().__init__(**kwargs) + self.mean_range = self._parse_range(mean, 'mean') + self.std_range = self._parse_range(std, 'std', min_constraint=0) + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + means_by_name: dict[str, float] = {} + stds_by_name: dict[str, float] = {} + seeds_by_name: dict[str, int] = {} + for image_name in images_dict: + mean, std, seed = self.get_params(self.mean_range, self.std_range) + means_by_name[image_name] = mean + stds_by_name[image_name] = std + seeds_by_name[image_name] = seed + transform = Noise( + mean=means_by_name, + std=stds_by_name, + seed=seeds_by_name, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + def get_params( + self, + mean_range: tuple[float, float], + std_range: tuple[float, float], + ) -> tuple[float, float, int]: + mean = self.sample_uniform(*mean_range) + std = self.sample_uniform(*std_range) + seed = self._get_random_seed() + return mean, std, seed + + +class Noise(IntensityTransform): + r"""Add Gaussian noise. + + Add noise sampled from a normal distribution. + + Args: + mean: Mean $\mu$ of the Gaussian distribution + from which the noise is sampled. + std: Standard deviation $\sigma$ of the Gaussian distribution + from which the noise is sampled. + seed: Seed for the random number generator. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + mean: float | dict[str, float], + std: float | dict[str, float], + seed: int | dict[str, int], + **kwargs, + ): + super().__init__(**kwargs) + self.noise_mean: float | dict[str, float] = mean + self.noise_std: float | dict[str, float] = std + self.seed: int | dict[str, int] = seed + self.invert_transform = False + self.args_names = ['mean', 'std', 'seed'] + + def apply_transform(self, subject: Subject) -> Subject: + for name, image in self.get_images_dict(subject).items(): + mean = self.get_parameter(self.noise_mean, name) + std = self.get_parameter(self.noise_std, name) + seed = self.get_parameter(self.seed, name) + with self._use_seed(seed): + noise = get_noise(image.data, mean, std) + if self.invert_transform: + noise *= -1 + image.set_data(image.data + noise) + return subject + + def _get_named_arguments(self) -> dict[str, object]: + return { + 'mean': self.noise_mean, + 'std': self.noise_std, + 'seed': self.seed, + } + + +def get_noise(tensor: torch.Tensor, mean: float, std: float) -> torch.Tensor: + return torch.randn(*tensor.shape) * std + mean diff --git a/src/torchio/transforms/augmentation/intensity/random_spike.py b/src/torchio/transforms/augmentation/intensity/random_spike.py new file mode 100644 index 000000000..ce6b272d7 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_spike.py @@ -0,0 +1,167 @@ +import numpy as np +import torch + +from ....data.subject import Subject +from ...fourier import FourierTransform +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + + +class RandomSpike(RandomTransform, IntensityTransform, FourierTransform): + r"""Add random MRI spike artifacts. + + Also known as [Herringbone artifact + ](https://radiopaedia.org/articles/herringbone-artifact), + crisscross artifact or corduroy artifact, it creates stripes in different + directions in image space due to spikes in k-space. + + Args: + num_spikes: Number of spikes $n$ present in k-space. + If a tuple $(a, b)$ is provided, then + $n \sim \mathcal{U}(a, b) \cap \mathbb{N}$. + If only one value $d$ is provided, + $n \sim \mathcal{U}(0, d) \cap \mathbb{N}$. + Larger values generate more distorted images. + intensity: Ratio $r$ between the spike intensity and the maximum + of the spectrum. + If a tuple $(a, b)$ is provided, then + $r \sim \mathcal{U}(a, b)$. + If only one value $d$ is provided, + $r \sim \mathcal{U}(-d, d)$. + Larger values generate more distorted images. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Note: + The execution time of this transform does not depend on the + number of spikes. + """ + + def __init__( + self, + num_spikes: int | tuple[int, int] = (1, 1), + intensity: float | tuple[float, float] = (1, 3), + **kwargs, + ): + super().__init__(**kwargs) + self.intensity_range = self._parse_range( + intensity, + 'intensity_range', + ) + self.num_spikes_range = self._parse_range( + num_spikes, + 'num_spikes', + min_constraint=0, + type_constraint=int, + ) + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + spikes_positions_by_name: dict[str, np.ndarray] = {} + intensity_by_name: dict[str, float] = {} + for image_name in images_dict: + spikes_positions_param, intensity_param = self.get_params( + self.num_spikes_range, + self.intensity_range, + ) + spikes_positions_by_name[image_name] = spikes_positions_param + intensity_by_name[image_name] = intensity_param + transform = Spike( + spikes_positions=spikes_positions_by_name, + intensity=intensity_by_name, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + def get_params( + self, + num_spikes_range: tuple[int, int], + intensity_range: tuple[float, float], + ) -> tuple[np.ndarray, float]: + ns_min, ns_max = num_spikes_range + num_spikes_param = int(torch.randint(ns_min, ns_max + 1, (1,)).item()) + intensity_param = self.sample_uniform(*intensity_range) + spikes_positions = torch.rand(num_spikes_param, 3).numpy() + return spikes_positions, intensity_param + + +class Spike(IntensityTransform, FourierTransform): + r"""Add MRI spike artifacts. + + Also known as [Herringbone artifact + ](https://radiopaedia.org/articles/herringbone-artifact), + crisscross artifact or corduroy artifact, it creates stripes in different + directions in image space due to spikes in k-space. + + Args: + spikes_positions: + intensity: Ratio $r$ between the spike intensity and the maximum + of the spectrum. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Note: + The execution time of this transform does not depend on the + number of spikes. + """ + + def __init__( + self, + spikes_positions: np.ndarray | dict[str, np.ndarray], + intensity: float | dict[str, float], + **kwargs, + ): + super().__init__(**kwargs) + self.spikes_positions = spikes_positions + self.intensity = intensity + self.args_names = ['spikes_positions', 'intensity'] + self.invert_transform = False + + def apply_transform(self, subject: Subject) -> Subject: + for image_name, image in self.get_images_dict(subject).items(): + spikes_positions = self.get_parameter(self.spikes_positions, image_name) + intensity = self.get_parameter(self.intensity, image_name) + transformed_tensors = [] + for channel in image.data: + transformed_tensor = self.add_artifact( + channel, + np.asarray(spikes_positions), + float(intensity), + ) + transformed_tensors.append(transformed_tensor) + image.set_data(torch.stack(transformed_tensors)) + return subject + + def add_artifact( + self, + tensor: torch.Tensor, + spikes_positions: np.ndarray, + intensity_factor: float, + ): + if intensity_factor == 0 or len(spikes_positions) == 0: + return tensor + spectrum = self.fourier_transform(tensor) + shape = np.array(spectrum.shape) + mid_shape = shape // 2 + indices = np.floor(spikes_positions * shape).astype(int) + for index in indices: + diff = index - mid_shape + i, j, k = mid_shape + diff + artifact = spectrum.cpu().abs().max() * intensity_factor + if self.invert_transform: + spectrum[i, j, k] -= artifact + else: + spectrum[i, j, k] += artifact + # If we wanted to add a pure cosine, we should add spikes to both + # sides of k-space. However, having only one is a better + # representation og the actual cause of the artifact in real + # scans. Therefore the next two lines have been removed. + # #i, j, k = mid_shape - diff + # #spectrum[i, j, k] = spectrum.max() * intensity_factor + result = self.inv_fourier_transform(spectrum).real.float() + return result diff --git a/src/torchio/transforms/augmentation/intensity/random_swap.py b/src/torchio/transforms/augmentation/intensity/random_swap.py new file mode 100644 index 000000000..627fa30b7 --- /dev/null +++ b/src/torchio/transforms/augmentation/intensity/random_swap.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np +import torch + +from ....data.subject import Subject +from ....types import TypeTripletInt +from ....types import TypeTuple +from ....utils import to_tuple +from ...intensity_transform import IntensityTransform +from .. import RandomTransform + +TypeLocations = Sequence[tuple[TypeTripletInt, TypeTripletInt]] + + +class RandomSwap(RandomTransform, IntensityTransform): + r"""Randomly swap patches within an image. + + This is typically used in [context restoration for self-supervised learning + ](https://www.sciencedirect.com/science/article/pii/S1361841518304699). + + Args: + patch_size: Tuple of integers $(w, h, d)$ to swap patches + of size $w \times h \times d$. + If a single number $n$ is provided, $w = h = d = n$. + num_iterations: Number of times that two patches will be swapped. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + patch_size: TypeTuple = 15, + num_iterations: int = 100, + **kwargs, + ): + super().__init__(**kwargs) + self.patch_size = np.array(to_tuple(patch_size)) + self.num_iterations = self._parse_num_iterations(num_iterations) + + @staticmethod + def _parse_num_iterations(num_iterations): + if not isinstance(num_iterations, int): + raise TypeError( + f'num_iterations must be an int,not {num_iterations}', + ) + if num_iterations < 0: + raise ValueError( + f'num_iterations must be positive,not {num_iterations}', + ) + return num_iterations + + @staticmethod + def get_params( + tensor: torch.Tensor, + patch_size: np.ndarray, + num_iterations: int, + ) -> list[tuple[TypeTripletInt, TypeTripletInt]]: + si, sj, sk = tensor.shape[-3:] + spatial_shape = si, sj, sk # for mypy + patch_size_values = [int(value) for value in patch_size.tolist()] + locations = [] + for _ in range(num_iterations): + first_ini, first_fin = get_random_indices_from_shape( + spatial_shape, + patch_size_values, + ) + while True: + second_ini, second_fin = get_random_indices_from_shape( + spatial_shape, + patch_size_values, + ) + larger_than_initial = np.all(second_ini >= first_ini) + less_than_final = np.all(second_fin <= first_fin) + if larger_than_initial and less_than_final: + continue # patches overlap + else: + break # patches don't overlap + first_values = [int(value) for value in first_ini.tolist()] + second_values = [int(value) for value in second_ini.tolist()] + first_location: TypeTripletInt = ( + first_values[0], + first_values[1], + first_values[2], + ) + second_location: TypeTripletInt = ( + second_values[0], + second_values[1], + second_values[2], + ) + location = first_location, second_location + locations.append(location) + return locations + + def apply_transform(self, subject: Subject) -> Subject: + images_dict = self.get_images_dict(subject) + if not images_dict: + return subject + + locations_by_name: dict[str, TypeLocations] = {} + patch_size_by_name: dict[str, TypeTripletInt] = {} + broadcast = np.broadcast_to(self.patch_size, 3) + patch_size = (int(broadcast[0]), int(broadcast[1]), int(broadcast[2])) + for name, image in images_dict.items(): + locations = self.get_params( + image.data, + self.patch_size, + self.num_iterations, + ) + locations_by_name[name] = locations + patch_size_by_name[name] = patch_size + transform = Swap( + patch_size=patch_size_by_name, + locations=locations_by_name, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + +class Swap(IntensityTransform): + r"""Swap patches within an image. + + This is typically used in [context restoration for self-supervised learning + ](https://www.sciencedirect.com/science/article/pii/S1361841518304699). + + Args: + patch_size: Tuple of integers $(w, h, d)$ to swap patches + of size $w \times h \times d$. + If a single number $n$ is provided, $w = h = d = n$. + num_iterations: Number of times that two patches will be swapped. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + patch_size: TypeTripletInt | dict[str, TypeTripletInt], + locations: TypeLocations | dict[str, TypeLocations], + **kwargs, + ): + super().__init__(**kwargs) + self.locations = locations + self.patch_size = patch_size + self.args_names = ['locations', 'patch_size'] + self.invert_transform = False + + def apply_transform(self, subject: Subject) -> Subject: + for name, image in self.get_images_dict(subject).items(): + locations = list(self.get_parameter(self.locations, name)) + patch_size = self.get_parameter(self.patch_size, name) + if self.invert_transform: + locations.reverse() + swapped = _swap(image.data, patch_size, locations) + image.set_data(swapped) + return subject + + +def _swap( + tensor: torch.Tensor, + patch_size: TypeTuple, + locations: TypeLocations, +) -> torch.Tensor: + # Note this function modifies the input in-place + tensor = tensor.clone() + patch_size_array = np.array(patch_size) + for first_ini, second_ini in locations: + first_ini_array = np.asarray(first_ini) + second_ini_array = np.asarray(second_ini) + first_fin = first_ini_array + patch_size_array + second_fin = second_ini_array + patch_size_array + first_patch = _crop(tensor, first_ini_array, first_fin) + second_patch = _crop(tensor, second_ini_array, second_fin).clone() + _insert(tensor, first_patch, second_ini_array) + _insert(tensor, second_patch, first_ini_array) + return tensor + + +def _insert( + tensor: torch.Tensor, + patch: torch.Tensor, + index_ini: np.ndarray, +) -> None: + index_fin = index_ini + np.array(patch.shape[-3:]) + i_ini, j_ini, k_ini = index_ini + i_fin, j_fin, k_fin = index_fin + tensor[:, i_ini:i_fin, j_ini:j_fin, k_ini:k_fin] = patch + + +def _crop( + image: torch.Tensor, + index_ini: np.ndarray, + index_fin: np.ndarray, +) -> torch.Tensor: + i_ini, j_ini, k_ini = index_ini + i_fin, j_fin, k_fin = index_fin + return image[:, i_ini:i_fin, j_ini:j_fin, k_ini:k_fin] + + +def get_random_indices_from_shape( + spatial_shape: Sequence[int], + patch_size: Sequence[int], +) -> tuple[np.ndarray, np.ndarray]: + assert len(spatial_shape) == 3 + assert len(patch_size) in (1, 3) + shape_array = np.array(spatial_shape) + patch_size_array = np.array(patch_size) + max_index_ini_unchecked = shape_array - patch_size_array + if (max_index_ini_unchecked < 0).any(): + message = ( + f'Patch size {patch_size} cannot be' + f' larger than image spatial shape {spatial_shape}' + ) + raise ValueError(message) + max_index_ini = max_index_ini_unchecked.astype(np.uint16) + coordinates = [] + for max_coordinate in max_index_ini.tolist(): + if max_coordinate == 0: + coordinate = 0 + else: + coordinate = int(torch.randint(max_coordinate, size=(1,)).item()) + coordinates.append(coordinate) + index_ini = np.array(coordinates, np.uint16) + index_fin = index_ini + patch_size_array + return index_ini, index_fin diff --git a/src/torchio/transforms/augmentation/random_transform.py b/src/torchio/transforms/augmentation/random_transform.py new file mode 100644 index 000000000..5c9c80d96 --- /dev/null +++ b/src/torchio/transforms/augmentation/random_transform.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import torch + +from ...types import TypeRangeFloat +from ...types import TypeSextetFloat +from ...types import TypeTripletFloat +from ..transform import Transform + + +class RandomTransform(Transform): + """Base class for stochastic augmentation transforms. + + Args: + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def parse_degrees( + self, + degrees: TypeRangeFloat, + ) -> tuple[float, float]: + return self._parse_range(degrees, 'degrees') + + def parse_translation( + self, + translation: TypeRangeFloat, + ) -> tuple[float, float]: + return self._parse_range(translation, 'translation') + + @staticmethod + def sample_uniform(a: float, b: float) -> float: + return torch.FloatTensor(1).uniform_(a, b).item() + + @staticmethod + def _get_random_seed() -> int: + """Generate a random seed. + + Returns: + A random seed as an int. + """ + return int(torch.randint(0, 2**31, (1,)).item()) + + @staticmethod + def sample_uniform_sextet(params: TypeSextetFloat) -> TypeTripletFloat: + results = [] + for a, b in zip(params[::2], params[1::2], strict=True): + results.append(RandomTransform.sample_uniform(a, b)) + sx, sy, sz = results + return sx, sy, sz diff --git a/src/torchio/transforms/augmentation/spatial/__init__.py b/src/torchio/transforms/augmentation/spatial/__init__.py new file mode 100644 index 000000000..800b2eedc --- /dev/null +++ b/src/torchio/transforms/augmentation/spatial/__init__.py @@ -0,0 +1,21 @@ +from .random_affine import Affine +from .random_affine import RandomAffine +from .random_affine_elastic_deformation import AffineElasticDeformation +from .random_affine_elastic_deformation import RandomAffineElasticDeformation +from .random_anisotropy import RandomAnisotropy +from .random_elastic_deformation import ElasticDeformation +from .random_elastic_deformation import RandomElasticDeformation +from .random_flip import Flip +from .random_flip import RandomFlip + +__all__ = [ + 'RandomFlip', + 'Flip', + 'RandomAffine', + 'Affine', + 'RandomAnisotropy', + 'RandomElasticDeformation', + 'ElasticDeformation', + 'RandomAffineElasticDeformation', + 'AffineElasticDeformation', +] diff --git a/src/torchio/transforms/augmentation/spatial/random_affine.py b/src/torchio/transforms/augmentation/spatial/random_affine.py new file mode 100644 index 000000000..ede10f6f7 --- /dev/null +++ b/src/torchio/transforms/augmentation/spatial/random_affine.py @@ -0,0 +1,506 @@ +from __future__ import annotations + +from collections.abc import Sequence +from numbers import Number +from typing import Protocol +from typing import cast + +import numpy as np +import SimpleITK as sitk +import torch + +from ....constants import INTENSITY +from ....constants import TYPE +from ....data.io import nib_to_sitk +from ....data.subject import Subject +from ....types import TypeRangeFloat +from ....types import TypeSextetFloat +from ....types import TypeTripletFloat +from ....utils import get_major_sitk_version +from ....utils import to_tuple +from ...spatial_transform import SpatialTransform +from .. import RandomTransform + +TypeOneToSixFloat = TypeRangeFloat | TypeTripletFloat | TypeSextetFloat + + +class _CompositeTransform(Protocol): + def AddTransform(self, transform: sitk.Transform) -> None: ... + + +class RandomAffine(RandomTransform, SpatialTransform): + r"""Apply a random affine transformation and resample the image. + + Args: + scales: Tuple $(a_1, b_1, a_2, b_2, a_3, b_3)$ defining the + scaling ranges. + The scaling values along each dimension are $(s_1, s_2, s_3)$, + where $s_i \sim \mathcal{U}(a_i, b_i)$. + If two values $(a, b)$ are provided, + then $s_i \sim \mathcal{U}(a, b)$. + If only one value $x$ is provided, + then $s_i \sim \mathcal{U}(1 - x, 1 + x)$. + If three values $(x_1, x_2, x_3)$ are provided, + then $s_i \sim \mathcal{U}(1 - x_i, 1 + x_i)$. + For example, using `scales=(0.5, 0.5)` will zoom out the image, + making the objects inside look twice as small while preserving + the physical size and position of the image bounds. + degrees: Tuple $(a_1, b_1, a_2, b_2, a_3, b_3)$ defining the + rotation ranges in degrees. + Rotation angles around each axis are + $(\theta_1, \theta_2, \theta_3)$, + where $\theta_i \sim \mathcal{U}(a_i, b_i)$. + If two values $(a, b)$ are provided, + then $\theta_i \sim \mathcal{U}(a, b)$. + If only one value $x$ is provided, + then $\theta_i \sim \mathcal{U}(-x, x)$. + If three values $(x_1, x_2, x_3)$ are provided, + then $\theta_i \sim \mathcal{U}(-x_i, x_i)$. + translation: Tuple $(a_1, b_1, a_2, b_2, a_3, b_3)$ defining the + translation ranges in mm. + Translation along each axis is $(t_1, t_2, t_3)$, + where $t_i \sim \mathcal{U}(a_i, b_i)$. + If two values $(a, b)$ are provided, + then $t_i \sim \mathcal{U}(a, b)$. + If only one value $x$ is provided, + then $t_i \sim \mathcal{U}(-x, x)$. + If three values $(x_1, x_2, x_3)$ are provided, + then $t_i \sim \mathcal{U}(-x_i, x_i)$. + For example, if the image is in RAS+ orientation (e.g., after + applying [`ToCanonical`][torchio.transforms.preprocessing.ToCanonical]) + and the translation is $(10, 20, 30)$, the sample will move + 10 mm to the right, 20 mm to the front, and 30 mm upwards. + If the image was in, e.g., PIR+ orientation, the sample will move + 10 mm to the back, 20 mm downwards, and 30 mm to the right. + isotropic: If `True`, only one scaling factor will be sampled for all dimensions, + i.e. $s_1 = s_2 = s_3$. + If one value $x$ is provided in `scales`, the scaling factor along all + dimensions will be $s \sim \mathcal{U}(1 - x, 1 + x)$. + If two values provided $(a, b)$ in `scales`, the scaling factor along all + dimensions will be $s \sim \mathcal{U}(a, b)$. + center: If `'image'`, rotations and scaling will be performed around + the image center. If `'origin'`, rotations and scaling will be + performed around the origin in world coordinates. + default_pad_value: As the image is rotated, some values near the + borders will be undefined. + If `'minimum'`, the fill value will be the image minimum. + If `'mean'`, the fill value is the mean of the border values. + If `'otsu'`, the fill value is the mean of the values at the + border that lie under an + [Otsu threshold ](https://ieeexplore.ieee.org/document/4310076). + If it is a number, that value will be used. + This parameter applies to intensity images only. + default_pad_label: As the label map is rotated, some values near the + borders will be undefined. This numeric value will be used to fill + those undefined regions. This parameter applies to label maps only. + image_interpolation: See Interpolation. + label_interpolation: See Interpolation. + check_shape: If `True` an error will be raised if the images are in + different physical spaces. If `False`, `center` should + probably not be `'image'` but `'center'`. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> image = tio.datasets.Colin27().t1 + >>> transform = tio.RandomAffine( + ... scales=(0.9, 1.2), + ... degrees=15, + ... ) + >>> transformed = transform(image) + + """ + + def __init__( + self, + scales: TypeOneToSixFloat = 0.1, + degrees: TypeOneToSixFloat = 10, + translation: TypeOneToSixFloat = 0, + isotropic: bool = False, + center: str = 'image', + default_pad_value: str | float = 'minimum', + default_pad_label: int | float = 0, + image_interpolation: str = 'linear', + label_interpolation: str = 'nearest', + check_shape: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + self.isotropic = isotropic + _parse_scales_isotropic(scales, isotropic) + self.scales = self.parse_params(scales, 1, 'scales', min_constraint=0) + self.degrees = self.parse_params(degrees, 0, 'degrees') + self.translation = self.parse_params(translation, 0, 'translation') + if center not in ('image', 'origin'): + message = f'Center argument must be "image" or "origin", not "{center}"' + raise ValueError(message) + self.center = center + self.default_pad_value = _parse_default_value(default_pad_value) + if not isinstance(default_pad_label, (int, float)): + message = 'default_pad_label must be a number, ' + message += f'but it is "{default_pad_label}"' + raise ValueError(message) + self.default_pad_label = float(default_pad_label) + self.image_interpolation = self.parse_interpolation( + image_interpolation, + ) + self.label_interpolation = self.parse_interpolation( + label_interpolation, + ) + self.check_shape = check_shape + + @staticmethod + def get_params( + scales: TypeSextetFloat, + degrees: TypeSextetFloat, + translation: TypeSextetFloat, + isotropic: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + scaling_params = torch.as_tensor( + RandomTransform.sample_uniform_sextet(scales), + dtype=torch.float64, + ) + if isotropic: + scaling_params.fill_(scaling_params[0]) + rotation_params = torch.as_tensor( + RandomTransform.sample_uniform_sextet(degrees), + dtype=torch.float64, + ) + translation_params = torch.as_tensor( + RandomTransform.sample_uniform_sextet(translation), + dtype=torch.float64, + ) + return scaling_params, rotation_params, translation_params + + def apply_transform(self, subject: Subject) -> Subject: + scaling_params, rotation_params, translation_params = self.get_params( + self.scales, + self.degrees, + self.translation, + self.isotropic, + ) + scaling_values = [float(value) for value in scaling_params.tolist()] + rotation_values = [float(value) for value in rotation_params.tolist()] + translation_values = [float(value) for value in translation_params.tolist()] + transform = Affine( + scales=(scaling_values[0], scaling_values[1], scaling_values[2]), + degrees=(rotation_values[0], rotation_values[1], rotation_values[2]), + translation=( + translation_values[0], + translation_values[1], + translation_values[2], + ), + center=self.center, + default_pad_value=self.default_pad_value, + default_pad_label=self.default_pad_label, + image_interpolation=self.image_interpolation, + label_interpolation=self.label_interpolation, + check_shape=self.check_shape, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + +class Affine(SpatialTransform): + r"""Apply affine transformation. + + Args: + scales: Tuple $(s_1, s_2, s_3)$ defining the + scaling values along each dimension. + degrees: Tuple $(\theta_1, \theta_2, \theta_3)$ defining the + rotation around each axis. + translation: Tuple $(t_1, t_2, t_3)$ defining the + translation in mm along each axis. + center: If `'image'`, rotations and scaling will be performed around + the image center. If `'origin'`, rotations and scaling will be + performed around the origin in world coordinates. + default_pad_value: As the image is rotated, some values near the + borders will be undefined. + If `'minimum'`, the fill value will be the image minimum. + If `'mean'`, the fill value is the mean of the border values. + If `'otsu'`, the fill value is the mean of the values at the + border that lie under an + [Otsu threshold ](https://ieeexplore.ieee.org/document/4310076). + If it is a number, that value will be used. + This parameter applies to intensity images only. + default_pad_label: As the label map is rotated, some values near the + borders will be undefined. This numeric value will be used to fill + those undefined regions. This parameter applies to label maps only. + image_interpolation: See Interpolation. + label_interpolation: See Interpolation. + check_shape: If `True` an error will be raised if the images are in + different physical spaces. If `False`, `center` should + probably not be `'image'` but `'center'`. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + scales: TypeTripletFloat, + degrees: TypeTripletFloat, + translation: TypeTripletFloat, + center: str = 'image', + default_pad_value: str | float = 'minimum', + default_pad_label: int | float = 0, + image_interpolation: str = 'linear', + label_interpolation: str = 'nearest', + check_shape: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + self.scales = self.parse_params( + scales, + None, + 'scales', + make_ranges=False, + min_constraint=0, + ) + self.degrees = self.parse_params( + degrees, + None, + 'degrees', + make_ranges=False, + ) + self.translation = self.parse_params( + translation, + None, + 'translation', + make_ranges=False, + ) + if center not in ('image', 'origin'): + message = f'Center argument must be "image" or "origin", not "{center}"' + raise ValueError(message) + self.center = center + self.use_image_center = center == 'image' + self.default_pad_value = _parse_default_value(default_pad_value) + if not isinstance(default_pad_label, (int, float)): + message = 'default_pad_label must be a number, ' + message += f'but it is "{default_pad_label}"' + raise ValueError(message) + self.default_pad_label = float(default_pad_label) + self.image_interpolation = self.parse_interpolation( + image_interpolation, + ) + self.label_interpolation = self.parse_interpolation( + label_interpolation, + ) + self.invert_transform = False + self.check_shape = check_shape + self.args_names = [ + 'scales', + 'degrees', + 'translation', + 'center', + 'default_pad_value', + 'default_pad_label', + 'image_interpolation', + 'label_interpolation', + 'check_shape', + ] + + @staticmethod + def _get_scaling_transform( + scaling_params: Sequence[float] | np.ndarray, + center_lps: TypeTripletFloat | None = None, + ) -> sitk.ScaleTransform: + # 1.5 means the objects look 1.5 times larger + transform = sitk.ScaleTransform(3) + scaling_params_array = np.array(scaling_params).astype(float) + transform.SetScale(scaling_params_array) + if center_lps is not None: + transform.SetCenter(center_lps) + return transform + + @staticmethod + def _get_rotation_transform( + degrees: Sequence[float] | np.ndarray, + translation: Sequence[float] | np.ndarray, + center_lps: TypeTripletFloat | None = None, + ) -> sitk.Euler3DTransform: + def ras_to_lps(triplet: Sequence[float] | np.ndarray) -> np.ndarray: + return np.array((-1, -1, 1), dtype=float) * np.asarray(triplet) + + transform = sitk.Euler3DTransform() + radians = np.asarray(np.radians(degrees), dtype=float) + + # SimpleITK uses LPS + radians_lps = ras_to_lps(radians) + translation_lps = ras_to_lps(translation) + + transform.SetRotation(*radians_lps) + transform.SetTranslation(translation_lps) + if center_lps is not None: + transform.SetCenter(center_lps) + return transform + + def get_affine_transform(self, image): + scaling = np.asarray(self.scales).copy() + rotation = np.asarray(self.degrees).copy() + translation = np.asarray(self.translation).copy() + + if image.is_2d(): + scaling[2] = 1 + rotation[:-1] = 0 + + if self.use_image_center: + center_lps = image.get_center(lps=True) + else: + center_lps = None + + scaling_transform = self._get_scaling_transform( + scaling, + center_lps=center_lps, + ) + rotation_transform = self._get_rotation_transform( + rotation, + translation, + center_lps=center_lps, + ) + + sitk_major_version = get_major_sitk_version() + if sitk_major_version == 1: + composite = cast( + _CompositeTransform, + sitk.Transform(3, sitk.sitkComposite), + ) + composite.AddTransform(scaling_transform) + composite.AddTransform(rotation_transform) + transform = cast(sitk.Transform, composite) + elif sitk_major_version == 2: + transforms = [scaling_transform, rotation_transform] + transform = sitk.CompositeTransform(transforms) + + # ResampleImageFilter expects the transform from the output space to + # the input space. Intuitively, the passed arguments should take us + # from the input space to the output space, so we need to invert the + # transform. + # More info at https://github.com/TorchIO-project/torchio/discussions/693 + transform = transform.GetInverse() + + if self.invert_transform: + transform = transform.GetInverse() + + return transform + + def get_default_pad_value( + self, tensor: torch.Tensor, sitk_image: sitk.Image + ) -> float: + default_value: float + if self.default_pad_value == 'minimum': + default_value = tensor.min().item() + elif self.default_pad_value == 'mean': + default_value = get_borders_mean( + sitk_image, + filter_otsu=False, + ) + elif self.default_pad_value == 'otsu': + default_value = get_borders_mean( + sitk_image, + filter_otsu=True, + ) + else: + assert isinstance(self.default_pad_value, Number) + default_value = float(self.default_pad_value) + return default_value + + def apply_transform(self, subject: Subject) -> Subject: + if self.check_shape: + subject.check_consistent_spatial_shape() + default_value: float + for image in self.get_images(subject): + transform = self.get_affine_transform(image) + transformed_tensors = [] + for tensor in image.data: + sitk_image = nib_to_sitk( + tensor[np.newaxis], + image.affine, + force_3d=True, + ) + if image[TYPE] != INTENSITY: + interpolation = self.label_interpolation + default_value = self.default_pad_label + else: + interpolation = self.image_interpolation + default_value = self.get_default_pad_value(tensor, sitk_image) + transformed_tensor = self.apply_affine_transform( + sitk_image, + transform, + interpolation, + default_value, + ) + transformed_tensors.append(transformed_tensor) + image.set_data(torch.stack(transformed_tensors)) + return subject + + def apply_affine_transform( + self, + sitk_image: sitk.Image, + transform: sitk.Transform, + interpolation: str, + default_value: float, + ) -> torch.Tensor: + floating = reference = sitk_image + + resampler = sitk.ResampleImageFilter() + resampler.SetInterpolator(self.get_sitk_interpolator(interpolation)) + resampler.SetReferenceImage(reference) + resampler.SetDefaultPixelValue(float(default_value)) + resampler.SetOutputPixelType(sitk.sitkFloat32) + resampler.SetTransform(transform) + resampled = resampler.Execute(floating) + + np_array = sitk.GetArrayFromImage(resampled) + np_array = np_array.transpose() # ITK to NumPy + tensor = torch.as_tensor(np_array) + return tensor + + +def get_borders_mean(image, filter_otsu=True): + array = sitk.GetArrayViewFromImage(image) + borders_tuple = ( + array[0, :, :], + array[-1, :, :], + array[:, 0, :], + array[:, -1, :], + array[:, :, 0], + array[:, :, -1], + ) + borders_flat = np.hstack([border.ravel() for border in borders_tuple]) + if not filter_otsu: + return borders_flat.mean() + borders_reshaped = borders_flat.reshape(1, 1, -1) + borders_image = sitk.GetImageFromArray(borders_reshaped) + otsu = sitk.OtsuThresholdImageFilter() + otsu.Execute(borders_image) + threshold = otsu.GetThreshold() + values = borders_flat[borders_flat < threshold] + if values.any(): + default_value = values.mean() + else: + default_value = borders_flat.mean() + return default_value + + +def _parse_scales_isotropic(scales, isotropic): + scales = to_tuple(scales) + if isotropic and len(scales) in (3, 6): + message = ( + 'If "isotropic" is True, the value for "scales" must have' + f' length 1 or 2, but "{scales}" was passed.' + ' If you want to set isotropic scaling, use a single value or two values as a range' + ' for the scaling factor. Refer to the documentation for more information.' + ) + raise ValueError(message) + + +def _parse_default_value(value: str | float) -> str | float: + if isinstance(value, Number) or value in ('minimum', 'otsu', 'mean'): + return value + message = ( + 'Value for default_pad_value must be "minimum", "otsu", "mean" or a number' + ) + raise ValueError(message) diff --git a/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py b/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py new file mode 100644 index 000000000..c1e282234 --- /dev/null +++ b/src/torchio/transforms/augmentation/spatial/random_affine_elastic_deformation.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import SimpleITK as sitk +import torch + +from ....constants import INTENSITY +from ....constants import TYPE +from ....data.io import nib_to_sitk +from ....data.subject import Subject +from ...spatial_transform import SpatialTransform +from .. import RandomTransform +from .random_affine import Affine +from .random_elastic_deformation import ElasticDeformation + + +class RandomAffineElasticDeformation(RandomTransform, SpatialTransform): + r"""Apply a RandomAffine and RandomElasticDeformation simultaneously. + + Optimization to use only a single SimpleITK resampling. For additional details on + the transformations, see [`RandomAffine`][torchio.transforms.RandomAffine] + and [`RandomElasticDeformation`][torchio.transforms.RandomElasticDeformation] + + Args: + affine_first: Apply affine before elastic deformation. + affine_kwargs: See [`RandomAffine`][torchio.transforms.RandomAffine] for kwargs. + elastic_kwargs: See [`RandomElasticDeformation`][torchio.transforms.RandomElasticDeformation] + for kwargs. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> image = tio.datasets.Colin27().t1 + >>> affine_kwargs = {'scales': (0.9, 1.2), 'degrees': 15} + >>> elastic_kwargs = {'max_displacement': (17, 12, 2)} + >>> transform = tio.RandomAffineElasticDeformation( + ... affine_kwargs, + ... elastic_kwargs + ... ) + >>> transformed = transform(image) + + """ + + def __init__( + self, + affine_first: bool = True, + affine_kwargs: dict[str, Any] | None = None, + elastic_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.affine_first = affine_first + + # Avoid circular imports + from .random_affine import RandomAffine + from .random_elastic_deformation import RandomElasticDeformation + + self.affine_kwargs = affine_kwargs or {} + self.random_affine = RandomAffine(**self.affine_kwargs) + + self.elastic_kwargs = elastic_kwargs or {} + self.random_elastic = RandomElasticDeformation(**self.elastic_kwargs) + + def get_params(self): + affine_params = self.random_affine.get_params( + self.random_affine.scales, + self.random_affine.degrees, + self.random_affine.translation, + self.random_affine.isotropic, + ) + elastic_params = self.random_elastic.get_params( + self.random_elastic.num_control_points, + self.random_elastic.max_displacement, + self.random_elastic.num_locked_borders, + ) + return affine_params, elastic_params + + def apply_transform(self, subject: Subject) -> Subject: + affine_params, elastic_params = self.get_params() + + scaling_params, rotation_params, translation_params = affine_params + affine_params = { + 'scales': scaling_params.tolist(), + 'degrees': rotation_params.tolist(), + 'translation': translation_params.tolist(), + 'center': self.random_affine.center, + 'default_pad_value': self.random_affine.default_pad_value, + 'image_interpolation': self.random_affine.image_interpolation, + 'label_interpolation': self.random_affine.label_interpolation, + 'check_shape': self.random_affine.check_shape, + } + + elastic_params = { + 'control_points': elastic_params, + 'max_displacement': self.random_elastic.max_displacement, + 'image_interpolation': self.random_elastic.image_interpolation, + 'label_interpolation': self.random_elastic.label_interpolation, + } + + transform = AffineElasticDeformation( + affine_first=self.affine_first, + affine_params=affine_params, + elastic_params=elastic_params, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + +class AffineElasticDeformation(SpatialTransform): + r"""Apply an Affine and ElasticDeformation simultaneously. + + Optimization to use only a single SimpleITK resampling. For additional details + on the transformations, see [`Affine`][torchio.transforms.augmentation.Affine] + and [`ElasticDeformation`][torchio.transforms.augmentation.ElasticDeformation] + + Args: + affine_first: Apply affine before elastic deformation. + affine_kwargs: See [`RandomAffine`][torchio.transforms.augmentation.RandomAffine] for kwargs. + elastic_kwargs: See + [`RandomElasticDeformation`][torchio.transforms.augmentation.RandomElasticDeformation] for kwargs. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + affine_first: bool, + affine_params: dict[str, Any], + elastic_params: dict[str, Any], + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.affine_first = affine_first + + self.affine_params = affine_params + self._affine = Affine( + **self.affine_params, + **kwargs, + ) + self.elastic_params = elastic_params + self._elastic = ElasticDeformation( + **self.elastic_params, + **kwargs, + ) + + self.args_names = ['affine_first', 'affine_params', 'elastic_params'] + + def apply_transform(self, subject: Subject) -> Subject: + if self._affine.check_shape: + subject.check_consistent_spatial_shape() + default_value: float + + for image in self.get_images(subject): + affine_transform = self._affine.get_affine_transform(image) + transformed_tensors = [] + for tensor in image.data: + sitk_image = nib_to_sitk( + tensor[np.newaxis], + image.affine, + force_3d=True, + ) + if image[TYPE] != INTENSITY: + interpolation = self._affine.label_interpolation + default_value = 0 + else: + interpolation = self._affine.image_interpolation + default_value = self._affine.get_default_pad_value( + tensor, sitk_image + ) + + bspline_transform = self._elastic.get_bspline_transform(sitk_image) + self._elastic.parse_free_form_transform( + bspline_transform, + self._elastic.max_displacement, + ) + + # stack: LIFO + if self.affine_first: + combined_transforms = [affine_transform, bspline_transform] + else: + combined_transforms = [bspline_transform, affine_transform] + composite_transform = sitk.CompositeTransform(combined_transforms) + + transformed_tensor = self.apply_composite_transform( + sitk_image, + composite_transform, + interpolation, + default_value, + ) + transformed_tensors.append(transformed_tensor) + image.set_data(torch.stack(transformed_tensors)) + return subject + + def apply_composite_transform( + self, + sitk_image: sitk.Image, + transform: sitk.Transform, + interpolation: str, + default_value: float, + ) -> torch.Tensor: + floating = reference = sitk_image + + resampler = sitk.ResampleImageFilter() + resampler.SetInterpolator(self.get_sitk_interpolator(interpolation)) + resampler.SetReferenceImage(reference) + resampler.SetDefaultPixelValue(float(default_value)) + resampler.SetOutputPixelType(sitk.sitkFloat32) + resampler.SetTransform(transform) + resampled = resampler.Execute(floating) + + np_array = sitk.GetArrayFromImage(resampled) + np_array = np_array.transpose() # ITK to NumPy + tensor = torch.as_tensor(np_array) + return tensor diff --git a/src/torchio/transforms/augmentation/spatial/random_anisotropy.py b/src/torchio/transforms/augmentation/spatial/random_anisotropy.py new file mode 100644 index 000000000..8b8541cdd --- /dev/null +++ b/src/torchio/transforms/augmentation/spatial/random_anisotropy.py @@ -0,0 +1,133 @@ +import warnings + +import torch + +from ....data.subject import Subject +from ....types import TypeRangeFloat +from ....utils import to_tuple +from ...preprocessing import Resample +from .. import RandomTransform + + +class RandomAnisotropy(RandomTransform): + r"""Downsample an image along an axis and upsample to initial space. + + This transform simulates an image that has been acquired using anisotropic + spacing and resampled back to its original spacing. + + Similar to the work by Billot et al.: [Partial Volume Segmentation of Brain + MRI Scans of any Resolution and + Contrast ](https://link.springer.com/chapter/10.1007/978-3-030-59728-3_18). + + Args: + axes: Axis or tuple of axes along which the image will be downsampled. + downsampling: Downsampling factor $m \gt 1$. If a tuple + $(a, b)$ is provided then $m \sim \mathcal{U}(a, b)$. + image_interpolation: Image interpolation used to upsample the image + back to its initial spacing. Downsampling is performed using + nearest neighbor interpolation. See Interpolation for + supported interpolation types. + scalars_only: Apply only to instances of [`torchio.ScalarImage`][torchio.ScalarImage]. + This is useful when the segmentation quality needs to be kept, + as in [Billot et al. ](https://link.springer.com/chapter/10.1007/978-3-030-59728-3_18). + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> transform = tio.RandomAnisotropy(axes=1, downsampling=2) + >>> transform = tio.RandomAnisotropy( + ... axes=(0, 1, 2), + ... downsampling=(2, 5), + ... ) # Multiply spacing of one of the 3 axes by a factor randomly chosen in [2, 5] + >>> colin = tio.datasets.Colin27() + >>> transformed = transform(colin) + """ + + def __init__( + self, + axes: int | tuple[int, ...] = (0, 1, 2), + downsampling: TypeRangeFloat = (1.5, 5), + image_interpolation: str = 'linear', + scalars_only: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + self.axes = self.parse_axes(axes) + self.downsampling_range = self._parse_range( + downsampling, + 'downsampling', + min_constraint=1, + ) + parsed_interpolation = self.parse_interpolation(image_interpolation) + self.image_interpolation = parsed_interpolation + self.scalars_only = scalars_only + + def get_params( + self, + axes: tuple[int, ...], + downsampling_range: tuple[float, float], + ) -> tuple[int, float]: + axis = axes[torch.randint(0, len(axes), (1,))] + downsampling = self.sample_uniform(*downsampling_range) + return axis, downsampling + + @staticmethod + def parse_axes(axes: int | tuple[int, ...]): + axes_tuple = to_tuple(axes) + for axis in axes_tuple: + is_int = isinstance(axis, int) + if not is_int or axis not in (0, 1, 2): + raise ValueError('All axes must be 0, 1 or 2') + return axes_tuple + + def apply_transform(self, subject: Subject) -> Subject: + is_2d = subject.get_first_image().is_2d() + axes = self.axes + if is_2d and 2 in self.axes: + warnings.warn( + f'Input image is 2D, but "2" is in axes: {self.axes}', + RuntimeWarning, + stacklevel=2, + ) + axes = tuple(axis for axis in self.axes if axis != 2) + axis, downsampling = self.get_params( + axes, + self.downsampling_range, + ) + target_spacing = list(subject.spacing) + target_spacing[axis] *= downsampling + + # NOTE: If copy=False, the underlying image data will be modified in place. + # We have to obtain the target spatial shape and affine before the transform + image = subject.get_first_image() + downsample = Resample( + target=( + float(target_spacing[0]), + float(target_spacing[1]), + float(target_spacing[2]), + ), + image_interpolation='nearest', + scalars_only=self.scalars_only, + copy=self.copy, + include=self.include, + exclude=self.exclude, + keep=self.keep, + parse_input=self.parse_input, + label_keys=self.label_keys, + ) + downsampled = downsample(subject) + upsample = Resample( + target=(image.spatial_shape, image.affine), + image_interpolation=self.image_interpolation, + scalars_only=self.scalars_only, + copy=self.copy, + include=self.include, + exclude=self.exclude, + keep=self.keep, + parse_input=self.parse_input, + label_keys=self.label_keys, + ) + upsampled = upsample(downsampled) + assert isinstance(upsampled, Subject) + return upsampled diff --git a/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py b/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py new file mode 100644 index 000000000..2ceff4ab7 --- /dev/null +++ b/src/torchio/transforms/augmentation/spatial/random_elastic_deformation.py @@ -0,0 +1,348 @@ +import warnings +from numbers import Number +from typing import cast + +import numpy as np +import SimpleITK as sitk +import torch + +from ....data.image import ScalarImage +from ....data.io import nib_to_sitk +from ....data.subject import Subject +from ....types import TypeTripletFloat +from ....types import TypeTripletInt +from ....utils import to_tuple +from ...spatial_transform import SpatialTransform +from .. import RandomTransform + +SPLINE_ORDER = 3 + + +class RandomElasticDeformation(RandomTransform, SpatialTransform): + r"""Apply dense random elastic deformation. + + A random displacement is assigned to a coarse grid of control points around + and inside the image. The displacement at each voxel is interpolated from + the coarse grid using cubic B-splines. + + The ['Deformable Registration' ](https://www.sciencedirect.com/topics/computer-science/deformable-registration) + topic on ScienceDirect contains useful articles explaining interpolation of + displacement fields using cubic B-splines. + + Warning: + This transform is slow as it requires expensive computations. + If your images are large you might want to use + [`RandomAffine`][torchio.transforms.RandomAffine] instead. + + Args: + num_control_points: Number of control points along each dimension of + the coarse grid $(n_x, n_y, n_z)$. + If a single value $n$ is passed, + then $n_x = n_y = n_z = n$. + Smaller numbers generate smoother deformations. + The minimum number of control points is `4` as this transform + uses cubic B-splines to interpolate displacement. + max_displacement: Maximum displacement along each dimension at each + control point $(D_x, D_y, D_z)$. + The displacement along dimension $i$ at each control point is + $d_i \sim \mathcal{U}(0, D_i)$. + If a single value $D$ is passed, + then $D_x = D_y = D_z = D$. + Note that the total maximum displacement would actually be + $D_{max} = \sqrt{D_x^2 + D_y^2 + D_z^2}$. + locked_borders: If `0`, all displacement vectors are kept. + If `1`, displacement of control points at the + border of the coarse grid will be set to `0`. + If `2`, displacement of control points at the border of the image + (red dots in the image below) will also be set to `0`. + image_interpolation: See Interpolation. + Note that this is the interpolation used to compute voxel + intensities when resampling using the dense displacement field. + The value of the dense displacement at each voxel is always + interpolated with cubic B-splines from the values at the control + points of the coarse grid. + label_interpolation: See Interpolation. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + [This gist ](https://gist.github.com/fepegar/b723d15de620cd2a3a4dbd71e491b59d) + can also be used to better understand the meaning of the parameters. + + This is an example from the + [3D Slicer registration FAQ ](https://www.slicer.org/wiki/Documentation/4.10/FAQ/Registration#What.27s_the_BSpline_Grid_Size.3F). + + ![B-spline example from 3D Slicer documentation](https://www.slicer.org/w/img_auth.php/6/6f/RegLib_BSplineGridModel.png) + + To generate a similar grid of control points with TorchIO, + the transform can be instantiated as follows: + + Examples: + >>> from torchio import RandomElasticDeformation + >>> transform = RandomElasticDeformation( + ... num_control_points=(7, 7, 7), # or just 7 + ... locked_borders=2, + ... ) + + Note that control points outside the image bounds are not showed in the + example image (they would also be red as we set `locked_borders` + to `2`). + + Warning: + Image folding may occur if the maximum displacement is larger + than half the coarse grid spacing. The grid spacing can be computed + using the image bounds in physical space and the number of control + points. + + Using a `max_displacement` larger than the computed + `potential_folding` will raise a `RuntimeWarning`. + + Technically, $2 \epsilon$ should be added to the + image bounds, where $\epsilon = 2^{-3}$ [according to ITK + source code](https://github.com/InsightSoftwareConsortium/ITK/blob/633f84548311600845d54ab2463d3412194690a8/Modules/Core/Transform/include/itkBSplineTransformInitializer.hxx#L116-L138). + + Examples: + >>> import numpy as np + >>> import torchio as tio + >>> image = tio.datasets.Slicer().MRHead.as_sitk() + >>> image.GetSize() # in voxels + (256, 256, 130) + >>> image.GetSpacing() # in mm + (1.0, 1.0, 1.2999954223632812) + >>> bounds = np.array(image.GetSize()) * np.array(image.GetSpacing()) + >>> bounds # mm + array([256. , 256. , 168.99940491]) + >>> num_control_points = np.array((7, 7, 6)) + >>> grid_spacing = bounds / (num_control_points - 2) + >>> grid_spacing + array([51.2 , 51.2 , 42.24985123]) + >>> potential_folding = grid_spacing / 2 + >>> potential_folding # mm + array([25.6 , 25.6 , 21.12492561]) + """ + + def __init__( + self, + num_control_points: int | TypeTripletInt = 7, + max_displacement: float | TypeTripletFloat = 7.5, + locked_borders: int = 2, + image_interpolation: str = 'linear', + label_interpolation: str = 'nearest', + **kwargs, + ): + super().__init__(**kwargs) + self._bspline_transformation = None + self.num_control_points = cast( + TypeTripletInt, + to_tuple(num_control_points, length=3), + ) + _parse_num_control_points(self.num_control_points) + self.max_displacement = cast( + TypeTripletFloat, + to_tuple(max_displacement, length=3), + ) + _parse_max_displacement(self.max_displacement) + self.num_locked_borders = locked_borders + if locked_borders not in (0, 1, 2): + raise ValueError('locked_borders must be 0, 1, or 2') + if locked_borders == 2 and 4 in self.num_control_points: + message = ( + 'Setting locked_borders to 2 and using less than 5 control' + 'points results in an identity transform. Lock fewer borders' + ' or use more control points.' + ) + raise ValueError(message) + self.image_interpolation = self.parse_interpolation( + image_interpolation, + ) + self.label_interpolation = self.parse_interpolation( + label_interpolation, + ) + + @staticmethod + def get_params( + num_control_points: TypeTripletInt, + max_displacement: tuple[float, float, float], + num_locked_borders: int, + ) -> np.ndarray: + grid_shape = num_control_points + num_dimensions = 3 + coarse_field = torch.rand(*grid_shape, num_dimensions) # [0, 1) + coarse_field -= 0.5 # [-0.5, 0.5) + coarse_field *= 2 # [-1, 1] + for dimension in range(3): + # [-max_displacement, max_displacement) + coarse_field[..., dimension] *= max_displacement[dimension] + + # Set displacement to 0 at the borders + for i in range(num_locked_borders): + coarse_field[i, :] = 0 + coarse_field[-1 - i, :] = 0 + coarse_field[:, i] = 0 + coarse_field[:, -1 - i] = 0 + + return coarse_field.numpy() + + def apply_transform(self, subject: Subject) -> Subject: + subject.check_consistent_spatial_shape() + control_points = self.get_params( + self.num_control_points, + self.max_displacement, + self.num_locked_borders, + ) + + transform = ElasticDeformation( + control_points=control_points, + max_displacement=self.max_displacement, + image_interpolation=self.image_interpolation, + label_interpolation=self.label_interpolation, + **self._get_base_args(), + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + +class ElasticDeformation(SpatialTransform): + r"""Apply dense elastic deformation. + + Args: + control_points: + max_displacement: + image_interpolation: See Interpolation. + label_interpolation: See Interpolation. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__( + self, + control_points: np.ndarray, + max_displacement: TypeTripletFloat, + image_interpolation: str = 'linear', + label_interpolation: str = 'nearest', + **kwargs, + ): + super().__init__(**kwargs) + self.control_points = control_points + self.max_displacement = max_displacement + self.image_interpolation = self.parse_interpolation( + image_interpolation, + ) + self.label_interpolation = self.parse_interpolation( + label_interpolation, + ) + self.invert_transform = False + self.args_names = [ + 'control_points', + 'image_interpolation', + 'label_interpolation', + 'max_displacement', + ] + + def get_bspline_transform( + self, + image: sitk.Image, + ) -> sitk.BSplineTransform: + control_points = self.control_points.copy() + if self.invert_transform: + control_points *= -1 + is_2d = image.GetSize()[2] == 1 + if is_2d: + control_points[..., -1] = 0 # no displacement in IS axis + num_control_points = control_points.shape[:-1] + mesh_shape = [n - SPLINE_ORDER for n in num_control_points] + bspline_transform = sitk.BSplineTransformInitializer(image, mesh_shape) + parameters = control_points.flatten(order='F').tolist() + bspline_transform.SetParameters(parameters) + return bspline_transform + + @staticmethod + def parse_free_form_transform( + transform: sitk.BSplineTransform, + max_displacement: TypeTripletFloat, + ) -> None: + """Issue a warning is possible folding is detected.""" + coefficient_images = transform.GetCoefficientImages() + grid_spacing = coefficient_images[0].GetSpacing() + conflicts = np.array(max_displacement) > np.array(grid_spacing) / 2 + if np.any(conflicts): + (where,) = np.where(conflicts) + message = ( + 'The maximum displacement is larger than the coarse grid' + f' spacing for dimensions: {where.tolist()}, so folding may' + ' occur. Choose fewer control points or a smaller' + ' maximum displacement' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + + def apply_transform(self, subject: Subject) -> Subject: + no_displacement = not any(self.max_displacement) + if no_displacement: + return subject + subject.check_consistent_spatial_shape() + for image in self.get_images(subject): + if not isinstance(image, ScalarImage): + interpolation = self.label_interpolation + else: + interpolation = self.image_interpolation + transformed = self.apply_bspline_transform( + image.data, + image.affine, + interpolation, + ) + image.set_data(transformed) + return subject + + def apply_bspline_transform( + self, + tensor: torch.Tensor, + affine: np.ndarray, + interpolation: str, + ) -> torch.Tensor: + assert tensor.dim() == 4 + results = [] + for component in tensor: + image = nib_to_sitk(component[np.newaxis], affine, force_3d=True) + floating = reference = image + bspline_transform = self.get_bspline_transform(image) + self.parse_free_form_transform( + bspline_transform, + self.max_displacement, + ) + interpolator = self.get_sitk_interpolator(interpolation) + resampler = sitk.ResampleImageFilter() + resampler.SetReferenceImage(reference) + resampler.SetTransform(bspline_transform) + resampler.SetInterpolator(interpolator) + resampler.SetDefaultPixelValue(component.min().item()) + resampler.SetOutputPixelType(sitk.sitkFloat32) + resampled = resampler.Execute(floating) + result, _ = self.sitk_to_nib(resampled) + results.append(torch.as_tensor(result)) + tensor = torch.cat(results) + return tensor + + +def _parse_num_control_points( + num_control_points: TypeTripletInt, +) -> None: + for axis, number in enumerate(num_control_points): + if not isinstance(number, int) or number < 4: + message = ( + f'The number of control points for axis {axis} must be' + f' an integer greater than 3, not {number}' + ) + raise ValueError(message) + + +def _parse_max_displacement( + max_displacement: tuple[float, float, float], +) -> None: + for axis, number in enumerate(max_displacement): + if not isinstance(number, Number) or number < 0: + message = ( + 'The maximum displacement at each control point' + f' for axis {axis} must be' + f' a number greater or equal to 0, not {number}' + ) + raise ValueError(message) diff --git a/src/torchio/transforms/augmentation/spatial/random_flip.py b/src/torchio/transforms/augmentation/spatial/random_flip.py new file mode 100644 index 000000000..896bdc8f2 --- /dev/null +++ b/src/torchio/transforms/augmentation/spatial/random_flip.py @@ -0,0 +1,131 @@ +import numpy as np +import torch + +from ....data.subject import Subject +from ....utils import to_tuple +from ...spatial_transform import SpatialTransform +from .. import RandomTransform + + +class RandomFlip(RandomTransform, SpatialTransform): + """Reverse the order of elements in an image along the given axes. + + Args: + axes: Index or tuple of indices of the spatial dimensions along which + the image might be flipped. If they are integers, they must be in + `(0, 1, 2)`. Anatomical labels may also be used, such as + `'Left'`, `'Right'`, `'Anterior'`, `'Posterior'`, + `'Inferior'`, `'Superior'`, `'Height'` and `'Width'`, + `'AP'` (antero-posterior), `'lr'` (lateral), `'w'` (width) or + `'i'` (inferior). Only the first letter of the string will be + used. If the image is 2D, `'Height'` and `'Width'` may be + used. + flip_probability: Probability that the image will be flipped. This is + computed on a per-axis basis. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> fpg = tio.datasets.FPG() + >>> flip = tio.RandomFlip(axes=('LR',)) # flip along lateral axis only + + Tip: + It is handy to specify the axes as anatomical labels when the + image orientation is not known. + """ + + def __init__( + self, + axes: int | tuple[int, ...] = 0, + flip_probability: float = 0.5, + **kwargs, + ): + super().__init__(**kwargs) + self.axes = _parse_axes(axes) + self.flip_probability = self.parse_probability(flip_probability) + + def apply_transform(self, subject: Subject) -> Subject: + potential_axes = _ensure_axes_indices(subject, self.axes) + axes_to_flip_hot = self.get_params(self.flip_probability) + for i in range(3): + if i not in potential_axes: + axes_to_flip_hot[i] = False + (axes,) = np.where(axes_to_flip_hot) + axes_list = axes.tolist() + if not axes_list: + return subject + + transform = Flip(axes=axes_list, **self._get_base_args()) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed + + @staticmethod + def get_params(probability: float) -> list[bool]: + return (probability > torch.rand(3)).tolist() + + +class Flip(SpatialTransform): + """Reverse the order of elements in an image along the given axes. + + Args: + axes: Index or tuple of indices of the spatial dimensions along which + the image will be flipped. See + [`RandomFlip`][torchio.transforms.augmentation.spatial.random_flip.RandomFlip] + for more information. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Tip: + It is handy to specify the axes as anatomical labels when the + image orientation is not known. + """ + + def __init__(self, axes, **kwargs): + super().__init__(**kwargs) + self.axes = _parse_axes(axes) + self.args_names = ['axes'] + + def apply_transform(self, subject: Subject) -> Subject: + axes = _ensure_axes_indices(subject, self.axes) + for image in self.get_images(subject): + _flip_image(image, axes) + return subject + + def is_invertible(self): + return True + + def inverse(self): + return self + + +def _parse_axes(axes: int | tuple[int, ...]): + axes_tuple = to_tuple(axes) + for axis in axes_tuple: + is_int = isinstance(axis, int) + is_string = isinstance(axis, str) + valid_number = is_int and axis in (0, 1, 2) + if not is_string and not valid_number: + message = ( + f'All axes must be 0, 1 or 2, but found "{axis}" with type {type(axis)}' + ) + raise ValueError(message) + return axes_tuple + + +def _ensure_axes_indices(subject, axes): + if any(isinstance(n, str) for n in axes): + subject.check_consistent_orientation() + image = subject.get_first_image() + axes = sorted(3 + image.axis_name_to_index(n) for n in axes) + return axes + + +def _flip_image(image, axes): + spatial_axes = tuple(int(axis) + 1 for axis in axes) + data = image.numpy() + data = np.flip(data, axis=spatial_axes) + data = np.ascontiguousarray(data) # remove negative strides + data = torch.as_tensor(data) + image.set_data(data) diff --git a/src/torchio/transforms/compose.py b/src/torchio/transforms/compose.py deleted file mode 100644 index 68d9a97be..000000000 --- a/src/torchio/transforms/compose.py +++ /dev/null @@ -1,362 +0,0 @@ -"""Transform composition: Compose, OneOf, SomeOf.""" - -from __future__ import annotations - -import contextlib -import copy -from collections.abc import Mapping -from collections.abc import Sequence -from typing import Any -from typing import cast - -import torch - -from .transform import Transform - - -@contextlib.contextmanager -def _disabled_copy(transforms: Sequence[Transform]): - """Temporarily set ``copy=False`` on each transform. - - The composing transform copies the input once at the top level, so - the children it applies must not copy again. The original `copy` - flags are restored on exit. - - Args: - transforms: The child transforms to apply without copying. - """ - saved = [t.copy for t in transforms] - for t in transforms: - t.copy = False - try: - yield - finally: - for t, previous in zip(transforms, saved, strict=True): - t.copy = previous - - -class Compose(Transform): - """Compose several transforms together. - - The input is deep-copied once before the pipeline runs (by - default), then each transform operates in-place on the copy. - This avoids redundant copies when chaining many transforms. - - Args: - transforms: Sequence of transforms to apply sequentially, or a - mapping whose values are the transforms (keys are used as - human-readable names and ignored at runtime). - copy: If `True` (default), deep-copy the input before - applying the pipeline. Set to `False` when this - `Compose` is nested inside another `Compose`. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> preprocessing = tio.Compose([ - ... tio.Flip(axes=(0,), p=0.5), - ... tio.Noise(std=(0.01, 0.1)), - ... ]) - >>> augmented = preprocessing(subject) - >>> named = tio.Compose({ - ... "flip": tio.Flip(axes=(0,), p=0.5), - ... "noise": tio.Noise(std=(0.01, 0.1)), - ... }) - """ - - def __init__( - self, - transforms: Sequence[Transform] | Mapping[str, Transform] | None = None, - *, - copy: bool = True, - **kwargs: Any, - ) -> None: - super().__init__(copy=copy, **kwargs) - if transforms is None: - self.transforms: list[Transform] = [] - elif isinstance(transforms, Mapping): - mapping = cast(Mapping[str, Transform], transforms) - self.transforms = list(mapping.values()) - else: - self.transforms = list(transforms) - - def forward(self, data): - if self.copy: - data = copy.deepcopy(data) - subject, unwrap = self._wrap(data) - for transform in self.transforms: - old_copy = transform.copy - transform.copy = False - subject = transform(subject) - transform.copy = old_copy - return unwrap(subject) - - def to_hydra(self) -> dict[str, Any]: - cfg = super().to_hydra() - cfg["transforms"] = [t.to_hydra() for t in self.transforms] - return cfg - - -class OneOf(Transform): - """Apply one of the given transforms, chosen at random. - - When applied to a batch with `per_instance=True` (the default), - each batch element independently chooses which transform to apply. - This requires shape- and schema-preserving transforms so the - elements can be re-stacked. Pass `per_instance=False` to choose a - single transform for the whole batch. - - Args: - transforms: Sequence of transforms, or a `dict` mapping - transforms to their relative weights. If a sequence is - given, all transforms have equal probability. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> augmentation = tio.OneOf({ - ... tio.Noise(std=0.1): 0.7, - ... tio.Flip(axes=(0,)): 0.3, - ... }) - """ - - def __init__( - self, - transforms: Sequence[Transform] | dict[Transform, float], - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - if isinstance(transforms, dict): - weight_dict = cast(dict[Transform, float], transforms) - self.transforms = list(weight_dict.keys()) - w_list: list[float] = list(weight_dict.values()) - total: float = sum(w_list) - self.weights = [w / total for w in w_list] - else: - self.transforms = list(transforms) - n = len(self.transforms) - self.weights = [1.0 / n] * n - - def forward(self, data): - if self.copy: - data = copy.deepcopy(data) - batch, unwrap = self._wrap(data) - # The input is copied once above, so children apply without copying. - with _disabled_copy(self.transforms): - if self.per_instance and batch.batch_size > 1: - return unwrap(self._forward_per_element(batch)) - if torch.rand(1).item() >= self.p: - return unwrap(batch) - idx = int( - torch.multinomial( - torch.tensor(self.weights), - num_samples=1, - ).item() - ) - batch = self.transforms[idx](batch) - return unwrap(batch) - - def _forward_per_element(self, batch): - """Apply an independently chosen transform to each batch element.""" - if self.p == 0: - return batch - weights = torch.tensor(self.weights) - out_subjects = [] - any_applied = False - for subject in batch.unbatch(): - if torch.rand(1).item() < self.p: - any_applied = True - idx = int(torch.multinomial(weights, num_samples=1).item()) - subject = _apply_to_element(subject, self.transforms[idx]) - out_subjects.append(subject) - if not any_applied: - return batch - return _rebatch_with_history(out_subjects, "OneOf") - - def to_hydra(self) -> dict[str, Any]: - cfg = super().to_hydra() - cfg["transforms"] = [t.to_hydra() for t in self.transforms] - return cfg - - -class SomeOf(Transform): - """Apply a random subset of the given transforms. - - When applied to a batch with `per_instance=True` (the default), - each batch element independently samples its own subset. This - requires shape- and schema-preserving transforms so the elements - can be re-stacked. Pass `per_instance=False` to sample a single - subset for the whole batch. - - Args: - transforms: Sequence of candidate transforms. - num_transforms: How many transforms to apply. An `int` for a - fixed count, or a `(min, max)` tuple to sample the count - uniformly from that range. - replace: If `True`, sample with replacement (the same - transform may be applied more than once). - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> augmentation = tio.SomeOf( - ... [tio.Noise(), tio.Flip(), tio.Noise(std=0.5)], - ... num_transforms=2, - ... ) - """ - - def __init__( - self, - transforms: Sequence[Transform] | None = None, - *, - num_transforms: int | tuple[int, int] = 1, - replace: bool = False, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.transforms = list(transforms) if transforms else [] - self.num_transforms = num_transforms - self.replace = replace - - @property - def _min_n(self) -> int: - if isinstance(self.num_transforms, int): - return self.num_transforms - return self.num_transforms[0] - - @property - def _max_n(self) -> int: - if isinstance(self.num_transforms, int): - return self.num_transforms - return self.num_transforms[1] - - def forward(self, data): - if self.copy: - data = copy.deepcopy(data) - batch, unwrap = self._wrap(data) - # The input is copied once above, so children apply without copying. - with _disabled_copy(self.transforms): - if self.per_instance and batch.batch_size > 1: - return unwrap(self._forward_per_element(batch)) - if torch.rand(1).item() >= self.p: - return unwrap(batch) - batch = self._apply_subset(batch) - return unwrap(batch) - - def _apply_subset(self, batch): - """Apply a randomly chosen subset of transforms to *batch*.""" - n = int(torch.randint(self._min_n, self._max_n + 1, size=(1,)).item()) - n_transforms = len(self.transforms) - if self.replace: - indices = torch.randint(0, n_transforms, (n,)) - else: - n = min(n, n_transforms) - indices = torch.randperm(n_transforms)[:n] - for idx in indices: - batch = self.transforms[idx](batch) - return batch - - def _forward_per_element(self, batch): - """Apply an independently chosen subset to each batch element.""" - if self.p == 0: - return batch - out_subjects = [] - any_applied = False - for subject in batch.unbatch(): - if torch.rand(1).item() < self.p: - any_applied = True - subject = _apply_to_element(subject, self._apply_subset) - out_subjects.append(subject) - if not any_applied: - return batch - return _rebatch_with_history(out_subjects, "SomeOf") - - def to_hydra(self) -> dict[str, Any]: - cfg = super().to_hydra() - cfg["transforms"] = [t.to_hydra() for t in self.transforms] - return cfg - - -def _apply_to_element(subject: Any, apply_fn: Any) -> Any: - """Apply a transform (or callable) to one element, preserving history. - - Wrapping a `Subject` directly into a batch discards its existing - history, so the element is wrapped into a one-element batch seeded - with its prior history; the transform then appends to that history. - - Args: - subject: The single subject to transform (carrying its history). - apply_fn: A transform or callable taking and returning a - `SubjectsBatch`. - - Returns: - The transformed subject, with its full history. - """ - from ..data.batch import SubjectsBatch - - element_batch = SubjectsBatch.from_subjects([subject]) - element_batch.applied_transforms = list(subject.applied_transforms) - element_batch = apply_fn(element_batch) - return element_batch.unbatch()[0] - - -def _rebatch_with_history(subjects: list[Any], transform_name: str) -> Any: - """Re-stack per-element subjects and freeze their distinct histories. - - Args: - subjects: The transformed subjects, one per batch element. - transform_name: Name of the branching transform, used for a - clearer error message when shapes or schemas diverge. - - Returns: - A `SubjectsBatch` whose `unbatch()` restores each element's own - transform history. - """ - from ..data.batch import SubjectsBatch - - _check_consistent_schema(subjects, transform_name) - try: - batch = SubjectsBatch.from_subjects(subjects) - except (RuntimeError, KeyError) as error: - msg = ( - f"Per-instance {transform_name} produced batch elements with" - " different shapes or schemas, which cannot be re-stacked. Use" - " only shape- and schema-preserving transforms with per-instance" - f" {transform_name}, or pass per_instance=False." - ) - raise RuntimeError(msg) from error - batch.set_per_element_history([s.applied_transforms for s in subjects]) - return batch - - -def _check_consistent_schema(subjects: list[Any], transform_name: str) -> None: - """Ensure all subjects share the same image names and classes. - - Per-element branching may apply different transforms to different - elements; if those change the set of images (or their type), the - elements can no longer be re-stacked into one batch. This raises a - clear error instead of silently dropping data. - - Args: - subjects: The subjects about to be re-stacked. - transform_name: Name of the branching transform for the message. - - Raises: - RuntimeError: If image names or classes differ across subjects. - """ - if not subjects: - return - reference = {name: type(image) for name, image in subjects[0].images.items()} - for subject in subjects[1:]: - current = {name: type(image) for name, image in subject.images.items()} - if current != reference: - msg = ( - f"Per-instance {transform_name} produced batch elements with" - " different image names or types, which cannot be re-stacked." - " Use only schema-preserving transforms with per-instance" - f" {transform_name}, or pass per_instance=False." - ) - raise RuntimeError(msg) diff --git a/src/torchio/transforms/cornucopia_adapter.py b/src/torchio/transforms/cornucopia_adapter.py deleted file mode 100644 index f8b0588d7..000000000 --- a/src/torchio/transforms/cornucopia_adapter.py +++ /dev/null @@ -1,168 +0,0 @@ -"""CornucopiaAdapter: wrap Cornucopia transforms for use in TorchIO pipelines.""" - -from __future__ import annotations - -import copy as _copy -from collections.abc import Callable -from typing import Any - -import torch - -from ..data.batch import SubjectsBatch -from ..data.image import Image -from ..data.image import LabelMap -from ..data.image import ScalarImage -from ..data.subject import Subject -from .transform import Transform - - -class CornucopiaAdapter(Transform): - """Wrap a Cornucopia transform for use in TorchIO pipelines. - - `Cornucopia `_ transforms - operate on `(C, I, J, K)` tensors and support passing multiple - tensors to share spatial parameters (e.g., the same elastic - deformation is applied to an image and its segmentation). - - The adapter extracts image tensors from the subject, passes them - to the Cornucopia transform as positional arguments (scalar images - first, then label maps), and writes the results back. - - Args: - cornucopia_transform: A Cornucopia transform (any callable - accepting one or more `(C, I, J, K)` tensors). - Requires `cornucopia` to be installed: - `pip install cornucopia`. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> import cornucopia as cc # doctest: +SKIP - >>> adapter = tio.CornucopiaAdapter( - ... cc.ElasticTransform(), - ... ) # doctest: +SKIP - >>> result = adapter(subject) # doctest: +SKIP - - Note: - `CornucopiaAdapter` does **not** record itself in the - subject's transform history, because Cornucopia transform - objects are not guaranteed to be serializable. - """ - - def __init__( - self, - cornucopia_transform: Callable, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - if not callable(cornucopia_transform): - msg = ( - "cornucopia_transform must be callable, " - f"got {type(cornucopia_transform).__name__}" - ) - raise TypeError(msg) - self.cornucopia_transform = cornucopia_transform - - def forward(self, data: Any) -> Any: - """Apply without recording history.""" - batch, unwrap = self._wrap(data) - if self.copy: - batch = _copy.deepcopy(batch) - if torch.rand(1).item() > self.p: - return unwrap(batch) - subjects = batch.unbatch() - for subject in subjects: - _apply_cornucopia(subject, self.cornucopia_transform, self) - from ..data.batch import SubjectsBatch - - result = SubjectsBatch.from_subjects(subjects) - result.adopt_history(batch, subjects) - return unwrap(result) - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Not used: CornucopiaAdapter overrides forward directly.""" - return batch - - def add_transform_to_subject_history(self, *args: Any) -> None: - """No-op: Cornucopia transforms are opaque.""" - - @property - def invertible(self) -> bool: - """Cornucopia transforms are not invertible through TorchIO.""" - return False - - -def _apply_cornucopia( - subject: Subject, - cornucopia_transform: Callable, - adapter: CornucopiaAdapter, -) -> None: - """Apply a Cornucopia transform to a subject's images. - - All images are passed as positional arguments so that spatial - transforms share the same random parameters. Scalar images are - passed first, then label maps. - - Args: - subject: The subject to transform. - cornucopia_transform: The Cornucopia callable. - adapter: The adapter instance (for include/exclude filtering). - """ - images = _get_filtered_images(subject, adapter) - if not images: - return - - names = list(images.keys()) - tensors = [images[n].data for n in names] - - # Cornucopia transforms accept multiple tensors as *args - # and return the same number of tensors. - results = cornucopia_transform(*tensors) - - # If only one image, result is a single tensor (not a tuple). - if len(names) == 1: - results = (results,) - - for name, result_tensor in zip(names, results, strict=True): - if isinstance(result_tensor, torch.Tensor): - images[name].set_data(result_tensor) - - -def _filter_images( - images: dict[str, Image], - include: list[str] | None, - exclude: list[str] | None, -) -> dict[str, Image]: - """Apply include/exclude filters to an image dict.""" - if include is not None: - images = {k: v for k, v in images.items() if k in include} - if exclude is not None: - images = {k: v for k, v in images.items() if k not in exclude} - return images - - -def _get_filtered_images( - subject: Subject, - adapter: CornucopiaAdapter, -) -> dict[str, Image]: - """Get images from a subject, filtered and ordered. - - Scalar images come first so that Cornucopia transforms that - treat the first argument specially (e.g., intensity-only noise) - apply to the right image. - - Args: - subject: The subject. - adapter: The adapter (for include/exclude). - - Returns: - Ordered dict: scalar images first, then label maps. - """ - filtered = _filter_images(subject.images, adapter.include, adapter.exclude) - scalars = {k: v for k, v in filtered.items() if isinstance(v, ScalarImage)} - labels = {k: v for k, v in filtered.items() if isinstance(v, LabelMap)} - return {**scalars, **labels} diff --git a/src/torchio/transforms/data_parser.py b/src/torchio/transforms/data_parser.py new file mode 100644 index 000000000..ff5a7a27a --- /dev/null +++ b/src/torchio/transforms/data_parser.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +from collections.abc import Mapping +from collections.abc import Sequence +from typing import Generic +from typing import TypeAlias +from typing import TypeVar +from typing import cast + +import nibabel as nib +import numpy as np +import SimpleITK as sitk +import torch + +from ..data.image import Image +from ..data.image import LabelMap +from ..data.image import ScalarImage +from ..data.io import nib_to_sitk +from ..data.io import sitk_to_nib +from ..data.subject import Subject +from ..types import TypeData +from ..types import TypeImageData + +TypeTransformInput: TypeAlias = ( + Subject + | Image + | torch.Tensor + | np.ndarray + | sitk.Image + | dict[str, object] + | nib.Nifti1Image +) +ParserInput = TypeVar('ParserInput', bound=TypeTransformInput) + + +class DataParser(Generic[ParserInput]): + def __init__( + self, + data: ParserInput, + keys: Sequence[str] | None = None, + label_keys: Sequence[str] | None = None, + ): + self.data = data + self.keys = keys + self.label_keys = label_keys + self.default_image_name = 'default_image_name' + self.is_tensor = False + self.is_array = False + self.is_dict = False + self.is_image = False + self.is_sitk = False + self.is_nib = False + + def get_subject(self) -> Subject: + if isinstance(self.data, nib.Nifti1Image): + tensor = self.data.get_fdata(dtype=np.float32) + if tensor.ndim == 3: + tensor = tensor[np.newaxis] + elif tensor.ndim == 5: + tensor = tensor.transpose(3, 4, 0, 1, 2) + # Assume a unique timepoint + tensor = tensor[0] + data = ScalarImage(tensor=tensor, affine=self.data.affine) + subject = self._get_subject_from_image(data) + self.is_nib = True + elif isinstance(self.data, (np.ndarray, torch.Tensor)): + subject = self._parse_tensor(self.data) + self.is_array = isinstance(self.data, np.ndarray) + self.is_tensor = True + elif isinstance(self.data, Image): + subject = self._get_subject_from_image(self.data) + self.is_image = True + elif isinstance(self.data, Subject): + subject = self.data + elif isinstance(self.data, sitk.Image): + subject = self._get_subject_from_sitk_image(self.data) + self.is_sitk = True + elif isinstance(self.data, dict): # e.g. Eisen or MONAI dicts + if self.keys is None: + message = ( + 'If the input is a dictionary, a value for "include" must' + ' be specified when instantiating the transform. See the' + ' docs for Transform:' + ' https://docs.torchio.org/transforms/transforms.html#torchio.transforms.Transform' + ) + raise RuntimeError(message) + data_dict = dict(self.data) + subject = self._get_subject_from_dict( + data_dict, + self.keys, + self.label_keys, + ) + self.is_dict = True + else: + raise ValueError(f'Input type not recognized: {type(self.data)}') + assert isinstance(subject, Subject) + return subject + + def get_output(self, transformed: Subject) -> ParserInput: + output: object + if self.is_tensor or self.is_sitk: + image = transformed.get_image(self.default_image_name) + output = image.data + if self.is_array: + output = output.numpy() + elif self.is_sitk: + output = nib_to_sitk(image.data, image.affine) + elif self.is_image: + output = transformed.get_image(self.default_image_name) + elif self.is_dict: + output_dict: dict[str, object] = dict(transformed) + for key, value in output_dict.items(): + if isinstance(value, Image): + output_dict[key] = value.data + output = output_dict + elif self.is_nib: + image = transformed.get_image(self.default_image_name) + data = image.data + output = nib.Nifti1Image(data[0].numpy(), image.affine) + else: + output = transformed + return cast(ParserInput, output) + + def _parse_tensor(self, data: TypeData) -> Subject: + if data.ndim != 4: + message = ( + 'The input must be a 4D tensor with dimensions' + f' (channels, x, y, z) but it has shape {tuple(data.shape)}.' + ' Tips: if it is a volume, please add the channels dimension;' + ' if it is 2D, also add a dimension of size 1 for the z axis' + ) + raise ValueError(message) + return self._get_subject_from_tensor(data) + + def _get_subject_from_tensor(self, tensor: TypeImageData) -> Subject: + image = ScalarImage(tensor=tensor) + return self._get_subject_from_image(image) + + def _get_subject_from_image(self, image: Image) -> Subject: + subject = Subject({self.default_image_name: image}) + return subject + + @staticmethod + def _get_subject_from_dict( + data: Mapping[str, object], + image_keys: Sequence[str], + label_keys: Sequence[str] | None = None, + ) -> Subject: + subject_dict: dict[str, object] = {} + label_keys = [] if label_keys is None else label_keys + for key, value in data.items(): + if key in image_keys: + if not isinstance(value, (np.ndarray, torch.Tensor)): + message = ( + 'Input dictionary values selected as images must be' + f' tensors or arrays, not {type(value)}' + ) + raise TypeError(message) + class_ = LabelMap if key in label_keys else ScalarImage + value = class_(tensor=value) + subject_dict[key] = value + return Subject(subject_dict) + + def _get_subject_from_sitk_image(self, image: sitk.Image) -> Subject: + tensor, affine = sitk_to_nib(image) + scalar_image = ScalarImage(tensor=tensor, affine=affine) + return self._get_subject_from_image(scalar_image) diff --git a/src/torchio/transforms/fourier.py b/src/torchio/transforms/fourier.py new file mode 100644 index 000000000..2c52e3d6d --- /dev/null +++ b/src/torchio/transforms/fourier.py @@ -0,0 +1,34 @@ +import numpy as np +import torch + + +class FourierTransform: + @staticmethod + def fourier_transform(tensor: torch.Tensor) -> torch.Tensor: + try: + import torch.fft + + transformed = torch.fft.fftn(tensor) + fshift = torch.fft.fftshift(transformed) + return fshift + except (ModuleNotFoundError, AttributeError): + import torch + + transformed = np.fft.fftn(tensor) + fshift = np.fft.fftshift(transformed) + return torch.from_numpy(fshift) + + @staticmethod + def inv_fourier_transform(tensor: torch.Tensor) -> torch.Tensor: + try: + import torch.fft + + f_ishift = torch.fft.ifftshift(tensor) + img_back = torch.fft.ifftn(f_ishift) + return img_back + except (ModuleNotFoundError, AttributeError): + import torch + + f_ishift = np.fft.ifftshift(tensor) + img_back = np.fft.ifftn(f_ishift) + return torch.from_numpy(img_back) diff --git a/src/torchio/transforms/intensity/__init__.py b/src/torchio/transforms/intensity/__init__.py deleted file mode 100644 index b37e240d5..000000000 --- a/src/torchio/transforms/intensity/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Intensity transforms.""" diff --git a/src/torchio/transforms/intensity/bias_field.py b/src/torchio/transforms/intensity/bias_field.py deleted file mode 100644 index 73d1fc4b3..000000000 --- a/src/torchio/transforms/intensity/bias_field.py +++ /dev/null @@ -1,341 +0,0 @@ -"""BiasField: simulate MRI B1 field inhomogeneity. - -Uses the SynthSeg approach (Billot et al.): sample a small random tensor, -trilinearly upsample to image size, exponentiate, and multiply. Pure -PyTorch, GPU-native, and differentiable. -""" - -from __future__ import annotations - -from typing import Any -from typing import cast - -import torch -import torch.nn.functional as functional -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ..parameter_range import to_nonneg_range -from ..transform import IntensityTransform - - -class BiasField(IntensityTransform): - r"""Corrupt an image with a smooth multiplicative bias field. - - The bias field is generated by: - - 1. Sampling a small 3D tensor from $\mathcal{N}(0, \sigma)$. - 2. Trilinearly upsampling it to the image spatial shape. - 3. Taking the voxel-wise exponential to make it strictly positive. - 4. Multiplying the image by the resulting field. - - This follows the approach used in [SynthSeg: Segmentation of brain MRI - scans of any contrast and resolution without - retraining](https://www.sciencedirect.com/science/article/pii/S1361841523000506). - - Args: - std: Standard deviation $\sigma$ of the normal distribution - used to sample the coarse bias field. Larger values produce - stronger inhomogeneity. If two values $(a, b)$ are provided, - $\sigma \sim \mathcal{U}(a, b)$. A - `torch.distributions.Distribution` may also be passed. - scale: Ratio between the coarse field size and the image spatial - shape. Smaller values produce smoother fields. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.BiasField() - >>> transform = tio.BiasField(std=0.8) - >>> transform = tio.BiasField(std=(0.0, 1.0)) - """ - - def __init__( - self, - *, - std: float | tuple[float, float] | torch.distributions.Distribution = 0.5, - scale: float = 0.025, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.std = to_nonneg_range(std) - if scale <= 0 or scale > 1: - msg = f"scale must be in (0, 1], got {scale}" - raise ValueError(msg) - self.scale = scale - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample the bias field standard deviation (per element when batched).""" - n = self._resolve_n(batch) - if n is None: - std = self.std.sample_1d() - seed = int(torch.randint(0, 2**31, (1,)).item()) - return { - "std": std, - "seed": seed, - "scale": self.scale, - } - keep = self._keep_mask(batch, n) - std = self._mask_identity(self.std.sample_1d(n), keep, identity=0.0) - # A per-element seed makes each element's field reproducible from its - # own recorded parameters, so it inverts correctly after unbatching. - seeds = [int(torch.randint(0, 2**31, (1,)).item()) for _ in range(n)] - params = { - "std": self._serialize_param(std), - "seed": seeds, - "scale": self.scale, - } - self._tag_batched(params, batch, n, keep, ["std", "seed"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Multiply each selected image by a smooth random bias field.""" - std = params["std"] - seed = params["seed"] - scale = params["scale"] - - per_instance = self._is_per_instance_params(params) - if not per_instance and std == 0: - return batch - - for _name, img_batch in self._get_images(batch).items(): - if per_instance: - img_batch.data = _apply_bias_per_element( - img_batch.data, - std, - seed, - scale, - divide=False, - ) - else: - field = _generate_bias_field( - img_batch.data.shape, - std=std, - scale=scale, - seed=seed, - device=img_batch.data.device, - ) - img_batch.data = img_batch.data * field - - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> _BiasFieldInverse: - """Build the inverse by dividing by the same bias field.""" - return _BiasFieldInverse( - std=params["std"], - seed=params["seed"], - scale=params["scale"], - copy=False, - ) - - -class _BiasFieldInverse(IntensityTransform): - """Inverse of BiasField for history replay.""" - - def __init__( - self, - *, - std: float | list[float], - seed: int | list[int], - scale: float, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self._std = std - self._seed = seed - self._scale = scale - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Return empty params; all state is in instance attributes.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Divide by the regenerated bias field.""" - per_element = isinstance(self._std, list) - if not per_element and self._std == 0: - return batch - - for _name, img_batch in self._get_images(batch).items(): - if per_element: - img_batch.data = _apply_bias_per_element( - img_batch.data, - self._std, - cast("list[int]", self._seed), - self._scale, - divide=True, - ) - else: - field = _generate_bias_field( - img_batch.data.shape, - std=cast("float", self._std), - scale=self._scale, - seed=cast("int", self._seed), - device=img_batch.data.device, - ) - img_batch.data = img_batch.data / field - - return batch - - -def _apply_bias_per_element( - data: Tensor, - std_per_element: list[float], - seed_per_element: list[int], - scale: float, - *, - divide: bool, -) -> Tensor: - """Multiply (or divide) each batch element by its own bias field. - - Each element's field is regenerated from its own `(std, seed)`, so - the result is reproducible from per-element parameters alone (and - therefore invertible after unbatching). - - Args: - data: `(B, C, I, J, K)` tensor. - std_per_element: Per-element standard deviations. - seed_per_element: Per-element random seeds. - scale: Ratio between the coarse field and the spatial shape. - divide: If `True`, divide by the field (inverse); otherwise - multiply. - - Returns: - The corrected `(B, C, I, J, K)` tensor. - """ - identity_rows = [std == 0 for std in std_per_element] - if all(identity_rows): - return data - - coarse_field = _sample_coarse_bias_fields( - data.shape, - std_per_element=std_per_element, - seed_per_element=seed_per_element, - scale=scale, - device=data.device, - ) - field = functional.interpolate( - coarse_field, - size=list(data.shape[2:]), - mode="trilinear", - align_corners=True, - ) - field = torch.exp(field) - corrected = data / field if divide else data * field - corrected = corrected.to(data.dtype) - - if any(identity_rows): - identity_mask = torch.tensor( - identity_rows, - dtype=torch.bool, - device=data.device, - ) - corrected[identity_mask] = data[identity_mask] - - return corrected - - -def _sample_coarse_bias_fields( - shape: tuple[int, ...], - *, - std_per_element: list[float], - seed_per_element: list[int], - scale: float, - device: torch.device, -) -> Tensor: - """Sample per-element coarse bias fields before batched upsampling. - - Args: - shape: `(B, C, I, J, K)` tensor shape. - std_per_element: Per-element standard deviations. - seed_per_element: Per-element random seeds. - scale: Ratio between the coarse field and the spatial shape. - device: Device to move the stacked coarse fields to. - - Returns: - `(B, C, small_I, small_J, small_K)` tensor sampled from each - element's own seeded CPU generator. - """ - channels = shape[1] - spatial = shape[2:] - small_shape = [max(round(s * scale), 4) for s in spatial] - coarse_fields = [] - for std, seed in zip(std_per_element, seed_per_element, strict=True): - generator = torch.Generator(device="cpu") - generator.manual_seed(seed) - coarse_field = torch.normal( - mean=0.0, - std=std, - size=(1, channels, *small_shape), - generator=generator, - ) - coarse_fields.append(coarse_field) - return torch.cat(coarse_fields, dim=0).to(device) - - -def _generate_bias_field( - shape: tuple[int, ...], - *, - std: float, - scale: float, - seed: int, - device: torch.device, -) -> Tensor: - """Generate a smooth multiplicative bias field. - - Args: - shape: `(B, C, I, J, K)` tensor shape. - std: Standard deviation of the normal distribution. - scale: Ratio between the coarse field and the spatial shape. - seed: Random seed for reproducibility. - device: Device to create the field on. - - Returns: - `(B, C, I, J, K)` bias field (strictly positive). - """ - batch_size, channels = shape[0], shape[1] - spatial = shape[2:] - # Compute the coarse field shape (at least 4 per axis for interpolation). - small_shape = [max(round(s * scale), 4) for s in spatial] - - generator = torch.Generator(device="cpu") - generator.manual_seed(seed) - - # Sample on CPU for reproducibility, then move to device. - small_field = torch.normal( - mean=0.0, - std=std, - size=(batch_size, channels, *small_shape), - generator=generator, - ).to(device) - - # Trilinearly upsample to full spatial size. - field = functional.interpolate( - small_field, - size=list(spatial), - mode="trilinear", - align_corners=True, - ) - - # Exponentiate to make strictly positive. - return torch.exp(field) diff --git a/src/torchio/transforms/intensity/blur.py b/src/torchio/transforms/intensity/blur.py deleted file mode 100644 index dc85b9819..000000000 --- a/src/torchio/transforms/intensity/blur.py +++ /dev/null @@ -1,328 +0,0 @@ -"""Blur: Gaussian smoothing augmentation.""" - -from __future__ import annotations - -from typing import Any - -import numpy as np -import torch -import torch.nn.functional as functional -from einops import rearrange -from einops import repeat - -from ...data.batch import ImagesBatch -from ...data.batch import SubjectsBatch -from ..parameter_range import to_nonneg_range -from ..transform import IntensityTransform - - -class Blur(IntensityTransform): - r"""Blur an image using a Gaussian filter. - - The standard deviations $(\sigma_1, \sigma_2, \sigma_3)$ of the - Gaussian kernel along each spatial axis are independently sampled - from the given range. Sigmas are specified in mm and internally - converted to voxels using the image spacing. - - Args: - std: Standard deviation of the Gaussian kernel in mm. - A scalar $x$ means $\sigma_i = x$ for every axis - (deterministic). - A 2-tuple $(a, b)$ means - $\sigma_i \sim \mathcal{U}(a, b)$. - A 6-tuple $(a_1, b_1, a_2, b_2, a_3, b_3)$ means - $\sigma_i \sim \mathcal{U}(a_i, b_i)$ independently. - A `Choice` or `Distribution` may also be passed. - The default `std=0` is a no-op (and warns). - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Blur(std=2.0) - >>> transform = tio.Blur(std=(0, 4)) - """ - - def __init__( - self, - *, - std: float | tuple[float, float] = 0.0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.std = to_nonneg_range(std) - self._warn_if_noop(is_noop=self.std.is_constant(0.0), hint="std=(0, 2)") - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample per-axis standard deviations (per element when batched).""" - n = self._resolve_n(batch) - if n is None: - return {"std": self.std.sample()} - keep = self._keep_mask(batch, n) - std = self.std.sample(n) - if keep is not None: - std[~keep] = 0.0 - params = {"std": self._serialize_param(std)} - self._tag_batched(params, batch, n, keep, ["std"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply Gaussian smoothing to each selected image.""" - per_instance = self._is_per_instance_params(params) - for _name, img_batch in self._get_images(batch).items(): - if per_instance: - img_batch.data = _blur_per_element(img_batch, params["std"]) - else: - spacing = np.asarray(img_batch.affines[0].spacing, dtype=np.float64) - sigmas_vox = _sigmas_mm_to_voxels(params["std"], spacing) - img_batch.data = _gaussian_smooth(img_batch.data, sigmas_vox) - return batch - - -def _sigmas_mm_to_voxels( - sigmas_mm: list[float], - spacing: np.ndarray, -) -> list[float]: - """Convert per-axis sigmas from mm to voxels.""" - return [s / sp if sp > 0 else 0.0 for s, sp in zip(sigmas_mm, spacing, strict=True)] - - -def _blur_per_element( - img_batch: ImagesBatch, - sigmas_mm_per_element: list[list[float]], -) -> torch.Tensor: - """Blur each batch element with its own per-axis sigmas. - - Args: - img_batch: The image batch to blur. - sigmas_mm_per_element: One `[si, sj, sk]` (in mm) per element. - - Returns: - The blurred `(B, C, I, J, K)` tensor. - """ - data = img_batch.data - sigmas_mm = np.asarray(sigmas_mm_per_element, dtype=np.float64) - spacings = np.asarray( - [affine.spacing for affine in img_batch.affines], - dtype=np.float64, - ) - sigmas_vox = np.divide( - sigmas_mm, - spacings, - out=np.zeros_like(sigmas_mm), - where=spacings > 0, - ) - return _gaussian_smooth(data, sigmas_vox) - - -def _gaussian_smooth( - data: torch.Tensor, - sigmas: list[float] | np.ndarray, -) -> torch.Tensor: - """Apply separable Gaussian smoothing to a 5D tensor. - - Args: - data: `(B, C, I, J, K)` tensor. - sigmas: Per-axis sigma in voxels, or per-element per-axis sigmas - with shape `(B, 3)`. Zero means skip that axis. - - Returns: - Smoothed tensor. - """ - sigmas_array = np.asarray(sigmas, dtype=np.float64) - if np.all(sigmas_array <= 0): - return data - if sigmas_array.ndim == 1: - # Sigmas shared across the batch: build one kernel set and apply it - # to every element, which is much cheaper than a grouped conv. - return _gaussian_smooth_shared(data, sigmas_array) - if np.all(sigmas_array == sigmas_array[0]): - # Per-element sigmas that happen to be identical collapse to the - # shared fast path. - return _gaussian_smooth_shared(data, sigmas_array[0]) - return _gaussian_smooth_per_element(data, sigmas_array) - - -def _gaussian_smooth_shared( - data: torch.Tensor, - sigmas: np.ndarray, -) -> torch.Tensor: - """Separable Gaussian smoothing with a single shared kernel set. - - Args: - data: `(B, C, I, J, K)` tensor. - sigmas: Per-axis sigma in voxels (length 3), shared by every - batch element. - - Returns: - Smoothed tensor (input dtype preserved). - """ - if np.all(sigmas <= 0): - return data - result = data.float() - b, c = result.shape[:2] - for axis_idx in range(3): - sigma = float(sigmas[axis_idx]) - if sigma <= 0: - continue - radius = max(int(np.ceil(3 * sigma)), 1) - kernel_size = 2 * radius + 1 - x = torch.arange(kernel_size, dtype=torch.float32, device=data.device) - radius - kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2) - kernel_1d = kernel_1d / kernel_1d.sum() - - kernel_patterns = ( - "k -> 1 1 k 1 1", - "k -> 1 1 1 k 1", - "k -> 1 1 1 1 k", - ) - kernel_3d = rearrange(kernel_1d, kernel_patterns[axis_idx]) - - pad = [0] * 6 - pad_idx = 2 * (2 - axis_idx) - pad[pad_idx] = radius - pad[pad_idx + 1] = radius - - padded = functional.pad(result, pad, mode="replicate") - result = functional.conv3d( - rearrange(padded, "b c i j k -> (b c) 1 i j k"), - kernel_3d, - padding=0, - ) - result = rearrange(result, "(b c) 1 i j k -> b c i j k", b=b, c=c) - return result.to(data.dtype) - - -def _gaussian_smooth_per_element( - data: torch.Tensor, - sigmas: np.ndarray, -) -> torch.Tensor: - """Apply Gaussian smoothing with per-element sigmas. - - Args: - data: `(B, C, I, J, K)` tensor. - sigmas: Per-element per-axis sigmas in voxels with shape `(B, 3)`. - - Returns: - Smoothed tensor with the same dtype as `data`. - """ - result = data.float() - b, c = result.shape[:2] - no_blur_rows = np.all(sigmas <= 0, axis=1) - for axis_idx in range(3): - axis_sigmas = sigmas[:, axis_idx] - if np.all(axis_sigmas <= 0): - continue - kernel_3d, radius = _make_grouped_axis_kernel( - axis_sigmas, - axis_idx, - c, - data.device, - ) - - # Replicate-pad along the target axis. - pad = [0] * 6 - pad_idx = 2 * (2 - axis_idx) - pad[pad_idx] = radius - pad[pad_idx + 1] = radius - - padded = functional.pad(result, pad, mode="replicate") - result = functional.conv3d( - rearrange(padded, "b c i j k -> 1 (b c) i j k"), - kernel_3d, - groups=b * c, - padding=0, - ) - result = rearrange(result, "1 (b c) i j k -> b c i j k", b=b, c=c) - result = result.to(data.dtype) - if no_blur_rows.any(): - no_blur_mask = torch.as_tensor(no_blur_rows, device=data.device) - result[no_blur_mask] = data[no_blur_mask] - return result - - -def _make_grouped_axis_kernel( - sigmas: np.ndarray, - axis_idx: int, - channels: int, - device: torch.device, -) -> tuple[torch.Tensor, int]: - """Build one grouped 3D convolution kernel for one spatial axis. - - Args: - sigmas: Per-element sigma in voxels for the target axis. - axis_idx: Spatial axis index, from 0 to 2. - channels: Number of image channels per batch element. - device: Device on which the kernel should be allocated. - - Returns: - The `(B * C, 1, kI, kJ, kK)` grouped convolution kernel and the - maximum radius used to pad the input. - """ - radii = np.zeros_like(sigmas, dtype=np.int64) - positive = sigmas > 0 - radii[positive] = np.maximum(np.ceil(3 * sigmas[positive]).astype(np.int64), 1) - max_radius = int(radii.max()) - kernel_1d = _make_stacked_1d_kernels(sigmas, radii, max_radius, device) - if axis_idx == 0: - kernel_3d = rearrange(kernel_1d, "b i -> b 1 i 1 1") - elif axis_idx == 1: - kernel_3d = rearrange(kernel_1d, "b j -> b 1 1 j 1") - else: - kernel_3d = rearrange(kernel_1d, "b k -> b 1 1 1 k") - kernel_3d = repeat( - kernel_3d, - "b one i j k -> (b c) one i j k", - c=channels, - ) - return kernel_3d, max_radius - - -def _make_stacked_1d_kernels( - sigmas: np.ndarray, - radii: np.ndarray, - max_radius: int, - device: torch.device, -) -> torch.Tensor: - """Build centered 1D Gaussian or identity kernels. - - Args: - sigmas: Per-element sigma in voxels for one axis. - radii: Per-element kernel radii. - max_radius: Maximum radius across all elements for the axis. - device: Device on which the kernels should be allocated. - - Returns: - A `(B, 2 * max_radius + 1)` tensor of normalized 1D kernels. - """ - kernel_size = 2 * max_radius + 1 - offsets = torch.arange(kernel_size, dtype=torch.float32, device=device) - max_radius - sigmas_tensor = torch.as_tensor(sigmas, dtype=torch.float32, device=device) - radii_tensor = torch.as_tensor(radii, device=device) - offsets_row = rearrange(offsets, "k -> 1 k") - sigmas_column = rearrange(sigmas_tensor, "b -> b 1") - radii_column = rearrange(radii_tensor, "b -> b 1") - safe_sigmas = torch.where( - sigmas_column > 0, - sigmas_column, - torch.ones_like(sigmas_column), - ) - kernels = torch.exp(-0.5 * (offsets_row / safe_sigmas) ** 2) - within_radius = torch.abs(offsets_row) <= radii_column - kernels = torch.where(within_radius, kernels, torch.zeros_like(kernels)) - - delta = torch.zeros_like(kernels) - delta[:, max_radius] = 1.0 - kernels = torch.where(sigmas_column > 0, kernels, delta) - return kernels / kernels.sum(dim=1, keepdim=True) diff --git a/src/torchio/transforms/intensity/clamp.py b/src/torchio/transforms/intensity/clamp.py deleted file mode 100644 index fac72ea90..000000000 --- a/src/torchio/transforms/intensity/clamp.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Clamp: clip intensity values to a range.""" - -from __future__ import annotations - -from typing import Any - -from ...data.batch import SubjectsBatch -from ..transform import IntensityTransform - - -class Clamp(IntensityTransform): - r"""Clamp intensity values into the range $[a, b]$. - - Wraps [`torch.clamp`][torch.clamp]. - - Args: - out_min: Minimum value $a$. `None` means no lower bound. - out_max: Maximum value $b$. `None` means no upper bound. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> # CT windowing: clip to [-1000, 1000] Hounsfield units - >>> clamp = tio.Clamp(out_min=-1000, out_max=1000) - >>> # Clip negative values only - >>> clamp = tio.Clamp(out_min=0) - """ - - def __init__( - self, - *, - out_min: float | None = None, - out_max: float | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - if out_min is not None and out_max is not None and out_min > out_max: - msg = f"out_min ({out_min}) must be <= out_max ({out_max})" - raise ValueError(msg) - self.out_min = out_min - self.out_max = out_max - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {"out_min": self.out_min, "out_max": self.out_max} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Clamp each selected image.""" - out_min = params["out_min"] - out_max = params["out_max"] - for _name, img_batch in self._get_images(batch).items(): - img_batch.data = img_batch.data.clamp(min=out_min, max=out_max) - return batch diff --git a/src/torchio/transforms/intensity/gamma.py b/src/torchio/transforms/intensity/gamma.py deleted file mode 100644 index 068490d4e..000000000 --- a/src/torchio/transforms/intensity/gamma.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Gamma: randomly change contrast via power-law transform.""" - -from __future__ import annotations - -import math -from typing import Any - -import torch -from einops import rearrange -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ..parameter_range import to_range -from ..transform import IntensityTransform - - -class Gamma(IntensityTransform): - r"""Change image contrast by raising values to the power $\gamma$. - - The exponent is computed as $\gamma = e^{\beta}$, where $\beta$ is - sampled from the specified range. Positive $\beta$ increases - contrast (gamma expansion), negative $\beta$ decreases it (gamma - compression). See the - [Gamma correction](https://en.wikipedia.org/wiki/Gamma_correction) - Wikipedia entry for more information. - - Note: - Fractional exponentiation of negative values is not - well-defined for non-complex numbers. If negative values are - found in the input image $I$, the applied transform is - $\text{sign}(I) \cdot |I|^{\gamma}$ instead of the usual - $I^{\gamma}$. Use [`Normalize`][torchio.Normalize] to ensure - all values are positive if needed. - - Args: - log_gamma: Range for $\beta$ in $\gamma = e^{\beta}$. - A scalar $x$ means $\beta = x$ (deterministic). - A 2-tuple $(a, b)$ means $\beta \sim \mathcal{U}(a, b)$. - A `Choice` or `Distribution` may also be passed. - The default `log_gamma=0` is a no-op (and warns). - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Gamma(log_gamma=0.5) - >>> transform = tio.Gamma(log_gamma=(-0.3, 0.3)) - """ - - def __init__( - self, - *, - log_gamma: float | tuple[float, float] = 0.0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.log_gamma = to_range(log_gamma) - self._warn_if_noop( - is_noop=self.log_gamma.is_constant(0.0), - hint="log_gamma=(-0.3, 0.3)", - ) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample the log-gamma value (per element when per-instance).""" - n = self._resolve_n(batch) - keep = self._keep_mask(batch, n) - log_gamma = self.log_gamma.sample_1d(n) - log_gamma = self._mask_identity(log_gamma, keep, identity=0.0) - params = {"log_gamma": self._serialize_param(log_gamma)} - self._tag_batched(params, batch, n, keep, ["log_gamma"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Raise each intensity image to the power gamma.""" - log_gamma = params["log_gamma"] - for _name, img_batch in self._get_images(batch).items(): - gamma = _gamma_from_log(log_gamma, img_batch.data) - data = img_batch.data - img_batch.data = data.sign() * data.abs().pow(gamma) - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> _GammaInverse: - """Invert by applying 1/gamma.""" - return _GammaInverse(log_gamma=params["log_gamma"], copy=False) - - -def _gamma_from_log( - log_gamma: float | list[float], - data: Tensor, -) -> float | Tensor: - """Compute the gamma exponent, broadcasting over the batch if needed. - - Args: - log_gamma: A scalar (batch-shared) or a per-element list. - data: The `(B, C, I, J, K)` tensor the exponent applies to. - - Returns: - A Python float for the scalar case, or a `(B, 1, 1, 1, 1)` - tensor for the per-element case. - """ - if isinstance(log_gamma, list): - values = torch.tensor(log_gamma, dtype=torch.float32, device=data.device) - return rearrange(torch.exp(values), "b -> b 1 1 1 1") - return math.exp(log_gamma) - - -class _GammaInverse(IntensityTransform): - """Inverse of Gamma for history replay.""" - - def __init__(self, *, log_gamma: float | list[float], **kwargs: Any) -> None: - super().__init__(**kwargs) - self._log_gamma = log_gamma - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - for _name, img_batch in self._get_images(batch).items(): - gamma = _gamma_from_log(_negate_log(self._log_gamma), img_batch.data) - data = img_batch.data - img_batch.data = data.sign() * data.abs().pow(gamma) - return batch - - -def _negate_log(log_gamma: float | list[float]) -> float | list[float]: - """Negate a scalar or per-element log-gamma value.""" - if isinstance(log_gamma, list): - return [-value for value in log_gamma] - return -log_gamma diff --git a/src/torchio/transforms/intensity/ghosting.py b/src/torchio/transforms/intensity/ghosting.py deleted file mode 100644 index 72273eef2..000000000 --- a/src/torchio/transforms/intensity/ghosting.py +++ /dev/null @@ -1,277 +0,0 @@ -"""Ghosting: simulate MRI ghosting artifacts along the phase-encode axis.""" - -from __future__ import annotations - -from typing import Any - -import torch -from einops import rearrange -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ..parameter_range import to_nonneg_range -from ..transform import IntensityTransform - - -class Ghosting(IntensityTransform): - r"""Add random MRI ghosting artifacts. - - Discrete "ghost" replicas of the imaged anatomy appear along the - phase-encode direction when signal intensity varies periodically - during acquisition. Common causes include pulsatile blood flow, - cardiac motion, and respiratory motion. - (See [mriquestions.com](http://mriquestions.com/why-discrete-ghosts.html).) - - The artifact is simulated by zeroing periodic planes in k-space - along a randomly chosen axis, then restoring a fraction of the - central k-space to avoid extreme artifacts. - - Args: - num_ghosts: Number of ghost replicas. A scalar $n$ is - deterministic; a 2-tuple $(a, b)$ samples - $n \sim \mathcal{U}(a, b) \cap \mathbb{N}$. - axes: Spatial axes along which ghosts may appear. One is - chosen at random per application. - intensity: Artifact strength relative to the k-space maximum. - A scalar is deterministic; a 2-tuple $(a, b)$ means - $s \sim \mathcal{U}(a, b)$. - The default `intensity=0` is a no-op (and warns). - restore: Fraction of central k-space to restore after - zeroing. `None` restores only the single central - slice. - **kwargs: See [`Transform`][torchio.Transform]. - - Note: - Execution time does not depend on the number of ghosts. - - Examples: - >>> import torchio as tio - >>> transform = tio.Ghosting(intensity=0.8) - >>> transform = tio.Ghosting(num_ghosts=6, intensity=0.8) - """ - - def __init__( - self, - *, - num_ghosts: int | tuple[int, int] = 4, - axes: tuple[int, ...] = (0, 1, 2), - intensity: float | tuple[float, float] = 0.0, - restore: float | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.num_ghosts = to_nonneg_range(num_ghosts) - self.axes = axes - self.intensity = to_nonneg_range(intensity) - self.restore = restore - self._warn_if_noop( - is_noop=self.intensity.is_constant(0.0) or self.num_ghosts.is_constant(0.0), - hint="intensity=(0.5, 1)", - ) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample ghosting parameters (per element when batched).""" - restore = self.restore if self.restore is not None else 0.0 - n = self._resolve_n(batch) - if n is None: - num_ghosts = max(1, round(self.num_ghosts.sample_1d())) - axis = self.axes[int(torch.randint(len(self.axes), (1,)).item())] - return { - "num_ghosts": num_ghosts, - "axis": axis, - "intensity": self.intensity.sample_1d(), - "restore": restore, - } - keep = self._keep_mask(batch, n) - num_ghosts_list: list[int] = [] - axis_list: list[int] = [] - intensity_list: list[float] = [] - for batch_index in range(n): - if keep is not None and not keep[batch_index]: - num_ghosts_list.append(0) - axis_list.append(self.axes[0]) - intensity_list.append(0.0) - continue - num_ghosts_list.append(max(1, round(self.num_ghosts.sample_1d()))) - axis_list.append(self.axes[int(torch.randint(len(self.axes), (1,)).item())]) - intensity_list.append(self.intensity.sample_1d()) - params = { - "num_ghosts": num_ghosts_list, - "axis": axis_list, - "intensity": intensity_list, - "restore": restore, - } - self._tag_batched( - params, - batch, - n, - keep, - ["num_ghosts", "axis", "intensity"], - ) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Add ghosting artifacts to each selected image.""" - per_instance = self._is_per_instance_params(params) - restore = params["restore"] - for _name, img_batch in self._get_images(batch).items(): - if per_instance: - img_batch.data = _add_ghosting_per_element( - img_batch.data, - num_ghosts=params["num_ghosts"], - axis=params["axis"], - intensity=params["intensity"], - restore=restore, - ) - else: - img_batch.data = _add_ghosting( - img_batch.data, - num_ghosts=params["num_ghosts"], - axis=params["axis"], - intensity=params["intensity"], - restore=restore, - ) - return batch - - -def _add_ghosting_per_element( - data: Tensor, - *, - num_ghosts: list[int], - axis: list[int], - intensity: list[float], - restore: float, -) -> Tensor: - """Add ghosting with independent parameters for each batch element. - - Args: - data: `(B, C, I, J, K)` image tensor. - num_ghosts: Number of ghost replicas per batch element. - axis: Spatial axis per batch element. - intensity: Artifact strength per batch element. - restore: Fraction of central k-space to restore. - - Returns: - Corrupted `(B, C, I, J, K)` tensor. - """ - result = data.float() - spectrum = torch.fft.fftshift( - torch.fft.fftn(result, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ) - mask = torch.ones( - data.shape[0], - 1, - *data.shape[2:], - dtype=result.dtype, - device=data.device, - ) - active = torch.zeros(data.shape[0], dtype=torch.bool, device=data.device) - for batch_index, (ghosts, phase_axis, strength) in enumerate( - zip(num_ghosts, axis, intensity, strict=True) - ): - if not ghosts or strength == 0: - continue - active[batch_index] = True - fft_dim = phase_axis + 2 - size = result.shape[fft_dim] - line_mask = torch.ones(size, dtype=result.dtype, device=data.device) - step = max(size // ghosts, 1) - line_mask[::step] = 1 - strength - if restore > 0: - mid = size // 2 - half_restore = max(int(size * restore / 2), 1) - lo, hi = mid - half_restore, mid + half_restore - line_mask[lo:hi] = 1 - line_patterns = { - 2: "s -> 1 1 s 1 1", - 3: "s -> 1 1 1 s 1", - 4: "s -> 1 1 1 1 s", - } - mask[batch_index : batch_index + 1] = rearrange( - line_mask, - line_patterns[fft_dim], - ) - - corrupted = spectrum * mask - result = torch.fft.ifftn( - torch.fft.ifftshift(corrupted, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ).real - result = result.to(data.dtype) - active = rearrange(active, "b -> b 1 1 1 1") - return torch.where(active, result, data) - - -def _add_ghosting( - data: Tensor, - *, - num_ghosts: int, - axis: int, - intensity: float, - restore: float, -) -> Tensor: - """Add ghosting artifacts to a 5D tensor via k-space manipulation. - - Args: - data: `(B, C, I, J, K)` image tensor. - num_ghosts: Number of ghost replicas. - axis: Spatial axis (0, 1, or 2) for the phase-encode direction. - intensity: Artifact strength (0 = none, 1 = strong). - restore: Fraction of central k-space to restore. - - Returns: - Corrupted `(B, C, I, J, K)` tensor. - """ - if not num_ghosts or intensity == 0: - return data - - result = data.float() - # FFT over spatial dims only: dims 2, 3, 4. - spectrum = torch.fft.fftshift( - torch.fft.fftn(result, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ) - - # Build mask along the chosen axis. - fft_dim = axis + 2 # map spatial axis to tensor dim - size = result.shape[fft_dim] - mask = torch.ones(size, device=data.device) - step = max(size // num_ghosts, 1) - mask[::step] = 1 - intensity - - # Reshape mask for broadcasting: (1, 1, 1, 1, 1) with size at fft_dim. - mask_patterns = { - 2: "s -> 1 1 s 1 1", - 3: "s -> 1 1 1 s 1", - 4: "s -> 1 1 1 1 s", - } - mask = rearrange(mask, mask_patterns[fft_dim]) - corrupted = spectrum * mask - - # Restore the center of k-space. - if restore > 0: - mid = size // 2 - half_restore = max(int(size * restore / 2), 1) - lo, hi = mid - half_restore, mid + half_restore - slices: list[slice] = [slice(None)] * 5 - slices[fft_dim] = slice(lo, hi) - corrupted[tuple(slices)] = spectrum[tuple(slices)] - - result = torch.fft.ifftn( - torch.fft.ifftshift(corrupted, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ).real - return result.to(data.dtype) diff --git a/src/torchio/transforms/intensity/histogram_standardization.py b/src/torchio/transforms/intensity/histogram_standardization.py deleted file mode 100644 index ab47a5275..000000000 --- a/src/torchio/transforms/intensity/histogram_standardization.py +++ /dev/null @@ -1,303 +0,0 @@ -"""Histogram standardization following Nyúl and Udupa (1999). - -This module provides: - -- `compute_histogram_landmarks`: a standalone function that computes - average percentile landmarks from a set of training images. -- [`HistogramStandardization`][torchio.HistogramStandardization]: - a transform that applies piecewise-linear histogram mapping using - precomputed landmarks. -""" - -from __future__ import annotations - -from collections.abc import Callable -from collections.abc import Sequence -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ...data.image import ScalarImage -from ..transform import IntensityTransform - -DEFAULT_CUTOFF: tuple[float, float] = (0.01, 0.99) -STANDARD_RANGE: tuple[float, float] = (0.0, 100.0) - -# The v1-compatible percentile set: cutoff endpoints + deciles + quartiles. -# After dedup and sort this gives 12 values; the piecewise map uses 11 segments. -_DEFAULT_QUANTILES: tuple[float, ...] = ( - 0.01, - 0.10, - 0.20, - 0.25, - 0.30, - 0.40, - 0.50, - 0.60, - 0.70, - 0.75, - 0.80, - 0.90, - 0.99, -) - - -def compute_histogram_landmarks( - images: Sequence[ScalarImage | Path | str], - *, - quantiles: Sequence[float] | None = None, - cutoff: tuple[float, float] = DEFAULT_CUTOFF, - masking_method: Callable[[Tensor], Tensor] | None = None, -) -> Tensor: - """Compute average histogram landmarks from training images. - - Implements the training phase of - `Nyúl and Udupa (1999) `_. - The returned landmarks tensor can be passed directly to - [`HistogramStandardization`][torchio.HistogramStandardization]. - - Args: - images: Training images. Each element can be a - [`ScalarImage`][torchio.ScalarImage], a file path, or a - string path. - quantiles: Quantile positions in `[0, 1]` used as control - points. Must be sorted and include the cutoff endpoints. - `None` uses the default 13-point scheme (v1-compatible). - cutoff: Lower and upper quantile bounds for the intensity - range of interest. Defaults to `(0.01, 0.99)`. - masking_method: Optional callable that takes a 4-D tensor - `(C, I, J, K)` and returns a boolean mask of the same - shape. Only `True` voxels are used for percentile - computation. `None` uses all voxels. - - Returns: - 1-D tensor of landmark values, one per quantile. - - Examples: - >>> import torchio as tio - >>> from torchio.transforms.histogram_standardization import ( - ... compute_histogram_landmarks, - ... ) - >>> landmarks = compute_histogram_landmarks([ # doctest: +SKIP - ... tio.ScalarImage("subject_a_t1.nii"), - ... tio.ScalarImage("subject_b_t1.nii"), - ... ]) - """ - if quantiles is None: - quantiles = _build_quantiles(cutoff) - else: - quantiles = tuple(sorted(set(quantiles))) - - _validate_quantiles(quantiles, cutoff) - percentiles = [100.0 * q for q in quantiles] - - all_percentile_values: list[np.ndarray] = [] - for img_source in images: - tensor = _load_tensor(img_source) - if masking_method is not None: - mask = masking_method(tensor) - else: - mask = torch.ones_like(tensor, dtype=torch.bool) - values = tensor[mask].numpy() - pv = np.percentile(values, percentiles) - all_percentile_values.append(pv) - - database = np.vstack(all_percentile_values) - landmarks = _compute_average_mapping(database) - return torch.as_tensor(landmarks, dtype=torch.float32) - - -def _build_quantiles(cutoff: tuple[float, float]) -> tuple[float, ...]: - """Build the default quantile set from cutoff and standard positions.""" - raw = set(_DEFAULT_QUANTILES) - raw.add(cutoff[0]) - raw.add(cutoff[1]) - return tuple(sorted(raw)) - - -def _validate_quantiles( - quantiles: tuple[float, ...], - cutoff: tuple[float, float], -) -> None: - """Validate quantile array.""" - if len(quantiles) < 2: - msg = f"Need at least 2 quantiles, got {len(quantiles)}" - raise ValueError(msg) - if any(q < 0 or q > 1 for q in quantiles): - msg = "All quantiles must be in [0, 1]" - raise ValueError(msg) - if cutoff[0] not in quantiles or cutoff[1] not in quantiles: - msg = ( - f"Cutoff values {cutoff} must be included in quantiles. " - f"Got quantiles: {quantiles}" - ) - raise ValueError(msg) - - -def _load_tensor(source: ScalarImage | Path | str) -> Tensor: - """Load a 4-D tensor from various source types.""" - if isinstance(source, ScalarImage): - return source.data - return ScalarImage(source).data - - -def _compute_average_mapping(database: np.ndarray) -> np.ndarray: - """Map percentile landmarks to the standard range via linear regression. - - Args: - database: `(N, P)` array of percentile values for *N* images - and *P* quantile positions. - - Returns: - `(P,)` array of averaged landmark values in the standard range. - """ - pc_low = database[:, 0] - pc_high = database[:, -1] - s_low, s_high = STANDARD_RANGE - slopes = (s_high - s_low) / (pc_high - pc_low) - slopes = np.nan_to_num(slopes) - intercept = float(np.mean(s_low - slopes * pc_low)) - n = len(database) - mapping = slopes @ database / n + intercept - return mapping - - -class HistogramStandardization(IntensityTransform): - r"""Apply piecewise-linear histogram standardization. - - Implementation of - [Nyúl and Udupa (1999)](https://ieeexplore.ieee.org/document/836373). - - Landmarks must be precomputed using - [`compute_histogram_landmarks`][torchio.transforms.histogram_standardization.compute_histogram_landmarks] - and are passed directly to this transform. Each instance targets - **one modality**; for multi-modal subjects, compose multiple - instances with the `include` parameter: - - ```python - tio.Compose([ - tio.HistogramStandardization(t1_landmarks, include=["t1"]), - tio.HistogramStandardization(t2_landmarks, include=["t2"]), - ]) - ``` - - Args: - landmarks: 1-D tensor (or path to a `.npy` / `.pt` file) - of standard-space landmark values, as returned by - [`compute_histogram_landmarks`][torchio.transforms.histogram_standardization.compute_histogram_landmarks]. - cutoff: Lower and upper quantile bounds. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> landmarks = torch.linspace(0, 100, 13) - >>> transform = tio.HistogramStandardization(landmarks) - """ - - def __init__( - self, - landmarks: Tensor | Path | str, - *, - cutoff: tuple[float, float] = DEFAULT_CUTOFF, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.landmarks = _load_landmarks(landmarks) - self.cutoff = cutoff - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply histogram standardization to each selected image.""" - for _name, img_batch in self._get_images(batch).items(): - for i in range(img_batch.batch_size): - img_batch.data[i] = _apply_histogram_standardization( - img_batch.data[i], - self.landmarks, - self.cutoff, - ) - return batch - - -def _load_landmarks(source: Tensor | Path | str) -> Tensor: - """Load landmarks from various sources.""" - if isinstance(source, Tensor): - return source.float() - path = Path(source) - if path.suffix == ".npy": - arr = np.load(path) - return torch.as_tensor(arr, dtype=torch.float32) - if path.suffix in (".pt", ".pth"): - data = torch.load(path, weights_only=True) - if isinstance(data, Tensor): - return data.float() - msg = f"Expected a Tensor in {path}, got {type(data).__name__}" - raise TypeError(msg) - msg = f"Unsupported landmarks file extension: {path.suffix}" - raise ValueError(msg) - - -def _apply_histogram_standardization( - tensor: Tensor, - landmarks: Tensor, - cutoff: tuple[float, float], -) -> Tensor: - """Apply piecewise-linear histogram mapping to a 4-D tensor. - - Args: - tensor: `(C, I, J, K)` image tensor. - landmarks: 1-D standard-space landmarks. - cutoff: `(low, high)` quantile cutoff. - - Returns: - Standardized `(C, I, J, K)` tensor. - """ - quantiles = _build_quantiles(cutoff) - percentiles = [100.0 * q for q in quantiles] - num_landmarks = len(landmarks) - if num_landmarks != len(percentiles): - msg = ( - f"Number of landmarks ({num_landmarks}) does not match " - f"the number of quantile positions ({len(percentiles)}). " - "Ensure the same quantile scheme was used for training." - ) - raise ValueError(msg) - - data = tensor.float() - flat = data.reshape(-1) - - # Compute input percentiles. - pv = np.percentile(flat.cpu().numpy(), percentiles) - input_landmarks = torch.as_tensor(pv, dtype=torch.float32, device=data.device) - - # Build piecewise-linear mapping. - diff_landmarks = torch.diff(landmarks.to(data.device)) - diff_input = torch.diff(input_landmarks) - - # Handle flat segments (constant regions). - eps = 1e-5 - diff_input = torch.where( - diff_input.abs() < eps, - torch.tensor(float("inf"), device=data.device), - diff_input, - ) - - slopes = diff_landmarks / diff_input - intercepts = landmarks[:-1].to(data.device) - slopes * input_landmarks[:-1] - - # Digitize: find which segment each voxel falls into. - bin_edges = input_landmarks[1:-1] - bin_ids = torch.bucketize(flat, bin_edges, right=False) - - result = slopes[bin_ids] * flat + intercepts[bin_ids] - return result.reshape(data.shape) diff --git a/src/torchio/transforms/intensity/labels_to_image.py b/src/torchio/transforms/intensity/labels_to_image.py deleted file mode 100644 index a17952535..000000000 --- a/src/torchio/transforms/intensity/labels_to_image.py +++ /dev/null @@ -1,290 +0,0 @@ -"""LabelsToImage: generate a synthetic image from a label map.""" - -from __future__ import annotations - -from collections.abc import Sequence -from typing import Any - -import torch -from einops import rearrange -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ...data.image import ScalarImage -from ..parameter_range import to_range -from ..transform import Transform - - -class LabelsToImage(Transform): - r"""Generate a synthetic image from a label map. - - For each label, Gaussian-distributed tissue is created with a - sampled mean and standard deviation, weighted by the label mask. - The per-label contributions are summed to produce the output - image. - - This is the building block behind - [SynthSeg](https://github.com/BBillot/SynthSeg)-style synthesis. - For best results, compose with - [`Blur`][torchio.Blur] and - [`BiasField`][torchio.BiasField]. - - The generated image is added to the subject under the key given - by *image_key*. Existing images are **not** modified. - - Only [`LabelMap`][torchio.LabelMap] images are used as input. - - Args: - label_key: Name of the label map to use. If `None`, the - first `LabelMap` found is used. - image_key: Name for the generated `ScalarImage`. - mean: Per-label mean ranges. If `None`, each label gets a - mean sampled from *default_mean*. - std: Per-label std ranges. If `None`, each label gets a - std sampled from *default_std*. - default_mean: Fallback range for label means. - default_std: Fallback range for label stds. - ignore_background: If `True`, label 0 is left as zero. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.LabelsToImage(label_key="seg") - >>> transform = tio.LabelsToImage( - ... label_key="seg", - ... mean=[(0.8, 1.0), (0.3, 0.5)], - ... std=[(0.01, 0.05), (0.02, 0.08)], - ... ) - """ - - def __init__( - self, - label_key: str | None = None, - *, - image_key: str = "image_from_labels", - mean: Sequence[float | tuple[float, float]] | None = None, - std: Sequence[float | tuple[float, float]] | None = None, - default_mean: float | tuple[float, float] = (0.1, 0.9), - default_std: float | tuple[float, float] = (0.01, 0.1), - ignore_background: bool = False, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.label_key = label_key - self.image_key = image_key - self.mean_ranges = [to_range(m) for m in mean] if mean is not None else None - self.std_ranges = [to_range(s) for s in std] if std is not None else None - self.default_mean = to_range(default_mean) - self.default_std = to_range(default_std) - self.ignore_background = ignore_background - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample per-label mean and std values (per element when batched).""" - label_batch = self._find_label_batch(batch) - # Discover unique labels from the first sample. - unique = sorted(int(v) for v in label_batch.data[0].unique().tolist()) - n = self._resolve_n(batch) - if n is None: - means, stds = self._sample_label_values(unique) - return {"means": means, "stds": stds} - means_list: list[dict[int, float]] = [] - stds_list: list[dict[int, float]] = [] - for _ in range(n): - means, stds = self._sample_label_values(unique) - means_list.append(means) - stds_list.append(stds) - params = {"means": means_list, "stds": stds_list} - self._tag_batched(params, batch, n, None, ["means", "stds"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - def _sample_label_values( - self, - unique: list[int], - ) -> tuple[dict[int, float], dict[int, float]]: - """Sample one mean and std per label. - - Args: - unique: Sorted list of unique label values. - - Returns: - A `(means, stds)` pair of per-label dictionaries. - """ - means: dict[int, float] = {} - stds: dict[int, float] = {} - for idx, label in enumerate(unique): - if self.ignore_background and label == 0: - means[label] = 0.0 - stds[label] = 0.0 - continue - if self.mean_ranges is not None and idx < len(self.mean_ranges): - means[label] = self.mean_ranges[idx].sample_1d() - else: - means[label] = self.default_mean.sample_1d() - if self.std_ranges is not None and idx < len(self.std_ranges): - stds[label] = self.std_ranges[idx].sample_1d() - else: - stds[label] = abs(self.default_std.sample_1d()) - return means, stds - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Generate a synthetic image and add it to the batch.""" - label_batch = self._find_label_batch(batch) - if self._is_per_instance_params(params): - generated = _generate_per_element( - label_batch.data, - params["means"], - params["stds"], - ) - else: - generated = _generate_from_labels( - label_batch.data, - params["means"], - params["stds"], - ) - # Create a new image batch entry. - from ...data.batch import ImagesBatch - - new_batch = ImagesBatch( - data=generated, - affines=label_batch.affines, - image_class=ScalarImage, - ) - batch.images[self.image_key] = new_batch - return batch - - def _find_label_batch(self, batch: SubjectsBatch) -> Any: - """Find the label map batch to use.""" - if self.label_key is not None: - if self.label_key not in batch.images: - msg = ( - f"Label key '{self.label_key}' not found. " - f"Available: {list(batch.images.keys())}" - ) - raise KeyError(msg) - return batch.images[self.label_key] - # Auto-detect first LabelMap. - for _name, img_batch in batch.images.items(): - if issubclass(img_batch._image_class, LabelMap): - return img_batch - msg = "No LabelMap found in the subject" - raise KeyError(msg) - - -def _generate_per_element( - label_data: Tensor, - means_per_element: list[dict[int, float]], - stds_per_element: list[dict[int, float]], -) -> Tensor: - """Generate synthetic tissue with per-element label statistics. - - Args: - label_data: `(B, C, I, J, K)` label tensor. - means_per_element: One per-label mean dict per batch element. - stds_per_element: One per-label std dict per batch element. - - Returns: - `(B, 1, I, J, K)` synthetic image tensor. - """ - b = label_data.shape[0] - spatial = label_data.shape[2:] - result = torch.zeros(b, 1, *spatial, device=label_data.device) - - for label_val in _label_values_from(means_per_element): - means = _broadcast_values( - means_per_element, - label_val, - result, - ) - stds = _broadcast_values( - stds_per_element, - label_val, - result, - ) - if _is_all_zero(means) and _is_all_zero(stds): - continue - mask = (label_data[:, 0:1] == label_val).to(dtype=result.dtype) - tissue = torch.randn_like(result) * stds + means - result += tissue * mask - - return result - - -def _label_values_from(values_per_element: list[dict[int, float]]) -> list[int]: - """Get sorted label values represented by per-element dictionaries. - - Args: - values_per_element: One value dictionary per batch element. - - Returns: - Sorted union of labels in the dictionaries. - """ - return sorted(set().union(*(values.keys() for values in values_per_element))) - - -def _broadcast_values( - values_per_element: list[dict[int, float]], - label_val: int, - reference: Tensor, -) -> Tensor: - """Convert per-element values for one label to a broadcastable tensor. - - Args: - values_per_element: One value dictionary per batch element. - label_val: Label whose values are needed. - reference: Tensor defining device and dtype. - - Returns: - Tensor of shape `(B, 1, 1, 1, 1)`, on the same device and dtype as - `reference`, broadcastable over `(B, 1, I, J, K)`. - """ - values = [values.get(label_val, 0.0) for values in values_per_element] - tensor = torch.as_tensor( - values, - device=reference.device, - dtype=reference.dtype, - ) - return rearrange(tensor, "b -> b 1 1 1 1") - - -def _is_all_zero(values: Tensor) -> bool: - """Return whether all tensor entries are zero.""" - return torch.count_nonzero(values).item() == 0 - - -def _generate_from_labels( - label_data: Tensor, - means: dict[int, float], - stds: dict[int, float], -) -> Tensor: - """Generate Gaussian tissue for each label. - - Args: - label_data: `(B, C, I, J, K)` label tensor. - means: Per-label mean values. - stds: Per-label std values. - - Returns: - `(B, 1, I, J, K)` synthetic image tensor. - """ - b = label_data.shape[0] - spatial = label_data.shape[2:] - result = torch.zeros(b, 1, *spatial, device=label_data.device) - - for label_val, mean in means.items(): - std = stds.get(label_val, 0.0) - if mean == 0.0 and std == 0.0: - continue - mask = (label_data[:, 0:1] == label_val).float() - tissue = torch.randn_like(result) * std + mean - result += tissue * mask - - return result diff --git a/src/torchio/transforms/intensity/mask.py b/src/torchio/transforms/intensity/mask.py deleted file mode 100644 index c24b07ed4..000000000 --- a/src/torchio/transforms/intensity/mask.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Mask: zero out voxels outside a mask region.""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import Any - -import torch -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import IntensityTransform - - -class Mask(IntensityTransform): - """Set voxels outside a mask to a constant value. - - Useful for brain extraction, region-of-interest cropping, or - zeroing out background before training. - - Args: - masking_method: Defines the mask. Can be: - - - A `str`: key to a [`LabelMap`][torchio.LabelMap] in - the subject. - - A callable: receives the image tensor and returns a - boolean mask. - outside_value: Value to assign to voxels outside the mask. - labels: If using a label map, which label values to include - in the mask. `None` means all nonzero values. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> # Use a brain mask to zero out non-brain voxels - >>> transform = tio.Mask(masking_method="brain") - >>> # Use a callable mask - >>> transform = tio.Mask(masking_method=lambda x: x > 0) - >>> # Keep only specific labels - >>> transform = tio.Mask(masking_method="seg", labels=[1, 2]) - """ - - def __init__( - self, - *, - masking_method: str | Callable = "brain", - outside_value: float = 0.0, - labels: list[int] | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.masking_method = masking_method - self.outside_value = outside_value - self.labels = labels - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply the mask to each selected image.""" - mask = self._resolve_mask(batch) - for _name, img_batch in self._get_images(batch).items(): - expanded = mask.expand_as(img_batch.data) - img_batch.data = torch.where(expanded, img_batch.data, self.outside_value) - return batch - - def _resolve_mask(self, batch: SubjectsBatch) -> Tensor: - """Build a boolean mask from the masking method.""" - if callable(self.masking_method) and not isinstance(self.masking_method, str): - first_img = next(iter(self._get_images(batch).values())) - return self.masking_method(first_img.data[0]).bool() - - if isinstance(self.masking_method, str): - key = self.masking_method - if key not in batch.images: - msg = ( - f'Masking method "{key}" not found in batch images.' - f" Available: {list(batch.images.keys())}" - ) - raise KeyError(msg) - mask_batch = batch.images[key] - if not issubclass(mask_batch._image_class, LabelMap): - msg = f'Masking method "{key}" must refer to a LabelMap.' - raise TypeError(msg) - mask_data = mask_batch.data[0] - if self.labels is not None: - mask = torch.zeros_like(mask_data, dtype=torch.bool) - for label in self.labels: - mask = mask | (mask_data == label) - return mask - return mask_data.bool() - - msg = ( - f"masking_method must be a str or callable, got {type(self.masking_method)}" - ) - raise TypeError(msg) diff --git a/src/torchio/transforms/intensity/motion.py b/src/torchio/transforms/intensity/motion.py deleted file mode 100644 index 06abe6450..000000000 --- a/src/torchio/transforms/intensity/motion.py +++ /dev/null @@ -1,561 +0,0 @@ -"""Motion: simulate MRI motion artifacts via k-space corruption.""" - -from __future__ import annotations - -from functools import partial -from operator import itemgetter -from typing import Any - -import torch -import torch.nn.functional as functional -from einops import rearrange -from einops import repeat -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ..parameter_range import to_range -from ..transform import IntensityTransform - -RigidTransform = dict[str, tuple[float, float, float]] -MotionTransforms = list[RigidTransform] -PerElementMotionTransforms = list[MotionTransforms] -MotionParameters = tuple[Tensor, Tensor] - -_IDENTITY_TRANSFORM: RigidTransform = { - "degrees": (0.0, 0.0, 0.0), - "translation": (0.0, 0.0, 0.0), -} - - -class Motion(IntensityTransform): - r"""Simulate MRI motion artifacts. - - Motion during MR acquisition corrupts different segments of - k-space with different rigid-body transforms, producing - characteristic ringing and blurring. This implementation follows - [Shaw et al., 2019](http://proceedings.mlr.press/v102/shaw19a.html). - - The simulation: - - 1. Splits k-space into *num_transforms* + 1 segments along a - random axis. - 2. For each segment, applies a random rigid-body transform to - the image and fills the corresponding k-space lines from the - transformed image. - 3. Reconstructs the corrupted image via inverse FFT. - - Args: - degrees: Rotation range in degrees. A scalar $d$ means - $\theta_i \sim \mathcal{U}(-d, d)$. A 2-tuple $(a, b)$ - means $\theta_i \sim \mathcal{U}(a, b)$. - translation: Translation range in voxels, same convention as - *degrees*. The translation is applied in normalized grid - coordinates (a voxel-space approximation), not in millimeters. - num_transforms: Number of inter-segment motion events. - More transforms produce more distortion. - **kwargs: See [`Transform`][torchio.Transform]. - - Warning: - Large numbers of transforms increase execution time - significantly for 3D volumes. - - Examples: - >>> import torchio as tio - >>> transform = tio.Motion() - >>> transform = tio.Motion(degrees=15, translation=10, num_transforms=4) - """ - - def __init__( - self, - *, - degrees: float | tuple[float, float] = 10.0, - translation: float | tuple[float, float] = 10.0, - num_transforms: int = 2, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.degrees = to_range(degrees) - self.translation = to_range(translation) - if not isinstance(num_transforms, int) or num_transforms < 1: - msg = f"num_transforms must be a positive int, got {num_transforms}" - raise ValueError(msg) - self.num_transforms = num_transforms - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample motion parameters (per element when batched).""" - n = self._resolve_n(batch) - if n is None: - transforms = self._sample_transforms() - return {"transforms": transforms} - keep = self._keep_mask(batch, n) - transforms_list: list[Any] = [] - for index in range(n): - if keep is not None and not keep[index]: - transforms_list.append([]) - continue - transforms_list.append(self._sample_transforms()) - params = {"transforms": transforms_list} - self._tag_batched(params, batch, n, keep, ["transforms"]) - return params - - def _sample_transforms(self) -> MotionTransforms: - """Sample one list of rigid sub-transforms.""" - return [ - { - "degrees": self.degrees.sample(), - "translation": self.translation.sample(), - } - for _ in range(self.num_transforms) - ] - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Corrupt each selected image with simulated motion.""" - per_instance = self._is_per_instance_params(params) - for _name, img_batch in self._get_images(batch).items(): - if per_instance: - img_batch.data = _apply_motion_per_instance( - img_batch.data, - params["transforms"], - ) - else: - img_batch.data = _apply_motion( - img_batch.data, - params["transforms"], - ) - return batch - - -def _apply_motion( - data: Tensor, - motion_transforms: MotionTransforms, -) -> Tensor: - """Apply motion corruption to a 5D tensor. - - Args: - data: `(B, C, I, J, K)` image tensor. - motion_transforms: List of dicts with `degrees` and - `translation` 3-tuples. - - Returns: - Motion-corrupted `(B, C, I, J, K)` tensor. - """ - if not motion_transforms: - return data - batch_size = data.shape[0] - segment_parameters = [ - _shared_motion_parameters(transform, batch_size, data.device) - for transform in motion_transforms - ] - return _apply_motion_segments(data, segment_parameters) - - -def _apply_motion_per_instance( - data: Tensor, - motion_transforms: PerElementMotionTransforms, -) -> Tensor: - """Apply motion corruption with per-element rigid parameters. - - Args: - data: `(B, C, I, J, K)` image tensor. - motion_transforms: One transform list per batch element. Empty - lists mark gated-out elements. - - Returns: - Motion-corrupted `(B, C, I, J, K)` tensor, with inactive rows - restored exactly from the input. - """ - _check_batch_size(data, motion_transforms) - active = _active_motion_mask(motion_transforms, data.device) - if not active.any().item(): - return data - segment_parameters = [ - _per_instance_motion_parameters(motion_transforms, segment_index, data.device) - for segment_index in range(_num_motion_transforms(motion_transforms)) - ] - transformed = _apply_motion_segments(data, segment_parameters) - active = rearrange(active, "b -> b 1 1 1 1") - return torch.where(active, transformed, data) - - -def _check_batch_size( - data: Tensor, - motion_transforms: PerElementMotionTransforms, -) -> None: - """Validate that parameter lists match the batch size. - - Args: - data: `(B, C, I, J, K)` image tensor. - motion_transforms: One transform list per batch element. - - Raises: - ValueError: If the parameter count differs from the batch size. - """ - if len(motion_transforms) == data.shape[0]: - return - msg = ( - f"Expected {data.shape[0]} motion parameter lists, got {len(motion_transforms)}" - ) - raise ValueError(msg) - - -def _active_motion_mask( - motion_transforms: PerElementMotionTransforms, - device: torch.device, -) -> Tensor: - """Return a boolean mask for elements with sampled motion. - - Args: - motion_transforms: One transform list per batch element. - device: Device where the mask will be allocated. - - Returns: - Boolean `(B,)` tensor. - """ - return torch.as_tensor( - tuple(map(bool, motion_transforms)), - dtype=torch.bool, - device=device, - ) - - -def _num_motion_transforms( - motion_transforms: PerElementMotionTransforms, -) -> int: - """Return the uniform number of transforms for active elements. - - Args: - motion_transforms: One transform list per batch element. - - Returns: - Number of rigid transforms per active element. - - Raises: - ValueError: If active elements have inconsistent transform counts. - """ - lengths = set(map(len, motion_transforms)) - lengths.discard(0) - if len(lengths) <= 1: - return max(lengths, default=0) - msg = f"Expected uniform motion transform counts, got {sorted(lengths)}" - raise ValueError(msg) - - -def _shared_motion_parameters( - motion_transform: RigidTransform, - batch_size: int, - device: torch.device, -) -> MotionParameters: - """Convert a shared rigid transform into batched tensors. - - Args: - motion_transform: Shared rigid transform parameters. - batch_size: Number of batch elements. - device: Device where tensors will be allocated. - - Returns: - Batched degrees and translation tensors of shape `(B, 3)`. - """ - degrees = _repeat_parameter(motion_transform["degrees"], batch_size, device) - translation = _repeat_parameter( - motion_transform["translation"], - batch_size, - device, - ) - return degrees, translation - - -def _per_instance_motion_parameters( - motion_transforms: PerElementMotionTransforms, - segment_index: int, - device: torch.device, -) -> MotionParameters: - """Collect one segment's per-element parameters as tensors. - - Args: - motion_transforms: One transform list per batch element. - segment_index: Segment transform index. - device: Device where tensors will be allocated. - - Returns: - Batched degrees and translation tensors of shape `(B, 3)`. - """ - get_segment_transform = partial( - _segment_transform_or_identity, - segment_index=segment_index, - ) - segment_transforms = tuple(map(get_segment_transform, motion_transforms)) - degrees = torch.as_tensor( - tuple(map(itemgetter("degrees"), segment_transforms)), - dtype=torch.float32, - device=device, - ) - translation = torch.as_tensor( - tuple(map(itemgetter("translation"), segment_transforms)), - dtype=torch.float32, - device=device, - ) - return degrees, translation - - -def _segment_transform_or_identity( - transforms: MotionTransforms, - *, - segment_index: int, -) -> RigidTransform: - """Return segment parameters or identity for inactive elements. - - Args: - transforms: Rigid transform list for one batch element. - segment_index: Segment transform index. - - Returns: - The segment's rigid transform or identity parameters. - """ - if not transforms: - return _IDENTITY_TRANSFORM - return transforms[segment_index] - - -def _repeat_parameter( - parameter: tuple[float, float, float], - batch_size: int, - device: torch.device, -) -> Tensor: - """Repeat a shared 3-vector parameter across the batch. - - Args: - parameter: Shared 3-vector parameter. - batch_size: Number of batch elements. - device: Device where the result will be allocated. - - Returns: - `(B, 3)` float tensor. - """ - tensor = torch.as_tensor(parameter, dtype=torch.float32, device=device) - return repeat(tensor, "component -> batch component", batch=batch_size) - - -def _apply_motion_segments( - data: Tensor, - segment_parameters: list[MotionParameters], -) -> Tensor: - """Apply k-space segment replacements for a whole batch. - - Args: - data: `(B, C, I, J, K)` image tensor. - segment_parameters: Per-segment batched degrees and translations. - - Returns: - Motion-corrupted `(B, C, I, J, K)` tensor. - """ - result = data.float() - num_segments = len(segment_parameters) + 1 - spatial_shape = result.shape[-3:] - segment_size = spatial_shape[0] // num_segments - if segment_size == 0: - msg = ( - f"Cannot split {spatial_shape[0]} k-space slices into" - f" {num_segments} motion segments; reduce num_transforms or use a" - " larger image along the first spatial axis." - ) - raise ValueError(msg) - spectrum = torch.fft.fftn(result, dim=(-3, -2, -1)) - for segment_index, (degrees, translation) in enumerate( - segment_parameters, - start=1, - ): - moved = _apply_rigid_transform(result, degrees, translation) - moved_spectrum = torch.fft.fftn(moved, dim=(-3, -2, -1)) - start, end = _segment_bounds( - segment_index, - num_segments, - segment_size, - spatial_shape[0], - ) - spectrum[:, :, start:end] = moved_spectrum[:, :, start:end] - - reconstructed = torch.fft.ifftn(spectrum, dim=(-3, -2, -1)).real - return reconstructed.to(data.dtype) - - -def _segment_bounds( - segment_index: int, - num_segments: int, - segment_size: int, - first_spatial_size: int, -) -> tuple[int, int]: - """Return start and end indices for a k-space segment. - - Args: - segment_index: One-based segment index. - num_segments: Total number of k-space segments. - segment_size: Size of every non-final segment. - first_spatial_size: Size of the first spatial axis. - - Returns: - Start and end indices along the first spatial axis. - """ - start = segment_index * segment_size - if segment_index == num_segments - 1: - return start, first_spatial_size - return start, (segment_index + 1) * segment_size - - -def _apply_rigid_transform( - tensor: Tensor, - degrees: Tensor, - translation: Tensor, -) -> Tensor: - """Apply per-element rigid-body transforms to a 5-D tensor. - - Each batch element gets its own affine grid, shared by all channels. - - Args: - tensor: `(B, C, I, J, K)` tensor. - degrees: Euler angles in degrees, with shape `(B, 3)`. - translation: Translation in voxels, with shape `(B, 3)`. - - Returns: - Transformed `(B, C, I, J, K)` tensor. - """ - batch_size, channels, *shape = tensor.shape - theta = _affine_matrices(degrees, translation, shape) - grid = functional.affine_grid( - theta, - [batch_size, 1, shape[0], shape[1], shape[2]], - align_corners=True, - ) - input_5d = rearrange(tensor, "b c i j k -> (b c) 1 i j k").float() - grid = repeat(grid, "b i j k xyz -> (b c) i j k xyz", c=channels) - output = functional.grid_sample( - input_5d, - grid, - mode="bilinear", - padding_mode="zeros", - align_corners=True, - ) - return rearrange(output, "(b c) 1 i j k -> b c i j k", b=batch_size) - - -def _affine_matrices( - degrees: Tensor, - translation: Tensor, - spatial_shape: list[int], -) -> Tensor: - """Build batched affine matrices for `affine_grid`. - - Args: - degrees: Euler angles in degrees, with shape `(B, 3)`. - translation: Translation in voxels, with shape `(B, 3)`. - spatial_shape: Spatial tensor shape `(I, J, K)`. - - Returns: - Batched affine matrices with shape `(B, 3, 4)`. - """ - theta = torch.zeros( - degrees.shape[0], - 3, - 4, - dtype=degrees.dtype, - device=degrees.device, - ) - theta[:, :3, :3] = _rotation_matrices(degrees) - theta[:, :3, 3] = _normalized_translation(translation, spatial_shape) - return theta - - -def _normalized_translation( - translation: Tensor, - spatial_shape: list[int], -) -> Tensor: - """Normalize voxel translations to `affine_grid` coordinates. - - Args: - translation: Translation in voxels, with shape `(B, 3)`. - spatial_shape: Spatial tensor shape `(I, J, K)`. - - Returns: - Normalized translations with shape `(B, 3)`. - """ - shape = torch.as_tensor( - spatial_shape, - dtype=translation.dtype, - device=translation.device, - ) - return translation / (shape / 2) - - -def _rotation_matrices(degrees: Tensor) -> Tensor: - """Build batched Euler rotation matrices. - - Args: - degrees: Euler angles in degrees, with shape `(B, 3)`. - - Returns: - Rotation matrices with shape `(B, 3, 3)`. - """ - radians = torch.deg2rad(degrees) - rx, ry, rz = radians.unbind(dim=-1) - r_x = _axis_rotation_matrices(rx, axis=0) - r_y = _axis_rotation_matrices(ry, axis=1) - r_z = _axis_rotation_matrices(rz, axis=2) - return r_z @ r_y @ r_x - - -def _axis_rotation_matrices(angles: Tensor, *, axis: int) -> Tensor: - """Build batched rotation matrices around one axis. - - Args: - angles: Rotation angles in radians, with shape `(B,)`. - axis: Rotation axis, where 0, 1 and 2 are x, y and z. - - Returns: - Rotation matrices with shape `(B, 3, 3)`. - - Raises: - ValueError: If `axis` is not 0, 1 or 2. - """ - cos = torch.cos(angles) - sin = torch.sin(angles) - matrices = torch.zeros( - angles.shape[0], - 3, - 3, - dtype=angles.dtype, - device=angles.device, - ) - if axis == 0: - matrices[:, 0, 0] = 1 - matrices[:, 1, 1] = cos - matrices[:, 1, 2] = -sin - matrices[:, 2, 1] = sin - matrices[:, 2, 2] = cos - return matrices - if axis == 1: - matrices[:, 0, 0] = cos - matrices[:, 0, 2] = sin - matrices[:, 1, 1] = 1 - matrices[:, 2, 0] = -sin - matrices[:, 2, 2] = cos - return matrices - if axis == 2: - matrices[:, 0, 0] = cos - matrices[:, 0, 1] = -sin - matrices[:, 1, 0] = sin - matrices[:, 1, 1] = cos - matrices[:, 2, 2] = 1 - return matrices - msg = f"Expected axis to be 0, 1 or 2, got {axis}" - raise ValueError(msg) diff --git a/src/torchio/transforms/intensity/noise.py b/src/torchio/transforms/intensity/noise.py deleted file mode 100644 index f2bb4798a..000000000 --- a/src/torchio/transforms/intensity/noise.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Noise transform: add Gaussian noise to intensity images.""" - -from __future__ import annotations - -from typing import Any - -import torch -from einops import rearrange -from torch import Tensor -from torch.distributions import Distribution - -from ...data.batch import SubjectsBatch -from ..parameter_range import to_nonneg_range -from ..parameter_range import to_range -from ..transform import IntensityTransform - - -class Noise(IntensityTransform): - r"""Add Gaussian or Rician noise with random parameters. - - Add noise sampled from a normal distribution with random - parameters. When `rician=True`, the magnitude of complex - Gaussian noise is used instead, producing - [Rician-distributed](https://en.wikipedia.org/wiki/Rice_distribution) - noise typical of MRI acquisitions: - - $$I_{\text{noisy}} = \sqrt{(I + n_1)^2 + n_2^2}$$ - - where $n_1, n_2 \sim \mathcal{N}(\mu, \sigma^2)$ independently. - - Args: - mean: Mean $\mu$ of the Gaussian distribution from which the - noise is sampled. If two values $(a, b)$ are provided, - then $\mu \sim \mathcal{U}(a, b)$. - If only one value $d$ is provided, $\mu = d$ - (deterministic). A `torch.distributions.Distribution` - may also be passed for custom sampling. - std: Standard deviation $\sigma$ of the Gaussian distribution - from which the noise is sampled. If two values $(a, b)$ - are provided, then $\sigma \sim \mathcal{U}(a, b)$. - If only one value $d$ is provided, $\sigma = d$ - (deterministic). Must be non-negative. - A `torch.distributions.Distribution` may also be passed. - rician: If `True`, add Rician noise instead of Gaussian. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> # Gaussian noise (default) - >>> transform = tio.Noise(std=0.1) - >>> # Rician noise (typical for MRI) - >>> transform = tio.Noise(std=0.1, rician=True) - >>> # Random std from a uniform range - >>> transform = tio.Noise(std=(0.05, 0.2)) - >>> # Custom distribution for std - >>> from torch.distributions import LogNormal - >>> transform = tio.Noise(std=LogNormal(loc=-2, scale=0.5)) - """ - - def __init__( - self, - *, - mean: float | tuple[float, float] | Distribution = 0.0, - std: float | tuple[float, float] | Distribution = 0.25, - rician: bool = False, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.mean = to_range(mean) - self.std = to_nonneg_range(std) - self.rician = rician - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - seed = int(torch.randint(0, 2**31, (1,)).item()) - n = self._resolve_n(batch) - keep = self._keep_mask(batch, n) - # Identity for a gated-out element is zero mean and zero std. - mean = self._mask_identity(self.mean.sample_1d(n), keep, identity=0.0) - std = self._mask_identity(self.std.sample_1d(n), keep, identity=0.0) - params = { - "mean": self._serialize_param(mean), - "std": self._serialize_param(std), - "seed": seed, - "rician": self.rician, - } - self._tag_batched(params, batch, n, keep, ["mean", "std"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - mean = params["mean"] - std = params["std"] - seed = params["seed"] - rician = params.get("rician", False) - keep = params.get("_keep") - generator = torch.Generator(device="cpu") - generator.manual_seed(seed) - for _name, img_batch in self._get_images(batch).items(): - data = img_batch.data - mean_b = _broadcast(mean, data) - std_b = _broadcast(std, data) - noise = _sample_noise(data, mean_b, std_b, generator) - if rician: - noise_2 = _sample_noise(data, mean_b, std_b, generator) - transformed = torch.sqrt((data + noise) ** 2 + noise_2**2) - else: - transformed = data + noise - # The Rician map is non-linear, so zero mean/std is not a true - # no-op (it returns |data|). Restore gated-out elements explicitly. - img_batch.data = _restore_gated_out(transformed, data, keep) - return batch - - -def _restore_gated_out( - transformed: Tensor, - original: Tensor, - keep: list[bool] | None, -) -> Tensor: - """Keep the original data for elements excluded by per-element gating. - - Args: - transformed: The augmented `(B, C, I, J, K)` tensor. - original: The input `(B, C, I, J, K)` tensor. - keep: Per-element keep mask, or `None` when gating is inactive. - - Returns: - A tensor equal to *transformed* for kept elements and to - *original* for gated-out elements. - """ - if keep is None: - return transformed - keep_mask = torch.tensor(keep, dtype=torch.bool, device=transformed.device) - keep_mask = rearrange(keep_mask, "b -> b 1 1 1 1") - return torch.where(keep_mask, transformed, original) - - -def _broadcast(value: float | list[float], data: Tensor) -> float | Tensor: - """Broadcast a scalar or per-element value over a `(B, C, I, J, K)` tensor. - - Args: - value: A scalar (batch-shared) or a per-element list. - data: The tensor the value will be combined with. - - Returns: - A Python float for the scalar case, or a `(B, 1, 1, 1, 1)` - tensor for the per-element case. - """ - if isinstance(value, list): - tensor = torch.tensor(value, dtype=torch.float32, device=data.device) - return rearrange(tensor, "b -> b 1 1 1 1") - return value - - -def _sample_noise( - data: Tensor, - mean: float | Tensor, - std: float | Tensor, - generator: torch.Generator, -) -> Tensor: - """Draw `mean + std * N(0, 1)` noise shaped like *data*. - - Sampling a standard normal and scaling keeps reproducibility while - supporting per-element `mean`/`std` via broadcasting. - """ - base = torch.randn(data.shape, generator=generator).to(data.device) - return mean + std * base diff --git a/src/torchio/transforms/intensity/normalize.py b/src/torchio/transforms/intensity/normalize.py deleted file mode 100644 index fddfae64a..000000000 --- a/src/torchio/transforms/intensity/normalize.py +++ /dev/null @@ -1,369 +0,0 @@ -"""Normalize: linearly map voxel intensities to a target range.""" - -from __future__ import annotations - -import warnings -from collections.abc import Callable -from typing import Any -from typing import cast - -import torch -from torch import Tensor - -from ...data.batch import ImagesBatch -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from .._statistics import compute_quantile -from ..parameter_range import Choice -from ..parameter_range import _ParameterRange -from ..transform import IntensityTransform - -TypeParameterValue = float | tuple | Choice | torch.distributions.Distribution - - -def _to_range( - value: TypeParameterValue, -) -> _ParameterRange: - """Convert a scalar, tuple, Choice, or Distribution to a _ParameterRange.""" - if isinstance(value, (torch.distributions.Distribution, Choice)): - return _ParameterRange(value) - if isinstance(value, (int, float)): - return _ParameterRange(float(value)) - return _ParameterRange(tuple(float(v) for v in value)) - - -class Normalize(IntensityTransform): - r"""Linearly rescale voxel intensities to a target range. - - The transform clips values to an input range, then applies the - affine map: - - $$v_{\text{out}} = \frac{v - m_{\min}}{m_{\max} - m_{\min}} - \cdot (n_{\max} - n_{\min}) + n_{\min}$$ - - All six numeric parameters are independently randomizable via - scalar, `(low, high)` range, or `torch.distributions.Distribution`. - - Args: - out_min: Lower bound of the output range. - out_max: Upper bound of the output range. - in_min: Lower bound of the input range. If `None`, determined - from *percentile_low* of the (masked) input data. - in_max: Upper bound of the input range. If `None`, determined - from *percentile_high* of the (masked) input data. - percentile_low: Lower percentile for auto input range. - percentile_high: Upper percentile for auto input range. - Use `(0.5, 99.5)` for the nn-UNet convention. - masking_method: Which voxels to include when computing - percentiles. `None` uses all voxels. A `str` is - interpreted as a key to a - [`LabelMap`][torchio.LabelMap] in the subject. A callable - receives the image tensor and returns a boolean mask. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> # Rescale to [-1, 1] (default) - >>> transform = tio.Normalize() - >>> # CT windowing - >>> transform = tio.Normalize( - ... out_min=0.0, out_max=1.0, - ... in_min=-1000.0, in_max=1000.0, - ... ) - >>> # nn-UNet percentile clipping - >>> transform = tio.Normalize( - ... percentile_low=0.5, percentile_high=99.5, - ... ) - >>> # Random output range - >>> transform = tio.Normalize( - ... out_min=(-1.0, 0.0), out_max=(0.5, 1.0), - ... ) - """ - - def __init__( - self, - *, - out_min: TypeParameterValue = -1.0, - out_max: TypeParameterValue = 1.0, - in_min: TypeParameterValue | None = None, - in_max: TypeParameterValue | None = None, - percentile_low: TypeParameterValue = 0.0, - percentile_high: TypeParameterValue = 100.0, - masking_method: str | Callable[[Tensor], Tensor] | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.out_min = _to_range(out_min) - self.out_max = _to_range(out_max) - self.in_min = _to_range(in_min) if in_min is not None else None - self.in_max = _to_range(in_max) if in_max is not None else None - self.percentile_low = _to_range(percentile_low) - self.percentile_high = _to_range(percentile_high) - self.masking_method = masking_method - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample random parameters and compute the input range. - - When per-instance augmentation is active, the output range is - sampled independently per batch element; the data-driven input - range stays batch-shared. - - Returns: - Dict with `out_min`, `out_max`, and either `in_min`/`in_max` - or `in_ranges` (per image name). - """ - n = self._resolve_n(batch) - out_min = self.out_min.sample_1d(n) - out_max = self.out_max.sample_1d(n) - pct_low = self.percentile_low.sample_1d() - pct_high = self.percentile_high.sample_1d() - - params: dict[str, Any] = { - "out_min": self._serialize_param(out_min), - "out_max": self._serialize_param(out_max), - } - # If explicit in_min/in_max are given, sample them directly. - if self.in_min is not None and self.in_max is not None: - params["in_min"] = self.in_min.sample_1d() - params["in_max"] = self.in_max.sample_1d() - else: - # Otherwise, compute per-image input range from percentiles. - in_ranges: dict[str, tuple[float, float]] = {} - for name, img_batch in self._get_images(batch).items(): - mask = self._get_mask(img_batch, batch) - in_ranges[name] = _percentile_range( - img_batch.data[0], - mask, - pct_low, - pct_high, - name, - ) - params["in_ranges"] = in_ranges - - if n is not None: - self._tag_batched(params, batch, n, None, ["out_min", "out_max"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Clip and linearly rescale each selected image.""" - for name, img_batch in self._get_images(batch).items(): - if "in_min" in params: - in_min = params["in_min"] - in_max = params["in_max"] - else: - in_ranges = params.get("in_ranges", {}) - if name not in in_ranges: - continue - in_min, in_max = in_ranges[name] - - in_range = in_max - in_min - if in_range == 0: - warnings.warn( - f'Cannot rescale "{name}": input range is zero.', - RuntimeWarning, - stacklevel=2, - ) - continue - - data = img_batch.data.float() - out_min, out_range = _out_min_and_range( - params["out_min"], - params["out_max"], - data, - ) - data = data.clamp(in_min, in_max) - data = (data - in_min) / in_range * out_range + out_min - img_batch.data = data - - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> _RescaleInverse: - """Build the inverse transform from recorded parameters.""" - return _RescaleInverse( - out_min=params["out_min"], - out_max=params["out_max"], - in_min=params.get("in_min"), - in_max=params.get("in_max"), - in_ranges=params.get("in_ranges"), - copy=False, - ) - - def _get_mask( - self, - img_batch: ImagesBatch, - batch: SubjectsBatch, - ) -> Tensor | None: - """Resolve masking_method to a boolean tensor or None.""" - if self.masking_method is None: - return None - if callable(self.masking_method) and not isinstance(self.masking_method, str): - return self.masking_method(img_batch.data[0]).bool() - # String key: look up a LabelMap in the batch. - if isinstance(self.masking_method, str): - key = self.masking_method - if key not in batch.images: - msg = ( - f'Masking method "{key}" not found in batch images.' - f" Available: {list(batch.images.keys())}" - ) - raise KeyError(msg) - mask_batch = batch.images[key] - if not issubclass(mask_batch._image_class, LabelMap): - msg = f'Masking method "{key}" must refer to a LabelMap.' - raise TypeError(msg) - return mask_batch.data[0].bool() - msg = ( - "masking_method must be None, str, or callable, got" - f" {type(self.masking_method)}" - ) - raise TypeError(msg) - - -class _RescaleInverse(IntensityTransform): - """Inverse of Normalize for history replay.""" - - def __init__( - self, - *, - out_min: float | list[float], - out_max: float | list[float], - in_min: float | None, - in_max: float | None, - in_ranges: dict[str, tuple[float, float]] | None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self._out_min = out_min - self._out_max = out_max - self._in_min = in_min - self._in_max = in_max - self._in_ranges = in_ranges - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Return empty params; all state is in instance attributes.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Reverse the linear rescaling.""" - for name, img_batch in self._get_images(batch).items(): - if self._in_min is not None and self._in_max is not None: - in_min = self._in_min - in_max = self._in_max - elif self._in_ranges is not None and name in self._in_ranges: - in_min, in_max = self._in_ranges[name] - else: - continue - - in_range = in_max - in_min - if in_range == 0: - continue - - data = img_batch.data.float() - out_min, out_range = _out_min_and_range( - self._out_min, - self._out_max, - data, - ) - if isinstance(out_range, float): - if out_range == 0: - continue - data = (data - out_min) / out_range * in_range + in_min - else: - # Per-element: leave zero-range elements (out_min == out_max) - # unchanged instead of dividing by zero. - out_range_t = cast("Tensor", out_range) - out_min_t = cast("Tensor", out_min) - zero = out_range_t == 0 - safe_range = torch.where( - zero, torch.ones_like(out_range_t), out_range_t - ) - reversed_data = (data - out_min_t) / safe_range * in_range + in_min - data = torch.where(zero, data, reversed_data) - img_batch.data = data - - return batch - - -def _out_min_and_range( - out_min: float | list[float], - out_max: float | list[float], - data: Tensor, -) -> tuple[float | Tensor, float | Tensor]: - """Resolve the output min and range, broadcasting per element if needed. - - Args: - out_min: Scalar (batch-shared) or per-element output minimum. - out_max: Scalar or per-element output maximum. - data: The `(B, C, I, J, K)` tensor being rescaled. - - Returns: - A `(out_min, out_range)` pair, each a float (scalar case) or a - `(B, 1, 1, 1, 1)` tensor (per-element case). - """ - if isinstance(out_min, list): - from einops import rearrange - - min_t = torch.tensor(out_min, dtype=torch.float32, device=data.device) - max_t = torch.tensor(out_max, dtype=torch.float32, device=data.device) - min_b = rearrange(min_t, "b -> b 1 1 1 1") - max_b = rearrange(max_t, "b -> b 1 1 1 1") - return min_b, max_b - min_b - low = out_min - high = cast("float", out_max) - return low, high - low - - -def _percentile_range( - tensor: Tensor, - mask: Tensor | None, - pct_low: float, - pct_high: float, - image_name: str, -) -> tuple[float, float]: - """Compute the input range from percentiles of (masked) data. - - Args: - tensor: `(C, I, J, K)` image tensor (first sample). - mask: Optional boolean mask with compatible shape, or `None`. - pct_low: Lower percentile (0-100). - pct_high: Upper percentile (0-100). - image_name: Used in warning messages. - - Returns: - `(in_min, in_max)` tuple. - """ - values = tensor[mask.expand_as(tensor)] if mask is not None else tensor.reshape(-1) - - if values.numel() == 0: - warnings.warn( - f'Cannot compute percentiles for "{image_name}": mask is empty.' - " Using full range.", - RuntimeWarning, - stacklevel=3, - ) - values = tensor.reshape(-1) - - low = float(compute_quantile(values.float(), pct_low / 100.0).item()) - high = float(compute_quantile(values.float(), pct_high / 100.0).item()) - return low, high - - -# Backwards-compatible alias. -RescaleIntensity = Normalize diff --git a/src/torchio/transforms/intensity/pca.py b/src/torchio/transforms/intensity/pca.py deleted file mode 100644 index 65547adfa..000000000 --- a/src/torchio/transforms/intensity/pca.py +++ /dev/null @@ -1,140 +0,0 @@ -"""PCA: dimensionality reduction of multi-channel images.""" - -from __future__ import annotations - -from typing import Any - -import torch -from einops import rearrange -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ..transform import IntensityTransform - - -class PCA(IntensityTransform): - r"""Apply PCA to reduce the channel dimension. - - Reshapes a $(C, I, J, K)$ image to $(N, I \cdot J \cdot K)$, - performs PCA, and reshapes back to - $(\text{num\_components}, I, J, K)$. - - This is useful for visualizing high-dimensional feature maps - (e.g., neural network embeddings) as RGB images. - - The implementation uses [`torch.pca_lowrank`][torch.pca_lowrank], so no - external dependencies are needed. - - Args: - num_components: Number of principal components to keep. - whiten: If `True`, normalize each component to unit - variance. - normalize: If `True`, divide all components by the - standard deviation of the first component. - values_range: Linear mapping range for normalization to - $[0, 1]$. The default $(-2.3, 2.3)$ covers - $\approx 99\%$ of a standard normal distribution. - clip: If `True`, clip output to $[0, 1]$ after - normalization. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.PCA(num_components=3) - """ - - def __init__( - self, - num_components: int = 3, - *, - whiten: bool = True, - normalize: bool = True, - values_range: tuple[float, float] = (-2.3, 2.3), - clip: bool = True, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - if num_components < 1: - msg = f"num_components must be >= 1, got {num_components}" - raise ValueError(msg) - self.num_components = num_components - self.whiten = whiten - self.normalize = normalize - self.values_range = values_range - self.clip = clip - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply PCA to each selected image.""" - for _name, img_batch in self._get_images(batch).items(): - results = [] - for i in range(img_batch.batch_size): - results.append(self._pca_single(img_batch.data[i])) - img_batch.data = torch.stack(results) - return batch - - def _pca_single(self, tensor: Tensor) -> Tensor: - """Apply PCA to a single `(C, I, J, K)` tensor. - - Args: - tensor: Input with *C* channels. - - Returns: - Tensor with *num_components* channels. - """ - c, si, sj, sk = tensor.shape - if c < self.num_components: - msg = ( - f"Image has {c} channels but num_components=" - f"{self.num_components}. Need at least as many " - "channels as components." - ) - raise ValueError(msg) - - # (C, I*J*K) → (voxels, channels) - flat = rearrange(tensor.float(), "c i j k -> (i j k) c") - - # Center. - mean = flat.mean(dim=0, keepdim=True) - centered = flat - mean - - # PCA via torch.pca_lowrank. - _u, s, v = torch.pca_lowrank(centered, q=self.num_components) - # Project: (voxels, channels) @ (channels, n_comp) → (voxels, n_comp) - projected = centered @ v - - if self.whiten: - # s are singular values; variance ≈ s² / (n - 1) - n = flat.shape[0] - denom = (n - 1) ** 0.5 if n > 1 else 1.0 - std = s / denom - std = std.clamp(min=1e-8) - projected = projected / std.unsqueeze(0) - - if self.normalize and projected.shape[1] > 0: - first_std = projected[:, 0].std().clamp(min=1e-8) - projected = projected / first_std - - # Map values_range to [0, 1]. - lo, hi = self.values_range - projected = (projected - lo) / (hi - lo) - - if self.clip: - projected = projected.clamp(0, 1) - - # Reshape back: (voxels, n_comp) → (n_comp, I, J, K) - result = rearrange( - projected, - "(i j k) c -> c i j k", - i=si, - j=sj, - k=sk, - ) - return result diff --git a/src/torchio/transforms/intensity/spike.py b/src/torchio/transforms/intensity/spike.py deleted file mode 100644 index fe7ad93a5..000000000 --- a/src/torchio/transforms/intensity/spike.py +++ /dev/null @@ -1,223 +0,0 @@ -"""Spike: simulate k-space spike (herringbone) artifacts.""" - -from __future__ import annotations - -from typing import Any - -import torch -from einops import rearrange -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ..parameter_range import to_nonneg_range -from ..parameter_range import to_range -from ..transform import IntensityTransform - - -class Spike(IntensityTransform): - r"""Add random MRI spike artifacts. - - Also known as - [herringbone artifact](https://radiopaedia.org/articles/herringbone-artifact), - crisscross artifact, or corduroy artifact. Spikes in k-space - create stripes in image space. - - The artifact is simulated by adding point impulses to the Fourier - spectrum of the image. All operations use `torch.fft` and run - on GPU. - - Args: - num_spikes: Number of spikes. A scalar $n$ is deterministic; - a 2-tuple $(a, b)$ samples - $n \sim \mathcal{U}(a, b) \cap \mathbb{N}$. - intensity: Ratio between the spike amplitude and the spectrum - maximum. A scalar is deterministic; a 2-tuple $(a, b)$ - means $r \sim \mathcal{U}(a, b)$. - The default `intensity=0` is a no-op (and warns). - **kwargs: See [`Transform`][torchio.Transform]. - - Note: - Execution time does not depend on the number of spikes. - - Examples: - >>> import torchio as tio - >>> transform = tio.Spike(intensity=2.0) - >>> transform = tio.Spike(num_spikes=3, intensity=2.0) - """ - - def __init__( - self, - *, - num_spikes: int | tuple[int, int] = 1, - intensity: float | tuple[float, float] = 0.0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.num_spikes = to_nonneg_range(num_spikes) - self.intensity = to_range(intensity) - self._warn_if_noop( - is_noop=self.intensity.is_constant(0.0) or self.num_spikes.is_constant(0.0), - hint="intensity=(1, 3)", - ) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample the number and positions of spikes (per element when batched).""" - n = self._resolve_n(batch) - if n is None: - num_spikes = max(1, round(self.num_spikes.sample_1d())) - positions = torch.rand(num_spikes, 3).tolist() - intensity = self.intensity.sample_1d() - return { - "positions": positions, - "intensity": intensity, - } - keep = self._keep_mask(batch, n) - positions_list: list[list[list[float]]] = [] - intensity_list: list[float] = [] - keep_values = [True] * n if keep is None else keep.tolist() - for should_keep in keep_values: - if not should_keep: - positions_list.append([]) - intensity_list.append(0.0) - continue - num_spikes = max(1, round(self.num_spikes.sample_1d())) - positions_list.append(torch.rand(num_spikes, 3).tolist()) - intensity_list.append(self.intensity.sample_1d()) - params = { - "positions": positions_list, - "intensity": intensity_list, - } - self._tag_batched(params, batch, n, keep, ["positions", "intensity"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Add spike artifacts to each selected image.""" - per_instance = self._is_per_instance_params(params) - for _name, img_batch in self._get_images(batch).items(): - if per_instance: - img_batch.data = _add_spikes_per_instance( - img_batch.data, - params["positions"], - params["intensity"], - ) - else: - img_batch.data = _add_spikes( - img_batch.data, - params["positions"], - params["intensity"], - ) - return batch - - -def _add_spikes( - data: Tensor, - positions: list[list[float]], - intensity: float, -) -> Tensor: - """Add point spikes to the k-space of a 5D tensor. - - Args: - data: `(B, C, I, J, K)` image tensor. - positions: List of `[pi, pj, pk]` in `[0, 1)` range. - intensity: Ratio between the spike amplitude and the spectrum - maximum. - - Returns: - Corrupted `(B, C, I, J, K)` tensor. - """ - if intensity == 0 or not positions: - return data - - result = data.float() - shape = result.shape[2:] # (I, J, K) - - # FFT over spatial dims. - spectrum = torch.fft.fftshift( - torch.fft.fftn(result, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ) - # Peak per (B, C): shape (B, C, 1, 1, 1) for broadcasting. - peak = spectrum.abs().amax(dim=(-3, -2, -1), keepdim=True) - - for pos in positions: - idx = [int(p * s) % s for p, s in zip(pos, shape, strict=True)] - spectrum[:, :, idx[0], idx[1], idx[2]] += peak[..., 0, 0, 0] * intensity - - result = torch.fft.ifftn( - torch.fft.ifftshift(spectrum, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ).real - return result.to(data.dtype) - - -def _add_spikes_per_instance( - data: Tensor, - positions: list[list[list[float]]], - intensities: list[float], -) -> Tensor: - """Add independently sampled point spikes to a batched 5D tensor. - - Args: - data: `(B, C, I, J, K)` image tensor. - positions: Per-element lists of `[pi, pj, pk]` positions in `[0, 1)`. - intensities: Per-element spike amplitude ratios. - - Returns: - Corrupted `(B, C, I, J, K)` tensor, with inactive elements unchanged. - """ - active = torch.as_tensor( - [ - bool(batch_positions) and batch_intensity != 0 - for batch_positions, batch_intensity in zip( - positions, - intensities, - strict=True, - ) - ], - device=data.device, - ) - if not active.any().item(): - return data - - result = data.float() - shape = result.shape[2:] # (I, J, K) - - spectrum = torch.fft.fftshift( - torch.fft.fftn(result, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ) - peak = spectrum.abs().amax(dim=(-3, -2, -1), keepdim=True) - - for batch_index, (batch_positions, batch_intensity) in enumerate( - zip( - positions, - intensities, - strict=True, - ) - ): - if not batch_positions or batch_intensity == 0: - continue - for pos in batch_positions: - idx = [int(p * s) % s for p, s in zip(pos, shape, strict=True)] - spectrum[batch_index, :, idx[0], idx[1], idx[2]] += ( - peak[batch_index, :, 0, 0, 0] * batch_intensity - ) - - transformed = torch.fft.ifftn( - torch.fft.ifftshift(spectrum, dim=(-3, -2, -1)), - dim=(-3, -2, -1), - ).real.to(data.dtype) - active = rearrange(active, "b -> b 1 1 1 1") - return torch.where(active, transformed, data) diff --git a/src/torchio/transforms/intensity/standardize.py b/src/torchio/transforms/intensity/standardize.py deleted file mode 100644 index 6a16b0a1b..000000000 --- a/src/torchio/transforms/intensity/standardize.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Standardize: subtract mean and divide by standard deviation.""" - -from __future__ import annotations - -import warnings -from collections.abc import Callable -from typing import Any - -from torch import Tensor - -from ...data.batch import ImagesBatch -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import IntensityTransform - - -class Standardize(IntensityTransform): - r"""Subtract mean and divide by standard deviation (z-score). - - $$v_{\text{out}} = \frac{v - \mu}{\sigma}$$ - - The statistics $\mu$ and $\sigma$ are computed from the (optionally - masked) voxels and applied to the entire image. - - Args: - masking_method: Which voxels to include when computing the - mean and standard deviation. `None` uses all voxels. - A `str` is interpreted as a key to a - [`LabelMap`][torchio.LabelMap] in the subject. - A callable receives the image tensor and returns a boolean - mask. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Standardize() - >>> # Use only brain voxels for statistics - >>> transform = tio.Standardize(masking_method="brain") - >>> # Use voxels above mean - >>> transform = tio.Standardize(masking_method=lambda x: x > x.mean()) - """ - - def __init__( - self, - *, - masking_method: str | Callable[[Tensor], Tensor] | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.masking_method = masking_method - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Compute per-image mean and std from the first sample. - - Returns: - Dict mapping image names to `(mean, std)` pairs. - """ - images = self._get_images(batch) - stats: dict[str, tuple[float, float]] = {} - for name, img_batch in images.items(): - mask = _get_mask(self.masking_method, img_batch, batch) - tensor = img_batch.data[0] - values = ( - tensor[mask.expand_as(tensor)] - if mask is not None - else tensor.reshape(-1) - ) - if values.numel() == 0: - warnings.warn( - f'Mask is empty for "{name}". Using all voxels.', - RuntimeWarning, - stacklevel=2, - ) - values = tensor.reshape(-1) - mean = float(values.float().mean().item()) - std = float(values.float().std().item()) - stats[name] = (mean, std) - return {"stats": stats} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Subtract mean and divide by std for each selected image.""" - stats = params["stats"] - for name, img_batch in self._get_images(batch).items(): - if name not in stats: - continue - mean, std = stats[name] - if std == 0: - msg = ( - f'Standard deviation is zero for masked values in "{name}".' - " Cannot standardize." - ) - raise RuntimeError(msg) - img_batch.data = (img_batch.data.float() - mean) / std - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> _StandardizeInverse: - """Build the inverse using the recorded mean and std.""" - return _StandardizeInverse(stats=params["stats"], copy=False) - - -class _StandardizeInverse(IntensityTransform): - """Inverse of Standardize for history replay.""" - - def __init__( - self, - *, - stats: dict[str, tuple[float, float]], - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self._stats = stats - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Return empty params; all state is in instance attributes.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Reverse the standardization: `data * std + mean`.""" - for name, img_batch in self._get_images(batch).items(): - if name not in self._stats: - continue - mean, std = self._stats[name] - if std == 0: - continue - img_batch.data = img_batch.data.float() * std + mean - return batch - - -def _get_mask( - masking_method: str | Callable[[Tensor], Tensor] | None, - img_batch: ImagesBatch, - batch: SubjectsBatch, -) -> Tensor | None: - """Resolve a masking method to a boolean tensor or `None`. - - Args: - masking_method: `None`, a string key, or a callable. - img_batch: The image batch being processed. - batch: The full subject batch (for string-key lookup). - - Returns: - Boolean mask tensor, or `None` for no masking. - """ - if masking_method is None: - return None - if callable(masking_method) and not isinstance(masking_method, str): - return masking_method(img_batch.data[0]).bool() - if isinstance(masking_method, str): - if masking_method not in batch.images: - msg = ( - f'Masking method "{masking_method}" not found in batch images.' - f" Available: {list(batch.images.keys())}" - ) - raise KeyError(msg) - mask_batch = batch.images[masking_method] - if not issubclass(mask_batch._image_class, LabelMap): - msg = f'Masking method "{masking_method}" must refer to a LabelMap.' - raise TypeError(msg) - return mask_batch.data[0].bool() - msg = f"masking_method must be None, str, or callable, got {type(masking_method)}" - raise TypeError(msg) - - -# Backwards-compatible alias. -ZNormalization = Standardize diff --git a/src/torchio/transforms/intensity/swap.py b/src/torchio/transforms/intensity/swap.py deleted file mode 100644 index 920987a72..000000000 --- a/src/torchio/transforms/intensity/swap.py +++ /dev/null @@ -1,364 +0,0 @@ -"""Swap: randomly swap patches within an image for self-supervised learning.""" - -from __future__ import annotations - -import warnings -from typing import Any - -import torch -from einops import rearrange -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..parameter_range import to_nonneg_range -from ..transform import IntensityTransform - -Origin = tuple[int, int, int] -SwapLocation = tuple[Origin, Origin] -PatchIndices = tuple[Tensor, Tensor, Tensor, Tensor, Tensor] - - -class Swap(IntensityTransform): - r"""Randomly swap patches within an image. - - This is typically used in - [context restoration for self-supervised learning](https://www.sciencedirect.com/science/article/pii/S1361841518304699). - Pairs of same-sized patches are selected at random and their - contents are exchanged. - - Warning: - This transform is intended for **self-supervised** or - **unsupervised** workflows. Because the spatial content is - rearranged, aligned label maps become inconsistent with the - swapped image. A warning is emitted if `LabelMap` images - are present in the subject. - - Args: - patch_size: Spatial size of the patches to swap. A single - integer $n$ means $(n, n, n)$. - num_iterations: Number of patch pairs to swap. A 2-tuple - $(a, b)$ samples $n \sim \mathcal{U}(a, b)$. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Swap(patch_size=15, num_iterations=100) - """ - - def __init__( - self, - *, - patch_size: int | tuple[int, int, int] = 15, - num_iterations: int | tuple[int, int] = 100, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - if isinstance(patch_size, int): - patch_size = (patch_size, patch_size, patch_size) - self.patch_size = patch_size - self.num_iterations = to_nonneg_range(num_iterations) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample swap locations (per element when batched).""" - # Warn if label maps are present. - for _name, img_batch in batch.images.items(): - if issubclass(img_batch._image_class, LabelMap): - warnings.warn( - "Swap is applied to a subject containing LabelMap " - "images. The spatial rearrangement will make labels " - "inconsistent with the swapped image. This transform " - "is intended for self-supervised learning.", - stacklevel=2, - ) - break - - any_img = next(iter(batch.images.values())) - spatial_shape = any_img.data.shape[2:] # (I, J, K) - - n = self._resolve_n(batch) - if n is None: - iterations = max(1, round(self.num_iterations.sample_1d())) - locations = _sample_swap_locations( - spatial_shape, - self.patch_size, - iterations, - ) - return {"locations": locations} - - keep = self._keep_mask(batch, n) - locations_list: list[Any] = [] - for index in range(n): - if keep is not None and not keep[index]: - locations_list.append([]) - continue - iterations = max(1, round(self.num_iterations.sample_1d())) - locations_list.append( - _sample_swap_locations(spatial_shape, self.patch_size, iterations) - ) - params = {"locations": locations_list} - self._tag_batched(params, batch, n, keep, ["locations"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Swap patches in each selected image.""" - per_instance = self._is_per_instance_params(params) - for _name, img_batch in self._get_images(batch).items(): - if per_instance: - img_batch.data = _apply_swaps_per_instance( - img_batch.data, - params["locations"], - self.patch_size, - ) - else: - img_batch.data = _apply_swaps( - img_batch.data, - params["locations"], - self.patch_size, - ) - return batch - - -def _sample_swap_locations( - spatial_shape: tuple[int, ...], - patch_size: tuple[int, int, int], - num_iterations: int, -) -> list[SwapLocation]: - """Sample pairs of non-overlapping patch origins. - - Args: - spatial_shape: `(I, J, K)` spatial dimensions. - patch_size: `(pi, pj, pk)` patch dimensions. - num_iterations: Number of pairs to sample. - - Returns: - List of `(origin_a, origin_b)` tuples. - """ - locations: list[SwapLocation] = [] - max_ini = [s - p for s, p in zip(spatial_shape, patch_size, strict=True)] - if any(m < 0 for m in max_ini): - msg = ( - f"Patch size {patch_size} cannot be larger than " - f"spatial shape {tuple(spatial_shape)}" - ) - raise ValueError(msg) - - for _ in range(num_iterations): - first = _random_origin(max_ini) - # Resample second until non-overlapping with first. - for _ in range(100): - second = _random_origin(max_ini) - if not _patches_overlap(first, second, patch_size): - break - locations.append((first, second)) - - return locations - - -def _random_origin( - max_ini: list[int], -) -> Origin: - """Sample a random patch origin.""" - coords = [] - for m in max_ini: - if m == 0: - coords.append(0) - else: - coords.append(int(torch.randint(m + 1, (1,)).item())) - return (coords[0], coords[1], coords[2]) - - -def _patches_overlap( - a: Origin, - b: Origin, - patch_size: tuple[int, int, int], -) -> bool: - """Check whether two axis-aligned patches overlap.""" - for ai, bi, p in zip(a, b, patch_size, strict=True): - if ai + p <= bi or bi + p <= ai: - return False - return True - - -def _apply_swaps( - data: Tensor, - locations: list[SwapLocation], - patch_size: tuple[int, int, int], -) -> Tensor: - """Swap patch pairs in a 5D tensor. - - Args: - data: `(B, C, I, J, K)` tensor. - locations: List of `(origin_a, origin_b)` pairs. - patch_size: `(pi, pj, pk)` patch dimensions. - - Returns: - Tensor with patches swapped. - """ - result = data.clone() - pi, pj, pk = patch_size - - for (ai, aj, ak), (bi, bj, bk) in locations: - patch_a = result[:, :, ai : ai + pi, aj : aj + pj, ak : ak + pk].clone() - patch_b = result[:, :, bi : bi + pi, bj : bj + pj, bk : bk + pk].clone() - result[:, :, ai : ai + pi, aj : aj + pj, ak : ak + pk] = patch_b - result[:, :, bi : bi + pi, bj : bj + pj, bk : bk + pk] = patch_a - - return result - - -def _apply_swaps_per_instance( - data: Tensor, - locations: list[list[SwapLocation]], - patch_size: tuple[int, int, int], -) -> Tensor: - """Swap per-element patch pairs in a 5D tensor. - - Args: - data: `(B, C, I, J, K)` tensor. - locations: One list of `(origin_a, origin_b)` pairs per batch element. - patch_size: `(pi, pj, pk)` patch dimensions. - - Returns: - Tensor with each element's patches swapped. - """ - result = data.clone() - num_swaps = max( - (len(element_locations) for element_locations in locations), default=0 - ) - if num_swaps == 0: - return result - - origins_a, origins_b = _get_batched_origins( - locations, - num_swaps, - data.device, - ) - patch_indices = _make_patch_indices(data, patch_size) - for swap_index in range(num_swaps): - _swap_batched_patches( - result, - origins_a[:, swap_index], - origins_b[:, swap_index], - patch_indices, - ) - - return result - - -def _get_batched_origins( - locations: list[list[SwapLocation]], - num_swaps: int, - device: torch.device, -) -> tuple[Tensor, Tensor]: - """Build origin tensors for batched indexed swapping. - - Args: - locations: One list of `(origin_a, origin_b)` pairs per batch element. - num_swaps: Number of sequential swap steps to encode. - device: Device on which the index tensors are created. - - Returns: - Two tensors of shape `(B, num_swaps, 3)` for the first and second patch - origins. - """ - batch_size = len(locations) - origins_a = torch.zeros(batch_size, num_swaps, 3, dtype=torch.long, device=device) - origins_b = torch.zeros_like(origins_a) - for batch_index, element_locations in enumerate(locations): - for swap_index, (origin_a, origin_b) in enumerate(element_locations): - origins_a[batch_index, swap_index] = torch.as_tensor( - origin_a, - dtype=torch.long, - device=device, - ) - origins_b[batch_index, swap_index] = torch.as_tensor( - origin_b, - dtype=torch.long, - device=device, - ) - return origins_a, origins_b - - -def _make_patch_indices( - data: Tensor, - patch_size: tuple[int, int, int], -) -> PatchIndices: - """Create shared batch, channel, and patch-offset index tensors. - - Args: - data: `(B, C, I, J, K)` tensor. - patch_size: `(pi, pj, pk)` patch dimensions. - - Returns: - Index tensors that broadcast to `(B, C, pi, pj, pk)`. - """ - batch_size, channels = data.shape[:2] - pi, pj, pk = patch_size - device = data.device - batch_index = rearrange( - torch.arange(batch_size, device=device), - "b -> b 1 1 1 1", - ) - channel_index = rearrange( - torch.arange(channels, device=device), - "c -> 1 c 1 1 1", - ) - i_offsets = rearrange(torch.arange(pi, device=device), "i -> 1 1 i 1 1") - j_offsets = rearrange(torch.arange(pj, device=device), "j -> 1 1 1 j 1") - k_offsets = rearrange(torch.arange(pk, device=device), "k -> 1 1 1 1 k") - return batch_index, channel_index, i_offsets, j_offsets, k_offsets - - -def _swap_batched_patches( - data: Tensor, - origins_a: Tensor, - origins_b: Tensor, - patch_indices: PatchIndices, -) -> None: - """Swap one patch pair per batch element using batched indexing. - - Args: - data: `(B, C, I, J, K)` tensor to update in place. - origins_a: Tensor of shape `(B, 3)` with first patch origins. - origins_b: Tensor of shape `(B, 3)` with second patch origins. - patch_indices: Broadcastable batch, channel, and offset indices. - """ - indices_a = _get_patch_indices(origins_a, patch_indices) - indices_b = _get_patch_indices(origins_b, patch_indices) - patch_a = data[indices_a].clone() - patch_b = data[indices_b].clone() - data[indices_a] = patch_b - data[indices_b] = patch_a - - -def _get_patch_indices( - origins: Tensor, - patch_indices: PatchIndices, -) -> PatchIndices: - """Build full tensor indices for per-element patch origins. - - Args: - origins: Tensor of shape `(B, 3)` with per-element patch origins. - patch_indices: Broadcastable batch, channel, and offset indices. - - Returns: - Index tensors that select one patch per batch element. - """ - batch_index, channel_index, i_offsets, j_offsets, k_offsets = patch_indices - i_index = rearrange(origins[:, 0], "b -> b 1 1 1 1") + i_offsets - j_index = rearrange(origins[:, 1], "b -> b 1 1 1 1") + j_offsets - k_index = rearrange(origins[:, 2], "b -> b 1 1 1 1") + k_offsets - return batch_index, channel_index, i_index, j_index, k_index diff --git a/src/torchio/transforms/intensity_transform.py b/src/torchio/transforms/intensity_transform.py new file mode 100644 index 000000000..7786c3401 --- /dev/null +++ b/src/torchio/transforms/intensity_transform.py @@ -0,0 +1,50 @@ +from collections.abc import Mapping +from typing import TypeVar +from typing import cast + +from ..data.image import ScalarImage +from ..data.subject import Subject +from .transform import Transform + +ValueT = TypeVar('ValueT') + + +class IntensityTransform(Transform): + """Transform that modifies voxel intensities only.""" + + def get_images_dict(self, subject: Subject) -> dict[str, ScalarImage]: + return subject.get_images_dict( + intensity_only=True, + include=self.include, + exclude=self.exclude, + ) + + def get_images(self, subject: Subject) -> list[ScalarImage]: + return subject.get_images( + intensity_only=True, + include=self.include, + exclude=self.exclude, + ) + + @staticmethod + def get_parameter(value: ValueT | Mapping[str, ValueT], name: str) -> ValueT: + if isinstance(value, Mapping): + mapping = cast(Mapping[str, ValueT], value) + return mapping[name] + return value + + def arguments_are_dict(self) -> bool: + """Check if main arguments are dict. + + Return `True` if the type of all attributes specified in the + `args_names` have `dict` type. + """ + args = list(self._get_named_arguments().values()) + are_dict = [isinstance(arg, dict) for arg in args] + if all(are_dict): + return True + elif not any(are_dict): + return False + else: + message = 'Either all or none of the arguments must be dicts' + raise ValueError(message) diff --git a/src/torchio/transforms/interpolation.py b/src/torchio/transforms/interpolation.py new file mode 100644 index 000000000..2060e2bf9 --- /dev/null +++ b/src/torchio/transforms/interpolation.py @@ -0,0 +1,60 @@ +import enum + +import SimpleITK as sitk + + +class Interpolation(enum.Enum): + """Interpolation techniques available in ITK. + + For a full quantitative comparison of interpolation methods, you can read + [Meijering et al. 1999, Quantitative Comparison of Sinc-Approximating Kernels for + Medical Image Interpolation ](https://link.springer.com/chapter/10.1007/10704282_23). + + Examples: + >>> import torchio as tio + >>> transform = tio.RandomAffine(image_interpolation='bspline') + """ + + NEAREST = 'sitkNearestNeighbor' + """Nearest neighbor interpolation.""" + + LINEAR = 'sitkLinear' + """Linear interpolation.""" + + BSPLINE = 'sitkBSpline' + """B-Spline of order 3 (cubic) interpolation.""" + + CUBIC = 'sitkBSpline' + """Same as `BSPLINE`.""" + + GAUSSIAN = 'sitkGaussian' + """Gaussian interpolation. Sigma is set to 0.8 input pixels and alpha is 4.""" + + LABEL_GAUSSIAN = 'sitkLabelGaussian' + """Smoothly interpolate multi-label images. Sigma is set to 1 input pixel and alpha is 1.""" + + HAMMING = 'sitkHammingWindowedSinc' + """Hamming windowed sinc kernel.""" + + COSINE = 'sitkCosineWindowedSinc' + """Cosine windowed sinc kernel.""" + + WELCH = 'sitkWelchWindowedSinc' + """Welch windowed sinc kernel.""" + + LANCZOS = 'sitkLanczosWindowedSinc' + """Lanczos windowed sinc kernel.""" + + BLACKMAN = 'sitkBlackmanWindowedSinc' + """Blackman windowed sinc kernel.""" + + +def get_sitk_interpolator(interpolation: str) -> int: + if not isinstance(interpolation, str): + message = ( + f'Interpolation must be a string, not "{interpolation}"' + f' of type {type(interpolation)}' + ) + raise ValueError(message) + string = getattr(Interpolation, interpolation.upper()).value + return getattr(sitk, string) diff --git a/src/torchio/transforms/inverse.py b/src/torchio/transforms/inverse.py deleted file mode 100644 index 56bff7ba2..000000000 --- a/src/torchio/transforms/inverse.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Inverse transform utilities.""" - -from __future__ import annotations - -import warnings -from typing import Any - -from .compose import Compose -from .transform import _TRANSFORM_REGISTRY -from .transform import AppliedTransform -from .transform import IntensityTransform -from .transform import Transform - - -def get_inverse_transform( - history: list[AppliedTransform], - *, - warn: bool = True, - ignore_intensity: bool = False, -) -> Compose: - """Build a Compose that inverts a list of applied transforms. - - Walks the history in reverse order. For each invertible transform, - calls `transform.inverse(params)` to get the inverse transform - instance. Non-invertible transforms are skipped. - - Args: - history: List of `AppliedTransform` records. - warn: Issue a warning for non-invertible transforms. - ignore_intensity: Skip all intensity transforms. - - Returns: - A `Compose` of inverse transforms. - """ - steps: list[Transform] = [] - for trace in reversed(history): - cls = _TRANSFORM_REGISTRY.get(trace.name) - if cls is None: - if warn: - warnings.warn( - f"Unknown transform {trace.name!r} in history, skipping", - stacklevel=2, - ) - continue - if ignore_intensity and issubclass(cls, IntensityTransform): - continue - instance = object.__new__(cls) - if not instance.invertible: - if warn: - warnings.warn( - f"{trace.name} is not invertible, skipping", - stacklevel=2, - ) - continue - inverse = instance.inverse(trace.params) - inverse.include = trace.include - inverse.exclude = trace.exclude - steps.append(inverse) - # copy=True (the default) so applying the inverse does not mutate the - # caller's data in place, matching every other transform. - return Compose(steps) - - -def apply_inverse_transform( - data: Any, - *, - warn: bool = True, - ignore_intensity: bool = False, -) -> Any: - """Apply the inverse of all recorded transforms on the data. - - Works on any type that has `applied_transforms` (Subject, - SubjectsBatch, Image, etc.). Non-invertible transforms are skipped. - - Args: - data: Transformed data with an `applied_transforms` attribute. - warn: Issue a warning for non-invertible transforms. - ignore_intensity: Skip all intensity transforms. - - Returns: - Data with transforms undone, same type as input. - """ - if not hasattr(data, "applied_transforms"): - return data - # Batches with per-element histories (from per-instance OneOf/SomeOf) - # know how to invert each element; delegate to their own method. - if getattr(data, "_per_element_history", None) is not None: - return data.apply_inverse_transform( - warn=warn, - ignore_intensity=ignore_intensity, - ) - inverse = get_inverse_transform( - data.applied_transforms, - warn=warn, - ignore_intensity=ignore_intensity, - ) - result = inverse(data) - if hasattr(result, "applied_transforms"): - result.applied_transforms = [] - return result diff --git a/src/torchio/transforms/label/__init__.py b/src/torchio/transforms/label/__init__.py deleted file mode 100644 index f9e172668..000000000 --- a/src/torchio/transforms/label/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Label transforms.""" diff --git a/src/torchio/transforms/label/contour.py b/src/torchio/transforms/label/contour.py deleted file mode 100644 index 0991329c5..000000000 --- a/src/torchio/transforms/label/contour.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Contour: extract label boundaries.""" - -from __future__ import annotations - -from typing import Any - -import torch -import torch.nn.functional as functional - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import Transform - - -class Contour(Transform): - r"""Keep only the boundary voxels of each label. - - A voxel is on the boundary if any of its 6-connected neighbors - has a different value. The result is a binary mask of the - contours. - - Only [`LabelMap`][torchio.LabelMap] images are affected. - - Args: - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Contour() - """ - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Replace each label map with its boundary voxels.""" - for _name, img_batch in batch.images.items(): - if not issubclass(img_batch._image_class, LabelMap): - continue - img_batch.data = _extract_contour(img_batch.data) - return batch - - -def _extract_contour(data: torch.Tensor) -> torch.Tensor: - """Extract boundaries using 3D erosion via max-pooling. - - A voxel is interior if all its 6-connected neighbors have the - same value as itself. We detect this by comparing the original - with a morphological erosion (min of neighbors). - - Args: - data: `(B, C, I, J, K)` label tensor. - - Returns: - Binary `(B, C, I, J, K)` tensor: 1 on boundaries, 0 inside. - """ - # Pad with -1 so boundary voxels at the edge are detected. - padded = functional.pad(data.float(), [1] * 6, mode="constant", value=-1) - # Min-pool with kernel 3 gives the morphological erosion. - eroded = -functional.max_pool3d(-padded, kernel_size=3, stride=1, padding=0) - # A voxel is on the contour if eroded != original (neighbor differs). - contour = (eroded != data.float()).float() - return contour diff --git a/src/torchio/transforms/label/keep_largest.py b/src/torchio/transforms/label/keep_largest.py deleted file mode 100644 index 7e6891f7e..000000000 --- a/src/torchio/transforms/label/keep_largest.py +++ /dev/null @@ -1,125 +0,0 @@ -"""KeepLargestComponent: keep only the largest connected component per label.""" - -from __future__ import annotations - -from collections.abc import Sequence -from typing import Any - -import SimpleITK as sitk -import torch -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import Transform - - -class KeepLargestComponent(Transform): - r"""Keep only the largest connected component of each label. - - For each specified label value, connected-component analysis is - performed and all but the largest component are removed (set to - the background value). This is useful for cleaning up noisy - segmentation predictions. - - Only single-channel [`LabelMap`][torchio.LabelMap] images are - affected. - - Args: - labels: Label values to filter. `None` means all non-zero - labels found in the data. - background_label: Value used for removed components. - fully_connected: If `True`, use 26-connectivity (voxels - sharing a corner are connected). If `False`, use - 6-connectivity (face-connected only). - **kwargs: See [`Transform`][torchio.Transform]. - - Raises: - RuntimeError: If a label map has more than one channel. - - Examples: - >>> import torchio as tio - >>> transform = tio.KeepLargestComponent() - >>> transform = tio.KeepLargestComponent(labels=[1, 2]) - """ - - def __init__( - self, - labels: Sequence[int] | None = None, - *, - background_label: int = 0, - fully_connected: bool = True, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.labels = list(labels) if labels is not None else None - self.background_label = background_label - self.fully_connected = fully_connected - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Keep only the largest connected component per label.""" - for _name, img_batch in batch.images.items(): - if not issubclass(img_batch._image_class, LabelMap): - continue - b, c = img_batch.data.shape[:2] - if c != 1: - msg = ( - "KeepLargestComponent requires single-channel" - f" label maps, got {c} channels" - ) - raise RuntimeError(msg) - for i in range(b): - img_batch.data[i, 0] = _keep_largest_per_label( - img_batch.data[i, 0], - labels=self.labels, - background_label=self.background_label, - fully_connected=self.fully_connected, - ) - return batch - - -def _keep_largest_per_label( - data: Tensor, - *, - labels: list[int] | None, - background_label: int, - fully_connected: bool, -) -> Tensor: - """Keep the largest connected component for each label. - - Args: - data: `(I, J, K)` label tensor. - labels: Which labels to filter. `None` means all - non-zero labels. - background_label: Value for removed voxels. - fully_connected: Whether to use 26- or 6-connectivity. - - Returns: - Filtered `(I, J, K)` tensor. - """ - result = data.clone() - if labels is None: - unique = data.unique().tolist() - labels = [int(v) for v in unique if int(v) != background_label] - - for label in labels: - binary = (data == label).cpu().numpy().astype("uint8") - if binary.sum() == 0: - continue - sitk_img = sitk.GetImageFromArray(binary) - cc = sitk.ConnectedComponent(sitk_img, fully_connected) - relabeled = sitk.RelabelComponent(cc, sortByObjectSize=True) - cc_array = sitk.GetArrayFromImage(relabeled) - # cc_array label 1 is the largest component; remove others. - mask = torch.from_numpy((cc_array >= 2).astype("uint8")).to(data.device) - result[mask.bool()] = background_label - - return result diff --git a/src/torchio/transforms/label/one_hot.py b/src/torchio/transforms/label/one_hot.py deleted file mode 100644 index af13a0780..000000000 --- a/src/torchio/transforms/label/one_hot.py +++ /dev/null @@ -1,97 +0,0 @@ -"""OneHot: one-hot encode label maps.""" - -from __future__ import annotations - -from typing import Any - -import torch.nn.functional as functional - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import Transform - - -class OneHot(Transform): - r"""One-hot encode label maps. - - Each label map with $K$ classes (including background) is converted - from shape $(1, I, J, K)$ to $(K, I, J, K)$, where channel $k$ - is 1 where the label equals $k$ and 0 elsewhere. - - Only [`LabelMap`][torchio.LabelMap] images are affected. - [`ScalarImage`][torchio.ScalarImage] instances are left unchanged. - - The inverse takes the argmax across channels, restoring the - original single-channel label map. - - Args: - num_classes: Total number of classes. `-1` (default) infers - from the data as `max_label + 1`. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.OneHot() - >>> transform = tio.OneHot(num_classes=5) - >>> # Invert back to single-channel - >>> restored = transformed.apply_inverse_transform() - """ - - def __init__( - self, - *, - num_classes: int = -1, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.num_classes = num_classes - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {"num_classes": self.num_classes} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """One-hot encode each label map in the batch.""" - num_classes = params["num_classes"] - for _name, img_batch in batch.images.items(): - if not issubclass(img_batch._image_class, LabelMap): - continue - # (B, 1, I, J, K) -> (B, num_classes, I, J, K) - data = img_batch.data.long() - flat = data[:, 0] # (B, I, J, K) - encoded = functional.one_hot(flat, num_classes=num_classes) - # (B, I, J, K, num_classes) -> (B, num_classes, I, J, K) - img_batch.data = encoded.permute(0, 4, 1, 2, 3).float() - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> _OneHotInverse: - """Invert by taking argmax.""" - return _OneHotInverse(copy=False) - - -class _OneHotInverse(Transform): - """Inverse of OneHot: argmax back to single-channel labels.""" - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - for _name, img_batch in batch.images.items(): - if not issubclass(img_batch._image_class, LabelMap): - continue - if img_batch.data.shape[1] > 1: - img_batch.data = img_batch.data.argmax(dim=1, keepdim=True).float() - return batch diff --git a/src/torchio/transforms/label/remap_labels.py b/src/torchio/transforms/label/remap_labels.py deleted file mode 100644 index 7e4ed4370..000000000 --- a/src/torchio/transforms/label/remap_labels.py +++ /dev/null @@ -1,69 +0,0 @@ -"""RemapLabels: reassign label values in a label map.""" - -from __future__ import annotations - -from typing import Any - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import Transform - - -class RemapLabels(Transform): - r"""Reassign label values in label maps. - - Each key in the *remapping* dict is replaced by its value. - Labels not mentioned in the dict are left unchanged. - - Only [`LabelMap`][torchio.LabelMap] images are affected. - - Args: - remapping: Dictionary mapping old labels to new labels. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> # Merge labels 2 and 3 into label 1 - >>> transform = tio.RemapLabels({2: 1, 3: 1}) - >>> # Swap labels 1 and 2 - >>> transform = tio.RemapLabels({1: 2, 2: 1}) - """ - - def __init__( - self, - remapping: dict[int, int], - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.remapping = remapping - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {"remapping": self.remapping} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Remap labels in each label map.""" - remapping = params["remapping"] - for _name, img_batch in batch.images.items(): - if not issubclass(img_batch._image_class, LabelMap): - continue - data = img_batch.data.clone() - for old, new in remapping.items(): - data[img_batch.data == old] = new - img_batch.data = data - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> RemapLabels: - """Invert by swapping keys and values.""" - remapping = params["remapping"] - inverse_remapping = {v: k for k, v in remapping.items()} - return RemapLabels(remapping=inverse_remapping, copy=False) diff --git a/src/torchio/transforms/label/remove_labels.py b/src/torchio/transforms/label/remove_labels.py deleted file mode 100644 index 6e580513a..000000000 --- a/src/torchio/transforms/label/remove_labels.py +++ /dev/null @@ -1,61 +0,0 @@ -"""RemoveLabels: set specified label values to background.""" - -from __future__ import annotations - -from collections.abc import Sequence -from typing import Any - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import Transform - - -class RemoveLabels(Transform): - """Set specified label values to a background value. - - This is a convenience wrapper that builds a remapping dict - internally. For more control, use - [`RemapLabels`][torchio.RemapLabels] directly. - - Only [`LabelMap`][torchio.LabelMap] images are affected. - - Args: - labels: Label values to remove. - background_label: Value to assign to removed labels. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.RemoveLabels([3, 4, 5]) - >>> transform = tio.RemoveLabels([2], background_label=0) - """ - - def __init__( - self, - labels: Sequence[int], - *, - background_label: int = 0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.labels = list(labels) - self.background_label = background_label - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Set specified labels to the background value.""" - for _name, img_batch in batch.images.items(): - if not issubclass(img_batch._image_class, LabelMap): - continue - data = img_batch.data.clone() - for label in self.labels: - data[img_batch.data == label] = self.background_label - img_batch.data = data - return batch diff --git a/src/torchio/transforms/label/sequential_labels.py b/src/torchio/transforms/label/sequential_labels.py deleted file mode 100644 index 8e1ed15dd..000000000 --- a/src/torchio/transforms/label/sequential_labels.py +++ /dev/null @@ -1,105 +0,0 @@ -"""SequentialLabels: renumber labels to consecutive integers.""" - -from __future__ import annotations - -from typing import Any - -import torch - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import Transform - - -class SequentialLabels(Transform): - r"""Renumber labels in label maps to consecutive integers starting from 0. - - For example, if a label map has values `{0, 5, 10}`, this - transform remaps them to `{0, 1, 2}`. - - Only [`LabelMap`][torchio.LabelMap] images are affected. - - Note: - The background (label 0) is always mapped to 0. Even if there - are no zeros in the input, zero will appear in the output. - - Args: - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.SequentialLabels() - """ - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Compute the remapping from the first sample's labels.""" - remappings: dict[str, dict[int, int]] = {} - for name, img_batch in batch.images.items(): - if not issubclass(img_batch._image_class, LabelMap): - continue - unique = sorted(int(v) for v in img_batch.data[0].unique().tolist()) - remappings[name] = {old: new for new, old in enumerate(unique)} - return {"remappings": remappings} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply sequential renumbering.""" - remappings = params["remappings"] - for name, img_batch in batch.images.items(): - if name not in remappings: - continue - remapping = remappings[name] - data = torch.zeros_like(img_batch.data) - for old, new in remapping.items(): - data[img_batch.data == old] = new - img_batch.data = data - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> _SequentialLabelsInverse: - """Invert by restoring original label values.""" - return _SequentialLabelsInverse( - remappings=params["remappings"], - copy=False, - ) - - -class _SequentialLabelsInverse(Transform): - """Inverse of SequentialLabels.""" - - def __init__( - self, - *, - remappings: dict[str, dict[int, int]], - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self._remappings = remappings - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - for name, img_batch in batch.images.items(): - if name not in self._remappings: - continue - inverse = {v: k for k, v in self._remappings[name].items()} - data = torch.zeros_like(img_batch.data) - for old, new in inverse.items(): - data[img_batch.data == old] = new - img_batch.data = data - return batch diff --git a/src/torchio/transforms/lambda_transform.py b/src/torchio/transforms/lambda_transform.py index 8b54c866e..6132a4fd0 100644 --- a/src/torchio/transforms/lambda_transform.py +++ b/src/torchio/transforms/lambda_transform.py @@ -1,76 +1,72 @@ -"""Lambda: apply a user-defined callable as a transform.""" - from __future__ import annotations -from collections.abc import Callable -from typing import Any +from collections.abc import Sequence -from torch import Tensor +import torch -from ..data.batch import SubjectsBatch -from ..data.image import LabelMap -from ..data.image import ScalarImage +from ..constants import TYPE +from ..data.subject import Subject +from ..types import TypeCallable from .transform import Transform class Lambda(Transform): - """Apply a user-defined function as a transform. - - The callable receives and returns a 4D tensor `(C, I, J, K)`. - Use *types_to_apply* to restrict which image types are affected. + """Applies a user-defined function as transform. Args: function: Callable that receives and returns a 4D [`torch.Tensor`][torch.Tensor]. - types_to_apply: Which image types the function applies to. - `"scalar"` for [`ScalarImage`][torchio.ScalarImage] only, - `"label"` for [`LabelMap`][torchio.LabelMap] only, - `None` for all images. - **kwargs: See [`Transform`][torchio.Transform]. + types_to_apply: List of strings corresponding to the image types to + which this transform should be applied. If `None`, the transform + will be applied to all images in the subject. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. Examples: >>> import torchio as tio - >>> invert = tio.Lambda(lambda x: -x, types_to_apply="scalar") - >>> double = tio.Lambda(lambda x: 2 * x) - >>> threshold = tio.Lambda(lambda x: (x > 0.5).float(), types_to_apply="label") + >>> invert_intensity = tio.Lambda(lambda x: -x, types_to_apply=[tio.INTENSITY]) + >>> invert_mask = tio.Lambda(lambda x: 1 - x, types_to_apply=[tio.LABEL]) + >>> def double(x): + ... return 2 * x + >>> double_transform = tio.Lambda(double) """ def __init__( self, - function: Callable[[Tensor], Tensor], - types_to_apply: str | None = None, - **kwargs: Any, - ) -> None: + function: TypeCallable, + types_to_apply: Sequence[str] | None = None, + **kwargs, + ): super().__init__(**kwargs) - if not callable(function): - msg = f"function must be callable, got {type(function).__name__}" - raise TypeError(msg) self.function = function self.types_to_apply = types_to_apply + self.args_names = ['function', 'types_to_apply'] - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply the callable to each matching image.""" - for _name, img_batch in batch.images.items(): - if not self._should_apply(img_batch._image_class): - continue - for i in range(img_batch.batch_size): - img_batch.data[i] = self.function(img_batch.data[i]) - return batch + def apply_transform(self, subject: Subject) -> Subject: + images = subject.get_images( + intensity_only=False, + include=self.include, + exclude=self.exclude, + ) + for image in images: + image_type = image[TYPE] + if self.types_to_apply is not None: + if image_type not in self.types_to_apply: + continue - def _should_apply(self, image_class: type) -> bool: - """Check whether this image type should be transformed.""" - if self.types_to_apply is None: - return True - if self.types_to_apply == "scalar": - return issubclass(image_class, ScalarImage) - if self.types_to_apply == "label": - return issubclass(image_class, LabelMap) - return True + function_arg = image.data + result = self.function(function_arg) + if not isinstance(result, torch.Tensor): + message = ( + 'The returned value from the callable argument must be' + f' of type {torch.Tensor}, not {type(result)}' + ) + raise ValueError(message) + if result.ndim != function_arg.ndim: + message = ( + 'The number of dimensions of the returned value must' + f' be {function_arg.ndim}, not {result.ndim}' + ) + raise ValueError(message) + image.set_data(result) + return subject diff --git a/src/torchio/transforms/monai_adapter.py b/src/torchio/transforms/monai_adapter.py index 64cb94b40..c2468cbf6 100644 --- a/src/torchio/transforms/monai_adapter.py +++ b/src/torchio/transforms/monai_adapter.py @@ -1,17 +1,13 @@ -"""MonaiAdapter: wrap MONAI transforms for use in TorchIO pipelines.""" - from __future__ import annotations -import copy as _copy import warnings from collections.abc import Callable from collections.abc import Mapping from types import ModuleType -from typing import Any +import numpy as np import torch -from ..data.affine import AffineMatrix from ..data.image import Image from ..data.image import ScalarImage from ..data.subject import Subject @@ -20,104 +16,218 @@ class MonaiAdapter(Transform): - """Wrap a MONAI transform for use in TorchIO pipelines. + """Wraps a MONAI transform for use in TorchIO pipelines. + + This adapter allows using + [MONAI transforms](https://docs.monai.io/en/stable/transforms.html) + within TorchIO workflows. Both **dictionary transforms** (e.g., + ``NormalizeIntensityd``) and **array transforms** (e.g., + ``NormalizeIntensity``) are supported. - Both **dictionary transforms** (subclasses of MONAI's - `MapTransform`, e.g., `NormalizeIntensityd`) and **array - transforms** (e.g., `NormalizeIntensity`) are supported. + The adapter handles conversion between TorchIO's + [`Subject`][torchio.Subject] (where values are + [`Image`][torchio.Image] objects) and MONAI's expected format + (where values are tensors or + [`MetaTensor`](https://docs.monai.io/en/stable/data.html#metatensor) + objects). Image tensors are passed as `MetaTensor` instances with the + affine matrix embedded, so spatial transforms (e.g., cropping, resizing) + correctly propagate affine changes. - Dictionary transforms operate on the full subject dictionary: - only the keys specified in the MONAI transform are modified. + **Dictionary transforms** (subclasses of MONAI's ``MapTransform``) + operate on the full subject dictionary — only the keys specified in + the MONAI transform are modified. - Array transforms are applied to each - [`ScalarImage`][torchio.ScalarImage] in the subject individually, - respecting the `include` / `exclude` parameters. + **Array transforms** (all other callables) are applied to each + image in the subject individually, respecting the ``include`` and + ``exclude`` parameters inherited from + [`Transform`][torchio.transforms.Transform]. Args: - monai_transform: A MONAI transform or any callable. Requires - MONAI to be installed: `pip install torchio[monai]`. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. + monai_transform: A MONAI transform (dictionary or array) or a + MONAI-compatible callable. Requires MONAI to be installed, + e.g., via ``pip install torchio[monai]``. + **kwargs: See [`Transform`][torchio.transforms.Transform] for + additional keyword arguments. Examples: + >>> import torch >>> import torchio as tio >>> from monai.transforms import NormalizeIntensity - >>> # Array transform: applied to each ScalarImage + >>> from monai.transforms import NormalizeIntensityd + >>> from monai.transforms import RandSpatialCropd + >>> subject = tio.Subject( + ... t1=tio.ScalarImage(tensor=torch.randn(1, 64, 64, 64)), + ... seg=tio.LabelMap(tensor=torch.ones(1, 64, 64, 64)), + ... ) + >>> # Array transform — applied to each image >>> adapter = tio.MonaiAdapter(NormalizeIntensity()) - >>> result = adapter(subject) + >>> transformed = adapter(subject) + >>> # Dictionary transform — applied to specified keys + >>> adapter = tio.MonaiAdapter(NormalizeIntensityd(keys=["t1"])) + >>> transformed = adapter(subject) + >>> # Spatial dict transform (affine is updated) + >>> adapter = tio.MonaiAdapter( + ... RandSpatialCropd(keys=["t1", "seg"], roi_size=[32, 32, 32]), + ... ) + >>> transformed = adapter(subject) >>> # Inside a Compose pipeline >>> pipeline = tio.Compose([ + ... tio.ToCanonical(), ... tio.MonaiAdapter(NormalizeIntensity()), - ... tio.Noise(std=0.1), + ... tio.RandomFlip(), ... ]) - - Note: - `MonaiAdapter` does **not** record itself in the subject's - transform history, because MONAI transform objects are not - serializable. + >>> transformed = pipeline(subject) """ - def __init__(self, monai_transform: Callable, **kwargs: Any) -> None: + def __init__( + self, + monai_transform: Callable, + **kwargs, + ): super().__init__(**kwargs) if not callable(monai_transform): - msg = ( - "monai_transform must be callable, " - f"got {type(monai_transform).__name__}" + message = ( + f'The monai_transform argument must be callable,' + f' but got {type(monai_transform)}' ) - raise TypeError(msg) + raise TypeError(message) self.monai_transform = monai_transform - def forward(self, data): - """Apply without recording history (MONAI transforms are opaque).""" - batch, unwrap = self._wrap(data) - if self.copy: - batch = _copy.deepcopy(batch) - if torch.rand(1).item() > self.p: - return unwrap(batch) - # MONAI transforms operate per-subject - monai = get_monai() - subjects = batch.unbatch() - for subject in subjects: - is_dict = isinstance( - self.monai_transform, - monai.transforms.MapTransform, - ) - if is_dict: - _apply_dict_transform(subject, self.monai_transform, monai) - else: - images = self._get_subject_images(subject) - _apply_array_transform(images, self.monai_transform, monai) - from ..data.batch import SubjectsBatch + def __repr__(self) -> str: + return f'{self.name}({self.monai_transform})' + + def _add_transform_to_subject_history(self, subject: Subject) -> None: + """Skip history recording for MONAI adapter transforms. + + The wrapped MONAI transform object cannot be reliably serialized or + reconstructed by `Subject.get_applied_transforms()`, so the adapter + is omitted from the subject history. This is similar to container + transforms such as `Compose` and `OneOf`, which rely on their inner + transforms being recorded instead. + """ + + def to_hydra_config(self) -> dict[str, object]: + """Not supported — MONAI transforms are not serializable. - result = SubjectsBatch.from_subjects(subjects) - result.adopt_history(batch, subjects) - return unwrap(result) + Raises: + NotImplementedError: Always. + """ + message = ( + 'MonaiAdapter cannot be exported to a Hydra config because' + ' the wrapped MONAI transform is not serializable' + ) + raise NotImplementedError(message) + + def apply_transform(self, subject: Subject) -> Subject: + """Apply the wrapped MONAI transform to a subject. + + Args: + subject: TorchIO subject to transform. + + Returns: + The transformed subject with updated tensor data and, for spatial + transforms, updated affine matrices. + """ + monai = get_monai() + is_dict_transform = isinstance( + self.monai_transform, + monai.transforms.MapTransform, + ) + if is_dict_transform: + self._apply_dict_transform(subject, monai) + else: + self._apply_array_transform(subject, monai) + return subject - def apply_transform(self, batch: Any, params: dict[str, Any]) -> Any: - # Not used: MonaiAdapter overrides forward directly - return batch + def _apply_dict_transform( + self, + subject: Subject, + monai: ModuleType, + ) -> None: + monai_dict = _subject_to_monai_dict(subject, monai) + result = self.monai_transform(monai_dict) + if not isinstance(result, Mapping): + message = ( + 'Expected the MONAI dict transform to return a single' + f' mapping, but got {type(result)}. Multi-sample' + ' transforms (returning a list of dicts) are not' + ' supported by MonaiAdapter.' + ) + raise TypeError(message) + _update_subject_from_monai_dict(subject, result, monai) - def _get_subject_images(self, subject: Subject) -> dict[str, Image]: - """Filter to ScalarImage, then apply include/exclude.""" - images: dict[str, Image] = { - k: v for k, v in subject.images.items() if isinstance(v, ScalarImage) - } - if self.include is not None: - images = {k: v for k, v in images.items() if k in self.include} - if self.exclude is not None: - images = {k: v for k, v in images.items() if k not in self.exclude} - return images + def _apply_array_transform( + self, + subject: Subject, + monai: ModuleType, + ) -> None: + images_dict = subject.get_images_dict( + intensity_only=False, + include=self.include, + exclude=self.exclude, + ) + if len(images_dict) > 1 and isinstance( + self.monai_transform, + monai.transforms.Randomizable, + ): + warnings.warn( + 'Applying a MONAI Randomizable array transform to a' + ' subject with multiple images. Each image will receive' + ' different random parameters, which may break spatial' + ' alignment. Consider using the dictionary version of' + ' this transform instead (e.g., RandFlipd instead of' + ' RandFlip).', + UserWarning, + stacklevel=4, + ) + for image in images_dict.values(): + meta_tensor = _image_to_meta_tensor(image, monai) + result = self.monai_transform(meta_tensor) + if not isinstance(result, torch.Tensor): + message = ( + 'Expected a torch.Tensor from the MONAI transform' + f' output, but got {type(result)}' + ) + raise TypeError(message) + _update_image_from_result(image, result, monai) -# ── Helpers ────────────────────────────────────────────────────────── +def _to_meta_tensor( + tensor: torch.Tensor, + affine: np.ndarray, + monai: ModuleType, +) -> torch.Tensor: + """Convert a tensor to a MONAI MetaTensor with affine.""" + MetaTensor = monai.data.MetaTensor + affine_tensor = torch.as_tensor(affine, dtype=torch.float64, device=tensor.device) + return MetaTensor(tensor, affine=affine_tensor) def _image_to_meta_tensor( image: Image, monai: ModuleType, ) -> torch.Tensor: - affine_tensor = image.affine.data.to(device=image.data.device) - return monai.data.MetaTensor(image.data, affine=affine_tensor) + """Convert a TorchIO Image to a MONAI MetaTensor.""" + return _to_meta_tensor(image.data, image.affine, monai) + + +def _unwrap_tensor(result: torch.Tensor, monai: ModuleType) -> torch.Tensor: + """Extract a plain tensor from a result, unwrapping MetaTensor if needed.""" + MetaTensor = monai.data.MetaTensor + if isinstance(result, MetaTensor): + return result.as_tensor() + return result + + +def _extract_affine( + result: torch.Tensor, + monai: ModuleType, +) -> np.ndarray | None: + """Extract the affine from a MetaTensor, or return None.""" + MetaTensor = monai.data.MetaTensor + if isinstance(result, MetaTensor): + return result.affine.cpu().numpy() + return None def _update_image_from_result( @@ -125,62 +235,67 @@ def _update_image_from_result( result: torch.Tensor, monai: ModuleType, ) -> None: - meta_tensor_cls = monai.data.MetaTensor - if isinstance(result, meta_tensor_cls): - image.set_data(result.as_tensor()) - new_affine = result.affine - if not torch.equal(new_affine.cpu().to(torch.float64), image.affine.data.cpu()): - image._affine = AffineMatrix(new_affine) - else: - image.set_data(result) - - -def _apply_array_transform( - images: dict[str, Image], - monai_transform: Callable, - monai: ModuleType, -) -> None: - if len(images) > 1 and isinstance( - monai_transform, - monai.transforms.Randomizable, - ): - warnings.warn( - "Applying a MONAI Randomizable array transform to multiple" - " images. Each image gets different random parameters." - " Use the dictionary version (e.g., RandFlipd) to keep" - " spatial alignment.", - UserWarning, - stacklevel=5, - ) - for image in images.values(): - meta_tensor = _image_to_meta_tensor(image, monai) - result = monai_transform(meta_tensor) - if not isinstance(result, torch.Tensor): - msg = ( - "Expected torch.Tensor from MONAI transform, " - f"got {type(result).__name__}" - ) - raise TypeError(msg) - _update_image_from_result(image, result, monai) + """Update a TorchIO Image from a MONAI transform result tensor.""" + image.set_data(_unwrap_tensor(result, monai)) + new_affine = _extract_affine(result, monai) + if new_affine is not None and not np.array_equal(new_affine, image.affine): + image.affine = new_affine -def _apply_dict_transform( +def _subject_to_monai_dict( subject: Subject, - monai_transform: Callable, monai: ModuleType, -) -> None: - monai_dict: dict[str, Any] = {} - for name, image in subject.images.items(): - monai_dict[name] = _image_to_meta_tensor(image, monai) - for key, value in subject.metadata.items(): - monai_dict[key] = value +) -> dict[str, object]: + """Convert a Subject to a MONAI-compatible dictionary. + + Image values are converted to MetaTensor instances. + Non-image values are passed through unchanged. + """ + monai_dict: dict[str, object] = {} + for key, value in subject.items(): + if isinstance(value, Image): + monai_dict[key] = _image_to_meta_tensor(value, monai) + else: + monai_dict[key] = value + return monai_dict - result = monai_transform(monai_dict) - if not isinstance(result, Mapping): - msg = f"Expected mapping from MONAI dict transform, got {type(result).__name__}" - raise TypeError(msg) +def _update_subject_from_monai_dict( + subject: Subject, + monai_dict: Mapping[str, object], + monai: ModuleType, +) -> None: + """Update a Subject in-place from the MONAI transform output. - for name, image in subject.images.items(): - if name in result and isinstance(result[name], torch.Tensor): - _update_image_from_result(image, result[name], monai) + For keys that correspond to Images in the subject, the tensor data + is updated. If the output is a MetaTensor with an updated affine, + the Image affine is also updated. + """ + for key, value in monai_dict.items(): + if key in subject and isinstance(subject[key], Image): + image = subject[key] + assert isinstance(image, Image) + if not isinstance(value, torch.Tensor): + message = ( + 'Expected a torch.Tensor from the MONAI transform' + f' output, but got {type(value)}' + ) + raise TypeError(message) + _update_image_from_result(image, value, monai) + else: + # New key or non-image key from MONAI output + if isinstance(value, torch.Tensor): + tensor = _unwrap_tensor(value, monai) + if tensor.ndim == 4: + new_affine = _extract_affine(value, monai) + if new_affine is None: + new_affine = np.eye(4) + subject[key] = ScalarImage( + tensor=tensor, + affine=new_affine, + ) + else: + subject[key] = tensor + else: + subject[key] = value + subject.update_attributes() diff --git a/src/torchio/transforms/parameter_range.py b/src/torchio/transforms/parameter_range.py deleted file mode 100644 index 7d2ef536c..000000000 --- a/src/torchio/transforms/parameter_range.py +++ /dev/null @@ -1,401 +0,0 @@ -"""_ParameterRange: scalar, range, choice, or distribution for transform parameters. - -Each axis can be specified independently when a 3-element tuple is -used, where each element can be a float, `(lo, hi)` range, -`Choice`, or `Distribution`. - -Examples: - >>> _ParameterRange(0.5) # deterministic - >>> _ParameterRange((5.0, 15.0)) # U(5, 15) all axes - >>> _ParameterRange((1.0, 2.0, 3.0)) # per-axis fixed - >>> _ParameterRange(Choice([-10, 0, 10])) # discrete choice - >>> _ParameterRange((0, 0, Choice([-90, 0, 90]))) # per-axis mix -""" - -from __future__ import annotations - -from collections.abc import Sequence -from typing import cast -from typing import overload - -import torch -from torch.distributions import Distribution - -# ── Choice ─────────────────────────────────────────────────────────── - - -class Choice: - """A discrete set of values to sample from. - - Args: - values: Sequence of numeric values to choose from. - probabilities: Optional per-value probabilities. - If `None`, all values are equally likely. - - Examples: - >>> Choice([-10, 0, 10]) - Choice([-10.0, 0.0, 10.0]) - >>> Choice([0.5, 1.0, 2.0], probabilities=[0.2, 0.6, 0.2]) - Choice([0.5, 1.0, 2.0], p=[0.20, 0.60, 0.20]) - """ - - def __init__( - self, - values: Sequence[float | int], - probabilities: Sequence[float] | None = None, - ) -> None: - if len(values) < 1: - msg = "Choice requires at least one value" - raise ValueError(msg) - self._values = torch.tensor([float(v) for v in values]) - if probabilities is not None: - if len(probabilities) != len(values): - msg = f"Expected {len(values)} probabilities, got {len(probabilities)}" - raise ValueError(msg) - self._probs = torch.tensor([float(p) for p in probabilities]) - else: - self._probs = torch.ones(len(values)) / len(values) - - def sample(self) -> float: - """Pick one value at random.""" - idx = torch.multinomial(self._probs, 1).item() - return float(self._values[int(idx)]) - - def sample_batched(self, n: int) -> torch.Tensor: - """Pick ``n`` values at random, with replacement. - - Args: - n: Number of independent draws. - - Returns: - A 1D tensor of shape ``(n,)``. - """ - idx = torch.multinomial(self._probs, n, replacement=True) - return self._values[idx] - - def __repr__(self) -> str: - vals = [f"{v:.1f}" if v == int(v) else f"{v}" for v in self._values.tolist()] - parts = f"[{', '.join(vals)}]" - if torch.allclose(self._probs, self._probs[0].expand_as(self._probs)): - return f"Choice({parts})" - probs = ", ".join(f"{p:.2f}" for p in self._probs.tolist()) - return f"Choice({parts}, p=[{probs}])" - - -# ── Per-axis sampler ───────────────────────────────────────────────── - -#: What each axis stores after parsing. -AxisSpec = float | tuple[float, float] | Choice | Distribution - - -def _sample_axis( - spec: AxisSpec, - *, - generator: torch.Generator | None = None, -) -> float: - """Sample a single float from an axis specification.""" - if isinstance(spec, (int, float)): - return float(spec) - if isinstance(spec, Choice): - return spec.sample() - if isinstance(spec, Distribution): - return spec.sample().item() - lo, hi = spec - if lo == hi: - return float(lo) - return torch.empty(1).uniform_(float(lo), float(hi), generator=generator).item() - - -def _sample_axis_batched( - spec: AxisSpec, - n: int, - *, - generator: torch.Generator | None = None, -) -> torch.Tensor: - """Sample ``n`` independent floats from an axis specification. - - Args: - spec: The axis specification (float, ``(lo, hi)`` range, - `Choice`, or `Distribution`). - n: Number of independent samples to draw. - generator: Optional generator for the uniform-range path. - Note that `Choice` and `Distribution` ignore it (see - [`_ParameterRange.sample`][]). - - Returns: - A 1D tensor of shape ``(n,)``. - """ - if isinstance(spec, (int, float)): - return torch.full((n,), float(spec)) - if isinstance(spec, Choice): - return spec.sample_batched(n) - if isinstance(spec, Distribution): - return spec.sample((n,)).reshape(n).to(torch.float32) - lo, hi = spec - if lo == hi: - return torch.full((n,), float(lo)) - return torch.empty(n).uniform_(float(lo), float(hi), generator=generator) - - -# ── _ParameterRange ─────────────────────────────────────────────────── - - -class _ParameterRange: - """Encapsulates the range-or-scalar pattern for transform params. - - Args: - value: Parameter specification. Accepted forms: - - - `float`: deterministic value broadcast to 3 axes. - - `(lo, hi)`: uniform range for all axes. - - `(a, b, c)`: deterministic per-axis (when all plain - numbers). - - `(lo0, hi0, lo1, hi1, lo2, hi2)`: per-axis ranges. - - `Choice`: discrete random, same for all axes. - - `Distribution`: sample from any distribution. - - 3-tuple of mixed specs, e.g., - `(0, Choice([-90, 0, 90]), (-10, 10))`. - """ - - def __init__( - self, - value: float | tuple | Distribution | Choice, - ) -> None: - self._original = value - self._axes: tuple[AxisSpec, AxisSpec, AxisSpec] - - if isinstance(value, (int, float)): - v = float(value) - self._axes = (v, v, v) - elif isinstance(value, (Choice, Distribution)): - self._axes = (value, value, value) - elif isinstance(value, tuple): - self._axes = _parse_tuple(value) - else: - msg = ( - "Expected float, tuple, Distribution, or Choice," - f" got {type(value).__name__}" - ) - raise TypeError(msg) - - @property - def is_deterministic(self) -> bool: - """Whether this range always returns the same values.""" - return all(isinstance(a, (int, float)) for a in self._axes) - - def is_constant(self, value: float) -> bool: - """Whether every axis deterministically equals `value`. - - Args: - value: The value to compare against on every axis. - - Returns: - `True` if each axis is a fixed number, or a degenerate - `(v, v)` range, equal to `value`. `Choice` and - `Distribution` axes are never constant. - """ - for axis in self._axes: - if isinstance(axis, (int, float)): - if float(axis) != float(value): - return False - elif isinstance(axis, tuple): - low, high = axis - if not (low == high == value): - return False - else: # Choice or Distribution - return False - return True - - @property - def _ranges( - self, - ) -> tuple[tuple[float, float], tuple[float, float], tuple[float, float]]: - """Legacy: per-axis (lo, hi) ranges for validation code.""" - result: list[tuple[float, float]] = [] - for a in self._axes: - if isinstance(a, (int, float)): - result.append((float(a), float(a))) - elif isinstance(a, tuple): - result.append(cast(tuple[float, float], a)) - else: - result.append((0.0, 0.0)) - return (result[0], result[1], result[2]) - - @property - def _distribution(self) -> Distribution | None: - """Legacy: shared Distribution if the first axis uses one.""" - if isinstance(self._axes[0], Distribution): - return self._axes[0] - return None - - @overload - def sample( - self, - n: None = ..., - *, - generator: torch.Generator | None = ..., - ) -> tuple[float, float, float]: ... - @overload - def sample( - self, - n: int, - *, - generator: torch.Generator | None = ..., - ) -> torch.Tensor: ... - def sample( - self, - n: int | None = None, - *, - generator: torch.Generator | None = None, - ) -> tuple[float, float, float] | torch.Tensor: - """Sample a 3-tuple of values, or a batch of them. - - Args: - n: If `None` (default), draw a single 3-tuple (legacy - behavior). If an integer, draw ``n`` independent - 3-tuples and return a tensor of shape ``(n, 3)``. - generator: Optional generator for the uniform-range path. - - Returns: - A tuple of three floats when ``n is None``, otherwise a - tensor of shape ``(n, 3)``. - """ - if n is None: - return ( - _sample_axis(self._axes[0], generator=generator), - _sample_axis(self._axes[1], generator=generator), - _sample_axis(self._axes[2], generator=generator), - ) - columns = [ - _sample_axis_batched(axis, n, generator=generator) for axis in self._axes - ] - return torch.stack(columns, dim=-1) - - @overload - def sample_1d( - self, - n: None = ..., - *, - generator: torch.Generator | None = ..., - ) -> float: ... - @overload - def sample_1d( - self, - n: int, - *, - generator: torch.Generator | None = ..., - ) -> torch.Tensor: ... - def sample_1d( - self, - n: int | None = None, - *, - generator: torch.Generator | None = None, - ) -> float | torch.Tensor: - """Sample a single float (from the first axis spec), or a batch. - - Args: - n: If `None` (default), draw a single float (legacy - behavior). If an integer, draw ``n`` independent values - and return a tensor of shape ``(n,)``. - generator: Optional generator for the uniform-range path. - - Returns: - A single float when ``n is None``, otherwise a tensor of - shape ``(n,)``. - """ - if n is None: - return _sample_axis(self._axes[0], generator=generator) - return _sample_axis_batched(self._axes[0], n, generator=generator) - - def __repr__(self) -> str: - v = self._original - if isinstance(v, (Distribution, Choice)): - return repr(v) - if isinstance(v, tuple): - inner = ", ".join(repr(x) for x in v) - return f"({inner})" - return str(v) - - -# ── Tuple parsing ──────────────────────────────────────────────────── - - -def _is_plain_number(x: object) -> bool: - return isinstance(x, (int, float)) - - -def _parse_tuple( - value: tuple, -) -> tuple[AxisSpec, AxisSpec, AxisSpec]: - """Parse a tuple into three per-axis specs.""" - n = len(value) - - # 3-element: per-axis fixed OR per-axis mixed - if n == 3: - if all(_is_plain_number(v) for v in value): - return (float(value[0]), float(value[1]), float(value[2])) - return ( - _parse_single(value[0]), - _parse_single(value[1]), - _parse_single(value[2]), - ) - - # Remaining forms require all plain numbers - if not all(_is_plain_number(v) for v in value): - msg = f"Mixed per-axis specs require exactly 3 elements, got {n}" - raise ValueError(msg) - - if n == 1: - v = float(value[0]) - return (v, v, v) - if n == 2: - r: tuple[float, float] = (float(value[0]), float(value[1])) - return (r, r, r) - if n == 6: - return ( - (float(value[0]), float(value[1])), - (float(value[2]), float(value[3])), - (float(value[4]), float(value[5])), - ) - msg = f"Tuple must have 1, 2, 3, or 6 elements, got {n}" - raise ValueError(msg) - - -def _parse_single(spec: object) -> AxisSpec: - """Parse a single axis specification.""" - if isinstance(spec, (int, float)): - return float(spec) - if isinstance(spec, (Choice, Distribution)): - return spec - if isinstance(spec, tuple) and len(spec) == 2: - lo, hi = spec - if isinstance(lo, (int, float)) and isinstance(hi, (int, float)): - return (float(lo), float(hi)) - msg = ( - "Per-axis spec must be a float, (lo, hi) tuple, Choice," - f" or Distribution, got {type(spec).__name__}" - ) - raise TypeError(msg) - - -# ── Convenience constructors ───────────────────────────────────────── - - -def to_range( - value: float | tuple | Distribution | Choice, -) -> _ParameterRange: - """Convert to a _ParameterRange.""" - return _ParameterRange(value) - - -def to_nonneg_range( - value: float | tuple | Distribution | Choice, -) -> _ParameterRange: - """Like `to_range`, but rejects negative values in tuple ranges.""" - pr = _ParameterRange(value) - if pr._distribution is None: - for lo, hi in pr._ranges: - if lo < 0 or hi < 0: - msg = f"Value must be non-negative, got {value}" - raise ValueError(msg) - return pr diff --git a/src/torchio/transforms/preprocessing/__init__.py b/src/torchio/transforms/preprocessing/__init__.py new file mode 100644 index 000000000..94ee2d0f1 --- /dev/null +++ b/src/torchio/transforms/preprocessing/__init__.py @@ -0,0 +1,51 @@ +from .intensity.clamp import Clamp +from .intensity.histogram_standardization import HistogramStandardization +from .intensity.mask import Mask +from .intensity.pca import PCA +from .intensity.rescale import RescaleIntensity +from .intensity.to import To +from .intensity.z_normalization import ZNormalization +from .label.contour import Contour +from .label.keep_largest_component import KeepLargestComponent +from .label.one_hot import OneHot +from .label.remap_labels import RemapLabels +from .label.remove_labels import RemoveLabels +from .label.sequential_labels import SequentialLabels +from .spatial.copy_affine import CopyAffine +from .spatial.crop import Crop +from .spatial.crop_or_pad import CropOrPad +from .spatial.ensure_shape_multiple import EnsureShapeMultiple +from .spatial.pad import Pad +from .spatial.resample import Resample +from .spatial.resize import Resize +from .spatial.to_canonical import ToCanonical +from .spatial.to_orientation import ToOrientation +from .spatial.to_reference_space import ToReferenceSpace +from .spatial.transpose import Transpose + +__all__ = [ + 'Pad', + 'Crop', + 'Resize', + 'Resample', + 'ToCanonical', + 'ToOrientation', + 'ToReferenceSpace', + 'Transpose', + 'CropOrPad', + 'CopyAffine', + 'EnsureShapeMultiple', + 'Mask', + 'PCA', + 'RescaleIntensity', + 'To', + 'Clamp', + 'ZNormalization', + 'HistogramStandardization', + 'OneHot', + 'Contour', + 'RemapLabels', + 'RemoveLabels', + 'SequentialLabels', + 'KeepLargestComponent', +] diff --git a/src/torchio/transforms/preprocessing/intensity/__init__.py b/src/torchio/transforms/preprocessing/intensity/__init__.py new file mode 100644 index 000000000..763a18477 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/__init__.py @@ -0,0 +1,5 @@ +from .normalization_transform import NormalizationTransform + +__all__ = [ + 'NormalizationTransform', +] diff --git a/src/torchio/transforms/preprocessing/intensity/clamp.py b/src/torchio/transforms/preprocessing/intensity/clamp.py new file mode 100644 index 000000000..e3f0686e5 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/clamp.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import torch + +from ....data.image import ScalarImage +from ....data.subject import Subject +from ...intensity_transform import IntensityTransform + + +class Clamp(IntensityTransform): + """Clamp intensity values into a range $[a, b]$. + + For more information, see `torch.clamp()`. + + Args: + out_min: Minimum value $a$ of the output image. If `None`, the + minimum of the image is used. + out_max: Maximum value $b$ of the output image. If `None`, the + maximum of the image is used. + + Examples: + >>> import torchio as tio + >>> ct = tio.datasets.Slicer('CTChest').CT_chest + >>> HOUNSFIELD_AIR, HOUNSFIELD_BONE = -1000, 1000 + >>> clamp = tio.Clamp(out_min=HOUNSFIELD_AIR, out_max=HOUNSFIELD_BONE) + >>> ct_clamped = clamp(ct) + + """ + + def __init__( + self, + out_min: float | None = None, + out_max: float | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.out_min, self.out_max = out_min, out_max + self.args_names = ['out_min', 'out_max'] + + def apply_transform(self, subject: Subject) -> Subject: + for image in self.get_images(subject): + assert isinstance(image, ScalarImage) + self.apply_clamp(image) + return subject + + def apply_clamp(self, image: ScalarImage) -> None: + image.set_data(self.clamp(image.data)) + + def clamp(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor.clamp(self.out_min, self.out_max) diff --git a/src/torchio/transforms/preprocessing/intensity/histogram_standardization.py b/src/torchio/transforms/preprocessing/intensity/histogram_standardization.py new file mode 100644 index 000000000..340b50d21 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/histogram_standardization.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +from collections.abc import Callable +from collections.abc import Iterable +from collections.abc import Sequence +from pathlib import Path +from typing import Union +from typing import cast + +import numpy as np +import torch +from tqdm.auto import tqdm + +from ....data.io import read_image +from ....data.subject import Subject +from ....types import TypePath +from .normalization_transform import NormalizationTransform +from .normalization_transform import TypeMaskingMethod + +DEFAULT_CUTOFF = 0.01, 0.99 +STANDARD_RANGE = 0, 100 +TypeLandmarks = Union[TypePath, dict[str, Union[TypePath, np.ndarray]]] + + +class HistogramStandardization(NormalizationTransform): + """Perform histogram standardization of intensity values. + + Implementation of [New variants of a method of MRI scale + standardization ](https://ieeexplore.ieee.org/document/836373). + + See example in `torchio.transforms.HistogramStandardization.train()`. + + Args: + landmarks: Dictionary (or path to a PyTorch file with `.pt` or `.pth` + extension in which a dictionary has been saved) whose keys are + image names in the subject and values are NumPy arrays or paths to + NumPy arrays defining the landmarks after training with + `torchio.transforms.HistogramStandardization.train()`. + masking_method: See + [`NormalizationTransform`][torchio.transforms.preprocessing.intensity.NormalizationTransform]. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torch + >>> import torchio as tio + >>> landmarks = { + ... 't1': 't1_landmarks.npy', + ... 't2': 't2_landmarks.npy', + ... } + >>> transform = tio.HistogramStandardization(landmarks) + >>> torch.save(landmarks, 'path_to_landmarks.pth') + >>> transform = tio.HistogramStandardization('path_to_landmarks.pth') + """ + + def __init__( + self, + landmarks: TypeLandmarks, + masking_method: TypeMaskingMethod = None, + **kwargs, + ): + super().__init__(masking_method=masking_method, **kwargs) + self.landmarks = landmarks + self.landmarks_dict = self._parse_landmarks(landmarks) + self.args_names = ['landmarks', 'masking_method'] + + @staticmethod + def _parse_landmarks(landmarks: TypeLandmarks) -> dict[str, np.ndarray]: + if isinstance(landmarks, (str, Path)): + path = Path(landmarks) + if path.suffix not in ('.pt', '.pth'): + message = ( + 'The landmarks file must have extension .pt or .pth,' + f' not "{path.suffix}"' + ) + raise ValueError(message) + landmarks_dict = torch.load(path) + else: + landmarks_dict = landmarks + parsed_landmarks: dict[str, np.ndarray] = {} + for key, value in landmarks_dict.items(): + if isinstance(value, (str, Path)): + parsed_landmarks[key] = np.load(value) + else: + parsed_landmarks[key] = value + return parsed_landmarks + + def apply_normalization( + self, + subject: Subject, + image_name: str, + mask: torch.Tensor, + ) -> None: + if image_name not in self.landmarks_dict: + keys = tuple(self.landmarks_dict.keys()) + message = ( + f'Image name "{image_name}" should be a key in the' + f' landmarks dictionary, whose keys are {keys}' + ) + raise KeyError(message) + image = subject.get_scalar_image(image_name) + landmarks = self.landmarks_dict[image_name] + normalized = _normalize(image.data, landmarks, mask=mask.numpy()) + image.set_data(normalized) + + @classmethod + def train( + cls, + images_paths: Sequence[TypePath], + cutoff: tuple[float, float] | None = None, + mask_path: Sequence[TypePath] | TypePath | None = None, + masking_function: Callable | None = None, + output_path: TypePath | None = None, + *, + progress: bool = True, + ) -> np.ndarray: + """Extract average histogram landmarks from images used for training. + + Args: + images_paths: List of image paths used to train. + cutoff: Optional minimum and maximum quantile values, + respectively, that are used to select a range of intensity of + interest. Equivalent to $pc_1$ and $pc_2$ in + [Nyúl and Udupa's paper ](https://pubmed.ncbi.nlm.nih.gov/10571928/). + mask_path: Path (or list of paths) to a binary image that will be + used to select the voxels use to compute the stats during + histogram training. If `None`, all voxels in the image will + be used. + masking_function: Function used to extract voxels used for + histogram training. + output_path: Optional file path with extension `.txt` or + `.npy`, where the landmarks will be saved. + + Examples: + >>> import torch + >>> import numpy as np + >>> from pathlib import Path + >>> from torchio.transforms import HistogramStandardization + >>> + >>> t1_paths = ['subject_a_t1.nii', 'subject_b_t1.nii.gz'] + >>> t2_paths = ['subject_a_t2.nii', 'subject_b_t2.nii.gz'] + >>> + >>> t1_landmarks_path = Path('t1_landmarks.npy') + >>> t2_landmarks_path = Path('t2_landmarks.npy') + >>> + >>> t1_landmarks = ( + ... t1_landmarks_path + ... if t1_landmarks_path.is_file() + ... else HistogramStandardization.train(t1_paths) + ... ) + >>> np.save(t1_landmarks_path, t1_landmarks) + >>> + >>> t2_landmarks = ( + ... t2_landmarks_path + ... if t2_landmarks_path.is_file() + ... else HistogramStandardization.train(t2_paths) + ... ) + >>> np.save(t2_landmarks_path, t2_landmarks) + >>> + >>> landmarks_dict = { + ... 't1': t1_landmarks, + ... 't2': t2_landmarks, + ... } + >>> + >>> transform = HistogramStandardization(landmarks_dict) + """ + mask_paths: Sequence[TypePath] | None = ( + mask_path + if isinstance(mask_path, Sequence) + and not isinstance(mask_path, (str, Path)) + else None + ) + if mask_paths is not None and len(mask_paths) != len(images_paths): + message = ( + f'Different number of images ({len(images_paths)})' + f' and mask ({len(mask_paths)}) paths found' + ) + raise ValueError(message) + quantiles_cutoff = DEFAULT_CUTOFF if cutoff is None else cutoff + percentiles_cutoff = 100 * np.array(quantiles_cutoff) + percentiles_database = [] + a, b = percentiles_cutoff # for mypy + percentiles = _get_percentiles((a, b)) + iterable = cast( + Iterable[TypePath], tqdm(images_paths) if progress else images_paths + ) + for i, image_file_path in enumerate(iterable): + tensor, _ = read_image(image_file_path) + if masking_function is not None: + mask = masking_function(tensor) + else: + if mask_path is None: + mask = np.ones_like(tensor, dtype=bool) + else: + path: TypePath + if mask_paths is not None: + path = mask_paths[i] + else: + assert isinstance(mask_path, (str, Path)) + path = mask_path + mask, _ = read_image(path) + mask = mask.numpy() > 0 + array = tensor.numpy() + percentile_values = np.percentile(array[mask], percentiles) + percentiles_database.append(percentile_values) + percentiles_database_array = np.vstack(percentiles_database) + mapping = _get_average_mapping(percentiles_database_array) + + if output_path is not None: + output_path = Path(output_path).expanduser() + extension = output_path.suffix + if extension == '.txt': + modality = 'image' + text = f'{modality} {" ".join(map(str, mapping))}' + output_path.write_text(text) + elif extension == '.npy': + np.save(output_path, mapping) + return mapping + + +def _standardize_cutoff(cutoff: Sequence[float]) -> np.ndarray: + """Standardize the cutoff values given in the configuration. + + Computes percentile landmark normalization by default. + """ + cutoff_array = np.asarray(cutoff) + cutoff_array[0] = max(0, cutoff_array[0]) + cutoff_array[1] = min(1, cutoff_array[1]) + cutoff_array[0] = np.min([cutoff_array[0], 0.09]) + cutoff_array[1] = np.max([cutoff_array[1], 0.91]) + return cutoff_array + + +def _get_average_mapping(percentiles_database: np.ndarray) -> np.ndarray: + """Map the landmarks of the database to the chosen range. + + Args: + percentiles_database: Percentiles database over which to perform the + averaging. + """ + # Assuming percentiles_database.shape == (num_data_points, num_percentiles) + pc1 = percentiles_database[:, 0] + pc2 = percentiles_database[:, -1] + s1, s2 = STANDARD_RANGE + slopes = (s2 - s1) / (pc2 - pc1) + slopes = np.nan_to_num(slopes) + intercepts = np.mean(s1 - slopes * pc1) + num_images = len(percentiles_database) + final_map = slopes.dot(percentiles_database) / num_images + intercepts + return final_map + + +def _get_percentiles(percentiles_cutoff: tuple[float, float]) -> np.ndarray: + quartiles = np.arange(25, 100, 25).tolist() + deciles = np.arange(10, 100, 10).tolist() + all_percentiles = list(percentiles_cutoff) + quartiles + deciles + percentiles = sorted(set(all_percentiles)) + return np.array(percentiles) + + +def _normalize( + tensor: torch.Tensor, + landmarks: np.ndarray, + mask: np.ndarray | None, + cutoff: tuple[float, float] | None = None, + epsilon: float = 1e-5, +) -> torch.Tensor: + cutoff_ = DEFAULT_CUTOFF if cutoff is None else cutoff + array = tensor.numpy() + mapping = landmarks + + data = array + shape = data.shape + data = data.reshape(-1).astype(np.float32) + + if mask is None: + mask = np.ones_like(data, bool) + mask = mask.reshape(-1) + + range_to_use = [0, 1, 2, 4, 5, 6, 7, 8, 10, 11, 12] + + quantiles_cutoff = _standardize_cutoff(cutoff_) + percentiles_cutoff = 100 * np.array(quantiles_cutoff) + a, b = percentiles_cutoff # for mypy + percentiles = _get_percentiles((a, b)) + percentile_values = np.percentile(data[mask], percentiles) + + # Apply linear histogram standardization + range_mapping = mapping[range_to_use] + range_perc = percentile_values[range_to_use] + diff_mapping = np.diff(range_mapping) + diff_perc = np.diff(range_perc) + + # Handling the case where two landmarks are the same + # for a given input image. This usually happens when + # image background is not removed from the image. + diff_perc[diff_perc < epsilon] = np.inf + + affine_map = np.zeros([2, len(range_to_use) - 1]) + + # Compute slopes of the linear models + affine_map[0] = diff_mapping / diff_perc + + # Compute intercepts of the linear models + affine_map[1] = range_mapping[:-1] - affine_map[0] * range_perc[:-1] + + bin_id = np.digitize(data, range_perc[1:-1], right=False) + lin_img = affine_map[0, bin_id] + aff_img = affine_map[1, bin_id] + new_img = lin_img * data + aff_img + new_img = new_img.reshape(shape) + new_img = new_img.astype(np.float32) + new_img = torch.as_tensor(new_img) + return new_img + + +# train_histogram kept for backward compatibility +train = train_histogram = HistogramStandardization.train diff --git a/src/torchio/transforms/preprocessing/intensity/mask.py b/src/torchio/transforms/preprocessing/intensity/mask.py new file mode 100644 index 000000000..770b05a10 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/mask.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import warnings +from collections.abc import Sequence + +import torch + +from ....data.image import ScalarImage +from ....data.subject import Subject +from ....transforms.transform import TypeMaskingMethod +from ...intensity_transform import IntensityTransform + + +class Mask(IntensityTransform): + """Set voxels outside of mask to a constant value. + + Args: + masking_method: See + [`NormalizationTransform`][torchio.transforms.preprocessing.intensity.NormalizationTransform]. + outside_value: Value to set for all voxels outside of the mask. + labels: If a label map is used to generate the mask, + sequence of labels to consider. If `None`, all values larger than + zero will be used for the mask. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Raises: + RuntimeWarning: If a 4D image is masked with a 3D mask, the mask will + be expanded along the channels (first) dimension, and a warning + will be raised. + + Examples: + >>> import torchio as tio + >>> subject = tio.datasets.Colin27() + >>> subject + Colin27(Keys: ('t1', 'head', 'brain'); images: 3) + >>> mask = tio.Mask(masking_method='brain') # Use "brain" image to mask + >>> transformed = mask(subject) # Set voxels outside of the brain to 0 + + """ + + def __init__( + self, + masking_method: TypeMaskingMethod, + outside_value: float = 0, + labels: Sequence[int] | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.masking_method = masking_method + self.masking_labels = labels + self.outside_value = outside_value + self.args_names = ['masking_method'] + + def apply_transform(self, subject: Subject) -> Subject: + for image in self.get_images(subject): + mask_data = self.get_mask_from_masking_method( + self.masking_method, + subject, + image.data, + self.masking_labels, + ) + assert isinstance(image, ScalarImage) + self.apply_masking(image, mask_data) + return subject + + def apply_masking( + self, + image: ScalarImage, + mask_data: torch.Tensor, + ) -> None: + masked = mask(image.data, mask_data, self.outside_value) + image.set_data(masked) + + +def mask( + tensor: torch.Tensor, + mask: torch.Tensor, + outside_value: float, +) -> torch.Tensor: + array = tensor.clone() + num_channels_array = array.shape[0] + num_channels_mask = mask.shape[0] + if num_channels_array != num_channels_mask: + assert num_channels_mask == 1 + message = ( + f'Expanding mask with shape {mask.shape}' + f' to match shape {array.shape} of input image' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + mask = mask.expand(*array.shape) + array[~mask] = outside_value + return array diff --git a/src/torchio/transforms/preprocessing/intensity/normalization_transform.py b/src/torchio/transforms/preprocessing/intensity/normalization_transform.py new file mode 100644 index 000000000..df25a2e69 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/normalization_transform.py @@ -0,0 +1,60 @@ +import torch + +from ....data.subject import Subject +from ....transforms.transform import TypeMaskingMethod +from ...intensity_transform import IntensityTransform + + +class NormalizationTransform(IntensityTransform): + """Base class for intensity preprocessing transforms. + + Args: + masking_method: Defines the mask used to compute the normalization statistics. It can be one of: + + - `None`: the mask image is all ones, i.e. all values in the image are used. + + - A string: key to a [`torchio.LabelMap`][torchio.LabelMap] in the subject which is used as a mask, + OR an anatomical label: `'Left'`, `'Right'`, `'Anterior'`, `'Posterior'`, + `'Inferior'`, `'Superior'` which specifies a side of the mask volume to be ones. + + - A function: the mask image is computed as a function of the intensity image. + The function must receive and return a [`torch.Tensor`][torch.Tensor] + + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> subject = tio.datasets.Colin27() + >>> subject + Colin27(Keys: ('t1', 'head', 'brain'); images: 3) + >>> transform = tio.ZNormalization() # ZNormalization is a subclass of NormalizationTransform + >>> transformed = transform(subject) # use all values to compute mean and std + >>> transform = tio.ZNormalization(masking_method='brain') + >>> transformed = transform(subject) # use only values within the brain + >>> transform = tio.ZNormalization(masking_method=lambda x: x > x.mean()) + >>> transformed = transform(subject) # use values above the image mean + """ + + def __init__(self, masking_method: TypeMaskingMethod = None, **kwargs): + super().__init__(**kwargs) + self.masking_method = masking_method + + def apply_transform(self, subject: Subject) -> Subject: + for image_name, image in self.get_images_dict(subject).items(): + mask = self.get_mask_from_masking_method( + self.masking_method, + subject, + image.data, + ) + self.apply_normalization(subject, image_name, mask) + return subject + + def apply_normalization( + self, + subject: Subject, + image_name: str, + mask: torch.Tensor, + ) -> None: + # There must be a nicer way of doing this + raise NotImplementedError diff --git a/src/torchio/transforms/preprocessing/intensity/pca.py b/src/torchio/transforms/preprocessing/intensity/pca.py new file mode 100644 index 000000000..a51c788f0 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/pca.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +from einops import rearrange + +from ....data.image import ScalarImage +from ....data.subject import Subject +from ....external.imports import get_sklearn +from ...intensity_transform import IntensityTransform + + +class PCA(IntensityTransform): + """Compute principal component analysis (PCA) of an image. + + PCA can be useful to visualize embeddings generated by a neural network. + See for example Figure 8 in [Cluster and Predict Latent Patches for + Improved Masked Image Modeling ](https://arxiv.org/abs/2502.08769). + + Args: + num_components: Number of components to compute. + whiten: If `True`, the components are normalized to have unit variance. + normalize: If `True`, all components are divided by the standard + deviation of the first component. + make_skewness_positive: If `True`, the skewness of each component is + made positive by multiplying the component by -1 if its skewness is + negative. + values_range: If not `None`, these values are linearly mapped to + $[0, 1]$. + clip: If `True`, the output values are clipped to $[0, 1]$. + pca_kwargs: Additional keyword arguments to pass to + [sklearn.decomposition.PCA][sklearn.decomposition.PCA]. + + Examples: + >>> import torchio as tio + >>> from torchio.visualization import build_image_from_reference + >>> ct = my_preprocessed_ct_image # Assume this is a preprocessed CT image + >>> ct + ScalarImage(shape: (1, 240, 480, 480); spacing: (1.50, 0.75, 0.75); orientation: SLP+; dtype: torch.FloatTensor; memory: 210.9 MiB) + >>> embedding_tensor = model(ct.data[None])[0] # `model` is some pre-trained neural network + >>> embedding_image = ToReferenceSpace(ct)(embedding_tensor) + >>> embedding_image + ScalarImage(shape: (512, 24, 24, 24); spacing: (15.00, 15.00, 15.00); orientation: SLP+; dtype: torch.FloatTensor; memory: 27.0 MiB) + >>> pca = tio.PCA()(embedding_image) + >>> pca + ScalarImage(shape: (3, 24, 24, 24); spacing: (15.00, 15.00, 15.00); orientation: SLP+; dtype: torch.FloatTensor; memory: 162.0 KiB) + """ + + def __init__( + self, + num_components: int = 3, + *, + whiten: bool = True, + normalize: bool = True, + make_skewness_positive: bool = True, + values_range: tuple[float, float] | None = (-2.3, 2.3), + clip: bool = True, + pca_kwargs: dict[str, Any] | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.num_components = num_components + self.whiten = whiten + self.normalize = normalize + self.make_skewness_positive = make_skewness_positive + self.values_range = values_range + self.clip = clip + self.pca_kwargs = pca_kwargs + self.args_names = [ + 'num_components', + 'whiten', + 'normalize', + 'make_skewness_positive', + 'values_range', + 'clip', + 'pca_kwargs', + ] + + def apply_transform(self, subject: Subject) -> Subject: + for image in self.get_images(subject): + kwargs = {} if self.pca_kwargs is None else self.pca_kwargs + pca_image = _compute_pca( + image, + num_components=self.num_components, + whiten=self.whiten, + normalize=self.normalize, + make_skewness_positive=self.make_skewness_positive, + values_range=self.values_range, + clip=self.clip, + **kwargs, + ) + image.set_data(pca_image.data) + return subject + + +def _compute_pca( + embeddings: ScalarImage, + num_components: int, + whiten: bool, + normalize: bool, + make_skewness_positive: bool, + values_range: tuple[float, float] | None, + clip: bool, + **pca_kwargs, +) -> ScalarImage: + # Adapted from https://github.com/facebookresearch/capi/blob/main/eval_visualizations.py + # 2.3 is roughly 2σ for a standard-normal variable, 99% of values map inside [0,1]. + sklearn = get_sklearn() + PCA = sklearn.decomposition.PCA + + data = embeddings.numpy() + _, size_x, size_y, size_z = data.shape + X = rearrange(data, 'c x y z -> (x y z) c') + pca = PCA(n_components=num_components, whiten=whiten, **pca_kwargs) + projected: np.ndarray = pca.fit_transform(X).T + if normalize: + projected /= projected[0].std() + if make_skewness_positive: + for component in projected: + third_cumulant = np.mean(component**3) + second_cumulant = np.mean(component**2) + skewness = third_cumulant / second_cumulant ** (3 / 2) + if skewness < 0: + component *= -1 + grid: np.ndarray = rearrange( + projected, + 'c (x y z) -> c x y z', + x=size_x, + y=size_y, + z=size_z, + ) + if values_range is not None: + vmin, vmax = values_range + else: + vmin, vmax = grid.min(), grid.max() + grid = (grid - vmin) / (vmax - vmin) + if clip: + grid = np.clip(grid, 0, 1) + return ScalarImage(tensor=grid, affine=embeddings.affine) diff --git a/src/torchio/transforms/preprocessing/intensity/rescale.py b/src/torchio/transforms/preprocessing/intensity/rescale.py new file mode 100644 index 000000000..73e4d81bc --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/rescale.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import warnings + +import numpy as np +import torch + +from ....data.subject import Subject +from ....types import TypeDoubleFloat +from .normalization_transform import NormalizationTransform +from .normalization_transform import TypeMaskingMethod + + +class RescaleIntensity(NormalizationTransform): + """Rescale intensity values to a certain range. + + Args: + out_min_max: Range $(n_{min}, n_{max})$ of output intensities. + If only one value $d$ is provided, + $(n_{min}, n_{max}) = (-d, d)$. + percentiles: Percentile values of the input image that will be mapped + to $(n_{min}, n_{max})$. They can be used for contrast + stretching, as in [this scikit-image example](https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_equalize.html#sphx-glr-auto-examples-color-exposure-plot-equalize-py). For example, + Isensee et al. use `(0.5, 99.5)` in their [nn-UNet paper](https://arxiv.org/abs/1809.10486). + If only one value $d$ is provided, + $(n_{min}, n_{max}) = (0, d)$. + masking_method: See + [`NormalizationTransform`][torchio.transforms.preprocessing.intensity.NormalizationTransform]. + in_min_max: Range $(m_{min}, m_{max})$ of input intensities that + will be mapped to $(n_{min}, n_{max})$. If `None`, the + minimum and maximum input intensities will be used. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> ct = tio.ScalarImage('ct_scan.nii.gz') + >>> ct_air, ct_bone = -1000, 1000 + >>> rescale = tio.RescaleIntensity( + ... out_min_max=(-1, 1), in_min_max=(ct_air, ct_bone)) + >>> ct_normalized = rescale(ct) + + """ + + def __init__( + self, + out_min_max: TypeDoubleFloat = (0, 1), + percentiles: TypeDoubleFloat = (0, 100), + masking_method: TypeMaskingMethod = None, + in_min_max: TypeDoubleFloat | None = None, + **kwargs, + ): + super().__init__(masking_method=masking_method, **kwargs) + self.out_min_max = out_min_max + self.in_min_max = in_min_max + self.out_min, self.out_max = self._parse_range( + out_min_max, + 'out_min_max', + ) + self.percentiles = self._parse_range( + percentiles, + 'percentiles', + min_constraint=0, + max_constraint=100, + ) + + if self.in_min_max is not None: + self.in_min_max = self._parse_range( + self.in_min_max, + 'in_min_max', + ) + + self.args_names = [ + 'out_min_max', + 'percentiles', + 'masking_method', + 'in_min_max', + ] + + def apply_normalization( + self, + subject: Subject, + image_name: str, + mask: torch.Tensor, + ) -> None: + image = subject.get_scalar_image(image_name) + image.set_data(self.rescale(image.data, mask, image_name)) + + def rescale( + self, + tensor: torch.Tensor, + mask: torch.Tensor, + image_name: str, + ) -> torch.Tensor: + # The tensor is cloned as in-place operations will be used + array = tensor.clone().float().numpy() + mask_array = mask.numpy() + if not mask_array.any(): + message = ( + f'Rescaling image "{image_name}" not possible' + ' because the mask to compute the statistics is empty' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + return tensor + + values = array[mask_array] + cutoff = np.percentile(values, self.percentiles) + np.clip(array, *cutoff, out=array) + + if self.in_min_max is None: + in_min, in_max = array.min(), array.max() + else: + in_min, in_max = self.in_min_max + in_range = in_max - in_min + if in_range == 0: # should this be compared using a tolerance? + message = ( + f'Rescaling image "{image_name}" not possible' + ' because all the intensity values are the same' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + return tensor + + out_range = self.out_max - self.out_min + + array -= in_min + array /= in_range + array *= out_range + array += self.out_min + return torch.as_tensor(array) diff --git a/src/torchio/transforms/preprocessing/intensity/to.py b/src/torchio/transforms/preprocessing/intensity/to.py new file mode 100644 index 000000000..a467ea691 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/to.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from typing import Any + +import torch + +from ....data.image import ScalarImage +from ....data.subject import Subject +from ...intensity_transform import IntensityTransform + + +class To(IntensityTransform): + """Convert the image tensor data type and/or device. + + This transform is a thin wrapper around `torch.Tensor.to()`. + + Args: + target: First argument to `torch.Tensor.to()`. + to_kwargs: Additional keyword arguments to pass to `torch.Tensor.to()`. + + Examples: + >>> import torchio as tio + >>> ct = tio.datasets.Slicer('CTChest').CT_chest + >>> clamp = tio.Clamp(out_min=-1000, out_max=1000) + >>> ct_clamped = clamp(ct) + >>> rescale = tio.RescaleIntensity(in_min_max=(-1000, 1000), out_min_max=(0, 255)) + >>> ct_rescaled = rescale(ct_clamped) + >>> to_uint8 = tio.To(torch.uint8) + >>> ct_uint8 = to_uint8(ct_rescaled) + """ + + def __init__( + self, + target: str | torch.dtype | torch.device, + to_kwargs: dict[str, Any] | None = None, + **kwargs, + ): + super().__init__(**kwargs) + self.target = target + if to_kwargs is None: + to_kwargs = {} + self.to_kwargs = to_kwargs + self.args_names = ['target', 'to_kwargs'] + + def apply_transform(self, subject: Subject) -> Subject: + for image in self.get_images(subject): + assert isinstance(image, ScalarImage) + image.set_data(image.data.to(self.target, **self.to_kwargs)) + return subject diff --git a/src/torchio/transforms/preprocessing/intensity/z_normalization.py b/src/torchio/transforms/preprocessing/intensity/z_normalization.py new file mode 100644 index 000000000..9940bec07 --- /dev/null +++ b/src/torchio/transforms/preprocessing/intensity/z_normalization.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import torch + +from ....data.subject import Subject +from .normalization_transform import NormalizationTransform +from .normalization_transform import TypeMaskingMethod + + +class ZNormalization(NormalizationTransform): + """Subtract mean and divide by standard deviation. + + Args: + masking_method: See + [`NormalizationTransform`][torchio.transforms.preprocessing.intensity.NormalizationTransform]. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__(self, masking_method: TypeMaskingMethod = None, **kwargs): + super().__init__(masking_method=masking_method, **kwargs) + self.args_names = ['masking_method'] + + def apply_normalization( + self, + subject: Subject, + image_name: str, + mask: torch.Tensor, + ) -> None: + image = subject.get_scalar_image(image_name) + standardized = self.znorm( + image.data, + mask, + ) + if standardized is None: + message = ( + 'Standard deviation is 0 for masked values' + f' in image "{image_name}" ({image.path})' + ) + raise RuntimeError(message) + image.set_data(standardized) + + @staticmethod + def znorm( + tensor: torch.Tensor, + mask: torch.Tensor, + ) -> torch.Tensor | None: + tensor = tensor.clone().float() + values = tensor[mask] + mean, std = values.mean(), values.std() + if std == 0: + return None + tensor -= mean + tensor /= std + return tensor diff --git a/src/torchio/transforms/preprocessing/label/__init__.py b/src/torchio/transforms/preprocessing/label/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/torchio/transforms/preprocessing/label/contour.py b/src/torchio/transforms/preprocessing/label/contour.py new file mode 100644 index 000000000..2e22816c5 --- /dev/null +++ b/src/torchio/transforms/preprocessing/label/contour.py @@ -0,0 +1,26 @@ +import SimpleITK as sitk + +from .label_transform import LabelTransform + + +class Contour(LabelTransform): + r"""Keep only the borders of each connected component in a binary image. + + Args: + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def apply_transform(self, subject): + for image in self.get_images(subject): + if image.num_channels > 1: + message = ( + 'The number of input channels must be 1,' + f' but it is {image.num_channels}' + ) + raise RuntimeError(message) + sitk_image = image.as_sitk() + contour = sitk.BinaryContour(sitk_image) + tensor, _ = self.sitk_to_nib(contour) + image.set_data(tensor) + return subject diff --git a/src/torchio/transforms/preprocessing/label/keep_largest_component.py b/src/torchio/transforms/preprocessing/label/keep_largest_component.py new file mode 100644 index 000000000..ea40d2c5c --- /dev/null +++ b/src/torchio/transforms/preprocessing/label/keep_largest_component.py @@ -0,0 +1,35 @@ +import SimpleITK as sitk + +from ....data.subject import Subject +from .label_transform import LabelTransform + + +class KeepLargestComponent(LabelTransform): + r"""Keep only the largest connected component in a binary label map. + + Args: + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Note: + For now, this transform only works for binary images, i.e., label + maps with a background and a foreground class. If you are interested in + extending this transform, please [open a new issue](https://github.com/TorchIO-project/torchio/issues/new?assignees=&labels=enhancement&template=feature_request.md&title=Improve%20KeepLargestComponent%20transform). + + """ + + def apply_transform(self, subject: Subject) -> Subject: + for image in self.get_images(subject): + if image.num_channels > 1: + message = ( + 'The number of input channels must be 1,' + f' but it is {image.num_channels}' + ) + raise RuntimeError(message) + sitk_image = image.as_sitk() + connected_components = sitk.ConnectedComponent(sitk_image) + labeled_cc = sitk.RelabelComponent(connected_components) + largest_cc = labeled_cc == 1 + tensor, _ = self.sitk_to_nib(largest_cc) + image.set_data(tensor) + return subject diff --git a/src/torchio/transforms/preprocessing/label/label_transform.py b/src/torchio/transforms/preprocessing/label/label_transform.py new file mode 100644 index 000000000..67ce59af6 --- /dev/null +++ b/src/torchio/transforms/preprocessing/label/label_transform.py @@ -0,0 +1,23 @@ +from ....data.image import LabelMap +from ....data.subject import Subject +from ...transform import Transform + + +class LabelTransform(Transform): + """Transform that modifies label maps.""" + + def get_images(self, subject: Subject) -> list[LabelMap]: + images = subject.get_images( + intensity_only=False, + include=self.include, + exclude=self.exclude, + ) + return [im for im in images if isinstance(im, LabelMap)] + + def get_images_dict(self, subject: Subject) -> dict[str, LabelMap]: + images = subject.get_images_dict( + intensity_only=False, + include=self.include, + exclude=self.exclude, + ) + return {k: v for (k, v) in images.items() if isinstance(v, LabelMap)} diff --git a/src/torchio/transforms/preprocessing/label/one_hot.py b/src/torchio/transforms/preprocessing/label/one_hot.py new file mode 100644 index 000000000..70c66b253 --- /dev/null +++ b/src/torchio/transforms/preprocessing/label/one_hot.py @@ -0,0 +1,45 @@ +import torch.nn.functional as F # noqa: N812 + +from ....data.image import Image +from .label_transform import LabelTransform + + +class OneHot(LabelTransform): + r"""Reencode label maps using one-hot encoding. + + Args: + num_classes: See [`torch.nn.functional.one_hot`][torch.nn.functional.one_hot]. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__(self, num_classes: int = -1, **kwargs): + super().__init__(**kwargs) + self.num_classes = num_classes + self.args_names = ['num_classes'] + self.invert_transform = False + + def apply_transform(self, subject): + for image in self.get_images(subject): + if self.invert_transform: + self.argmax(image) + else: + self.one_hot(image) + return subject + + @staticmethod + def argmax(image: Image) -> None: + data = image.data.argmax(dim=0, keepdim=True) + image.set_data(data) + + def one_hot(self, image: Image) -> None: + if image.num_channels > 1: + message = ( + 'The number of input channels must be 1,' + f' but it is {image.num_channels}' + ) + raise RuntimeError(message) + data = image.data[0] + num_classes = -1 if self.num_classes is None else self.num_classes + one_hot = F.one_hot(data.long(), num_classes=num_classes) + image.set_data(one_hot.permute(3, 0, 1, 2).type(data.type())) diff --git a/src/torchio/transforms/preprocessing/label/remap_labels.py b/src/torchio/transforms/preprocessing/label/remap_labels.py new file mode 100644 index 000000000..ba9e6894a --- /dev/null +++ b/src/torchio/transforms/preprocessing/label/remap_labels.py @@ -0,0 +1,137 @@ +from ...transform import TypeMaskingMethod +from .label_transform import LabelTransform + + +class RemapLabels(LabelTransform): + r"""Modify labels in a label map. + + Masking can be used to split the label into two during + the [inverse transformation ](../index.md#invertibility). + + Args: + remapping: Dictionary that specifies how labels should be remapped. + The keys are the old labels, and the corresponding values replace + them. + masking_method: Defines a mask for where the label remapping is applied. It can be one of: + + - `None`: the mask image is all ones, i.e. all values in the image are used. + + - A string: key to a [`torchio.LabelMap`][torchio.LabelMap] in the subject which is used as a mask, + OR an anatomical label: `'Left'`, `'Right'`, `'Anterior'`, `'Posterior'`, + `'Inferior'`, `'Superior'` which specifies a half of the mask volume to be ones. + + - A function: the mask image is computed as a function of the intensity image. + The function must receive and return a 4D [`torch.Tensor`][torch.Tensor]. + + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + + Examples: + >>> import torch + >>> import torchio as tio + >>> def get_image(*labels): + ... tensor = torch.as_tensor(labels).reshape(1, 1, 1, -1) + ... image = tio.LabelMap(tensor=tensor) + ... return image + ... + >>> image = get_image(0, 1, 2, 3, 4) + >>> remapping = {1: 2, 2: 1, 3: 1, 4: 7} + >>> transform = tio.RemapLabels(remapping) + >>> transform(image).data + tensor([[[[0, 2, 1, 1, 7]]]]) + + Warning: + The transform will not be correctly inverted if one of the values in + `remapping` is also in the input image: + + ```pycon + >>> tensor = torch.as_tensor([0, 1]).reshape(1, 1, 1, -1) + >>> subject = tio.Subject(label=tio.LabelMap(tensor=tensor)) + >>> mapping = {3: 1} # the value 1 is in the input image + >>> transform = tio.RemapLabels(mapping) + >>> transformed = transform(subject) + >>> back = transformed.apply_inverse_transform() + >>> original_label_set = set(subject.label.data.unique().tolist()) + >>> back_label_set = set(back.label.data.unique().tolist()) + >>> original_label_set + {0, 1} + >>> back_label_set + {0, 3} + ``` + + Examples: + >>> import torchio as tio + >>> # Target label map has the following labels: + >>> # { + >>> # 'left_ventricle': 1, 'right_ventricle': 2, + >>> # 'left_caudate': 3, 'right_caudate': 4, + >>> # 'left_putamen': 5, 'right_putamen': 6, + >>> # 'left_thalamus': 7, 'right_thalamus': 8, + >>> # } + >>> transform = tio.RemapLabels({2:1, 4:3, 6:5, 8:7}) + >>> # Merge right side labels with left side labels + >>> transformed = transform(subject) + >>> # Undesired behavior: The inverse transform will remap ALL left side labels to right side labels + >>> # so the label map only has right side labels. + >>> inverse_transformed = transformed.apply_inverse_transform() + >>> # Here's the *right* way to do it with masking: + >>> transform = tio.RemapLabels({2:1, 4:3, 6:5, 8:7}, masking_method="Right") + >>> # Remap the labels on the right side only (no difference yet). + >>> transformed = transform(subject) + >>> # Apply the inverse on the right side only. The labels are correctly split into left/right. + >>> inverse_transformed = transformed.apply_inverse_transform() + + """ + + def __init__( + self, + remapping: dict[int, int], + masking_method: TypeMaskingMethod = None, + **kwargs, + ): + super().__init__(**kwargs) + self.kwargs = kwargs + self.remapping = remapping + self.masking_method = masking_method + self.args_names = ['remapping', 'masking_method'] + + def apply_transform(self, subject): + for image in self.get_images(subject): + original_label_set = set(image.data.unique().tolist()) + source_label_set = set(self.remapping.keys()) + # Do nothing if no keys in the mapping are found in the image + if not source_label_set.intersection(original_label_set): + continue + new_data = image.data.clone() + mask = self.get_mask_from_masking_method( + self.masking_method, + subject, + new_data, + ) + for old_id, new_id in self.remapping.items(): + new_data[mask & (image.data == old_id)] = new_id + image.set_data(new_data) + + return subject + + def is_invertible(self): + # Not always, as explained in the docstring + return True + + def inverse(self): + targets = self.remapping.values() + unique_targets = set(targets) + if len(unique_targets) < len(targets): + message = ( + 'Labels mapping cannot be inverted because original values' + f' are not unique: {self.remapping}' + ) + raise RuntimeError(message) + inverse_remapping = {v: k for k, v in self.remapping.items()} + inverse_transform = RemapLabels( + inverse_remapping, + masking_method=self.masking_method, + **self.kwargs, + ) + return inverse_transform diff --git a/src/torchio/transforms/preprocessing/label/remove_labels.py b/src/torchio/transforms/preprocessing/label/remove_labels.py new file mode 100644 index 000000000..e1d3fe555 --- /dev/null +++ b/src/torchio/transforms/preprocessing/label/remove_labels.py @@ -0,0 +1,39 @@ +from collections.abc import Sequence + +from ...transform import TypeMaskingMethod +from .remap_labels import RemapLabels + + +class RemoveLabels(RemapLabels): + r"""Remove labels from a label map. + + The removed labels are remapped to the background label. + + This transformation is not [invertible ](../index.md#invertibility). + + Args: + labels: A sequence of label integers that will be removed. + background_label: integer that specifies which label is considered to + be background (typically, `0`). + masking_method: See [`RemapLabels`][torchio.transforms.RemapLabels]. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + """ + + def __init__( + self, + labels: Sequence[int], + background_label: int = 0, + masking_method: TypeMaskingMethod = None, + **kwargs, + ): + remapping = {label: background_label for label in labels} + super().__init__(remapping, masking_method, **kwargs) + self.labels = labels + self.background_label = background_label + self.masking_method = masking_method + self.args_names = ['labels', 'background_label', 'masking_method'] + + def is_invertible(self): + return False diff --git a/src/torchio/transforms/preprocessing/label/sequential_labels.py b/src/torchio/transforms/preprocessing/label/sequential_labels.py new file mode 100644 index 000000000..2675a75af --- /dev/null +++ b/src/torchio/transforms/preprocessing/label/sequential_labels.py @@ -0,0 +1,63 @@ +import torch + +from ...transform import TypeMaskingMethod +from .label_transform import LabelTransform +from .remap_labels import RemapLabels + + +class SequentialLabels(LabelTransform): + r"""Remap labels in a label map so they become consecutive. + + For example, if a label map has labels `(0, 3, 5)`, then this will apply + a [`RemapLabels`][torchio.RemapLabels] transform with `remapping={3: 1, 5: 2}`, + and therefore the output image will have labels `(0, 1, 2)`. + + Examples: + >>> import torch + >>> import torchio as tio + >>> def get_image(*labels): + ... tensor = torch.as_tensor(labels).reshape(1, 1, 1, -1) + ... image = tio.LabelMap(tensor=tensor) + ... return image + ... + >>> img_with_bg = get_image(0, 5, 10) + >>> transform = tio.SequentialLabels() + >>> transform(img_with_bg).data + tensor([[[[0, 1, 2]]]]) + >>> img_without_bg = get_image(7, 11, 99) + >>> transform(img_without_bg).data + tensor([[[[0, 1, 2]]]]) + + Note: + This transformation is always [fully invertible ](../index.md#invertibility). + + Warning: + The background is typically represented with the label `0`. There + will be zeros in the output image even if they are none in the input. + + Args: + masking_method: See [`RemapLabels`][torchio.transforms.RemapLabels]. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__(self, masking_method: TypeMaskingMethod = None, **kwargs): + super().__init__(**kwargs) + self.masking_method = masking_method + + def apply_transform(self, subject): + for name, image in self.get_images_dict(subject).items(): + unique_labels = torch.unique(image.data) + remapping = { + unique_labels[i].item(): i for i in range(0, len(unique_labels)) + } + init_kwargs = self._get_base_args() + init_kwargs['include'] = [name] + + transform = RemapLabels( + remapping=remapping, + masking_method=self.masking_method, + **init_kwargs, + ) + subject = transform(subject) + return subject diff --git a/src/torchio/transforms/preprocessing/spatial/__init__.py b/src/torchio/transforms/preprocessing/spatial/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/torchio/transforms/preprocessing/spatial/bounds_transform.py b/src/torchio/transforms/preprocessing/spatial/bounds_transform.py new file mode 100644 index 000000000..efe76e750 --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/bounds_transform.py @@ -0,0 +1,20 @@ +from ....transforms.transform import TypeBounds +from ...spatial_transform import SpatialTransform + + +class BoundsTransform(SpatialTransform): + """Base class for transforms that change image bounds. + + Args: + bounds_parameters: The meaning of this argument varies according to the + child class. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + """ + + def __init__(self, bounds_parameters: TypeBounds, **kwargs): + super().__init__(**kwargs) + self.bounds_parameters = self.parse_bounds(bounds_parameters) + + def is_invertible(self): + return True diff --git a/src/torchio/transforms/preprocessing/spatial/copy_affine.py b/src/torchio/transforms/preprocessing/spatial/copy_affine.py new file mode 100644 index 000000000..64ca8a343 --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/copy_affine.py @@ -0,0 +1,88 @@ +import copy + +from ....data.subject import Subject +from ...spatial_transform import SpatialTransform + + +class CopyAffine(SpatialTransform): + """Copy the spatial metadata from a reference image in the subject. + + Small unexpected differences in spatial metadata across different images + of a subject can arise due to rounding errors while converting formats. + + If the `shape` and `orientation` of the images are the same and their + `affine` attributes are different but very similar, this transform can be + used to avoid errors during safety checks in other transforms and samplers. + + Args: + target: Name of the image within the subject whose affine matrix will + be used. + + Examples: + >>> import torch + >>> import torchio as tio + >>> import numpy as np + >>> np.random.seed(0) + >>> affine = np.diag((*(np.random.rand(3) + 0.5), 1)) + >>> t1 = tio.ScalarImage(tensor=torch.rand(1, 100, 100, 100), affine=affine) + >>> # Let's simulate a loss of precision + >>> # (caused for example by NIfTI storing spatial metadata in single precision) + >>> bad_affine = affine.astype(np.float16) + >>> t2 = tio.ScalarImage(tensor=torch.rand(1, 100, 100, 100), affine=bad_affine) + >>> subject = tio.Subject(t1=t1, t2=t2) + >>> resample = tio.Resample(0.5) + >>> resample(subject).shape # error as images are in different spaces + Traceback (most recent call last): + File "", line 1, in + File "/Users/fernando/git/torchio/torchio/data/subject.py", line 101, in shape + self.check_consistent_attribute('shape') + File "/Users/fernando/git/torchio/torchio/data/subject.py", line 229, in check_consistent_attribute + raise RuntimeError(message) + RuntimeError: More than one shape found in subject images: + {'t1': (1, 210, 244, 221), 't2': (1, 210, 243, 221)} + >>> transform = tio.CopyAffine('t1') + >>> fixed = transform(subject) + >>> resample(fixed).shape + (1, 210, 244, 221) + + + Warning: + This transform should be used with caution. Modifying the + spatial metadata of an image manually can lead to incorrect processing + of the position of anatomical structures. For example, a machine + learning algorithm might incorrectly predict that a lesion on the right + lung is on the left lung. + + Note: + For more information, see some related discussions on GitHub: + + * https://github.com/TorchIO-project/torchio/issues/354 + * https://github.com/TorchIO-project/torchio/discussions/489 + * https://github.com/TorchIO-project/torchio/pull/584 + * https://github.com/TorchIO-project/torchio/issues/430 + * https://github.com/TorchIO-project/torchio/issues/382 + * https://github.com/TorchIO-project/torchio/pull/592 + """ + + def __init__(self, target: str, **kwargs): + super().__init__(**kwargs) + if not isinstance(target, str): + message = f'The target must be a string, but "{type(target)}" was found' + raise ValueError(message) + self.target = target + self.args_names = ['target'] + + def apply_transform(self, subject: Subject) -> Subject: + if self.target not in subject: + message = f'Target image "{self.target}" not found in subject' + raise RuntimeError(message) + reference = subject.get_image(self.target) + affine = copy.deepcopy(reference.affine) + for image in self.get_images(subject): + if image is reference: + continue + # We load the image to avoid complications + # https://github.com/TorchIO-project/torchio/issues/1071#issuecomment-1511814720 + image.load() + image.affine = affine + return subject diff --git a/src/torchio/transforms/preprocessing/spatial/crop.py b/src/torchio/transforms/preprocessing/spatial/crop.py new file mode 100644 index 000000000..dd1d5e06b --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/crop.py @@ -0,0 +1,129 @@ +from copy import deepcopy + +import numpy as np +from nibabel.affines import apply_affine + +from ....data.image import Image +from ....data.subject import Subject +from .bounds_transform import BoundsTransform +from .bounds_transform import TypeBounds + + +class Crop(BoundsTransform): + r"""Crop an image. + + Args: + cropping: Tuple + $(w_{ini}, w_{fin}, h_{ini}, h_{fin}, d_{ini}, d_{fin})$ + defining the number of values cropped from the edges of each axis. + If the initial shape of the image is + $W \times H \times D$, the final shape will be + $(- w_{ini} + W - w_{fin}) \times (- h_{ini} + H - h_{fin}) + \times (- d_{ini} + D - d_{fin})$. + If only three values $(w, h, d)$ are provided, then + $w_{ini} = w_{fin} = w$, + $h_{ini} = h_{fin} = h$ and + $d_{ini} = d_{fin} = d$. + If only one value $n$ is provided, then + $w_{ini} = w_{fin} = h_{ini} = h_{fin} + = d_{ini} = d_{fin} = n$. + copy: If `True`, each image will be cropped and the patch copied to a new + subject. If `False`, each image will be cropped in place. This transform + overwrites the copy argument of the base transform and copies only the + cropped patch instead of the whole image. This can provide a significant + speedup when cropping small patches from large images. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + See also: If you want to pass the output shape instead, please use + [`CropOrPad`][torchio.transforms.CropOrPad] instead. + """ + + def __init__(self, cropping: TypeBounds, copy=True, **kwargs): + self._copy_patch = copy + # Transform base class deepcopies whole subject by default + # We want to copy only the cropped patch, so we overwrite the functionality + super().__init__(cropping, copy=False, **kwargs) + self.cropping = cropping + self.args_names = ['cropping'] + + def apply_transform(self, subject: Subject) -> Subject: + assert self.bounds_parameters is not None + low = self.bounds_parameters[::2] + high = self.bounds_parameters[1::2] + index_ini = low + index_fin = np.array(subject.spatial_shape) - high + + if self._copy_patch: + # Create a clean new subject to copy the images into + # We do this __new__ to avoid calling __init__ so we don't have to specify images immediately + cropped_subject = subject.__class__.__new__(subject.__class__) + image_keys_to_crop = subject.get_images_dict( + intensity_only=False, + include=self.include, + exclude=self.exclude, + ).keys() + keys_to_expose = subject.keys() + # Copy all attributes we don't want to crop + # __dict__ returns all attributes, instead of just the images + for key, value in subject.__dict__.items(): + if key not in image_keys_to_crop: + copied_value = deepcopy(value) + # Setting __dict__ does not allow key indexing the attribute + # so we set it explicitly if we want to expose it + if key in keys_to_expose: + cropped_subject[key] = copied_value + cropped_subject.__dict__[str(key)] = copied_value + else: + # Images are always exposed, so we don't worry about setting __dict__ + cropped_subject[key] = self._crop_image( + value, + index_ini, + index_fin, + copy_patch=self._copy_patch, + ) + + # Update the __dict__ attribute to include the cropped images + cropped_subject.update_attributes() + return cropped_subject + else: + # Crop in place + for image in self.get_images(subject): + self._crop_image( + image, + index_ini, + index_fin, + copy_patch=self._copy_patch, + ) + return subject + + @staticmethod + def _crop_image( + image: Image, index_ini: tuple, index_fin: tuple, *, copy_patch: bool + ) -> Image: + new_origin = apply_affine(image.affine, index_ini) + new_affine = image.affine.copy() + new_affine[:3, 3] = new_origin + i0, j0, k0 = index_ini + i1, j1, k1 = index_fin + + # Crop the image data + if copy_patch: + # Create a new image with the cropped data + cropped_data = image.data[:, i0:i1, j0:j1, k0:k1].clone() + new_image = type(image)( + tensor=cropped_data, + affine=new_affine, + type=image.type, + path=image.path, + ) + return new_image + else: + image.set_data(image.data[:, i0:i1, j0:j1, k0:k1].clone()) + image.affine = new_affine + return image + + def inverse(self): + from .pad import Pad + + return Pad(self.cropping) diff --git a/src/torchio/transforms/preprocessing/spatial/crop_or_pad.py b/src/torchio/transforms/preprocessing/spatial/crop_or_pad.py new file mode 100644 index 000000000..5ae211fbd --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/crop_or_pad.py @@ -0,0 +1,337 @@ +from __future__ import annotations + +import warnings +from collections.abc import Sequence + +import numpy as np + +from ....data.subject import Subject +from ....utils import parse_spatial_shape +from ...spatial_transform import SpatialTransform +from ...transform import TypeSixBounds +from ...transform import TypeTripletInt +from .crop import Crop +from .pad import Pad + + +class CropOrPad(SpatialTransform): + """Modify the field of view by cropping or padding to match a target shape. + + This transform modifies the affine matrix associated to the volume so that + physical positions of the voxels are maintained. + + Args: + target_shape: Tuple $(W, H, D)$. If a single value $N$ is + provided, then $W = H = D = N$. If `None`, the shape will + be computed from the `mask_name` (and the `labels`, if + `labels` is not `None`). + padding_mode: Same as `padding_mode` in + [`Pad`][torchio.transforms.Pad]. + mask_name: If `None`, the centers of the input and output volumes + will be the same. + If a string is given, the output volume center will be the center + of the bounding box of non-zero values in the image named + `mask_name`. + labels: If a label map is used to generate the mask, sequence of labels + to consider. + only_crop: If `True`, padding will not be applied, only cropping will + be done. `only_crop` and `only_pad` cannot both be `True`. + only_pad: If `True`, cropping will not be applied, only padding will + be done. `only_crop` and `only_pad` cannot both be `True`. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> subject = tio.Subject( + ... chest_ct=tio.ScalarImage('subject_a_ct.nii.gz'), + ... heart_mask=tio.LabelMap('subject_a_heart_seg.nii.gz'), + ... ) + >>> subject.chest_ct.shape + torch.Size([1, 512, 512, 289]) + >>> transform = tio.CropOrPad( + ... (120, 80, 180), + ... mask_name='heart_mask', + ... ) + >>> transformed = transform(subject) + >>> transformed.chest_ct.shape + torch.Size([1, 120, 80, 180]) + + Warning: + If `target_shape` is `None`, subjects in the dataset + will probably have different shapes. This is probably fine if you are + using [patch-based training ](https://docs.torchio.org/patches/index.html). + If you are using full volumes for training and a batch size larger than + one, an error will be raised by the [`DataLoader`][torch.utils.data.DataLoader] + while trying to collate the batches. + + """ + + def __init__( + self, + target_shape: int | TypeTripletInt | None = None, + padding_mode: str | float = 0, + mask_name: str | None = None, + labels: Sequence[int] | None = None, + only_crop: bool = False, + only_pad: bool = False, + **kwargs, + ): + if target_shape is None and mask_name is None: + message = 'If mask_name is None, a target shape must be passed' + raise ValueError(message) + super().__init__(**kwargs) + if target_shape is None: + self.target_shape = None + else: + self.target_shape = parse_spatial_shape(target_shape) + self.padding_mode = padding_mode + if mask_name is not None and not isinstance(mask_name, str): + message = ( + f'If mask_name is not None, it must be a string, not {type(mask_name)}' + ) + raise ValueError(message) + if mask_name is None: + if labels is not None: + message = ( + 'If mask_name is None, labels should be None,' + f' but "{labels}" was passed' + ) + raise ValueError(message) + self.compute_crop_or_pad = self._compute_center_crop_or_pad + else: + if not isinstance(mask_name, str): + message = ( + 'If mask_name is not None, it must be a string,' + f' not {type(mask_name)}' + ) + raise ValueError(message) + self.compute_crop_or_pad = self._compute_mask_center_crop_or_pad + self.mask_name = mask_name + self.labels = labels + + if only_pad and only_crop: + message = 'only_crop and only_pad cannot both be True' + raise ValueError(message) + self.only_crop = only_crop + self.only_pad = only_pad + self.args_names = [ + 'target_shape', + 'padding_mode', + 'mask_name', + 'labels', + 'only_crop', + 'only_pad', + ] + + @staticmethod + def _bbox_mask(mask_volume: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Return 6 coordinates of a 3D bounding box from a given mask. + + Taken from [this SO question ](https://stackoverflow.com/questions/31400769/bounding-box-of-numpy-array). + + Args: + mask_volume: 3D NumPy array. + """ + i_any = np.any(mask_volume, axis=(1, 2)) + j_any = np.any(mask_volume, axis=(0, 2)) + k_any = np.any(mask_volume, axis=(0, 1)) + i_min, i_max = np.where(i_any)[0][[0, -1]] + j_min, j_max = np.where(j_any)[0][[0, -1]] + k_min, k_max = np.where(k_any)[0][[0, -1]] + bb_min = np.array([i_min, j_min, k_min]) + bb_max = np.array([i_max, j_max, k_max]) + 1 + return bb_min, bb_max + + @staticmethod + def _get_six_bounds_parameters( + parameters: np.ndarray, + ) -> TypeSixBounds: + r"""Compute bounds parameters for ITK filters. + + Args: + parameters: Tuple $(w, h, d)$ with the number of voxels to be + cropped or padded. + + Returns: + Tuple $(w_{ini}, w_{fin}, h_{ini}, h_{fin}, d_{ini}, d_{fin})$, + where $n_{ini} = \left \lceil \frac{n}{2} \right \rceil$ and + $n_{fin} = \left \lfloor \frac{n}{2} \right \rfloor$. + + Examples: + >>> p = np.array((4, 0, 7)) + >>> CropOrPad._get_six_bounds_parameters(p) + (2, 2, 0, 0, 4, 3) + """ + parameters = parameters / 2 + result = [] + for number in parameters: + ini, fin = int(np.ceil(number)), int(np.floor(number)) + result.extend([ini, fin]) + i1, i2, j1, j2, k1, k2 = result + return i1, i2, j1, j2, k1, k2 + + def _compute_cropping_padding_from_shapes( + self, + source_shape: TypeTripletInt, + ) -> tuple[TypeSixBounds | None, TypeSixBounds | None]: + diff_shape = np.array(self.target_shape) - source_shape + + cropping = -np.minimum(diff_shape, 0) + if cropping.any(): + cropping_params = self._get_six_bounds_parameters(cropping) + else: + cropping_params = None + + padding = np.maximum(diff_shape, 0) + if padding.any(): + padding_params = self._get_six_bounds_parameters(padding) + else: + padding_params = None + + return padding_params, cropping_params + + def _compute_center_crop_or_pad( + self, + subject: Subject, + ) -> tuple[TypeSixBounds | None, TypeSixBounds | None]: + source_shape = subject.spatial_shape + parameters = self._compute_cropping_padding_from_shapes(source_shape) + padding_params, cropping_params = parameters + return padding_params, cropping_params + + def _compute_mask_center_crop_or_pad( + self, + subject: Subject, + ) -> tuple[TypeSixBounds | None, TypeSixBounds | None]: + assert self.mask_name is not None + if self.mask_name not in subject: + message = ( + f'Mask name "{self.mask_name}"' + f' not found in subject keys "{tuple(subject.keys())}".' + ' Using volume center instead' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + return self._compute_center_crop_or_pad(subject=subject) + + mask_image = subject.get_image(self.mask_name) + mask_data = self.get_mask_from_masking_method( + self.mask_name, + subject, + mask_image.data, + self.labels, + ).numpy() + + if not np.any(mask_data): + message = ( + f'All values found in the mask "{self.mask_name}"' + ' are zero. Using volume center instead' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + return self._compute_center_crop_or_pad(subject=subject) + + # Let's assume that the center of first voxel is at coordinate 0.5 + # (which is typically not the case) + subject_shape = subject.spatial_shape + bb_min, bb_max = self._bbox_mask(mask_data[0]) + center_mask = np.mean((bb_min, bb_max), axis=0) + padding = [] + cropping = [] + + if self.target_shape is None: + target_shape = bb_max - bb_min + else: + target_shape = self.target_shape + + for dim in range(3): + target_dim = target_shape[dim] + center_dim = center_mask[dim] + subject_dim = subject_shape[dim] + + center_on_index = not (center_dim % 1) + target_even = not (target_dim % 2) + + # Approximation when the center cannot be computed exactly + # The output will be off by half a voxel, but this is just an + # implementation detail + if target_even ^ center_on_index: + center_dim -= 0.5 + + begin = center_dim - target_dim / 2 + if begin >= 0: + crop_ini = begin + pad_ini = 0 + else: + crop_ini = 0 + pad_ini = -begin + + end = center_dim + target_dim / 2 + if end <= subject_dim: + crop_fin = subject_dim - end + pad_fin = 0 + else: + crop_fin = 0 + pad_fin = end - subject_dim + + padding.extend([pad_ini, pad_fin]) + cropping.extend([crop_ini, crop_fin]) + # Conversion for SimpleITK compatibility + padding_array = np.asarray(padding, dtype=int) + cropping_array = np.asarray(cropping, dtype=int) + if padding_array.any(): + padding_values = [int(value) for value in padding_array.tolist()] + padding_params = ( + padding_values[0], + padding_values[1], + padding_values[2], + padding_values[3], + padding_values[4], + padding_values[5], + ) + else: + padding_params = None + if cropping_array.any(): + cropping_values = [int(value) for value in cropping_array.tolist()] + cropping_params = ( + cropping_values[0], + cropping_values[1], + cropping_values[2], + cropping_values[3], + cropping_values[4], + cropping_values[5], + ) + else: + cropping_params = None + return padding_params, cropping_params + + def apply_transform(self, subject: Subject) -> Subject: + subject.check_consistent_space() + padding_params, cropping_params = self.compute_crop_or_pad(subject) + if padding_params is not None and not self.only_crop: + pad = Pad( + padding_params, + padding_mode=self.padding_mode, + copy=self.copy, + include=self.include, + exclude=self.exclude, + keep=self.keep, + parse_input=self.parse_input, + label_keys=self.label_keys, + ) + transformed = pad(subject) + assert isinstance(transformed, Subject) + subject = transformed + if cropping_params is not None and not self.only_pad: + crop = Crop( + cropping_params, + copy=self.copy, + include=self.include, + exclude=self.exclude, + keep=self.keep, + parse_input=self.parse_input, + label_keys=self.label_keys, + ) + transformed = crop(subject) + assert isinstance(transformed, Subject) + subject = transformed + return subject diff --git a/src/torchio/transforms/preprocessing/spatial/ensure_shape_multiple.py b/src/torchio/transforms/preprocessing/spatial/ensure_shape_multiple.py new file mode 100644 index 000000000..5a5f69af5 --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/ensure_shape_multiple.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import numpy as np + +from ....data.subject import Subject +from ....types import TypeTripletInt +from ....utils import to_tuple +from ...spatial_transform import SpatialTransform +from .crop_or_pad import CropOrPad + + +class EnsureShapeMultiple(SpatialTransform): + """Ensure that all values in the image shape are divisible by $n$. + + Some convolutional neural network architectures need that the size of the + input across all spatial dimensions is a power of $2$. + + For example, the canonical 3D U-Net from + [Çiçek et al. ](https://link.springer.com/chapter/10.1007/978-3-319-46723-8_49) + includes three downsampling (pooling) and upsampling operations: + + ![3D U-Net](https://www.researchgate.net/profile/Olaf-Ronneberger/publication/304226155/figure/fig1/AS:375619658502144@1466566113191/The-3D-u-net-architecture-Blue-boxes-represent-feature-maps-The-number-of-channels-is.png) + + Pooling operations in PyTorch round down the output size: + + >>> import torch + >>> x = torch.rand(3, 10, 20, 31) + >>> x_down = torch.nn.functional.max_pool3d(x, 2) + >>> x_down.shape + torch.Size([3, 5, 10, 15]) + + If we upsample this tensor, the original shape is lost: + + >>> x_down_up = torch.nn.functional.interpolate(x_down, scale_factor=2) + >>> x_down_up.shape + torch.Size([3, 10, 20, 30]) + >>> x.shape + torch.Size([3, 10, 20, 31]) + + If we try to concatenate `x_down` and `x_down_up` (to create skip + connections), we will get an error. It is therefore good practice to ensure + that the size of our images is such that concatenations will be safe. + + Note: + In these examples, it's assumed that all convolutions in the + U-Net use padding so that the output size is the same as the input + size. + + The image above shows $3$ downsampling operations, so the input size + along all dimensions should be a multiple of $2^3 = 8$. + + Example (assuming `pip install unet` has been run before): + + >>> import torchio as tio + >>> import unet + >>> net = unet.UNet3D(padding=1) + >>> t1 = tio.datasets.Colin27().t1 + >>> tensor_bad = t1.data.unsqueeze(0) + >>> tensor_bad.shape + torch.Size([1, 1, 181, 217, 181]) + >>> net(tensor_bad).shape + Traceback (most recent call last): + File "", line 1, in + File "/home/fernando/miniconda3/envs/resseg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl + result = self.forward(*input, **kwargs) + File "/home/fernando/miniconda3/envs/resseg/lib/python3.7/site-packages/unet/unet.py", line 122, in forward + x = self.decoder(skip_connections, encoding) + File "/home/fernando/miniconda3/envs/resseg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl + result = self.forward(*input, **kwargs) + File "/home/fernando/miniconda3/envs/resseg/lib/python3.7/site-packages/unet/decoding.py", line 61, in forward + x = decoding_block(skip_connection, x) + File "/home/fernando/miniconda3/envs/resseg/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl + result = self.forward(*input, **kwargs) + File "/home/fernando/miniconda3/envs/resseg/lib/python3.7/site-packages/unet/decoding.py", line 131, in forward + x = torch.cat((skip_connection, x), dim=CHANNELS_DIMENSION) + RuntimeError: Sizes of tensors must match except in dimension 1. Got 45 and 44 in dimension 2 (The offending index is 1) + >>> num_poolings = 3 + >>> fix_shape_unet = tio.EnsureShapeMultiple(2**num_poolings) + >>> t1_fixed = fix_shape_unet(t1) + >>> tensor_ok = t1_fixed.data.unsqueeze(0) + >>> tensor_ok.shape + torch.Size([1, 1, 184, 224, 184]) # as expected + + Args: + target_multiple: Tuple $(n_w, n_h, n_d)$, so that the size of the + output along axis $i$ is a multiple of $n_i$. If a + single value $n$ is provided, then + $n_w = n_h = n_d = n$. + method: Either `'crop'` or `'pad'`. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torchio as tio + >>> image = tio.datasets.Colin27().t1 + >>> image.shape + (1, 181, 217, 181) + >>> transform = tio.EnsureShapeMultiple(8, method='pad') + >>> transformed = transform(image) + >>> transformed.shape + (1, 184, 224, 184) + >>> transform = tio.EnsureShapeMultiple(8, method='crop') + >>> transformed = transform(image) + >>> transformed.shape + (1, 176, 216, 176) + >>> image_2d = image.data[..., :1] + >>> image_2d.shape + torch.Size([1, 181, 217, 1]) + >>> transformed = transform(image_2d) + >>> transformed.shape + torch.Size([1, 176, 216, 1]) + """ + + def __init__( + self, + target_multiple: int | TypeTripletInt, + *, + method: str = 'pad', + **kwargs, + ): + super().__init__(**kwargs) + self.target_multiple = np.array(to_tuple(target_multiple, 3)) + if method not in ('crop', 'pad'): + raise ValueError('Method must be "crop" or "pad"') + self.method = method + + def apply_transform(self, subject: Subject) -> Subject: + source_shape = np.array(subject.spatial_shape, np.uint16) + if self.method == 'crop': + integer_ratio = np.floor(source_shape / self.target_multiple) + else: + integer_ratio = np.ceil(source_shape / self.target_multiple) + target_shape = integer_ratio * self.target_multiple + target_shape = np.maximum(target_shape, 1) + target_shape_values = [ + int(value) for value in target_shape.astype(int).tolist() + ] + transform = CropOrPad( + ( + target_shape_values[0], + target_shape_values[1], + target_shape_values[2], + ), + copy=self.copy, + include=self.include, + exclude=self.exclude, + keep=self.keep, + parse_input=self.parse_input, + label_keys=self.label_keys, + ) + transformed = transform(subject) + assert isinstance(transformed, Subject) + return transformed diff --git a/src/torchio/transforms/preprocessing/spatial/pad.py b/src/torchio/transforms/preprocessing/spatial/pad.py new file mode 100644 index 000000000..9a0d95ad0 --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/pad.py @@ -0,0 +1,154 @@ +import warnings +from numbers import Number +from typing import Literal +from typing import cast + +import numpy as np +import torch +from nibabel.affines import apply_affine + +from ....data.image import Image +from ....data.subject import Subject +from .bounds_transform import BoundsTransform +from .bounds_transform import TypeBounds + +NumpyPadMode = Literal[ + 'empty', + 'edge', + 'wrap', + 'constant', + 'linear_ramp', + 'maximum', + 'mean', + 'median', + 'minimum', + 'reflect', + 'symmetric', +] + + +class Pad(BoundsTransform): + r"""Pad an image. + + Args: + padding: Tuple + $(w_{ini}, w_{fin}, h_{ini}, h_{fin}, d_{ini}, d_{fin})$ + defining the number of values padded to the edges of each axis. + If the initial shape of the image is + $W \times H \times D$, the final shape will be + $(w_{ini} + W + w_{fin}) \times (h_{ini} + H + h_{fin}) + \times (d_{ini} + D + d_{fin})$. + If only three values $(w, h, d)$ are provided, then + $w_{ini} = w_{fin} = w$, + $h_{ini} = h_{fin} = h$ and + $d_{ini} = d_{fin} = d$. + If only one value $n$ is provided, then + $w_{ini} = w_{fin} = h_{ini} = h_{fin} = + d_{ini} = d_{fin} = n$. + padding_mode: See possible modes in [NumPy docs](https://numpy.org/doc/stable/reference/generated/numpy.pad.html). If it is a number, + the mode will be set to `'constant'`. If it is `'mean'`, + `'maximum'`, `'median'` or `'minimum'`, the statistic will be + computed from the whole volume, unlike in NumPy, which computes it + along the padded axis. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + See also: If you want to pass the output shape instead, please use + [`CropOrPad`][torchio.transforms.CropOrPad] instead. + + """ + + PADDING_MODES = ( + 'empty', + 'edge', + 'wrap', + 'constant', + 'linear_ramp', + 'maximum', + 'mean', + 'median', + 'minimum', + 'reflect', + 'symmetric', + ) + + def __init__( + self, + padding: TypeBounds, + padding_mode: str | float = 0, + **kwargs, + ): + super().__init__(padding, **kwargs) + self.padding = padding + self.check_padding_mode(padding_mode) + self.padding_mode = padding_mode + self.args_names = ['padding', 'padding_mode'] + + @classmethod + def check_padding_mode(cls, padding_mode): + is_number = isinstance(padding_mode, Number) + is_callable = callable(padding_mode) + if not (padding_mode in cls.PADDING_MODES or is_number or is_callable): + message = ( + f'Padding mode "{padding_mode}" not valid. Valid options are' + f' {list(cls.PADDING_MODES)}, a number or a function' + ) + raise KeyError(message) + + def _check_truncation(self, image: Image, mode: str | float) -> None: + if mode not in ('mean', 'median'): + return + if torch.is_floating_point(image.data): + return + message = ( + f'The constant value computed for padding mode "{mode}" might be truncated ' + ' in the output, as the data type of the input image is not float.' + ' Consider converting the image to a floating point type' + ' before applying this transform.' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + + def apply_transform(self, subject: Subject) -> Subject: + assert self.bounds_parameters is not None + low = self.bounds_parameters[::2] + for image in self.get_images(subject): + self._check_truncation(image, self.padding_mode) + new_origin = apply_affine(image.affine, -np.array(low)) + new_affine = image.affine.copy() + new_affine[:3, 3] = new_origin + + mode: NumpyPadMode = 'constant' + constant: int | float | None = None + if isinstance(self.padding_mode, Number): + constant = float(self.padding_mode) + elif self.padding_mode == 'maximum': + constant = image.data.max().item() + elif self.padding_mode == 'mean': + constant = image.data.float().mean().item() + elif self.padding_mode == 'median': + constant = torch.quantile(image.data.float(), 0.5).item() + elif self.padding_mode == 'minimum': + constant = image.data.min().item() + else: + constant = None + mode = cast(NumpyPadMode, self.padding_mode) + + pad_params = self.bounds_parameters + paddings = (0, 0), pad_params[:2], pad_params[2:4], pad_params[4:] + if constant is not None: + padded = np.pad( + image.data.numpy(), + paddings, + mode='constant', + constant_values=constant, + ) + else: + padded = np.pad(image.data.numpy(), paddings, mode=mode) + image.set_data(torch.as_tensor(padded)) + image.affine = new_affine + return subject + + def inverse(self): + from .crop import Crop + + return Crop(self.padding) diff --git a/src/torchio/transforms/preprocessing/spatial/resample.py b/src/torchio/transforms/preprocessing/spatial/resample.py new file mode 100644 index 000000000..b0b49eda1 --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/resample.py @@ -0,0 +1,455 @@ +from __future__ import annotations + +from collections.abc import Iterable +from collections.abc import Sequence +from collections.abc import Sized +from pathlib import Path +from typing import TypeAlias +from typing import Union + +import numpy as np +import SimpleITK as sitk +import torch + +from ....data.image import Image +from ....data.image import ScalarImage +from ....data.io import get_sitk_metadata_from_ras_affine +from ....data.io import sitk_to_nib +from ....data.subject import Subject +from ....types import TypePath +from ....types import TypeSpacing +from ....types import TypeTripletFloat +from ...spatial_transform import SpatialTransform + +TypeShapeAffine: TypeAlias = tuple[Sequence[int], np.ndarray] +TypeTarget = Union[TypeSpacing, str, Path, Image, TypeShapeAffine, None] +ONE_MILLIMITER_ISOTROPIC = 1 + + +class Resample(SpatialTransform): + """Resample image to a different physical space. + + This is a powerful transform that can be used to change the image shape + or spatial metadata, or to apply a spatial transformation. + + Args: + target: Argument to define the output space. Can be one of: + + - Output spacing $(s_w, s_h, s_d)$, in mm. If only one value + $s$ is specified, then $s_w = s_h = s_d = s$. + + - Path to an image that will be used as reference. + + - Instance of [`Image`][torchio.Image]. + + - Name of an image key in the subject. + + - Tuple `(spatial_shape, affine)` defining the output space. + + pre_affine_name: Name of the *image key* (not subject key) storing an + affine matrix that will be applied to the image header before + resampling. If `None`, the image is resampled with an identity + transform. See usage in the example below. + image_interpolation: See Interpolation. + label_interpolation: See Interpolation. + scalars_only: Apply only to instances of [`ScalarImage`][torchio.ScalarImage]. + Used internally by [`RandomAnisotropy`][torchio.transforms.RandomAnisotropy]. + antialias: If `True`, apply Gaussian smoothing before + downsampling along any dimension that will be downsampled. For example, + if the input image has spacing (0.5, 0.5, 4) and the target + spacing is (1, 1, 1), the image will be smoothed along the first two + dimensions before resampling. Label maps are not smoothed. + The standard deviations of the Gaussian kernels are computed according to + the method described in Cardoso et al., + [Scale factor point spread function matching: beyond aliasing in image + resampling + ](https://link.springer.com/chapter/10.1007/978-3-319-24571-3_81), + MICCAI 2015. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Examples: + >>> import torch + >>> import torchio as tio + >>> transform = tio.Resample() # resample all images to 1mm isotropic + >>> transform = tio.Resample(2) # resample all images to 2mm isotropic + >>> transform = tio.Resample('t1') # resample all images to 't1' image space + >>> # Example: using a precomputed transform to MNI space + >>> ref_path = tio.datasets.Colin27().t1.path # this image is in the MNI space, so we can use it as reference/target + >>> affine_matrix = tio.io.read_matrix('transform_to_mni.txt') # from a NiftyReg registration. Would also work with e.g. .tfm from SimpleITK + >>> image = tio.ScalarImage(tensor=torch.rand(1, 256, 256, 180), to_mni=affine_matrix) # 'to_mni' is an arbitrary name + >>> transform = tio.Resample(colin.t1.path, pre_affine_name='to_mni') # nearest neighbor interpolation is used for label maps + >>> transformed = transform(image) # "image" is now in the MNI space + + Note: + The `antialias` option is recommended when large (e.g. > 2×) downsampling + factors are expected, particularly for offline (before training) preprocessing, + when run times are not a concern. + + """ + + def __init__( + self, + target: TypeTarget = ONE_MILLIMITER_ISOTROPIC, + image_interpolation: str = 'linear', + label_interpolation: str = 'nearest', + pre_affine_name: str | None = None, + scalars_only: bool = False, + antialias: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.target = target + self.image_interpolation = self.parse_interpolation( + image_interpolation, + ) + self.label_interpolation = self.parse_interpolation( + label_interpolation, + ) + self.pre_affine_name = pre_affine_name + self.scalars_only = scalars_only + self.antialias = antialias + self.args_names = [ + 'target', + 'image_interpolation', + 'label_interpolation', + 'pre_affine_name', + 'scalars_only', + 'antialias', + ] + + @staticmethod + def _parse_spacing(spacing: TypeSpacing) -> tuple[float, float, float]: + if isinstance(spacing, (int, float)): + result = (float(spacing), float(spacing), float(spacing)) + elif isinstance(spacing, np.ndarray): + flat = list(spacing.flat) + if len(flat) != 3: + message = ( + 'Target must be a string, a positive number' + f' or a sequence of positive numbers, not {type(spacing)}' + ) + raise ValueError(message) + result = (float(flat[0]), float(flat[1]), float(flat[2])) + elif isinstance(spacing, Sequence): + if len(spacing) != 3: + message = ( + 'Target must be a string, a positive number' + f' or a sequence of positive numbers, not {type(spacing)}' + ) + raise ValueError(message) + values = [] + for value in spacing: + if not isinstance(value, (int, float)): + message = ( + 'Target must be a string, a positive number' + f' or a sequence of positive numbers, not {type(spacing)}' + ) + raise ValueError(message) + values.append(float(value)) + result = (values[0], values[1], values[2]) + else: + message = ( + 'Target must be a string, a positive number' + f' or a sequence of positive numbers, not {type(spacing)}' + ) + raise ValueError(message) + if any(value <= 0 for value in result): + message = f'Spacing must be strictly positive, not "{spacing}"' + raise ValueError(message) + return result + + @staticmethod + def check_affine(affine_name: str, image: Image): + if not isinstance(affine_name, str): + message = f'Affine name argument must be a string, not {type(affine_name)}' + raise TypeError(message) + if affine_name in image: + matrix = image[affine_name] + if not isinstance(matrix, (np.ndarray, torch.Tensor)): + message = ( + 'The affine matrix must be a NumPy array or PyTorch' + f' tensor, not {type(matrix)}' + ) + raise TypeError(message) + if matrix.shape != (4, 4): + message = f'The affine matrix shape must be (4, 4), not {matrix.shape}' + raise ValueError(message) + + @staticmethod + def check_affine_key_presence(affine_name: str, subject: Subject): + for image in subject.get_images(intensity_only=False): + if affine_name in image: + return + message = ( + f'An affine name was given ("{affine_name}"), but it was not found' + ' in any image in the subject' + ) + raise ValueError(message) + + def apply_transform(self, subject: Subject) -> Subject: + use_pre_affine = self.pre_affine_name is not None + if use_pre_affine: + assert self.pre_affine_name is not None # for mypy + self.check_affine_key_presence(self.pre_affine_name, subject) + + for image in self.get_images(subject): + # If the current image is the reference, don't resample it + if self.target is image: + continue + + # If the target is not a string, or is not an image in the subject, + # do nothing + if isinstance(self.target, str): + try: + target_image = subject.get_image(self.target) + except KeyError: + pass + else: + if target_image is image: + continue + + # Choose interpolation + if not isinstance(image, ScalarImage): + if self.scalars_only: + continue + interpolation = self.label_interpolation + else: + interpolation = self.image_interpolation + interpolator = self.get_sitk_interpolator(interpolation) + + # Apply given affine matrix if found in image + if use_pre_affine and self.pre_affine_name in image: + assert self.pre_affine_name is not None # for mypy + self.check_affine(self.pre_affine_name, image) + matrix = image[self.pre_affine_name] + if isinstance(matrix, torch.Tensor): + matrix = matrix.numpy() + image.affine = matrix @ image.affine + + floating_sitk = image.as_sitk(force_3d=True) + + resampler = self._get_resampler( + interpolator, + floating_sitk, + subject, + self.target, + ) + if self.antialias and isinstance(image, ScalarImage): + downsampling_factor = self._get_downsampling_factor( + floating_sitk, + resampler, + ) + sigmas = self._get_sigmas( + downsampling_factor, + floating_sitk.GetSpacing(), + ) + floating_sitk = self._smooth(floating_sitk, sigmas) + resampled = resampler.Execute(floating_sitk) + + array, affine = sitk_to_nib(resampled) + image.set_data(torch.as_tensor(array)) + image.affine = affine + return subject + + @staticmethod + def _smooth( + image: sitk.Image, + sigmas: np.ndarray, + epsilon: float = 1e-9, + ) -> sitk.Image: + """Smooth the image with a Gaussian kernel. + + Args: + image: Image to be smoothed. + sigmas: Standard deviations of the Gaussian kernel for each + dimension. If a value is NaN, no smoothing is applied in that + dimension. + epsilon: Small value to replace NaN values in sigmas, to avoid + division-by-zero errors. + """ + + sigmas[np.isnan(sigmas)] = epsilon # no smoothing in that dimension + gaussian = sitk.SmoothingRecursiveGaussianImageFilter() + gaussian.SetSigma(sigmas.tolist()) + smoothed = gaussian.Execute(image) + return smoothed + + @staticmethod + def _get_downsampling_factor( + floating: sitk.Image, + resampler: sitk.ResampleImageFilter, + ) -> np.ndarray: + """Get the downsampling factor for each dimension. + + The downsampling factor is the ratio between the output spacing and + the input spacing. If the output spacing is smaller than the input + spacing, the factor is set to NaN, meaning downsampling is not applied + in that dimension. + + Args: + floating: The input image to be resampled. + resampler: The resampler that will be used to resample the image. + """ + input_spacing = np.array(floating.GetSpacing()) + output_spacing = np.array(resampler.GetOutputSpacing()) + factors = output_spacing / input_spacing + no_downsampling = factors <= 1 + factors[no_downsampling] = np.nan + return factors + + def _get_resampler( + self, + interpolator: int, + floating: sitk.Image, + subject: Subject, + target: TypeTarget, + ) -> sitk.ResampleImageFilter: + """Instantiate a SimpleITK resampler.""" + if target is None: + raise RuntimeError('Target cannot be None') + resampler = sitk.ResampleImageFilter() + resampler.SetInterpolator(interpolator) + self._set_resampler_reference( + resampler, + target, + floating, + subject, + ) + return resampler + + def _set_resampler_reference( + self, + resampler: sitk.ResampleImageFilter, + target: TypeSpacing | TypePath | Image | TypeShapeAffine, + floating_sitk, + subject, + ): + # Target can be: + # 1) An instance of torchio.Image + # 2) An instance of pathlib.Path + # 3) A string, which could be a path or an image in subject + # 4) A number or sequence of numbers for spacing + # 5) A tuple of shape, affine + # The fourth case is the different one + if isinstance(target, (str, Path, Image)): + if isinstance(target, Image): + # It's a TorchIO image + image = target + elif Path(target).is_file(): + # It's an existing file + path = target + image = ScalarImage(path) + else: # assume it's the name of an image in the subject + try: + image = subject.get_image(target) + except KeyError as error: + message = ( + f'Image name "{target}" not found in subject.' + f' If "{target}" is a path, it does not exist or' + ' permission has been denied' + ) + raise ValueError(message) from error + self._set_resampler_from_shape_affine( + resampler, + image.spatial_shape, + image.affine, + ) + elif isinstance(target, (int, float)): # one number for target was passed + self._set_resampler_from_spacing(resampler, target, floating_sitk) + elif isinstance(target, tuple) and len(target) == 2: + shape = target[0] + affine = target[1] + if not (isinstance(shape, Sized) and len(shape) == 3): + message = ( + 'Target shape must be a sequence of three integers, but' + f' "{shape}" was passed' + ) + raise RuntimeError(message) + if not isinstance(affine, np.ndarray) or affine.shape != (4, 4): + message = ( + 'Target affine must have shape (4, 4) but the following' + f' was passed:\n{shape}' + ) + raise RuntimeError(message) + self._set_resampler_from_shape_affine( + resampler, + shape, + affine, + ) + elif ( + isinstance(target, Sized) + and isinstance(target, Iterable) + and len(target) == 3 + ): + self._set_resampler_from_spacing(resampler, target, floating_sitk) + else: + raise RuntimeError(f'Target not understood: "{target}"') + + def _set_resampler_from_shape_affine(self, resampler, shape, affine): + origin, spacing, direction = get_sitk_metadata_from_ras_affine(affine) + resampler.SetOutputDirection(direction) + resampler.SetOutputOrigin(origin) + resampler.SetOutputSpacing(spacing) + resampler.SetSize(shape) + + def _set_resampler_from_spacing(self, resampler, target, floating_sitk): + target_spacing = self._parse_spacing(target) + reference_image = self.get_reference_image( + floating_sitk, + target_spacing, + ) + resampler.SetReferenceImage(reference_image) + + @staticmethod + def get_reference_image( + floating_sitk: sitk.Image, + spacing: TypeTripletFloat, + ) -> sitk.Image: + old_spacing = np.array(floating_sitk.GetSpacing(), dtype=float) + new_spacing = np.array(spacing, dtype=float) + old_size = np.array(floating_sitk.GetSize()) + old_last_index = old_size - 1 + old_last_index_lps = np.array( + floating_sitk.TransformIndexToPhysicalPoint(old_last_index.tolist()), + dtype=float, + ) + old_origin_lps = np.array(floating_sitk.GetOrigin(), dtype=float) + center_lps = (old_last_index_lps + old_origin_lps) / 2 + # We use floor to avoid extrapolation by keeping the extent of the + # new image the same or smaller than the original. + new_size = np.floor(old_size * old_spacing / new_spacing) + # We keep singleton dimensions to avoid e.g. making 2D images 3D + new_size[old_size == 1] = 1 + direction = np.asarray(floating_sitk.GetDirection(), dtype=float).reshape(3, 3) + half_extent = (new_size - 1) / 2 * new_spacing + new_origin_lps = (center_lps - direction @ half_extent).tolist() + reference = sitk.Image( + new_size.astype(int).tolist(), + floating_sitk.GetPixelID(), + floating_sitk.GetNumberOfComponentsPerPixel(), + ) + reference.SetDirection(floating_sitk.GetDirection()) + reference.SetSpacing(new_spacing.tolist()) + reference.SetOrigin(new_origin_lps) + return reference + + @staticmethod + def _get_sigmas(downsampling_factor: np.ndarray, spacing: np.ndarray) -> np.ndarray: + """Compute optimal standard deviation for Gaussian kernel. + + From Cardoso et al., [Scale factor point spread function matching: + beyond aliasing in image resampling + ](https://link.springer.com/chapter/10.1007/978-3-319-24571-3_81), + MICCAI 2015. + + Args: + downsampling_factor: Array with the downsampling factor for each + dimension. + spacing: Array with the spacing of the input image in mm. + """ + k = downsampling_factor + # Equation from top of page 678 of proceedings (4/9 in the PDF) + variance = (k**2 - 1) * (2 * np.sqrt(2 * np.log(2))) ** (-2) + sigma = spacing * np.sqrt(variance) + return sigma diff --git a/src/torchio/transforms/preprocessing/spatial/resize.py b/src/torchio/transforms/preprocessing/spatial/resize.py new file mode 100644 index 000000000..b63eeed05 --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/resize.py @@ -0,0 +1,97 @@ +import warnings + +import numpy as np + +from ....data.subject import Subject +from ....types import TypeSpatialShape +from ....utils import to_tuple +from ...spatial_transform import SpatialTransform +from .crop_or_pad import CropOrPad +from .resample import Resample + + +class Resize(SpatialTransform): + """Resample images so the output shape matches the given target shape. + + The field of view remains the same. + + Warning: + In most medical image applications, this transform should not + be used as it will deform the physical object by scaling anisotropically + along the different dimensions. The solution to change an image size is + typically applying [`Resample`][torchio.transforms.Resample] and + [`CropOrPad`][torchio.transforms.CropOrPad]. + + Args: + target_shape: Tuple $(W, H, D)$. If a single value $N$ is + provided, then $W = H = D = N$. The size of dimensions set to + -1 will be kept. + image_interpolation: See Interpolation. + label_interpolation: See Interpolation. + """ + + def __init__( + self, + target_shape: TypeSpatialShape, + image_interpolation: str = 'linear', + label_interpolation: str = 'nearest', + **kwargs, + ): + super().__init__(**kwargs) + self.target_shape = np.asarray(to_tuple(target_shape, length=3)) + self.image_interpolation = self.parse_interpolation( + image_interpolation, + ) + self.label_interpolation = self.parse_interpolation( + label_interpolation, + ) + self.args_names = [ + 'target_shape', + 'image_interpolation', + 'label_interpolation', + ] + + def apply_transform(self, subject: Subject) -> Subject: + shape_in = np.asarray(subject.spatial_shape) + shape_out = self.target_shape + negative_mask = shape_out == -1 + shape_out[negative_mask] = shape_in[negative_mask] + spacing_in = np.asarray(subject.spacing) + spacing_out = shape_in / shape_out * spacing_in + resample = Resample( + spacing_out, + image_interpolation=self.image_interpolation, + label_interpolation=self.label_interpolation, + copy=self.copy, + include=self.include, + exclude=self.exclude, + keep=self.keep, + parse_input=self.parse_input, + label_keys=self.label_keys, + ) + resampled = resample(subject) + assert isinstance(resampled, Subject) + # Sometimes, the output shape is one voxel too large + if not resampled.spatial_shape == tuple(shape_out): + message = ( + f'Output shape {resampled.spatial_shape}' + f' != target shape {tuple(shape_out)}. Fixing with CropOrPad' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + target_shape_values = [int(value) for value in shape_out.tolist()] + crop_pad = CropOrPad( + ( + target_shape_values[0], + target_shape_values[1], + target_shape_values[2], + ), + copy=self.copy, + include=self.include, + exclude=self.exclude, + keep=self.keep, + parse_input=self.parse_input, + label_keys=self.label_keys, + ) + resampled = crop_pad(resampled) + assert isinstance(resampled, Subject) + return resampled diff --git a/src/torchio/transforms/preprocessing/spatial/to_canonical.py b/src/torchio/transforms/preprocessing/spatial/to_canonical.py new file mode 100644 index 000000000..6498cc68d --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/to_canonical.py @@ -0,0 +1,24 @@ +from .to_orientation import ToOrientation + + +class ToCanonical(ToOrientation): + """Reorder the data to be closest to canonical (RAS+) orientation. + + This transform reorders the voxels and modifies the affine matrix so that + the voxel orientations are nearest to: + + 1. First voxel axis goes from left to Right + 2. Second voxel axis goes from posterior to Anterior + 3. Third voxel axis goes from inferior to Superior + + See [NiBabel docs about image orientation](https://nipy.org/nibabel/image_orientation.html) for more information. + + Args: + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + Note: + The reorientation is performed using + [`ToOrientation()`](../ToOrientation/#torchio.transforms.ToOrientation). + + """ diff --git a/src/torchio/transforms/preprocessing/spatial/to_orientation.py b/src/torchio/transforms/preprocessing/spatial/to_orientation.py new file mode 100644 index 000000000..139a5136e --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/to_orientation.py @@ -0,0 +1,109 @@ +import nibabel as nib +import numpy as np +import torch +from einops import rearrange +from nibabel import orientations + +from ....data.subject import Subject +from ...spatial_transform import SpatialTransform + + +class ToOrientation(SpatialTransform): + """Reorient the data to a specified orientation. + + This transform reorders the voxels and modifies the affine matrix to match + the specified orientation code. + The image intensity values are not modified, and the sample locations in + the scanner space are preserved. + + Common orientation codes include: + + - `'RAS'` (neurological convention): + - The first axis goes from Left to Right (R). + - The second axis goes from Posterior to Anterior (A). + - The third axis goes from Inferior to Superior (S). + - `'LAS'` (radiological convention): + - The first axis goes from Right to Left (L). + - The second axis goes from Posterior to Anterior (A). + - The third axis goes from Inferior to Superior (S). + + See [NiBabel docs about image orientation](https://nipy.org/nibabel/image_orientation.html) for more information. + + Args: + orientation: A three-letter orientation code. Examples: `'RAS'`, + `'LAS'`, `'LPS'`, `'PLS'`, `'SLP'`. The code must contain + one character for each axis direction: R or L, A or P, and S or I. + **kwargs: See [`Transform`][torchio.transforms.Transform] for additional + keyword arguments. + + """ + + def __init__(self, orientation: str = 'RAS', **kwargs): + super().__init__(**kwargs) + if not isinstance(orientation, str) or len(orientation) != 3: + message = f'Orientation must be a 3-letter string, got "{orientation}"' + raise ValueError(message) + + valid_codes = set('RLAPIS') + orientation = orientation.upper() + all_valid = all(axis in valid_codes for axis in orientation) + if not all_valid: + message = ( + 'Orientation code must be composed of three distinct characters' + f' in {valid_codes} but got "{orientation}"' + ) + raise ValueError(message) + + # Check for valid axis directions + has_sagittal = 'R' in orientation or 'L' in orientation + has_coronal = 'A' in orientation or 'P' in orientation + has_axial = 'S' in orientation or 'I' in orientation + has_all = has_sagittal and has_coronal and has_axial + if not has_all: + message = ( + 'Orientation code must include one character for each axis direction:' + f' R or L, A or P, and S or I, but got "{orientation}"' + ) + raise ValueError(message) + + self.orientation = orientation + self.args_names = ['orientation'] + + def apply_transform(self, subject: Subject) -> Subject: + for image in subject.get_images(intensity_only=False): + current_orientation = ''.join(nib.orientations.aff2axcodes(image.affine)) + + # If the image is already in the target orientation, skip it + if current_orientation == self.orientation: + continue + + # NIfTI images should have channels in 5th dimension + array = rearrange(image.numpy(), 'C W H D -> W H D 1 C') + + nii = nib.nifti1.Nifti1Image(array, image.affine) + + # Compute transform from current orientation to target orientation + current_orientation = orientations.io_orientation(nii.affine) + target_orientation = orientations.axcodes2ornt(tuple(self.orientation)) + transform = orientations.ornt_transform( + current_orientation, + target_orientation, + ) + + # Reorder voxels + reoriented_array = orientations.apply_orientation(nii.dataobj, transform) + reoriented_array = rearrange(reoriented_array, 'W H D 1 C -> C W H D') + + # Calculate the new affine matrix reflecting the reorientation + reoriented_affine = nii.affine @ orientations.inv_ornt_aff( + transform, + nii.shape, + ) + + # Update the image data and affine + reoriented_array = np.ascontiguousarray(reoriented_array) + tensor = torch.from_numpy(reoriented_array) + image.set_data(tensor) + image.affine = reoriented_affine + + return subject diff --git a/src/torchio/transforms/preprocessing/spatial/to_reference_space.py b/src/torchio/transforms/preprocessing/spatial/to_reference_space.py new file mode 100644 index 000000000..e959d7a64 --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/to_reference_space.py @@ -0,0 +1,53 @@ +import numpy as np +import torch + +from ....data.image import Image +from ....data.subject import Subject +from ...spatial_transform import SpatialTransform +from .resample import Resample + + +class ToReferenceSpace(SpatialTransform): + """Modify the spatial metadata so it matches a reference space. + + This is useful, for example, to set meaningful spatial metadata of a neural + network embedding, for visualization or further processing such as + resampling a segmentation output. + + Examples: + >>> import torchio as tio + >>> image = tio.datasets.FPG().t1 + >>> embedding_tensor = my_network(image.tensor) # we lose metadata here + >>> embedding_image = tio.ToReferenceSpace.from_tensor(embedding_tensor, image) + """ + + def __init__(self, reference: Image, **kwargs): + super().__init__(**kwargs) + if not isinstance(reference, Image): + raise TypeError('The reference must be a TorchIO image') + self.reference = reference + + def apply_transform(self, subject: Subject) -> Subject: + for image in self.get_images(subject): + new_image = build_image_from_reference(image.data, self.reference) + image.set_data(new_image.data) + image.affine = new_image.affine + return subject + + @staticmethod + def from_tensor(tensor: torch.Tensor, reference: Image) -> Image: + """Build a TorchIO image from a tensor and a reference image.""" + return build_image_from_reference(tensor, reference) + + +def build_image_from_reference(tensor: torch.Tensor, reference: Image) -> Image: + input_shape = np.array(reference.spatial_shape) + output_shape = np.array(tensor.shape[-3:]) + downsampling_factor = input_shape / output_shape + input_spacing = np.array(reference.spacing) + output_spacing = input_spacing * downsampling_factor + downsample = Resample(output_spacing, image_interpolation='nearest') + reference = downsample(reference) + class_ = reference.__class__ + result = class_(tensor=tensor, affine=reference.affine) + return result diff --git a/src/torchio/transforms/preprocessing/spatial/transpose.py b/src/torchio/transforms/preprocessing/spatial/transpose.py new file mode 100644 index 000000000..541d0c0cc --- /dev/null +++ b/src/torchio/transforms/preprocessing/spatial/transpose.py @@ -0,0 +1,37 @@ +from ....data.subject import Subject +from ...spatial_transform import SpatialTransform +from .to_orientation import ToOrientation + + +class Transpose(SpatialTransform): + """Swap the first and last spatial dimensions of the image. + + The spatial metadata is updated accordingly, so the world coordinates of + all voxels in the input and output spaces match. + + Examples: + >>> import torchio as tio + >>> image = tio.datasets.FPG().t1 + >>> image + ScalarImage(shape: (1, 256, 256, 176); spacing: (1.00, 1.00, 1.00); orientation: PIR+; path: "/home/fernando/.cache/torchio/fpg/t1.nii.gz") + >>> transpose = tio.Transpose() + >>> transposed = transpose(image) + >>> transposed + ScalarImage(shape: (1, 176, 256, 256); spacing: (1.00, 1.00, 1.00); orientation: RIP+; dtype: torch.IntTensor; memory: 44.0 MiB) + """ + + def apply_transform(self, subject: Subject) -> Subject: + for image in self.get_images(subject): + old_orientation = image.orientation_str + new_orientation = old_orientation[::-1] + transform = ToOrientation(new_orientation) + transposed = transform(image) + image.set_data(transposed.data) + image.affine = transposed.affine + return subject + + def is_invertible(self): + return True + + def inverse(self): + return self diff --git a/src/torchio/transforms/spatial/__init__.py b/src/torchio/transforms/spatial/__init__.py deleted file mode 100644 index 282f01792..000000000 --- a/src/torchio/transforms/spatial/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Spatial transforms.""" diff --git a/src/torchio/transforms/spatial/_padding.py b/src/torchio/transforms/spatial/_padding.py deleted file mode 100644 index cb98ff491..000000000 --- a/src/torchio/transforms/spatial/_padding.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Shared spatial padding helpers.""" - -from __future__ import annotations - -import warnings -from typing import Literal -from typing import TypeGuard -from typing import get_args - -import torch -from torch import Tensor - -from ...types import TypeSixInts -from .._statistics import compute_quantile - -#: Accepted padding modes. -PaddingMode = Literal[ - "constant", - "reflect", - "replicate", - "circular", - "mean", - "median", - "minimum", -] - -_PADDING_MODES: tuple[PaddingMode, ...] = get_args(PaddingMode) -_STATISTIC_PADDING_MODES = "mean", "median", "minimum" - - -def _is_padding_mode(value: str) -> TypeGuard[PaddingMode]: - return value in _PADDING_MODES - - -def parse_padding_mode(padding_mode: str) -> PaddingMode: - """Validate and return a padding mode.""" - if not _is_padding_mode(padding_mode): - msg = f"padding_mode must be one of {_PADDING_MODES}, got {padding_mode!r}" - raise ValueError(msg) - return padding_mode - - -def _compute_padding_statistic( - data: Tensor, - padding_mode: PaddingMode, -) -> Tensor: - """Compute one whole-volume padding statistic per batch element.""" - flat = data.flatten(start_dim=1) - if padding_mode == "minimum": - return flat.amin(dim=1) - - if not torch.is_floating_point(data): - warnings.warn( - f'The constant value computed for padding mode "{padding_mode}"' - " might be truncated in the output, as the data type of the input" - " image is not float. Consider converting the image to a floating" - " point type before applying this transform.", - RuntimeWarning, - stacklevel=4, - ) - - float_flat = flat if data.dtype in (torch.float32, torch.float64) else flat.float() - if padding_mode == "mean": - statistic = float_flat.mean(dim=1) - else: - statistic = torch.stack( - [compute_quantile(values, 0.5) for values in float_flat], - ) - return statistic.to(data.dtype) - - -def pad_tensor( - data: Tensor, - padding: TypeSixInts, - padding_mode: PaddingMode, - fill: float, -) -> Tensor: - """Pad a 4D image tensor or 5D image batch.""" - if data.ndim not in (4, 5): - msg = f"Expected a 4D or 5D image tensor, got {data.ndim}D" - raise ValueError(msg) - i0, i1, j0, j1, k0, k1 = padding - pad_arg = k0, k1, j0, j1, i0, i1 - if padding_mode not in _STATISTIC_PADDING_MODES: - return torch.nn.functional.pad( - data, - pad_arg, - mode=padding_mode, - value=fill, - ) - - is_unbatched = data.ndim == 4 - batch = data.unsqueeze(0) if is_unbatched else data - statistic = _compute_padding_statistic(batch, padding_mode) - padded = torch.nn.functional.pad(batch, pad_arg) - interior = torch.ones( - (1, 1, *batch.shape[-3:]), - dtype=torch.bool, - device=batch.device, - ) - interior = torch.nn.functional.pad(interior, pad_arg) - fill_values = statistic.reshape(-1, 1, 1, 1, 1) - result = torch.where(interior, padded, fill_values) - return result[0] if is_unbatched else result diff --git a/src/torchio/transforms/spatial/anisotropy.py b/src/torchio/transforms/spatial/anisotropy.py deleted file mode 100644 index adc45ab06..000000000 --- a/src/torchio/transforms/spatial/anisotropy.py +++ /dev/null @@ -1,392 +0,0 @@ -"""Anisotropy: simulate low-resolution acquisition along one axis.""" - -from __future__ import annotations - -from typing import Any - -import torch -import torch.nn.functional as functional -from einops import rearrange - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..parameter_range import to_nonneg_range -from ..transform import Transform - - -class Anisotropy(Transform): - r"""Simulate an anisotropic acquisition. - - Downsample along a randomly chosen axis and then upsample back - to the original shape, emulating the through-plane blur seen in - clinical MRI when one axis has coarser resolution. - - This is useful as a data augmentation for super-resolution - training. - - Args: - axes: Spatial axes eligible for downsampling. One is chosen - at random per application. - downsampling: Downsampling factor $m \geq 1$. A scalar is - deterministic; a 2-tuple $(a, b)$ samples - $m \sim \mathcal{U}(a, b)$. The default `downsampling=1` - is a no-op (and warns). - image_interpolation: Interpolation mode used when upsampling - scalar images back to the original shape. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Anisotropy(downsampling=4) - >>> transform = tio.Anisotropy( - ... axes=(2,), - ... downsampling=(1.5, 5), - ... ) - """ - - def __init__( - self, - *, - axes: tuple[int, ...] = (0, 1, 2), - downsampling: float | tuple[float, float] = 1.0, - image_interpolation: str = "linear", - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.axes = axes - self.downsampling = to_nonneg_range(downsampling) - self.image_interpolation = image_interpolation - self._validate_downsampling() - self._warn_if_noop( - is_noop=self.downsampling.is_constant(1.0), - hint="downsampling=(1.5, 5)", - ) - - def _validate_downsampling(self) -> None: - """Ensure the range produces factors >= 1.""" - _lo, hi = self.downsampling._ranges[0] - if hi < 1.0: - msg = f"downsampling range upper bound must be >= 1, got {hi}" - raise ValueError(msg) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample axis and downsampling factor (per element when batched).""" - n = self._resolve_n(batch) - if n is None: - axis = self.axes[int(torch.randint(len(self.axes), (1,)).item())] - factor = max(1.0, self.downsampling.sample_1d()) - return {"axis": axis, "factor": factor} - keep = self._keep_mask(batch, n) - axis_list: list[int] = [] - factor_list: list[float] = [] - for index in range(n): - if keep is not None and not keep[index]: - axis_list.append(self.axes[0]) - factor_list.append(1.0) - continue - axis_list.append(self.axes[int(torch.randint(len(self.axes), (1,)).item())]) - factor_list.append(max(1.0, self.downsampling.sample_1d())) - params = {"axis": axis_list, "factor": factor_list} - self._tag_batched(params, batch, n, keep, ["axis", "factor"]) - return params - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Downsample then upsample along the chosen axis.""" - per_instance = self._is_per_instance_params(params) - for _name, img_batch in batch.images.items(): - is_label = issubclass(img_batch._image_class, LabelMap) - mode = "nearest" if is_label else self.image_interpolation - if per_instance: - data = img_batch.data - img_batch.data = _simulate_anisotropy_per_instance( - data, - axes=params["axis"], - factors=params["factor"], - mode=mode, - ) - else: - factor = params["factor"] - if factor <= 1.0: - continue - img_batch.data = _simulate_anisotropy( - img_batch.data, - axis=params["axis"], - factor=factor, - mode=mode, - ) - return batch - - -def _simulate_anisotropy_per_instance( - data: torch.Tensor, - *, - axes: list[int], - factors: list[float], - mode: str, -) -> torch.Tensor: - """Downsample then upsample each batch element with its own parameters. - - Args: - data: `(B, C, I, J, K)` tensor. - axes: Spatial axis per batch element. - factors: Downsampling factor per batch element. - mode: Interpolation mode for upsampling (`"nearest"` or - `"linear"`). - - Returns: - Degraded `(B, C, I, J, K)` tensor with original shape. - """ - axes_tensor = torch.as_tensor(axes, dtype=torch.long, device=data.device) - factors_tensor = torch.as_tensor( - factors, - dtype=torch.float64, - device=data.device, - ) - active = factors_tensor > 1.0 - if not bool(active.any()): - return data.to(data.dtype) - - active_axes = axes_tensor[active] - if bool(((active_axes < 0) | (active_axes > 2)).any()): - msg = f"Anisotropy axis must be in {{0, 1, 2}}, got {sorted(set(axes))}" - raise ValueError(msg) - - output = data.clone() - for axis in range(3): - axis_mask = active & (axes_tensor == axis) - if not bool(axis_mask.any()): - continue - output[axis_mask] = _simulate_anisotropy_fixed_axis( - data[axis_mask], - factors=factors_tensor[axis_mask], - axis=axis, - mode=mode, - ) - return output.to(data.dtype) - - -def _simulate_anisotropy_fixed_axis( - data: torch.Tensor, - *, - factors: torch.Tensor, - axis: int, - mode: str, -) -> torch.Tensor: - """Downsample then upsample one axis with per-element factors. - - Args: - data: `(B, C, I, J, K)` tensor whose elements share `axis`. - factors: Downsampling factors for the batch elements. - axis: Spatial axis (0, 1, or 2) to degrade. - mode: Interpolation mode for upsampling (`"nearest"` or - `"linear"`). - - Returns: - Degraded `(B, C, I, J, K)` tensor with original shape. - """ - length = data.shape[axis + 2] - down_sizes = _downsample_sizes(length, factors) - if mode == "nearest": - indices = _nearest_source_indices(length, down_sizes, data.device) - return _gather_axis(data.float(), indices, axis).to(data.dtype) - lower, upper, weights = _linear_source_indices(length, down_sizes, data.device) - lower_values = _gather_axis(data.float(), lower, axis) - upper_values = _gather_axis(data.float(), upper, axis) - weights = _broadcast_axis_weights(weights, axis) - degraded = lower_values * (1.0 - weights) + upper_values * weights - return degraded.to(data.dtype) - - -def _downsample_sizes(length: int, factors: torch.Tensor) -> torch.Tensor: - """Return PyTorch-compatible nearest-downsampled sizes. - - Args: - length: Original length along the degraded axis. - factors: Per-element downsampling factors. - - Returns: - Downsampled sizes matching `round(length / factor)`. - """ - sizes = torch.round(length / factors).clamp_min(1) - return sizes.to(torch.long) - - -def _nearest_source_indices( - length: int, - down_sizes: torch.Tensor, - device: torch.device, -) -> torch.Tensor: - """Return original-axis indices for nearest downsample and upsample. - - Args: - length: Original length along the degraded axis. - down_sizes: Downsampled sizes per batch element. - device: Device where the indices should be created. - - Returns: - A `(B, length)` tensor of source indices. - """ - positions = torch.arange(length, dtype=torch.long, device=device) - lowres_indices = torch.div( - positions * rearrange(down_sizes, "b -> b 1"), - length, - rounding_mode="floor", - ) - return _downsample_source_indices(length, down_sizes, lowres_indices) - - -def _linear_source_indices( - length: int, - down_sizes: torch.Tensor, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Return source indices and weights for trilinear upsampling. - - Args: - length: Original length along the degraded axis. - down_sizes: Downsampled sizes per batch element. - device: Device where the indices should be created. - - Returns: - Lower source indices, upper source indices and upper weights, - each shaped `(B, length)`. - """ - positions = torch.arange(length, dtype=torch.float32, device=device) - if length == 1: - lowres_positions = torch.zeros( - len(down_sizes), - 1, - dtype=torch.float32, - device=device, - ) - else: - scale = (down_sizes.to(torch.float32) - 1.0) / (length - 1) - lowres_positions = positions * rearrange(scale, "b -> b 1") - lower_lowres = lowres_positions.floor().to(torch.long) - upper_lowres = torch.minimum( - lower_lowres + 1, - rearrange(down_sizes, "b -> b 1") - 1, - ) - weights = lowres_positions - lower_lowres.to(torch.float32) - lower = _downsample_source_indices(length, down_sizes, lower_lowres) - upper = _downsample_source_indices(length, down_sizes, upper_lowres) - return lower, upper, weights - - -def _downsample_source_indices( - length: int, - down_sizes: torch.Tensor, - lowres_indices: torch.Tensor, -) -> torch.Tensor: - """Map low-resolution indices to nearest-downsampled source indices. - - Args: - length: Original length along the degraded axis. - down_sizes: Downsampled sizes per batch element. - lowres_indices: Low-resolution indices. - - Returns: - Original source indices. - """ - source = torch.div( - lowres_indices * length, - rearrange(down_sizes, "b -> b 1"), - rounding_mode="floor", - ) - return source.clamp(max=length - 1) - - -def _gather_axis( - data: torch.Tensor, - source_indices: torch.Tensor, - axis: int, -) -> torch.Tensor: - """Gather `data` along one spatial axis with per-element indices. - - Args: - data: `(B, C, I, J, K)` tensor. - source_indices: `(B, N)` indices for the selected spatial axis. - axis: Spatial axis (0, 1, or 2) to gather. - - Returns: - Gathered `(B, C, I, J, K)` tensor. - """ - if axis == 0: - indices = rearrange(source_indices, "b n -> b 1 n 1 1") - elif axis == 1: - indices = rearrange(source_indices, "b n -> b 1 1 n 1") - else: - indices = rearrange(source_indices, "b n -> b 1 1 1 n") - indices = indices.expand_as(data) - return torch.gather(data, dim=axis + 2, index=indices) - - -def _broadcast_axis_weights(weights: torch.Tensor, axis: int) -> torch.Tensor: - """Broadcast interpolation weights over channels and untouched axes. - - Args: - weights: `(B, N)` interpolation weights. - axis: Spatial axis (0, 1, or 2) corresponding to `N`. - - Returns: - Weights broadcast-compatible with `(B, C, I, J, K)`. - """ - if axis == 0: - return rearrange(weights, "b n -> b 1 n 1 1") - if axis == 1: - return rearrange(weights, "b n -> b 1 1 n 1") - return rearrange(weights, "b n -> b 1 1 1 n") - - -def _simulate_anisotropy( - data: torch.Tensor, - *, - axis: int, - factor: float, - mode: str, -) -> torch.Tensor: - """Downsample then upsample one axis of a 5-D tensor. - - Args: - data: `(B, C, I, J, K)` tensor. - axis: Spatial axis (0, 1, or 2) to degrade. - factor: Downsampling factor (> 1). - mode: Interpolation mode for upsampling (`"nearest"` or - `"linear"`). - - Returns: - Degraded `(B, C, I, J, K)` tensor with original shape. - """ - original_shape = list(data.shape[2:]) - down_shape = list(original_shape) - down_shape[axis] = max(1, round(original_shape[axis] / factor)) - - torch_mode_down = "nearest" - torch_mode_up = "nearest" if mode == "nearest" else "trilinear" - - # Downsample. - downsampled = functional.interpolate( - data.float(), - size=down_shape, - mode=torch_mode_down, - ) - # Upsample back. - upsampled = functional.interpolate( - downsampled, - size=original_shape, - mode=torch_mode_up, - align_corners=None if torch_mode_up == "nearest" else True, - ) - return upsampled.to(data.dtype) diff --git a/src/torchio/transforms/spatial/copy_affine.py b/src/torchio/transforms/spatial/copy_affine.py deleted file mode 100644 index 3e6b23b7e..000000000 --- a/src/torchio/transforms/spatial/copy_affine.py +++ /dev/null @@ -1,57 +0,0 @@ -"""CopyAffine: copy the affine matrix from a reference image.""" - -from __future__ import annotations - -import copy -from typing import Any - -from ...data.batch import SubjectsBatch -from ..transform import SpatialTransform - - -class CopyAffine(SpatialTransform): - """Copy the affine matrix from one image to all others. - - This is useful when slight numerical differences between affine - matrices cause downstream errors (e.g., in - [`Resample`][torchio.Resample]). NIfTI stores affines in - single precision, so saving and reloading can introduce - rounding errors. - - Args: - target: Name of the image whose affine will be copied to - all other images in the subject. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.CopyAffine(target="t1") - """ - - def __init__(self, target: str, **kwargs: Any) -> None: - super().__init__(**kwargs) - self.target = target - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Copy the reference affine to all other images.""" - if self.target not in batch.images: - msg = ( - f"Reference image '{self.target}' not found. " - f"Available: {list(batch.images.keys())}" - ) - raise KeyError(msg) - ref_affines = batch.images[self.target].affines - for name, img_batch in batch.images.items(): - if name == self.target: - continue - for i, affine in enumerate(img_batch.affines): - affine._matrix = copy.deepcopy(ref_affines[i]._matrix) - return batch diff --git a/src/torchio/transforms/spatial/crop.py b/src/torchio/transforms/spatial/crop.py deleted file mode 100644 index da977c259..000000000 --- a/src/torchio/transforms/spatial/crop.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Crop transform: remove voxels from the borders.""" - -from __future__ import annotations - -from typing import Any - -from ...data.batch import SubjectsBatch -from ...types import TypeSixInts -from ...types import TypeThreeInts -from ..transform import SpatialTransform - -#: Accepted cropping specifications. -#: `int` → same amount from each side of each axis. -#: 3-tuple → symmetric per axis `(i, j, k)`. -#: 6-tuple → per-side `(i_ini, i_fin, j_ini, j_fin, k_ini, k_fin)`. -CroppingParam = int | TypeThreeInts | TypeSixInts - - -def _parse_cropping(cropping: CroppingParam) -> TypeSixInts: - """Normalise cropping to a 6-tuple (i_ini, i_fin, j_ini, j_fin, k_ini, k_fin).""" - if isinstance(cropping, int): - return (cropping, cropping, cropping, cropping, cropping, cropping) - values = list(cropping) - n = len(values) - if n == 3: - i, j, k = values - return (i, i, j, j, k, k) - if n == 6: - return (values[0], values[1], values[2], values[3], values[4], values[5]) - msg = f"Cropping must have 1, 3, or 6 values, got {n}" - raise ValueError(msg) - - -class Crop(SpatialTransform): - r"""Remove a border of voxels from each side of the volume. - - Args: - cropping: Tuple - $(i_\text{ini}, i_\text{fin}, j_\text{ini}, j_\text{fin}, - k_\text{ini}, k_\text{fin})$ - defining the number of voxels cropped from the edges of - each axis. If the initial shape of the image is - $I \times J \times K$, the final shape will be - $(I - i_\text{ini} - i_\text{fin}) \times - (J - j_\text{ini} - j_\text{fin}) \times - (K - k_\text{ini} - k_\text{fin})$. - If only three values $(i, j, k)$ are provided, then - $i_\text{ini} = i_\text{fin} = i$, - $j_\text{ini} = j_\text{fin} = j$ and - $k_\text{ini} = k_\text{fin} = k$. - If only one value $n$ is provided, then - $i_\text{ini} = i_\text{fin} = j_\text{ini} = - j_\text{fin} = k_\text{ini} = k_\text{fin} = n$. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> transform = tio.Crop(cropping=10) - >>> # Equivalent to - >>> transform = tio.Crop(cropping=(10, 10, 10)) - >>> # Equivalent to - >>> transform = tio.Crop(cropping=(10, 10, 10, 10, 10, 10)) - """ - - def __init__( - self, - *, - cropping: CroppingParam, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.cropping = _parse_cropping(cropping) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - return {"cropping": self.cropping} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - i0, i1, j0, j1, k0, k1 = params["cropping"] - for _name, img_batch in self._get_images(batch).items(): - data = img_batch.data - # data shape: (B, C, I, J, K) - si = data.shape[-3] - sj = data.shape[-2] - sk = data.shape[-1] - img_batch.data = data[ - ..., - i0 : si - i1 or None, - j0 : sj - j1 or None, - k0 : sk - k1 or None, - ] - # Update each affine's origin - for affine in img_batch.affines: - origin_shift = affine.data[:3, :3] @ affine.data.new_tensor( - [float(i0), float(j0), float(k0)], - ) - affine._matrix[:3, 3] += origin_shift - return batch - - @property - def invertible(self) -> bool: - return True - - def inverse(self, params: dict[str, Any]) -> Any: - """Inverse of Crop is Pad.""" - from .pad import Pad - - return Pad(padding=params["cropping"], copy=False) diff --git a/src/torchio/transforms/spatial/crop_or_pad.py b/src/torchio/transforms/spatial/crop_or_pad.py deleted file mode 100644 index e958c48ba..000000000 --- a/src/torchio/transforms/spatial/crop_or_pad.py +++ /dev/null @@ -1,635 +0,0 @@ -"""CropOrPad transform: crop and/or pad to a target shape.""" - -from __future__ import annotations - -import copy as _copy -import math -from typing import Any -from typing import Literal - -import numpy as np -import torch -from loguru import logger -from torch import Tensor - -from ...data.affine import AffineMatrix -from ...data.backends import ImageDataBackend -from ...data.backends import normalize_index -from ...data.batch import SubjectsBatch -from ...data.image import Image -from ...data.subject import Subject -from ...types import SliceIndex -from ...types import TypeAffineMatrix -from ...types import TypeSixInts -from ...types import TypeSpacing -from ...types import TypeTensorShape -from ...types import TypeThreeInts -from ..compose import Compose -from ..transform import AppliedTransform -from ..transform import SpatialTransform -from ._padding import PaddingMode -from ._padding import pad_tensor -from ._padding import parse_padding_mode -from .crop import Crop -from .pad import Pad - -#: Accepted target shape specifications. -#: `int` or `float` → same size for each axis. -#: 3-tuple → per axis; use `None` to leave an axis unchanged. -TargetShapeParam = ( - int | float | tuple[int | float | None, int | float | None, int | float | None] -) - -#: Accepted unit values. -Units = Literal["voxels", "mm", "cm"] - -#: Accepted crop location strategies. -Location = Literal["center", "random"] - - -def _parse_target_shape( - target_shape: TargetShapeParam, -) -> tuple[float | None, float | None, float | None]: - """Normalise target_shape to a 3-tuple of floats or None.""" - if isinstance(target_shape, (int, float)): - return (float(target_shape), float(target_shape), float(target_shape)) - values = list(target_shape) - n = len(values) - if n == 3: - a, b, c = values - return ( - None if a is None else float(a), - None if b is None else float(b), - None if c is None else float(c), - ) - msg = f"target_shape must have 1 or 3 values, got {n}" - raise ValueError(msg) - - -def _to_voxels( - target: tuple[float | None, float | None, float | None], - units: Units, - spacing: TypeSpacing, - current_shape: TypeThreeInts, -) -> TypeThreeInts: - """Convert a target shape from the given units to integer voxels. - - `None` entries are replaced with the current size along that axis. - """ - result: list[int] = [] - for t, sp, cur in zip(target, spacing, current_shape, strict=True): - if t is None: - result.append(cur) - elif units == "voxels": - result.append(round(t)) - else: - factor = 10.0 if units == "cm" else 1.0 - result.append(round(t * factor / sp)) - return (result[0], result[1], result[2]) - - -def _split_per_axis( - diff: int, - location: Location, -) -> tuple[tuple[int, int], tuple[int, int]]: - """Compute (pad_ini, pad_fin) and (crop_ini, crop_fin) for one axis.""" - if diff > 0: - ini = math.ceil(diff / 2) - fin = math.floor(diff / 2) - return (ini, fin), (0, 0) - if diff < 0: - amount = -diff - if location == "random": - ini = int(torch.randint(0, amount + 1, (1,)).item()) - else: - ini = math.ceil(amount / 2) - return (0, 0), (ini, amount - ini) - return (0, 0), (0, 0) - - -def _compute_crop_and_pad( - current_shape: TypeThreeInts, - target_shape: TypeThreeInts, - *, - only_crop: bool, - only_pad: bool, - location: Location = "center", -) -> tuple[TypeSixInts | None, TypeSixInts | None]: - """Compute per-side crop and pad amounts to go from current to target. - - Args: - location: `"center"` splits evenly; `"random"` picks a - random crop start position for axes that need cropping. - - Returns: - `(padding_six, cropping_six)`: either may be `None` when no - padding or cropping is needed (or when `only_crop` / `only_pad` - suppress it). - """ - pad_values: list[int] = [] - crop_values: list[int] = [] - for cur, tgt in zip(current_shape, target_shape, strict=True): - pad, crop = _split_per_axis(tgt - cur, location) - pad_values.extend(pad) - crop_values.extend(crop) - - has_padding = any(v > 0 for v in pad_values) - has_cropping = any(v > 0 for v in crop_values) - - padding: TypeSixInts | None = None - if has_padding and not only_crop: - padding = ( - pad_values[0], - pad_values[1], - pad_values[2], - pad_values[3], - pad_values[4], - pad_values[5], - ) - - cropping: TypeSixInts | None = None - if has_cropping and not only_pad: - cropping = ( - crop_values[0], - crop_values[1], - crop_values[2], - crop_values[3], - crop_values[4], - crop_values[5], - ) - - return padding, cropping - - -class _CroppedBackend: - """Backend wrapper that defers spatial cropping until data is accessed.""" - - __slots__ = ("_affine", "_shape", "_source", "_spatial_slices") - - def __init__( - self, - source: ImageDataBackend, - spatial_slices: tuple[slice, slice, slice], - cropped_shape: tuple[int, int, int, int], - affine: TypeAffineMatrix, - ) -> None: - self._source = source - self._spatial_slices = spatial_slices - self._shape = cropped_shape - self._affine = affine - - @property - def shape(self) -> TypeTensorShape: - return self._shape - - @property - def affine(self) -> TypeAffineMatrix: - # The cropped affine (shifted origin), so it stays consistent with the - # owning image's affine rather than the uncropped source affine. - return self._affine - - @property - def dtype(self) -> np.dtype: - return self._source.dtype - - def to_tensor(self) -> Tensor: - slices = (slice(None), *self._spatial_slices) - return self._source[slices].to(dtype=torch.float32, copy=True) - - def __getitem__(self, slices: SliceIndex) -> Tensor: - return self.to_tensor()[normalize_index(slices)] - - -class _PaddedBackend: - """Backend wrapper that defers spatial padding until data is accessed.""" - - __slots__ = ("_affine", "_fill", "_padding", "_padding_mode", "_shape", "_source") - - def __init__( - self, - source: ImageDataBackend, - padding: TypeSixInts, - padded_shape: tuple[int, int, int, int], - affine: TypeAffineMatrix, - padding_mode: PaddingMode = "constant", - fill: float = 0, - ) -> None: - self._source = source - self._padding = padding - self._shape = padded_shape - self._affine = affine - self._padding_mode = padding_mode - self._fill = fill - - @property - def shape(self) -> TypeTensorShape: - return self._shape - - @property - def affine(self) -> TypeAffineMatrix: - # The padded affine (shifted origin), so it stays consistent with the - # owning image's affine rather than the unpadded source affine. - return self._affine - - @property - def dtype(self) -> np.dtype: - return self._source.dtype - - def to_tensor(self) -> Tensor: - base = self._source.to_tensor() - return pad_tensor( - base, - self._padding, - self._padding_mode, - self._fill, - ) - - def __getitem__(self, slices: SliceIndex) -> Tensor: - return self.to_tensor()[normalize_index(slices)] - - -def _get_images( - subject: Subject, - include: list[str] | None, - exclude: list[str] | None, -) -> dict[str, Image]: - """Filter subject images by include/exclude.""" - images = subject.images - if include is not None: - images = {k: v for k, v in images.items() if k in include} - if exclude is not None: - images = {k: v for k, v in images.items() if k not in exclude} - return images - - -def _crop_image_lazy(image: Image, cropping: TypeSixInts) -> Image: - """Crop an image lazily (data is only loaded when accessed).""" - i0, i1, j0, j1, k0, k1 = cropping - c, si, sj, sk = image.shape - - i_slice = slice(i0, si - i1 or None) - j_slice = slice(j0, sj - j1 or None) - k_slice = slice(k0, sk - k1 or None) - - # Compute new affine - affine_matrix = image.affine.data.clone() - start_voxel = torch.tensor( - [float(i0), float(j0), float(k0)], - dtype=torch.float64, - ) - affine_matrix[:3, 3] += affine_matrix[:3, :3] @ start_voxel - new_affine = AffineMatrix(affine_matrix) - - if image.is_loaded: - new_data = image.data[:, i_slice, j_slice, k_slice] - return image.new_like(data=new_data, affine=new_affine) - - # Install a cropped backend on a new Image from the same source - image._ensure_backend() - if image._backend is not None and ( - image._path is not None or image._zarr_store is not None - ): - cropped_shape = ( - c, - len(range(*i_slice.indices(si))), - len(range(*j_slice.indices(sj))), - len(range(*k_slice.indices(sk))), - ) - source = image._path if image._path is not None else image._zarr_store - new = type(image)( - source, - reader=image._reader, - reader_kwargs=dict(image._reader_kwargs), - affine=new_affine, - **dict(image._metadata), - ) - new._backend = _CroppedBackend( - image._backend, - (i_slice, j_slice, k_slice), - cropped_shape, - affine=new_affine.data, - ) - return new - - # No backend (custom reader) → fall back to eager crop - new_data = image.data[:, i_slice, j_slice, k_slice] - return image.new_like(data=new_data, affine=new_affine) - - -def _pad_image_lazy( - image: Image, - padding: TypeSixInts, - padding_mode: PaddingMode, - fill: float, -) -> Image: - """Pad an image lazily (data is only loaded when accessed).""" - i0, i1, j0, j1, k0, k1 = padding - c, si, sj, sk = image.shape - - # Compute new affine - affine_matrix = image.affine.data.clone() - origin_shift = affine_matrix[:3, :3] @ affine_matrix.new_tensor( - [-float(i0), -float(j0), -float(k0)], - ) - affine_matrix[:3, 3] += origin_shift - new_affine = AffineMatrix(affine_matrix) - - padded_shape = (c, si + i0 + i1, sj + j0 + j1, sk + k0 + k1) - - if image.is_loaded: - new_data = pad_tensor( - image.data, - padding, - padding_mode, - fill, - ) - return image.new_like(data=new_data, affine=new_affine) - - # Install a padded backend on a new Image from the same source - image._ensure_backend() - if image._backend is not None and ( - image._path is not None or image._zarr_store is not None - ): - source = image._path if image._path is not None else image._zarr_store - new = type(image)( - source, - reader=image._reader, - reader_kwargs=dict(image._reader_kwargs), - affine=new_affine, - **dict(image._metadata), - ) - new._backend = _PaddedBackend( - image._backend, - padding, - padded_shape, - affine=new_affine.data, - padding_mode=padding_mode, - fill=fill, - ) - return new - - # No backend → fall back to eager pad - new_data = pad_tensor( - image.data, - padding, - padding_mode, - fill, - ) - return image.new_like(data=new_data, affine=new_affine) - - -class CropOrPad(SpatialTransform): - r"""Crop and/or pad to a target spatial shape. - - If the current spatial size along an axis is larger than the target, that - axis is cropped symmetrically from both sides. If it is smaller, it is - padded symmetrically. The affine matrix is updated so that physical - positions of the voxels are maintained. - - The target shape can be specified in voxels (the default), millimetres, - or centimetres. When physical units are used, the target is converted to - voxels at transform time using the image spacing. - - When the input is a `Subject` or `Image`, the transform operates - lazily (data is not loaded from disk until it is actually accessed). - - Args: - target_shape: Desired spatial shape. A single `int` broadcasts - to all three axes. When `units` is `"mm"` or `"cm"`, - values may be floats representing the physical extent along - each axis. Use `None` for an axis to leave it unchanged, - e.g., `(256, 256, None)`. - units: Coordinate system for `target_shape`. One of - `"voxels"` (default), `"mm"`, or `"cm"`. - padding_mode: One of `'constant'`, `'reflect'`, - `'replicate'`, `'circular'`, `'mean'`, `'median'`, or - `'minimum'`. Statistical modes use one value computed from - the whole image volume. For integer inputs, `'mean'` and - `'median'` may be truncated to the input dtype and emit a - warning. - fill: Fill value when `padding_mode='constant'`. - only_crop: If `True`, padding is never applied. Mutually - exclusive with `only_pad`. - only_pad: If `True`, cropping is never applied. Mutually - exclusive with `only_crop`. - location: Where to place the crop window when the image is - larger than the target. `"center"` (default) centres the - window; `"random"` picks a uniformly random position. - Padding is always centred regardless of this parameter. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> transform = tio.CropOrPad(target_shape=(120, 80, 180)) - >>> transform = tio.CropOrPad(target_shape=256) - >>> transform = tio.CropOrPad(target_shape=(150.0, 200.0, 180.0), units='mm') - >>> transform = tio.CropOrPad(target_shape=(15.0, 20.0, 18.0), units='cm') - >>> transform = tio.CropOrPad(target_shape=256, only_pad=True) - >>> transform = tio.CropOrPad(target_shape=(256, 256, None)) # keep depth - >>> transform = tio.CropOrPad(target_shape=96, location='random') - >>> transform = tio.CropOrPad(target_shape=256, padding_mode='mean') - """ - - def __init__( - self, - target_shape: TargetShapeParam, - *, - units: Units = "voxels", - padding_mode: PaddingMode = "constant", - fill: float = 0, - only_crop: bool = False, - only_pad: bool = False, - location: Location = "center", - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - if only_crop and only_pad: - msg = "only_crop and only_pad cannot both be True" - raise ValueError(msg) - if units not in ("voxels", "mm", "cm"): - msg = f"units must be 'voxels', 'mm', or 'cm', got {units!r}" - raise ValueError(msg) - if location not in ("center", "random"): - msg = f"location must be 'center' or 'random', got {location!r}" - raise ValueError(msg) - self.target_shape = _parse_target_shape(target_shape) - self.units: Units = units - self.padding_mode = parse_padding_mode(padding_mode) - self.fill = fill - self.only_crop = only_crop - self.only_pad = only_pad - self.location: Location = location - - def forward(self, data): - """Apply the transform. - - For `Subject` and `Image` inputs, operates lazily per-image - without loading data from disk. For batched inputs, falls back - to the standard `SubjectsBatch` path. - """ - if isinstance(data, (Subject, Image)): - return self._forward_lazy(data) - return super().forward(data) - - def _forward_lazy(self, data: Subject | Image) -> Subject | Image: - is_image = isinstance(data, Image) - if is_image: - subject = Subject(tio_default_image=data) - else: - assert isinstance(data, Subject) - subject = data - - if self.copy: - subject = _copy.deepcopy(subject) - - if torch.rand(1).item() > self.p: - return subject.tio_default_image if is_image else subject - - first_image = next(iter(subject.images.values())) - current_shape: TypeThreeInts = first_image.spatial_shape - target_voxels = _to_voxels( - self.target_shape, - self.units, - first_image.affine.spacing, - current_shape, - ) - - padding, cropping = _compute_crop_and_pad( - current_shape, - target_voxels, - only_crop=self.only_crop, - only_pad=self.only_pad, - location=self.location, - ) - - self._apply_lazy_ops(subject, padding, cropping) - - return subject.tio_default_image if is_image else subject - - def _apply_lazy_ops( - self, - subject: Subject, - padding: TypeSixInts | None, - cropping: TypeSixInts | None, - ) -> None: - """Apply lazy pad/crop and record history.""" - images = _get_images(subject, self.include, self.exclude) - include = None if self.include is None else list(self.include) - exclude = None if self.exclude is None else list(self.exclude) - - if padding is not None: - for name, image in images.items(): - subject._images[name] = _pad_image_lazy( - image, - padding, - self.padding_mode, - self.fill, - ) - subject.applied_transforms.append( - AppliedTransform( - name="Pad", - params={ - "padding": padding, - "padding_mode": self.padding_mode, - "fill": self.fill, - }, - include=include, - exclude=exclude, - ), - ) - - if cropping is not None: - images = _get_images(subject, self.include, self.exclude) - for name, image in images.items(): - subject._images[name] = _crop_image_lazy(image, cropping) - subject.applied_transforms.append( - AppliedTransform( - name="Crop", - params={"cropping": cropping}, - include=include, - exclude=exclude, - ), - ) - - subject.applied_transforms.append( - AppliedTransform( - name="CropOrPad", - params={"padding": padding, "cropping": cropping}, - include=include, - exclude=exclude, - ), - ) - - # --- Standard batch path (for SubjectsBatch, Tensor, etc.) --- - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - first_images = next(iter(batch.images.values())) - spacing = first_images.affines[0].spacing - - data_tensor = first_images.data - current_shape: TypeThreeInts = ( - data_tensor.shape[-3], - data_tensor.shape[-2], - data_tensor.shape[-1], - ) - - target_voxels = _to_voxels( - self.target_shape, - self.units, - spacing, - current_shape, - ) - - if self.units != "voxels": - logger.debug( - "CropOrPad target {} {} → {} voxels (spacing {} mm)", - self.target_shape, - self.units, - target_voxels, - spacing, - ) - - padding, cropping = _compute_crop_and_pad( - current_shape, - target_voxels, - only_crop=self.only_crop, - only_pad=self.only_pad, - location=self.location, - ) - - return {"padding": padding, "cropping": cropping} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - padding: TypeSixInts | None = params["padding"] - cropping: TypeSixInts | None = params["cropping"] - - transforms: list[SpatialTransform] = [] - if padding is not None: - transforms.append( - Pad( - padding=padding, - padding_mode=self.padding_mode, - fill=self.fill, - include=self.include, - exclude=self.exclude, - ) - ) - if cropping is not None: - transforms.append( - Crop( - cropping=cropping, - include=self.include, - exclude=self.exclude, - ) - ) - - if transforms: - pipeline = Compose(transforms, copy=False) - batch = pipeline(batch) - - return batch diff --git a/src/torchio/transforms/spatial/ensure_shape_multiple.py b/src/torchio/transforms/spatial/ensure_shape_multiple.py deleted file mode 100644 index d2d6e61e4..000000000 --- a/src/torchio/transforms/spatial/ensure_shape_multiple.py +++ /dev/null @@ -1,178 +0,0 @@ -"""EnsureShapeMultiple: pad or crop so spatial dims are divisible by n.""" - -from __future__ import annotations - -import math -from typing import Any - -from ...data.batch import SubjectsBatch -from ...data.image import Image -from ...data.subject import Subject -from ...types import TypeThreeInts -from ..transform import SpatialTransform -from ._padding import PaddingMode -from ._padding import parse_padding_mode -from .crop_or_pad import CropOrPad - -#: Accepted target_multiple specifications. -#: `int` → same value for all axes. -#: 3-tuple → per-axis values. -TargetMultipleParam = int | TypeThreeInts - - -def _parse_target_multiple(value: TargetMultipleParam) -> TypeThreeInts: - """Normalise target_multiple to a 3-tuple of positive ints.""" - if isinstance(value, int): - if value < 1: - msg = f"target_multiple must be >= 1, got {value}" - raise ValueError(msg) - return (value, value, value) - values = tuple(value) - if len(values) != 3: - msg = f"target_multiple must have 1 or 3 values, got {len(values)}" - raise ValueError(msg) - for v in values: - if v < 1: - msg = f"All target_multiple values must be >= 1, got {v}" - raise ValueError(msg) - return (values[0], values[1], values[2]) - - -def _compute_target_shape( - current_shape: TypeThreeInts, - target_multiple: TypeThreeInts, - method: str, -) -> TypeThreeInts: - """Compute the target shape so each axis is a multiple of target_multiple.""" - result: list[int] = [] - for size, multiple in zip(current_shape, target_multiple, strict=True): - if method == "pad": - target = math.ceil(size / multiple) * multiple - else: - target = math.floor(size / multiple) * multiple - target = max(target, 1) - result.append(target) - return (result[0], result[1], result[2]) - - -class EnsureShapeMultiple(SpatialTransform): - r"""Ensure that all values in the image shape are divisible by $n$. - - Some convolutional neural network architectures need the size of the - input across all spatial dimensions to be a power of 2. - - For example, a 3D U-Net with 3 downsampling (pooling) operations - needs all spatial dimensions to be multiples of $2^3 = 8$. - - This transform computes the nearest valid shape and delegates to - [`CropOrPad`][torchio.CropOrPad] to reach it. - - Args: - target_multiple: Tuple $(n_i, n_j, n_k)$ so that the output - size along axis $d$ is a multiple of $n_d$. If a single - value $n$ is provided, then $n_i = n_j = n_k = n$. - method: Either `'pad'` (default) to pad up to the next - multiple, or `'crop'` to crop down to the previous - multiple. - padding_mode: Padding mode forwarded to `CropOrPad` when - `method='pad'`. One of `'constant'`, `'reflect'`, - `'replicate'`, `'circular'`, `'mean'`, `'median'`, or - `'minimum'`. - fill: Fill value when `padding_mode='constant'`. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> transform = tio.EnsureShapeMultiple(8) - >>> transform = tio.EnsureShapeMultiple(2**3, method='pad') - >>> transform = tio.EnsureShapeMultiple(16, method='crop') - >>> transform = tio.EnsureShapeMultiple((4, 8, 16)) - """ - - def __init__( - self, - target_multiple: TargetMultipleParam, - *, - method: str = "pad", - padding_mode: PaddingMode = "constant", - fill: float = 0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.target_multiple = _parse_target_multiple(target_multiple) - if method not in ("crop", "pad"): - msg = f"method must be 'crop' or 'pad', got {method!r}" - raise ValueError(msg) - self.method = method - self.padding_mode = parse_padding_mode(padding_mode) - self.fill = fill - - def forward(self, data: Any) -> Any: - """Apply the transform. - - For `Subject` and `Image` inputs, delegates to `CropOrPad` - for lazy operation without loading data from disk. - """ - if isinstance(data, (Subject, Image)): - return self._build_crop_or_pad(data).forward(data) - return super().forward(data) - - def _build_crop_or_pad(self, data: Subject | Image) -> CropOrPad: - """Build a CropOrPad targeting the nearest valid shape.""" - if isinstance(data, Image): - current_shape = data.spatial_shape - else: - current_shape = data.spatial_shape - target_shape = _compute_target_shape( - current_shape, - self.target_multiple, - self.method, - ) - return CropOrPad( - target_shape=target_shape, - padding_mode=self.padding_mode, - fill=self.fill, - only_crop=self.method == "crop", - only_pad=self.method == "pad", - p=self.p, - copy=self.copy, - include=self.include, - exclude=self.exclude, - ) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - first_images = next(iter(batch.images.values())) - data_tensor = first_images.data - current_shape: TypeThreeInts = ( - data_tensor.shape[-3], - data_tensor.shape[-2], - data_tensor.shape[-1], - ) - target_shape = _compute_target_shape( - current_shape, - self.target_multiple, - self.method, - ) - return {"target_shape": target_shape} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - target_shape = params["target_shape"] - crop_or_pad = CropOrPad( - target_shape=target_shape, - padding_mode=self.padding_mode, - fill=self.fill, - only_crop=self.method == "crop", - only_pad=self.method == "pad", - copy=False, - include=self.include, - exclude=self.exclude, - ) - return crop_or_pad.apply_transform( - batch, - crop_or_pad.make_params(batch), - ) diff --git a/src/torchio/transforms/spatial/flip.py b/src/torchio/transforms/spatial/flip.py deleted file mode 100644 index 1469321f7..000000000 --- a/src/torchio/transforms/spatial/flip.py +++ /dev/null @@ -1,263 +0,0 @@ -"""Flip transform: reverse voxel order along spatial axes.""" - -from __future__ import annotations - -from collections.abc import Sequence -from typing import Any - -import torch -from einops import rearrange - -from ...data.batch import ImagesBatch -from ...data.batch import SubjectsBatch -from ..transform import SpatialTransform - -# Map anatomical labels to axis indices. -# Only the first letter (uppercased) is used. -_LABEL_TO_AXIS: dict[str, tuple[str, str]] = { - "L": ("L", "R"), - "R": ("L", "R"), - "A": ("A", "P"), - "P": ("A", "P"), - "I": ("I", "S"), - "S": ("I", "S"), -} - - -def _resolve_axes( - axes: int | str | Sequence[int | str], - orientation: tuple[str, str, str] | None = None, -) -> tuple[int, ...]: - """Normalise axes to a tuple of ints (0, 1, 2). - - Accepts ints, anatomical strings (`'L'`, `'Right'`, `'AP'`, - etc.), or a mix. String axes are resolved using the image - orientation. - """ - if isinstance(axes, (int, str)): - axes = (axes,) - result: list[int] = [] - for axis in axes: - if isinstance(axis, int): - if axis not in (0, 1, 2): - msg = f"Axis must be 0, 1, or 2; got {axis}" - raise ValueError(msg) - result.append(axis) - elif isinstance(axis, str): - letter = axis[0].upper() - if letter not in _LABEL_TO_AXIS: - msg = ( - f"Unknown anatomical label {axis!r}." - " Use L, R, A, P, I, S or full names" - " like 'Left', 'Right', etc." - ) - raise ValueError(msg) - if orientation is None: - msg = ( - "Cannot resolve anatomical axis label" - f" {axis!r} without image orientation" - ) - raise ValueError(msg) - pair = _LABEL_TO_AXIS[letter] - for dim, code in enumerate(orientation): - if code in pair: - result.append(dim) - break - else: - msg = f"Axis must be int or str, got {type(axis).__name__}" - raise TypeError(msg) - return tuple(sorted(set(result))) - - -class Flip(SpatialTransform): - r"""Reverse the order of elements in an image along the given axes. - - Args: - axes: Index or tuple of indices of the spatial dimensions along - which the image might be flipped. Integers must be in - `{0, 1, 2}`. Anatomical labels may also be used, such as - `'Left'`, `'Right'`, `'Anterior'`, `'Posterior'`, - `'Inferior'`, `'Superior'`. Only the first letter of - the string is used. Anatomical labels are resolved using - the image orientation. - flip_probability: Probability that each axis will be flipped - (per-axis coin flip). This is independent of the `p` - parameter, which gates the entire transform. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Tip: - Specifying the axes as anatomical labels is useful when the - image orientation is not known. - - Examples: - >>> import torchio as tio - >>> # Flip along the first spatial axis - >>> transform = tio.Flip(axes=0) - >>> # Flip along the lateral axis (anatomical label) - >>> transform = tio.Flip(axes='LR') - >>> # Random per-axis flip with 50% chance each - >>> transform = tio.Flip(axes=(0, 1, 2), flip_probability=0.5) - """ - - def __init__( - self, - *, - axes: int | str | Sequence[int | str] = 0, - flip_probability: float = 1.0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.axes = axes - if not 0 <= flip_probability <= 1: - msg = f"flip_probability must be in [0, 1], got {flip_probability}" - raise ValueError(msg) - self.flip_probability = flip_probability - - def make_params( - self, - batch: SubjectsBatch, - ) -> dict[str, Any]: - images = self._get_images(batch) - if not images: - return {"axes": ()} - first_img = next(iter(images.values())) - - n = self._resolve_n(batch) - if n is None: - orientation = None - if first_img.batch_size > 0: - orientation = first_img[0].orientation - resolved = _resolve_axes(self.axes, orientation) - flip_mask = torch.rand(3) < self.flip_probability - axes_to_flip = tuple(a for a in resolved if flip_mask[a].item()) - return {"axes": axes_to_flip} - - keep = self._keep_mask(batch, n) - axes_list = self._sample_per_element_axes(n, first_img, keep) - params = {"axes": axes_list} - self._tag_batched(params, batch, n, keep, ["axes"]) - return params - - def _sample_per_element_axes( - self, - n: int, - first_img: ImagesBatch, - keep: torch.Tensor | None, - ) -> list[list[int]]: - """Sample the flip axes for each batch element. - - Gated-out elements (``keep[index]`` is false) get an empty axis - list so they are left unflipped. - - Args: - n: Number of batch elements. - first_img: First selected image batch, used for per-element - orientation. - keep: Per-element keep mask, or ``None`` to keep all elements. - - Returns: - One list of spatial axes (in ``{0, 1, 2}``) per element. - """ - axes_list: list[list[int]] = [] - for index in range(n): - if keep is not None and not keep[index]: - axes_list.append([]) - continue - # Resolve anatomical axes per element: each sample may have its - # own orientation in a batch with per-sample affines. - resolved = _resolve_axes(self.axes, first_img[index].orientation) - flip_mask = torch.rand(3) < self.flip_probability - axes_list.append([a for a in resolved if flip_mask[a].item()]) - return axes_list - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - return True - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - axes = params["axes"] - if self._is_per_instance_params(params): - for _name, img_batch in self._get_images(batch).items(): - img_batch.data = _flip_per_element(img_batch.data, axes) - return batch - if not axes: - return batch - dims = [a - 3 for a in axes] - for _name, img_batch in self._get_images(batch).items(): - img_batch.data = torch.flip(img_batch.data, dims) - return batch - - @property - def invertible(self) -> bool: - return True - - def inverse(self, params: dict[str, Any]) -> Flip | _FlipInverse: - """Flip is its own inverse.""" - if self._is_per_instance_params(params): - return _FlipInverse(axes_per_element=params["axes"], copy=False) - return Flip(axes=params["axes"], copy=False) - - -def _flip_per_element( - data: torch.Tensor, axes_per_element: list[list[int]] -) -> torch.Tensor: - """Flip each batch element along its own axes. - - Args: - data: `(B, C, I, J, K)` tensor. - axes_per_element: One list of spatial axes (in `{0, 1, 2}`) per - element. - - Returns: - The flipped `(B, C, I, J, K)` tensor. - """ - batch_size = data.shape[0] - result = data - # Flip the whole batch along each spatial axis once, then select per - # element with a boolean mask. Flips along distinct axes commute, so - # composing them sequentially matches flipping an element's axes at once. - for spatial_axis in range(3): - flip_flags = torch.tensor( - [spatial_axis in axes_per_element[index] for index in range(batch_size)], - device=data.device, - ) - if not bool(flip_flags.any()): - continue - flipped = torch.flip(result, [spatial_axis - 3]) - mask = rearrange(flip_flags, "b -> b 1 1 1 1") - result = torch.where(mask, flipped, result) - return result - - -class _FlipInverse(SpatialTransform): - """Inverse of a per-instance [`Flip`][torchio.Flip] for history replay.""" - - def __init__( - self, - *, - axes_per_element: list[list[int]], - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self._axes_per_element = axes_per_element - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - for _name, img_batch in self._get_images(batch).items(): - img_batch.data = _flip_per_element(img_batch.data, self._axes_per_element) - return batch diff --git a/src/torchio/transforms/spatial/pad.py b/src/torchio/transforms/spatial/pad.py deleted file mode 100644 index 92b8abe05..000000000 --- a/src/torchio/transforms/spatial/pad.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Pad transform: add voxels to the borders.""" - -from __future__ import annotations - -from typing import Any - -from ...data.batch import SubjectsBatch -from ...types import TypeSixInts -from ...types import TypeThreeInts -from ..transform import SpatialTransform -from ._padding import PaddingMode -from ._padding import pad_tensor -from ._padding import parse_padding_mode - -#: Accepted padding specifications. -#: `int` → same amount on each side of each axis. -#: 3-tuple → symmetric per axis `(i, j, k)`. -#: 6-tuple → per-side `(i_ini, i_fin, j_ini, j_fin, k_ini, k_fin)`. -PaddingParam = int | TypeThreeInts | TypeSixInts - - -def _parse_padding(padding: PaddingParam) -> TypeSixInts: - """Normalise padding to a 6-tuple.""" - if isinstance(padding, int): - return (padding, padding, padding, padding, padding, padding) - values = list(padding) - n = len(values) - if n == 3: - i, j, k = values - return (i, i, j, j, k, k) - if n == 6: - return (values[0], values[1], values[2], values[3], values[4], values[5]) - msg = f"Padding must have 1, 3, or 6 values, got {n}" - raise ValueError(msg) - - -class Pad(SpatialTransform): - r"""Add a border of voxels to each side of the volume. - - Args: - padding: Tuple - $(i_\text{ini}, i_\text{fin}, j_\text{ini}, j_\text{fin}, - k_\text{ini}, k_\text{fin})$ - defining the number of voxels added to the edges of - each axis. If the initial shape of the image is - $I \times J \times K$, the final shape will be - $(I + i_\text{ini} + i_\text{fin}) \times - (J + j_\text{ini} + j_\text{fin}) \times - (K + k_\text{ini} + k_\text{fin})$. - If only three values $(i, j, k)$ are provided, then - $i_\text{ini} = i_\text{fin} = i$, etc. - If only one value $n$ is provided, all six values are $n$. - padding_mode: One of `'constant'`, `'reflect'`, - `'replicate'`, `'circular'`, `'mean'`, `'median'`, or - `'minimum'`. Statistical modes use one value computed from - the whole image volume. For integer inputs, `'mean'` and - `'median'` may be truncated to the input dtype and emit a - warning. - fill: Fill value when `padding_mode='constant'`. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> transform = tio.Pad(padding=10) - >>> transform = tio.Pad(padding=(5, 10, 0)) - >>> transform = tio.Pad(padding=10, padding_mode='reflect') - >>> transform = tio.Pad(padding=10, padding_mode='minimum') - """ - - def __init__( - self, - *, - padding: PaddingParam, - padding_mode: PaddingMode = "constant", - fill: float = 0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.padding = _parse_padding(padding) - self.padding_mode = parse_padding_mode(padding_mode) - self.fill = fill - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - return { - "padding": self.padding, - "padding_mode": self.padding_mode, - "fill": self.fill, - } - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - i0, i1, j0, j1, k0, k1 = params["padding"] - mode = params["padding_mode"] - fill = params["fill"] - for _name, img_batch in self._get_images(batch).items(): - img_batch.data = pad_tensor( - img_batch.data, - (i0, i1, j0, j1, k0, k1), - mode, - fill, - ) - # Update each affine's origin (shift back) - for affine in img_batch.affines: - origin_shift = affine.data[:3, :3] @ affine.data.new_tensor( - [-float(i0), -float(j0), -float(k0)], - ) - affine._matrix[:3, 3] += origin_shift - return batch - - @property - def invertible(self) -> bool: - return True - - def inverse(self, params: dict[str, Any]) -> Any: - """Inverse of Pad is Crop.""" - from .crop import Crop - - return Crop(cropping=params["padding"], copy=False) diff --git a/src/torchio/transforms/spatial/reorient.py b/src/torchio/transforms/spatial/reorient.py deleted file mode 100644 index 96aa730c3..000000000 --- a/src/torchio/transforms/spatial/reorient.py +++ /dev/null @@ -1,179 +0,0 @@ -"""Reorient transform: reorder voxel axes to a target orientation.""" - -from __future__ import annotations - -from typing import Any - -import nibabel as nib -import numpy as np -import torch -from nibabel import orientations -from torch import Tensor - -from ...data.batch import SubjectsBatch -from ..transform import SpatialTransform - - -def _validate_orientation(orientation: str) -> str: - """Validate and normalise a 3-letter orientation code.""" - if not isinstance(orientation, str) or len(orientation) != 3: - msg = f'Orientation must be a 3-letter string, got "{orientation}"' - raise ValueError(msg) - - orientation = orientation.upper() - valid_codes = set("RLAPIS") - if not all(c in valid_codes for c in orientation): - msg = ( - "Orientation code must be composed of three distinct characters" - f' in {valid_codes} but got "{orientation}"' - ) - raise ValueError(msg) - - _check_axis_coverage(orientation) - return orientation - - -def _check_axis_coverage(orientation: str) -> None: - """Ensure the orientation code covers all three axis pairs.""" - pairs = [{"R", "L"}, {"A", "P"}, {"S", "I"}] - codes = set(orientation) - if not all(codes & pair for pair in pairs): - msg = ( - "Orientation code must include one character for each axis" - f' direction: R or L, A or P, and S or I, but got "{orientation}"' - ) - raise ValueError(msg) - - -def _compute_reorientation( - current_affine: np.ndarray, - target_codes: str, -) -> np.ndarray: - """Compute the ornt_transform from current affine to target codes. - - Returns: - (3, 2) array where column 0 is the input axis index and - column 1 is the flip direction (1 or -1). - """ - current_ornt = orientations.io_orientation(current_affine) - target_ornt = orientations.axcodes2ornt(tuple(target_codes)) - return orientations.ornt_transform(current_ornt, target_ornt) - - -def _apply_reorientation( - data: Tensor, - ornt: np.ndarray, -) -> Tensor: - """Apply an orientation transform to a tensor. - - Mirrors nibabel's `apply_orientation`: flip axes first, then - transpose. Works on both 4D `(C, I, J, K)` and 5D - `(B, C, I, J, K)` tensors. Spatial axes are always the last 3. - - Args: - data: Input tensor. - ornt: `(3, 2)` orientation transform from nibabel. - """ - n_leading = data.ndim - 3 - - # Step 1: flip axes where direction is -1 (on original axis indices). - for ax in range(3): - if ornt[ax, 1] == -1: - data = torch.flip(data, [ax + n_leading]) - - # Step 2: permute using argsort of the source-axis column, matching - # nibabel's `arr.transpose(np.argsort(ornt[:, 0]))`. - perm = np.argsort(ornt[:, 0]).astype(int) - leading = list(range(n_leading)) - spatial_perm = [int(p) + n_leading for p in perm] - data = data.permute(leading + spatial_perm) - - return data.contiguous() - - -class Reorient(SpatialTransform): - r"""Reorder voxel axes to match a target orientation. - - The voxels are permuted and/or flipped so that the image axes - align with the specified anatomical directions. The affine matrix - is updated to preserve the physical positions of the voxels. - - Common orientation codes: - - - `'RAS'`: Left→\ **R**\ ight, - Posterior→\ **A**\ nterior, Inferior→\ **S**\ uperior. - - `'LPS'`: Right→\ **L**\ eft, - Anterior→\ **P**\ osterior, Inferior→\ **S**\ uperior. - - See the `NiBabel docs on image orientation - `__ for details. - - Args: - orientation: Target three-letter orientation code. Must - contain one letter for each axis pair: R/L, A/P, S/I. - **kwargs: See [`Transform`][torchio.Transform] for additional - keyword arguments. - - Examples: - >>> import torchio as tio - >>> transform = tio.Reorient() # default: RAS - >>> transform = tio.Reorient(orientation='LPS') - >>> transform = tio.Reorient(orientation='SPL') - """ - - def __init__( - self, - orientation: str = "RAS", - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.orientation = _validate_orientation(orientation) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - first_images = next(iter(batch.images.values())) - affine_np = first_images.affines[0].numpy() - current_codes = "".join( - nib.orientations.aff2axcodes(affine_np), - ) - ornt = _compute_reorientation(affine_np, self.orientation) - return { - "ornt": ornt.tolist(), - "original_orientation": current_codes, - } - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - ornt = np.asarray(params["ornt"]) - - # No-op when already in target orientation - is_identity = np.array_equal(ornt[:, 0], [0, 1, 2]) and np.all(ornt[:, 1] == 1) - if is_identity: - return batch - - for _name, img_batch in self._get_images(batch).items(): - original_shape = img_batch.data.shape[-3:] - img_batch.data = _apply_reorientation(img_batch.data, ornt) - - for affine in img_batch.affines: - inv_aff = orientations.inv_ornt_aff(ornt, original_shape) - new_matrix = affine.numpy() @ inv_aff - affine._matrix = torch.as_tensor( - new_matrix, - dtype=torch.float64, - ) - - return batch - - @property - def invertible(self) -> bool: - return True - - def inverse(self, params: dict[str, Any]) -> Reorient: - """Inverse reorients back to the original orientation.""" - return Reorient( - orientation=params["original_orientation"], - copy=False, - ) diff --git a/src/torchio/transforms/spatial/resize.py b/src/torchio/transforms/spatial/resize.py deleted file mode 100644 index 22eef2488..000000000 --- a/src/torchio/transforms/spatial/resize.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Resize: resample to a target shape (not spacing).""" - -from __future__ import annotations - -from typing import Any - -import torch.nn.functional as functional - -from ...data.batch import SubjectsBatch -from ...data.image import LabelMap -from ..transform import SpatialTransform - - -class Resize(SpatialTransform): - r"""Resize images to a target spatial shape. - - The field of view is preserved; voxel spacing is adjusted to fit - the new shape. - - Warning: - In most medical image applications, this transform should - **not** be used as it scales anisotropically. Prefer - [`Resample`][torchio.Resample] (change spacing) combined with - [`CropOrPad`][torchio.CropOrPad] (change shape) instead. - - Args: - target_shape: Target spatial shape $(I, J, K)$. A single - integer $N$ means $(N, N, N)$. - image_interpolation: `"linear"` (default) for intensity - images. - label_interpolation: `"nearest"` (default) for label maps. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Resize(128) - >>> transform = tio.Resize((256, 256, 64)) - """ - - def __init__( - self, - target_shape: int | tuple[int, int, int], - *, - image_interpolation: str = "linear", - label_interpolation: str = "nearest", - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - if isinstance(target_shape, int): - target_shape = (target_shape, target_shape, target_shape) - self.target_shape = target_shape - self.image_interpolation = image_interpolation - self.label_interpolation = label_interpolation - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {"target_shape": self.target_shape} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Resize each image to the target shape.""" - target = list(params["target_shape"]) - for _name, img_batch in batch.images.items(): - is_label = issubclass(img_batch._image_class, LabelMap) - mode = self.label_interpolation if is_label else self.image_interpolation - torch_mode = "nearest" if mode == "nearest" else "trilinear" - old_shape = img_batch.data.shape[2:] - img_batch.data = functional.interpolate( - img_batch.data.float(), - size=target, - mode=torch_mode, - align_corners=None if torch_mode == "nearest" else True, - ).to(img_batch.data.dtype) - # Update affines: spacing changes to fit new shape in same FOV. - for affine in img_batch.affines: - for axis in range(3): - factor = old_shape[axis] / target[axis] - affine._matrix[:3, axis] *= factor - return batch diff --git a/src/torchio/transforms/spatial/spatial.py b/src/torchio/transforms/spatial/spatial.py deleted file mode 100644 index 117c46c3c..000000000 --- a/src/torchio/transforms/spatial/spatial.py +++ /dev/null @@ -1,2762 +0,0 @@ -"""Unified spatial transforms. - -Combines resampling, affine motion, and elastic deformation into a single -`grid_sample` call. The public API consists of four classes: - -- [`Spatial`][torchio.Spatial]: the unified transform. -- [`Resample`][torchio.Resample]: resampling-only convenience wrapper. -- [`Affine`][torchio.Affine]: affine-only convenience wrapper. -- [`ElasticDeformation`][torchio.ElasticDeformation]: elastic-only convenience wrapper. - -The module-level helpers handle coordinate math, grid construction, -serialization for history replay, and parameter parsing/validation. -""" - -from __future__ import annotations - -import warnings -from collections.abc import Sequence -from dataclasses import dataclass -from numbers import Number -from pathlib import Path -from typing import Any -from typing import Literal -from typing import TypeAlias -from typing import TypeGuard -from typing import cast - -import numpy as np -import numpy.typing as npt -import torch -import torch.nn.functional as functional -from einops import rearrange -from torch import Tensor -from torch.distributions import Distribution - -from ...data.affine import AffineMatrix -from ...data.batch import ImagesBatch -from ...data.batch import SubjectsBatch -from ...data.image import Image -from ...data.image import LabelMap -from ...data.image import ScalarImage -from ...types import TypeSpacing -from ...types import TypeThreeInts -from ..parameter_range import Choice -from ..parameter_range import _ParameterRange -from ..transform import SpatialTransform - - -@dataclass -class _PerSampleGrids: - """Per-element spatial parameters for per-instance augmentation. - - Each list has one entry per batch element. A `None` affine matrix - and `None` control-point field produce an identity grid for that - element (used for elements gated out by per-element probability). - - Attributes: - affine_matrices: Per-element world-space affine matrices. - control_points: Per-element elastic control-point fields. - max_displacements: Per-element maximum displacements. - """ - - affine_matrices: list[np.ndarray | None] - control_points: list[Tensor | None] - max_displacements: list[tuple[float, float, float] | None] - - -TypeParameterValue: TypeAlias = ( - int - | float - | tuple[int | float] - | tuple[int | float, int | float] - | tuple[int | float, int | float, int | float] - | tuple[ - int | float, - int | float, - int | float, - int | float, - int | float, - int | float, - ] - | tuple # per-axis mixed specs like (0, Choice([...]), (-10, 10)) - | Choice - | Distribution -) -TypeTarget: TypeAlias = ( - int - | float - | TypeSpacing - | tuple[float, float] # uniform spacing range, per axis - | tuple[float, float, float, float, float, float] # per-axis spacing ranges - | Choice - | Distribution - | str - | Path - | Image - | tuple[Sequence[int], AffineMatrix | Tensor | npt.ArrayLike] - | None -) -#: The spacing forms of `target` that resolve to a (possibly random) spacing. -TypeSpacingSpec: TypeAlias = ( - int | float | tuple | list | np.ndarray | Choice | Distribution -) -TypeControlPoints: TypeAlias = Tensor | npt.ArrayLike -TypeTargetSpace: TypeAlias = tuple[TypeThreeInts, AffineMatrix] -TypeImageInterpolation: TypeAlias = Literal[ - "nearest", - "linear", - "quadratic", - "cubic", - "fourth", - "fifth", - "sixth", - "seventh", -] -#: Label maps additionally accept the partial-volume-aware `"label"` mode. -TypeLabelInterpolation: TypeAlias = TypeImageInterpolation | Literal["label"] -#: Broad alias accepting any interpolation mode, including `"label"`. Used by -#: internal helpers that handle both image and label interpolation. -TypeInterpolation: TypeAlias = TypeLabelInterpolation -TypeCenter: TypeAlias = Literal["image", "origin"] -TypePadValue: TypeAlias = Literal["minimum", "mean", "otsu"] - -#: Partial-volume label interpolation mode. Only valid for label maps: the map -#: is one-hot encoded, each channel is resampled linearly, and the per-voxel -#: argmax recovers the discrete labels (see `_resample_label_partial_volume`). -LABEL_INTERPOLATION = "label" - -_SUPPORTED_INTERPOLATIONS = ( - "nearest", - "linear", - "quadratic", - "cubic", - "fourth", - "fifth", - "sixth", - "seventh", - LABEL_INTERPOLATION, -) -_INTERPOLATION_TO_ORDER: dict[str, int] = { - "nearest": 0, - "linear": 1, - "quadratic": 2, - "cubic": 3, - "fourth": 4, - "fifth": 5, - "sixth": 6, - "seventh": 7, -} -_TORCH_INTERPOLATION_MODE = { - "nearest": "nearest", - "linear": "bilinear", -} -_SUPPORTED_PAD_VALUES = ("minimum", "mean", "otsu") -_SPLINE_ORDER = 3 - - -class Spatial(SpatialTransform): - r"""Apply resampling, affine motion, and elastic deformation together. - - This transform can: - - 1. resample to a new space, - 2. apply a global affine mapping, and - 3. apply a dense elastic field, - - using a single sampling grid. - - The convenience wrappers [`Resample`][torchio.Resample], - [`Affine`][torchio.Affine], and - [`ElasticDeformation`][torchio.ElasticDeformation] expose subsets - of these parameters with sensible defaults. - - Args: - target: Output space. Can be one of: - - - A scalar or 3-tuple of floats: output voxel spacing in mm. - E.g., `1` for 1 mm isotropic, `(0.5, 0.5, 2.0)` for - anisotropic. - - A random spacing spec, sampled once per call at apply time: a - 2-tuple `(lo, hi)` (uniform, per axis), a 6-tuple - `(lo1, hi1, lo2, hi2, lo3, hi3)` (per-axis ranges), a `Choice`, - or a `torch.distributions.Distribution`. E.g., - `target=(1, 2)` or `target=(2, 4, 2, 4, 3, 6)`. - - A `str`: either a path to an image file, or the name of - an image in the subject (e.g., `"t1"`). - - An [`Image`][torchio.Image] instance. - - A `(spatial_shape, affine)` pair. - - `None` (default): the output grid matches the input grid. - scales: Scale factors $(s_1, s_2, s_3)$ for each axis. - If a single value $x$ is given, all axes use $x$. - If two values $(a, b)$ are given, - $s_i \sim \mathcal{U}(a, b)$. - If six values $(a_1, b_1, a_2, b_2, a_3, b_3)$ are given, - $s_i \sim \mathcal{U}(a_i, b_i)$ independently. - A `torch.distributions.Distribution` may also be passed. - For example, `scales=0.5` halves the apparent object - size (zoom out), and `scales=2` doubles it (zoom in). - degrees: Euler rotation angles $(\theta_1, \theta_2, \theta_3)$ - in degrees, following the same value/range/distribution - convention as *scales*. - translation: Translation $(t_1, t_2, t_3)$ in mm, following - the same convention. The direction depends on the image - orientation: in RAS+, `translation=(10, 0, 0)` shifts - 10 mm to the right. - isotropic: If `True`, sample a single scale factor and - reuse it for all three axes. *scales* must then be a - scalar or 2-value range. - center: Pivot point for rotation and scaling. - `"image"` (default) uses the image center; - `"origin"` uses the world-coordinate origin. - control_points: Optional pre-computed coarse displacement - field with shape `(n_i, n_j, n_k, 3)` in mm. If given, - *num_control_points*, *max_displacement*, and - *locked_borders* are ignored. - num_control_points: Number of control points along each - dimension of the coarse grid. Can be a single `int` - (isotropic) or a 3-tuple. Minimum is 4. Smaller values - produce smoother deformations. - max_displacement: Maximum displacement at each control point, - in mm. Follows the same value/range/distribution - convention as *scales*. Zero (default) disables elastic - deformation. - locked_borders: Number of outer control-point layers whose - displacement is forced to zero. `0` keeps all - displacements; `1` zeros the outermost layer; `2` - (default) zeros the two outermost layers. - affine_first: If `True` (default), apply the affine mapping - before the elastic field. If `False`, apply the elastic - field first. The difference is significant for large - transforms. - image_interpolation: Interpolation for - [`ScalarImage`][torchio.ScalarImage] instances. `"linear"` - (default) or `"nearest"` use a fast path; higher-order - B-spline modes `"quadratic"`, `"cubic"`, `"fourth"`, - `"fifth"`, `"sixth"`, and `"seventh"` are also supported, as - are the equivalent integer orders `0`-`7`. - label_interpolation: Interpolation for - [`LabelMap`][torchio.LabelMap] instances. Accepts the same - values as *image_interpolation* (`"nearest"` is the default), - plus the special `"label"` mode. The `"label"` mode performs - partial-volume-aware resampling: the label map is one-hot - encoded, each channel is resampled with linear interpolation, - and the per-voxel argmax recovers the discrete labels. - Compared with `"nearest"`, this reduces staircase artifacts - and yields more accurate label volumes, which is especially - useful when downsampling. For single-channel label maps it - never invents intermediate label values that were absent from - the input (the only new value that can appear is - `default_pad_label`, used for out-of-bounds voxels). A - multi-channel (already one-hot or probabilistic) label map is - instead resampled per channel, so its output can contain - fractional partial volumes. - - The `"label"` mode is more memory- and compute-intensive - because it processes one channel per label, so memory scales - with the number of distinct labels. For maps with many labels - (e.g. hundreds of brain structures) prefer processing patch by - patch with [`GridSampler`][torchio.GridSampler] and - [`PatchAggregator`][torchio.PatchAggregator] to bound memory. - one_hot_label_interpolation: Interpolation used for the one-hot - channels of the `"label"` mode. Accepts the same values as - *image_interpolation* (`"nearest"`, `"linear"`, or the - higher-order B-spline modes / integer orders 0-7), but **not** - `"label"`. Defaults to `"linear"`. Higher-order modes give - smoother label boundaries, which is particularly useful when - upsampling. Ignored unless `label_interpolation="label"`. - antialias: If `True`, apply Gaussian smoothing before - downsampling intensity images. Label maps are smoothed only - when `label_interpolation="label"` (the one-hot channels are - blurred before downsampling); otherwise they are left - untouched. The standard deviations follow - [Cardoso et al., MICCAI 2015](https://link.springer.com/chapter/10.1007/978-3-319-24571-3_81). - default_pad_value: Fill rule for out-of-bounds intensity - voxels. `"minimum"` (default), `"mean"`, `"otsu"`, - or a numeric value. - default_pad_label: Numeric fill value for out-of-bounds label - voxels. - **kwargs: See [`Transform`][torchio.Transform]. - - Note: - The randomizable parameters (`scales`, `degrees`, `translation`, - `max_displacement`, and the spacing form of `target`) follow a common - value/range/distribution convention: a scalar is deterministic, a - 2-tuple $(a, b)$ samples uniformly, a 3-tuple sets per-axis values, a - 6-tuple sets per-axis ranges, and a `Choice` or - `torch.distributions.Distribution` samples from a discrete set or a - distribution. Structural parameters (`center`, interpolation, padding) - are not randomizable. - - Examples: - >>> import torchio as tio - >>> # Resample to 1 mm isotropic with a random rotation - >>> transform = tio.Spatial( - ... target=1, - ... degrees=(-10, 10), - ... translation=(-5, 5), - ... ) - >>> # Elastic deformation only - >>> transform = tio.Spatial( - ... max_displacement=7.5, - ... num_control_points=7, - ... ) - >>> transformed = transform(subject) - """ - - def __init__( - self, - *, - target: TypeTarget = None, - scales: TypeParameterValue = 1.0, - degrees: TypeParameterValue = 0.0, - translation: TypeParameterValue = 0.0, - isotropic: bool = False, - center: TypeCenter = "image", - control_points: TypeControlPoints | None = None, - num_control_points: int | TypeThreeInts = 7, - max_displacement: TypeParameterValue = 0.0, - locked_borders: int = 2, - affine_first: bool = True, - image_interpolation: TypeImageInterpolation | int = "linear", - label_interpolation: TypeLabelInterpolation | int = "nearest", - one_hot_label_interpolation: TypeImageInterpolation | int = "linear", - antialias: bool = False, - default_pad_value: TypePadValue | float = "minimum", - default_pad_label: int | float = 0, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.target = target - _validate_isotropic(scales, isotropic) - self.scales = _to_positive_range(scales) - self.degrees = _to_parameter_range(degrees) - self.translation = _to_parameter_range(translation) - self.isotropic = isotropic - self.center = _parse_center(center) - self.control_points = ( - _parse_control_points(control_points) - if control_points is not None - else None - ) - self.num_control_points = _parse_num_control_points(num_control_points) - self.max_displacement = _to_nonnegative_parameter_range(max_displacement) - self.locked_borders = _parse_locked_borders(locked_borders) - if self.locked_borders == 2 and 4 in self.num_control_points: - msg = ( - "locked_borders=2 with 4 control points along any axis yields an" - " identity elastic field" - ) - raise ValueError(msg) - self.affine_first = affine_first - parsed_image_interpolation = _parse_interpolation(image_interpolation) - if parsed_image_interpolation == LABEL_INTERPOLATION: - msg = ( - f'image_interpolation cannot be "{LABEL_INTERPOLATION}"; that mode' - " is only valid for label_interpolation" - ) - raise ValueError(msg) - self.image_interpolation = parsed_image_interpolation - self.label_interpolation = _parse_interpolation(label_interpolation) - self.one_hot_label_interpolation = _parse_one_hot_label_interpolation( - one_hot_label_interpolation, - ) - self.antialias = antialias - self.default_pad_value = _parse_default_pad_value(default_pad_value) - if not isinstance(default_pad_label, Number): - msg = f"default_pad_label must be numeric, got {type(default_pad_label)}" - raise TypeError(msg) - self.default_pad_label = float(default_pad_label) - - @property - def supports_per_instance_params(self) -> bool: - return True - - @property - def supports_per_instance_p(self) -> bool: - # Per-element gating requires a shape-preserving transform so the - # gated-out element is a true no-op. A resampling target changes - # the grid, so keep batch-wide gating in that case. - return self.target is None - - def _sample_one( - self, - shape: TypeThreeInts, - affine: AffineMatrix, - ) -> tuple[ - np.ndarray | None, - Tensor | None, - tuple[float, float, float] | None, - bool, - ]: - """Sample one set of affine and elastic parameters. - - Args: - shape: Spatial shape used to build the affine and to suppress - out-of-plane components for 2D inputs. - affine: Affine of the reference image. - - Returns: - A tuple `(forward_affine, control_points, max_displacement, - has_geometry)` where `has_geometry` is `True` when either an - affine or an elastic component is present. - """ - sampled_scales = _sample_scales(self.scales, self.isotropic) - sampled_degrees = self.degrees.sample() - sampled_translation = self.translation.sample() - has_affine = _has_affine_component( - sampled_scales, - sampled_degrees, - sampled_translation, - ) - control_points, max_displacement = _resolve_control_points( - self.control_points, - self.num_control_points, - self.max_displacement, - self.locked_borders, - ) - has_elastic = control_points is not None - forward_affine = None - if has_affine: - forward_affine = _build_forward_affine( - scales=sampled_scales, - degrees=sampled_degrees, - translation=sampled_translation, - center=self.center, - shape=shape, - affine=affine, - ) - return ( - forward_affine, - control_points, - max_displacement, - (has_affine or has_elastic), - ) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample random parameters and resolve the output space. - - Scales, degrees, translation, and control-point displacements are - sampled per batch element when per-instance augmentation is - active (the default for batches), and once otherwise. - - Returns: - Dict of serializable parameters for `apply_transform` and - history replay. - """ - images = self._get_images(batch) - if not images: - return {"selected_images": []} - - _, first_batch = next(iter(images.items())) - first_shape = _get_spatial_shape(first_batch) - first_affine = first_batch.affines[0] - - params: dict[str, Any] = { - "selected_images": list(images.keys()), - "original": _serialize_space((first_shape, first_affine)), - "affine_first": self.affine_first, - "image_interpolation": self.image_interpolation, - "label_interpolation": self.label_interpolation, - "one_hot_label_interpolation": self.one_hot_label_interpolation, - "antialias": self.antialias, - "default_pad_value": self.default_pad_value, - "default_pad_label": self.default_pad_label, - } - - n = self._resolve_n(batch) - if n is None: - forward_affine, control_points, max_displacement, has_geometry = ( - self._sample_one(first_shape, first_affine) - ) - if has_geometry: - _check_shared_space(images, first_shape, first_affine) - # Resolve the (possibly random) target after sampling the - # geometry so the RNG stream matches the batch-shared path. - target_space = _resolve_target_space( - self.target, - batch, - first_shape, - first_affine, - ) - params["target"] = _serialize_space(target_space) - params["affine_matrix"] = _serialize_matrix(forward_affine) - params["control_points"] = _serialize_control_points(control_points) - params["max_displacement"] = ( - list(max_displacement) if max_displacement else None - ) - return params - - keep = self._keep_mask(batch, n) - affine_list, control_points_list, displacement_list, any_geometry = ( - self._sample_per_element_geometry(n, first_shape, first_affine, keep) - ) - if any_geometry: - _check_shared_space(images, first_shape, first_affine) - target_space = _resolve_target_space( - self.target, - batch, - first_shape, - first_affine, - ) - params["target"] = _serialize_space(target_space) - params["affine_matrix"] = affine_list - params["control_points"] = control_points_list - params["max_displacement"] = displacement_list - self._tag_batched( - params, - batch, - n, - keep, - ["affine_matrix", "control_points", "max_displacement"], - ) - return params - - def _sample_per_element_geometry( - self, - n: int, - first_shape: TypeThreeInts, - first_affine: AffineMatrix, - keep: Tensor | None, - ) -> tuple[list[Any], list[Any], list[Any], bool]: - """Sample serializable geometry for each batch element. - - Gated-out elements (``keep[index]`` is false) contribute ``None`` - placeholders so the per-element lists stay aligned with the batch. - - Args: - n: Number of batch elements. - first_shape: Spatial shape of the first image. - first_affine: Affine of the first image. - keep: Per-element keep mask, or ``None`` to keep all elements. - - Returns: - The affine, control-point and max-displacement lists plus a - flag indicating whether any element produced a geometry - transform. - """ - affine_list: list[Any] = [] - control_points_list: list[Any] = [] - displacement_list: list[Any] = [] - any_geometry = False - for index in range(n): - kept = keep is None or bool(keep[index]) - if not kept: - affine_list.append(None) - control_points_list.append(None) - displacement_list.append(None) - continue - forward_affine, control_points, max_displacement, has_geometry = ( - self._sample_one(first_shape, first_affine) - ) - any_geometry = any_geometry or has_geometry - affine_list.append(_serialize_matrix(forward_affine)) - control_points_list.append(_serialize_control_points(control_points)) - displacement_list.append( - list(max_displacement) if max_displacement else None - ) - return affine_list, control_points_list, displacement_list, any_geometry - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply the spatial mapping to every selected image in *batch*. - - One sampling grid is built per batch element when per-instance - parameters are present, and a single shared grid otherwise. - """ - selected_images = params.get("selected_images", []) - if not selected_images: - return batch - - target_space = _deserialize_space(params["target"]) - - affine_matrix, control_points, max_displacement, per_sample = ( - _resolve_spatial_params(params) - ) - is_noop = ( - target_space is None - and affine_matrix is None - and control_points is None - and max_displacement is None - and per_sample is None - ) - if is_noop: - # A true no-op (no resampling and no geometry, e.g. every - # element gated out) must leave the data and the per-sample - # affines untouched instead of rebuilding an identity grid. - return batch - _apply_spatial_to_batch( - batch=batch, - image_names=selected_images, - target_space=target_space, - affine_matrix=affine_matrix, - control_points=control_points, - max_displacement=max_displacement, - affine_first=params["affine_first"], - image_interpolation=params["image_interpolation"], - label_interpolation=params["label_interpolation"], - one_hot_label_interpolation=params.get( - "one_hot_label_interpolation", - "linear", - ), - antialias=params.get("antialias", False), - default_pad_value=params["default_pad_value"], - default_pad_label=float(params["default_pad_label"]), - per_sample=per_sample, - ) - return batch - - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return True - - def inverse(self, params: dict[str, Any]) -> _SpatialInverse: - """Build the inverse transform from recorded parameters. - - The affine component is inverted exactly. The elastic component - is approximated by negating the sampled displacement field. The - `affine_first` flag is flipped so that the inverse operations - run in the opposite order. Per-instance parameters are inverted - element by element. - - Args: - params: The parameter dict produced by `make_params`. - - Returns: - A `_SpatialInverse` that resamples back to the original grid. - """ - original_space = _deserialize_space(params["original"]) - if original_space is None: - msg = "Spatial inverse needs the original output space" - raise RuntimeError(msg) - - common: dict[str, Any] = { - "target": original_space, - "affine_first": not params["affine_first"], - "image_interpolation": params["image_interpolation"], - "label_interpolation": params["label_interpolation"], - "one_hot_label_interpolation": params.get( - "one_hot_label_interpolation", - "linear", - ), - "default_pad_value": params["default_pad_value"], - "default_pad_label": float(params["default_pad_label"]), - "copy": False, - # Invert only the images the forward pass actually transformed, - # so excluded images (e.g. label maps) are not resampled. - "include": params["selected_images"], - } - - batched_keys = params.get("_batched_keys") or [] - if "affine_matrix" in batched_keys: - per_sample = _invert_per_sample(params) - return _SpatialInverse( - affine_matrix=None, - control_points=None, - per_sample=per_sample, - **common, - ) - - affine_matrix = _deserialize_matrix(params["affine_matrix"]) - inverse_affine = None - if affine_matrix is not None: - inverse_affine = np.linalg.inv(affine_matrix) - control_points = _deserialize_control_points(params["control_points"]) - inverse_control_points = None - if control_points is not None: - inverse_control_points = -control_points - return _SpatialInverse( - affine_matrix=inverse_affine, - control_points=inverse_control_points, - **common, - ) - - -class _SpatialInverse(SpatialTransform): - """Concrete inverse of [`Spatial`][torchio.Spatial], used for history replay. - - Stores the exact inverse affine matrix, the negated elastic field, - and the original output space so that `apply_inverse_transform` - can restore the geometry of images transformed by [`Spatial`][torchio.Spatial]. - """ - - def __init__( - self, - *, - target: TypeTargetSpace, - affine_matrix: npt.ArrayLike | None, - control_points: TypeControlPoints | None, - affine_first: bool, - image_interpolation: TypeImageInterpolation, - label_interpolation: TypeLabelInterpolation, - one_hot_label_interpolation: TypeImageInterpolation = "linear", - default_pad_value: TypePadValue | float, - default_pad_label: float, - per_sample: _PerSampleGrids | None = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.target = target - self.affine_matrix = ( - np.asarray(affine_matrix, dtype=np.float64).copy() - if affine_matrix is not None - else None - ) - self.control_points = ( - _parse_control_points(control_points) - if control_points is not None - else None - ) - self.per_sample = per_sample - self.affine_first = affine_first - self.image_interpolation = cast( - TypeImageInterpolation, - _parse_interpolation(image_interpolation), - ) - self.label_interpolation = _parse_interpolation(label_interpolation) - self.one_hot_label_interpolation = _parse_one_hot_label_interpolation( - one_hot_label_interpolation, - ) - self.default_pad_value = _parse_default_pad_value(default_pad_value) - self.default_pad_label = float(default_pad_label) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Return empty params; all state is stored in instance attributes.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Resample back to the recorded original space.""" - max_displacement = None - if self.per_sample is None and self.control_points is not None: - max_displacement = _max_abs_displacement(self.control_points) - _apply_spatial_to_batch( - batch=batch, - image_names=list(self._get_images(batch)), - target_space=self.target, - affine_matrix=self.affine_matrix, - control_points=self.control_points, - max_displacement=max_displacement, - affine_first=self.affine_first, - image_interpolation=self.image_interpolation, - label_interpolation=self.label_interpolation, - one_hot_label_interpolation=self.one_hot_label_interpolation, - antialias=False, - default_pad_value=self.default_pad_value, - default_pad_label=self.default_pad_label, - per_sample=self.per_sample, - ) - return batch - - -class Resample(Spatial): - r"""Resample images to a different space. - - Convenience wrapper around [`Spatial`][torchio.Spatial] exposing - only the resampling parameters. - - Args: - target: Output space (see [`Spatial`][torchio.Spatial]). - Defaults to 1 mm isotropic. - image_interpolation: See [`Spatial`][torchio.Spatial]. - label_interpolation: See [`Spatial`][torchio.Spatial]. - one_hot_label_interpolation: See [`Spatial`][torchio.Spatial]. - antialias: See [`Spatial`][torchio.Spatial]. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Resample(2) # 2 mm isotropic - >>> transform = tio.Resample("t1") # match "t1" space - >>> transform = tio.Resample((1, 1, 3)) # anisotropic - >>> # Partial-volume-aware label resampling - >>> transform = tio.Resample( - ... 2, - ... label_interpolation="label", - ... antialias=True, - ... ) - """ - - def __init__( - self, - target: TypeTarget = 1, - image_interpolation: TypeImageInterpolation | int = "linear", - label_interpolation: TypeLabelInterpolation | int = "nearest", - one_hot_label_interpolation: TypeImageInterpolation | int = "linear", - antialias: bool = False, - **kwargs: Any, - ) -> None: - super().__init__( - target=target, - image_interpolation=image_interpolation, - label_interpolation=label_interpolation, - one_hot_label_interpolation=one_hot_label_interpolation, - antialias=antialias, - **kwargs, - ) - - -class Affine(Spatial): - r"""Apply a random or fixed affine transform. - - Convenience wrapper around [`Spatial`][torchio.Spatial] exposing - only the affine parameters. The affine matrix data structure is - available as [`AffineMatrix`][torchio.AffineMatrix]. - - Args: - scales: See [`Spatial`][torchio.Spatial]. - Default: `1.0` (no scaling). - degrees: See [`Spatial`][torchio.Spatial]. - Default: `0.0` (no rotation). - translation: See [`Spatial`][torchio.Spatial]. - isotropic: See [`Spatial`][torchio.Spatial]. - center: See [`Spatial`][torchio.Spatial]. - default_pad_value: See [`Spatial`][torchio.Spatial]. - default_pad_label: See [`Spatial`][torchio.Spatial]. - image_interpolation: See [`Spatial`][torchio.Spatial]. - label_interpolation: See [`Spatial`][torchio.Spatial]. - one_hot_label_interpolation: See [`Spatial`][torchio.Spatial]. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Affine(degrees=(-15, 15)) - >>> transform = tio.Affine(scales=1.0, degrees=(0, 0, 90)) - """ - - def __init__( - self, - *, - scales: TypeParameterValue = 1.0, - degrees: TypeParameterValue = 0.0, - translation: TypeParameterValue = 0.0, - isotropic: bool = False, - center: TypeCenter = "image", - default_pad_value: TypePadValue | float = "minimum", - default_pad_label: int | float = 0, - image_interpolation: TypeImageInterpolation | int = "linear", - label_interpolation: TypeLabelInterpolation | int = "nearest", - one_hot_label_interpolation: TypeImageInterpolation | int = "linear", - **kwargs: Any, - ) -> None: - super().__init__( - scales=scales, - degrees=degrees, - translation=translation, - isotropic=isotropic, - center=center, - default_pad_value=default_pad_value, - default_pad_label=default_pad_label, - image_interpolation=image_interpolation, - label_interpolation=label_interpolation, - one_hot_label_interpolation=one_hot_label_interpolation, - **kwargs, - ) - self._warn_if_noop( - is_noop=( - self.scales.is_constant(1.0) - and self.degrees.is_constant(0.0) - and self.translation.is_constant(0.0) - ), - hint="degrees=(-15, 15)", - ) - - -class ElasticDeformation(Spatial): - r"""Apply a dense random elastic deformation. - - Convenience wrapper around [`Spatial`][torchio.Spatial] exposing - only the elastic parameters. - - A random displacement is assigned to a coarse grid of control - points and trilinearly upsampled to the image resolution. - - Args: - control_points: See [`Spatial`][torchio.Spatial]. - num_control_points: See [`Spatial`][torchio.Spatial]. - max_displacement: See [`Spatial`][torchio.Spatial]. - Default: `7.5` mm. - locked_borders: See [`Spatial`][torchio.Spatial]. - image_interpolation: See [`Spatial`][torchio.Spatial]. - label_interpolation: See [`Spatial`][torchio.Spatial]. - one_hot_label_interpolation: See [`Spatial`][torchio.Spatial]. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.ElasticDeformation() - >>> transform = tio.ElasticDeformation( - ... max_displacement=10, - ... num_control_points=5, - ... ) - """ - - def __init__( - self, - *, - control_points: TypeControlPoints | None = None, - num_control_points: int | TypeThreeInts = 7, - max_displacement: TypeParameterValue = 7.5, - locked_borders: int = 2, - image_interpolation: TypeImageInterpolation | int = "linear", - label_interpolation: TypeLabelInterpolation | int = "nearest", - one_hot_label_interpolation: TypeImageInterpolation | int = "linear", - **kwargs: Any, - ) -> None: - super().__init__( - control_points=control_points, - num_control_points=num_control_points, - max_displacement=max_displacement, - locked_borders=locked_borders, - image_interpolation=image_interpolation, - label_interpolation=label_interpolation, - one_hot_label_interpolation=one_hot_label_interpolation, - **kwargs, - ) - - -def _invert_per_sample(params: dict[str, Any]) -> _PerSampleGrids: - """Build per-element inverse grid inputs from per-instance params. - - Affine matrices are inverted exactly and elastic fields are negated, - matching the batch-shared inverse. Identity (`None`) elements stay - identity. - - Args: - params: The per-instance parameter dict from `make_params`. - - Returns: - A `_PerSampleGrids` with inverted per-element parameters. - """ - inverse_affines: list[np.ndarray | None] = [] - inverse_control_points: list[Tensor | None] = [] - inverse_displacements: list[tuple[float, float, float] | None] = [] - for matrix, field in zip( - params["affine_matrix"], - params["control_points"], - strict=True, - ): - affine = _deserialize_matrix(matrix) - inverse_affines.append(np.linalg.inv(affine) if affine is not None else None) - control_points = _deserialize_control_points(field) - if control_points is not None: - control_points = -control_points - inverse_displacements.append(_max_abs_displacement(control_points)) - else: - inverse_displacements.append(None) - inverse_control_points.append(control_points) - return _PerSampleGrids( - inverse_affines, - inverse_control_points, - inverse_displacements, - ) - - -def _resolve_spatial_params( - params: dict[str, Any], -) -> tuple[ - np.ndarray | None, - Tensor | None, - tuple[float, float, float] | None, - _PerSampleGrids | None, -]: - """Deserialize spatial params into shared or per-sample grid inputs. - - Returns: - A tuple `(affine_matrix, control_points, max_displacement, - per_sample)`. For batch-shared params the first three are - populated and `per_sample` is `None`. For per-instance params - with at least one geometric element, only `per_sample` is - populated. Per-instance params with no geometry collapse to the - shared no-op path (all `None`). - """ - batched_keys = params.get("_batched_keys") or [] - if "affine_matrix" not in batched_keys: - affine_matrix = _deserialize_matrix(params["affine_matrix"]) - control_points = _deserialize_control_points(params["control_points"]) - max_displacement = _deserialize_max_displacement(params["max_displacement"]) - return affine_matrix, control_points, max_displacement, None - - affine_matrices = [_deserialize_matrix(m) for m in params["affine_matrix"]] - control_points_list = [ - _deserialize_control_points(c) for c in params["control_points"] - ] - displacements = [ - _deserialize_max_displacement(d) for d in params["max_displacement"] - ] - has_geometry = any(m is not None for m in affine_matrices) or any( - c is not None for c in control_points_list - ) - if not has_geometry: - return None, None, None, None - return ( - None, - None, - None, - _PerSampleGrids( - affine_matrices, - control_points_list, - displacements, - ), - ) - - -def _build_grid_per_sample_checked( - *, - per_sample: _PerSampleGrids, - first_img_batch: ImagesBatch, - input_shape: TypeThreeInts, - input_affine: AffineMatrix, - output_shape: TypeThreeInts, - output_affine: AffineMatrix, - affine_first: bool, -) -> Tensor: - """Build per-element sampling grids after validating the batch size. - - Args: - per_sample: The per-element geometry recorded at sampling time. - first_img_batch: First selected image batch (provides device and - batch size). - input_shape: Spatial shape of the input grid. - input_affine: Affine of the input grid. - output_shape: Spatial shape of the output grid. - output_affine: Affine of the output grid. - affine_first: Whether the affine is applied before the elastic - displacement. - - Returns: - The stacked per-element sampling grid. - - Raises: - RuntimeError: If the number of recorded elements does not match - the batch size. - """ - if len(per_sample.affine_matrices) != first_img_batch.batch_size: - msg = ( - "Per-instance spatial parameters were recorded for" - f" {len(per_sample.affine_matrices)} elements but the batch" - f" has {first_img_batch.batch_size}" - ) - raise RuntimeError(msg) - return _build_per_sample_grids( - input_shape=input_shape, - input_affine=input_affine, - output_shape=output_shape, - output_affine=output_affine, - affine_matrices=per_sample.affine_matrices, - control_points_list=per_sample.control_points, - max_displacements=per_sample.max_displacements, - affine_first=affine_first, - device=first_img_batch.data.device, - ) - - -def _passthrough_indices(per_sample: _PerSampleGrids) -> list[int]: - """Indices of batch elements with no geometry (identity / gated out). - - Args: - per_sample: The per-element geometry recorded at sampling time. - - Returns: - Indices whose affine and control points are both ``None``, i.e. - elements that should pass through a no-target resample unchanged. - """ - return [ - index - for index, affine in enumerate(per_sample.affine_matrices) - if affine is None and per_sample.control_points[index] is None - ] - - -def _restore_spatial_output_affines( - img_batch: ImagesBatch, - *, - original_data: Tensor, - original_affines: list[AffineMatrix], - output_affine: AffineMatrix, - passthrough: list[int], -) -> None: - """Set the output affines and restore no-op (passthrough) elements. - - Every resampled element adopts *output_affine*. Passthrough elements - (no geometry, no target) instead keep their original data and affine - so that per-element gating is an exact no-op for them, undoing the - tiny perturbation introduced by the float32 identity resample. - - Args: - img_batch: The image batch whose data/affines are updated in place. - original_data: The input data, before resampling. - original_affines: The input affines, before resampling. - output_affine: The affine shared by every resampled element. - passthrough: Indices of elements to restore exactly. - """ - new_affines = [output_affine.clone() for _ in img_batch.affines] - if passthrough: - data = img_batch.data - for index in passthrough: - data[index] = original_data[index] - new_affines[index] = original_affines[index] - img_batch.data = data - img_batch.affines[:] = new_affines - - -def _apply_spatial_to_batch( - *, - batch: SubjectsBatch, - image_names: list[str], - target_space: TypeTargetSpace | None, - affine_matrix: np.ndarray | None, - control_points: Tensor | None, - max_displacement: tuple[float, float, float] | None, - affine_first: bool, - image_interpolation: TypeImageInterpolation, - label_interpolation: TypeLabelInterpolation, - one_hot_label_interpolation: TypeImageInterpolation = "linear", - antialias: bool, - default_pad_value: TypePadValue | float, - default_pad_label: float, - per_sample: _PerSampleGrids | None = None, -) -> None: - """Apply the spatial mapping to all selected images in *batch*. - - By default a single sampling grid is built from the first image's - geometry and reused for every image and every sample in the batch. - When *per_sample* is given, one grid is built per batch element so - each element is augmented independently. Only the interpolation mode - and fill value change per image type. - """ - if not image_names: - return - - first_img_batch = batch.images[image_names[0]] - input_shape = _get_spatial_shape(first_img_batch) - input_affine = first_img_batch.affines[0] - output_shape = target_space[0] if target_space is not None else input_shape - output_affine = target_space[1] if target_space is not None else input_affine - - if per_sample is None: - grid = _build_sampling_grid( - input_shape=input_shape, - input_affine=input_affine, - output_shape=output_shape, - output_affine=output_affine, - affine_matrix=affine_matrix, - control_points=control_points, - max_displacement=max_displacement, - affine_first=affine_first, - device=first_img_batch.data.device, - ) - else: - grid = _build_grid_per_sample_checked( - per_sample=per_sample, - first_img_batch=first_img_batch, - input_shape=input_shape, - input_affine=input_affine, - output_shape=output_shape, - output_affine=output_affine, - affine_first=affine_first, - ) - - # Elements with no geometry and no target resampling must be exact - # no-ops, but the batched per-sample resample computes through an - # identity grid in float32, so restore those rows afterwards. - passthrough = ( - _passthrough_indices(per_sample) - if per_sample is not None and target_space is None - else [] - ) - - for name in image_names: - _resample_image_batch( - batch.images[name], - grid=grid, - per_sample=per_sample, - input_shape=input_shape, - input_affine=input_affine, - output_affine=output_affine, - image_interpolation=image_interpolation, - label_interpolation=label_interpolation, - one_hot_label_interpolation=one_hot_label_interpolation, - antialias=antialias, - default_pad_value=default_pad_value, - default_pad_label=default_pad_label, - passthrough=passthrough, - ) - - -def _resample_image_batch( - img_batch: ImagesBatch, - *, - grid: Tensor, - per_sample: _PerSampleGrids | None, - input_shape: TypeThreeInts, - input_affine: AffineMatrix, - output_affine: AffineMatrix, - image_interpolation: TypeImageInterpolation, - label_interpolation: TypeLabelInterpolation, - one_hot_label_interpolation: TypeImageInterpolation = "linear", - antialias: bool, - default_pad_value: TypePadValue | float, - default_pad_label: float, - passthrough: list[int], -) -> None: - """Resample a single image batch and restore no-op elements. - - Args: - img_batch: The image batch to resample in place. - grid: The shared or per-sample sampling grid. - per_sample: Per-element geometry, or ``None`` for the shared grid. - input_shape: Spatial shape of the input volume. - input_affine: Affine of the input volume. - output_affine: Affine shared by every resampled element. - image_interpolation: Interpolation mode for scalar images. - label_interpolation: Interpolation mode for label maps. - one_hot_label_interpolation: Per-channel interpolation used by the - ``"label"`` partial-volume mode. - antialias: Whether to blur scalar images before downsampling. - default_pad_value: Fill value for out-of-bounds scalar samples. - default_pad_label: Fill value for out-of-bounds label samples. - passthrough: Indices of no-op elements to restore exactly. - """ - original_data = img_batch.data - original_affines = list(img_batch.affines) - is_label = issubclass(img_batch._image_class, LabelMap) - interpolation = _interpolation_for_batch( - img_batch, - image_interpolation=image_interpolation, - label_interpolation=label_interpolation, - ) - if is_label and interpolation == LABEL_INTERPOLATION: - img_batch.data = _resample_label_partial_volume( - img_batch.data, - grid, - input_shape=input_shape, - input_affine=input_affine, - output_affine=output_affine, - antialias=antialias, - one_hot_label_interpolation=one_hot_label_interpolation, - per_sample=per_sample is not None, - default_pad_label=float(default_pad_label), - ) - else: - fill_value = _batch_fill_value( - img_batch, - default_pad_value=default_pad_value, - default_pad_label=default_pad_label, - ) - data = img_batch.data - # Antialias: blur ScalarImages before downsampling. - if antialias and not is_label: - data = _antialias_batch(data, input_affine, output_affine) - sampler = _sample_batch if per_sample is None else _sample_batch_per_sample - img_batch.data = sampler( - data, - grid, - input_shape=input_shape, - interpolation=interpolation, - fill_value=fill_value, - ) - _restore_spatial_output_affines( - img_batch, - original_data=original_data, - original_affines=original_affines, - output_affine=output_affine, - passthrough=passthrough, - ) - - -def _resample_label_partial_volume( - data: Tensor, - grid: Tensor, - *, - input_shape: TypeThreeInts, - input_affine: AffineMatrix, - output_affine: AffineMatrix, - antialias: bool, - one_hot_label_interpolation: TypeImageInterpolation = "linear", - per_sample: bool = False, - default_pad_label: float, -) -> Tensor: - r"""Resample a discrete label map in a partial-volume-aware way. - - Nearest-neighbor resampling of a label map produces staircase artifacts - and can drop thin structures, especially when downsampling. This helper - instead: - - 1. one-hot encodes the label map using the unique label values present - (robust to non-contiguous labels such as $\{0, 2, 5\}$), - 2. optionally Gaussian-smooths each channel before downsampling when - *antialias* is `True` (using the same `_antialias_sigmas` as the - intensity antialiasing path, from Cardoso et al., MICCAI 2015), - 3. resamples every channel with *one_hot_label_interpolation* - (`"linear"` by default; higher-order B-spline modes give smoother - boundaries, which is useful when upsampling), and - 4. takes the per-voxel argmax to recover the discrete labels. - - This single-channel pipeline (`C == 1`) fills out-of-bounds voxels with - *default_pad_label*. - - Multi-channel inputs (`C > 1`, an already one-hot or probabilistic map) - take a different path: their channels are resampled with - *one_hot_label_interpolation* **without** re-encoding or an argmax, - out-of-bounds voxels are filled with `0` (*default_pad_label* is **not** - applied), and the partial-volume result is returned as floating point so - the interpolated fractions are not truncated. - - Args: - data: `(B, C, I, J, K)` label batch. When `C == 1` the full one-hot - pipeline is used. When `C > 1` the channels are resampled - without re-encoding (see above). - grid: `(I_out, J_out, K_out, 3)` sampling grid in input voxel - coordinates. - input_shape: Spatial shape of the input volume `(I, J, K)`. - input_affine: Affine of the input grid. - output_affine: Affine of the output grid. - antialias: Whether to Gaussian-smooth the one-hot channels before - downsampling. - one_hot_label_interpolation: Interpolation used to resample the - one-hot channels (`"nearest"`, `"linear"`, or a higher-order - B-spline mode). Defaults to `"linear"`. - per_sample: If `True`, *grid* carries a batch dimension - `(B, I_out, J_out, K_out, 3)` and each element is resampled with - its own coordinates (per-instance augmentation); otherwise a - single shared grid is used. - default_pad_label: Value assigned to out-of-bounds voxels. Only - applied for single-channel (`C == 1`) inputs; multi-channel - inputs use `0` for out-of-bounds. - - Returns: - Resampled `(B, 1, I_out, J_out, K_out)` label batch with the input - dtype when `C == 1`. When `C > 1`, a `(B, C, ...)` partial-volume - batch is returned: the input dtype is preserved for floating-point - inputs, otherwise the output is `float32` to avoid truncating the - interpolated values. - """ - sampler = _sample_batch_per_sample if per_sample else _sample_batch - if data.shape[1] > 1: - smoothed = data.float() - if antialias: - smoothed = _antialias_batch(smoothed, input_affine, output_affine) - sampled = sampler( - smoothed, - grid, - input_shape=input_shape, - interpolation=one_hot_label_interpolation, - fill_value=0.0, - ) - if data.dtype.is_floating_point: - return sampled.to(data.dtype) - return sampled - - labels = torch.unique(data) - values = rearrange(data[:, 0], "b i j k -> b 1 i j k") - targets = rearrange(labels, "n -> 1 n 1 1 1") - one_hot = (values == targets).float() - - if antialias: - one_hot = _antialias_batch(one_hot, input_affine, output_affine) - - sampled = sampler( - one_hot, - grid, - input_shape=input_shape, - interpolation=one_hot_label_interpolation, - fill_value=0.0, - ) - - winners = sampled.argmax(dim=1) - resampled = labels[winners] - # In-bounds voxels keep partition of unity (channels sum to ~1). Near the - # border, the channel sum equals the in-bounds fraction of the sampling - # neighborhood. A voxel sampled mostly (>50%) from outside the input is - # treated as out-of-bounds and set to default_pad_label, matching the - # `mask > 0.5` fill convention used for intensity images in - # `_sample_batch_grid_sample`. - in_bounds = sampled.sum(dim=1) > 0.5 - resampled = torch.where( - in_bounds, - resampled, - torch.full_like(resampled, default_pad_label), - ) - out = rearrange(resampled, "b i j k -> b 1 i j k") - return out.to(data.dtype) - - -def _resolve_target_space( - target: TypeTarget, - batch: SubjectsBatch, - first_shape: TypeThreeInts, - first_affine: AffineMatrix, -) -> TypeTargetSpace | None: - """Convert the user-facing *target* specification to `(shape, affine)`. - - Accepts scalars (isotropic spacing), 3-tuples (per-axis spacing), image - names in the subject, file paths, `Image` instances, or explicit - `(shape, affine)` pairs. Returns `None` when the output grid should - match the input grid. - """ - if target is None: - return None - if isinstance(target, Image): - return target.spatial_shape, target.affine.clone() - if isinstance(target, (str, Path)): - return _target_from_string_or_path(target, batch, first_shape, first_affine) - if _is_target_space_tuple(target): - shape, affine = target - return _parse_target_space_tuple(shape, affine) - # Remaining cases: int, float, tuple of numbers, list, ndarray, or a - # random spacing spec (range, Choice, or Distribution). - if not isinstance( - target, (int, float, tuple, list, np.ndarray, Choice, Distribution) - ): - msg = f'Target not understood: "{target}"' - raise ValueError(msg) - spacing = _resolve_target_spacing(target) - return _compute_new_shape_affine(first_shape, first_affine, spacing) - - -def _target_from_string_or_path( - target: str | Path, - batch: SubjectsBatch, - first_shape: TypeThreeInts, - first_affine: AffineMatrix, -) -> TypeTargetSpace: - """Resolve a string or Path target (file, image name, or spacing).""" - path = Path(target) - if path.is_file(): - image = ScalarImage(path) - return image.spatial_shape, image.affine.clone() - if isinstance(target, str) and target in batch.images: - reference = batch.images[target] - return _get_spatial_shape(reference), reference.affines[0].clone() - msg = ( - f'Unknown target "{target}". Pass a file path, an image name in the' - " subject, an Image, or a spacing specification" - ) - raise ValueError(msg) - - -def _resolve_target_spacing( - value: TypeSpacingSpec, -) -> TypeSpacing: - """Resolve a (possibly random) spacing spec to a positive 3-tuple. - - Scalars and 3-number sequences are deterministic. Ranges (a 2-tuple - `(lo, hi)` or a 6-tuple of per-axis ranges), a `Choice`, or a - `torch.distributions.Distribution` are sampled once, here, at apply time, - using the same convention as the other spatial parameters. NumPy arrays are - treated as deterministic 3-element spacings. - - Args: - value: A scalar, sequence, range, `Choice`, or `Distribution`. - - Returns: - A strictly-positive `(sx, sy, sz)` spacing tuple. - """ - if isinstance(value, np.ndarray): - return _parse_spacing(tuple(float(v) for v in value.flat)) - if isinstance(value, (int, float)): - return _parse_spacing(float(value)) - spec = tuple(value) if isinstance(value, list) else value - sampled = _ParameterRange(spec).sample() - return _parse_spacing(sampled) - - -def _compute_new_shape_affine( - shape: TypeThreeInts, - affine: AffineMatrix, - spacing: TypeSpacing, -) -> TypeTargetSpace: - """Compute the output shape and affine for a target voxel spacing. - - The output grid is centered on the same physical center as the input - grid. Axes with a single voxel (size 1) are left unchanged. - """ - old_spacing = np.asarray(affine.spacing, dtype=np.float64) - new_spacing = np.asarray(spacing, dtype=np.float64) - old_shape = np.asarray(shape, dtype=np.float64) - - new_shape = np.floor(old_shape * old_spacing / new_spacing) - new_shape[old_shape == 1] = 1 - - # Keep the physical center of the volume fixed when changing spacing. - rotation = affine.direction.cpu().numpy() - old_origin = np.asarray(affine.origin, dtype=np.float64) - old_center = old_origin + rotation @ (((old_shape - 1) / 2) * old_spacing) - new_origin = old_center - rotation @ (((new_shape - 1) / 2) * new_spacing) - - new_affine = np.eye(4, dtype=np.float64) - new_affine[:3, :3] = rotation * new_spacing - new_affine[:3, 3] = new_origin - return ( - (int(new_shape[0]), int(new_shape[1]), int(new_shape[2])), - AffineMatrix(new_affine), - ) - - -def _build_sampling_grid( - *, - input_shape: TypeThreeInts, - input_affine: AffineMatrix, - output_shape: TypeThreeInts, - output_affine: AffineMatrix, - affine_matrix: np.ndarray | None, - control_points: Tensor | None, - max_displacement: tuple[float, float, float] | None, - affine_first: bool, - device: torch.device, -) -> Tensor: - """Build a sampling grid in **input voxel coordinates**. - - The grid maps each output voxel to its source location in the input - volume. The mapping is: - - .. code-block:: text - - output voxel → world → (optional inverse affine) → input voxel - - When elastic control points are provided, the dense displacement - field (in mm) is converted to voxel offsets and added to the mapped - coordinates. The `affine_first` flag controls whether the affine - mapping or the elastic displacement is applied first. - - Returns: - Grid tensor with shape `(I_out, J_out, K_out, 3)` in input - voxel coordinates. - """ - mapping = _output_to_input_voxel_matrix( - input_affine=input_affine, - output_affine=output_affine, - affine_matrix=affine_matrix, - device=device, - ) - output_coords = _output_voxel_coordinates(output_shape, device) - - if control_points is None: - return _apply_voxel_mapping(output_coords, mapping) - - output_spacing = np.asarray(output_affine.spacing, dtype=np.float64) - if max_displacement is None: - max_displacement = _max_abs_displacement(control_points) - _check_folding( - control_points.cpu().numpy(), - max_displacement, - output_shape, - output_spacing, - ) - displacement = _upsample_displacement_field( - control_points.to(device=device, dtype=torch.float32), - output_shape, - ) - # Convert mm displacements to voxel offsets using the spacing. - output_spacing_t = torch.as_tensor( - output_affine.spacing, - dtype=torch.float32, - device=device, - ) - input_spacing_t = torch.as_tensor( - input_affine.spacing, - dtype=torch.float32, - device=device, - ) - - if affine_first: - # Affine first: map to input space, then add elastic offset. - input_voxels = _apply_voxel_mapping(output_coords, mapping) - input_voxels = input_voxels + displacement / input_spacing_t - else: - # Elastic first: deform in output space, then map to input. - deformed_output = output_coords + displacement / output_spacing_t - input_voxels = _apply_voxel_mapping(deformed_output, mapping) - - return input_voxels - - -def _output_to_input_voxel_matrix( - *, - input_affine: AffineMatrix, - output_affine: AffineMatrix, - affine_matrix: np.ndarray | None, - device: torch.device, -) -> Tensor: - """Compute the 4x4 matrix mapping output voxels to input voxels. - - The composition is `A_in^{-1} @ T^{-1} @ A_out` where *T* is the - optional world-space affine transform. - """ - input_affine_inv = np.linalg.inv(input_affine.numpy()) - transform_inv = ( - np.eye(4, dtype=np.float64) - if affine_matrix is None - else np.linalg.inv(np.asarray(affine_matrix, dtype=np.float64)) - ) - matrix = input_affine_inv @ transform_inv @ output_affine.numpy() - return torch.as_tensor(matrix, dtype=torch.float32, device=device) - - -def _output_voxel_coordinates( - shape: TypeThreeInts, - device: torch.device, -) -> Tensor: - """Create an `(I, J, K, 3)` meshgrid of output voxel indices.""" - i = torch.arange(shape[0], dtype=torch.float32, device=device) - j = torch.arange(shape[1], dtype=torch.float32, device=device) - k = torch.arange(shape[2], dtype=torch.float32, device=device) - grid_i, grid_j, grid_k = torch.meshgrid(i, j, k, indexing="ij") - return torch.stack([grid_i, grid_j, grid_k], dim=-1) - - -def _apply_voxel_mapping( - coords: Tensor, - matrix: Tensor, -) -> Tensor: - """Apply a 4x4 homogeneous matrix to an `(..., 3)` coordinate tensor.""" - ones = torch.ones(*coords.shape[:-1], 1, dtype=coords.dtype, device=coords.device) - homogeneous = torch.cat([coords, ones], dim=-1) - mapped = homogeneous @ matrix.T - return mapped[..., :3] - - -def _voxel_coordinates_to_grid( - coords: Tensor, - input_shape: TypeThreeInts, -) -> Tensor: - """Normalize `(I, J, K, 3)` voxel coords to the `[-1, 1]` grid. - - `F.grid_sample` expects coordinates in `[-1, 1]` where `-1` - maps to the first voxel and `+1` to the last. The output is - rearranged to `(1, K, J, I, 3)` to match PyTorch's `(D, H, W)` - convention. - """ - size_i = max(input_shape[0] - 1, 1) - size_j = max(input_shape[1] - 1, 1) - size_k = max(input_shape[2] - 1, 1) - sizes = torch.tensor( - [size_i, size_j, size_k], - dtype=torch.float32, - device=coords.device, - ) - grid = 2.0 * coords / sizes - 1.0 - # (I, J, K, 3) -> (1, K, J, I, 3) for grid_sample - return rearrange(grid, "i j k d -> 1 k j i d") - - -def _sample_batch( - data: Tensor, - voxel_grid: Tensor, - *, - input_shape: TypeThreeInts, - interpolation: str, - fill_value: float | Tensor, -) -> Tensor: - """Resample a 5D batch using a shared sampling grid. - - For interpolation orders 0-1 (nearest, linear), the fast - `F.grid_sample` path is used. For orders 2+ (quadratic, - cubic, ...), `interpol.grid_pull` provides high-order B-spline - interpolation. - - Args: - data: `(B, C, I, J, K)` image batch. - voxel_grid: `(I_out, J_out, K_out, 3)` sampling grid in - input voxel coordinates. - input_shape: Spatial shape of the input volume `(I, J, K)`. - interpolation: Interpolation mode name. - fill_value: Scalar or per-channel fill for out-of-bounds - samples. - - Returns: - Resampled `(B, C, I_out, J_out, K_out)` tensor. - """ - order = _INTERPOLATION_TO_ORDER[interpolation] - if order <= 1: - return _sample_batch_grid_sample( - data, - voxel_grid, - input_shape=input_shape, - mode=_TORCH_INTERPOLATION_MODE[interpolation], - fill_value=fill_value, - ) - return _sample_batch_interpol( - data, - voxel_grid, - order=order, - fill_value=fill_value, - ) - - -def _sample_batch_grid_sample( - data: Tensor, - voxel_grid: Tensor, - *, - input_shape: TypeThreeInts, - mode: str, - fill_value: float | Tensor, -) -> Tensor: - """Fast path: resample with F.grid_sample (orders 0-1).""" - batch_size = data.shape[0] - # Normalize voxel coords to [-1, 1] for grid_sample. - grid = _voxel_coordinates_to_grid(voxel_grid, input_shape) - # (B, C, I, J, K) -> (B, C, K, J, I) for grid_sample - input_5d = rearrange(data, "b c i j k -> b c k j i").float() - # Expand grid from (1, ...) to (B, ...) - grid_b = grid.expand(batch_size, -1, -1, -1, -1) - sampled = functional.grid_sample( - input_5d, - grid_b, - mode=mode, - padding_mode="zeros", - align_corners=True, - ) - - fill_tensor = _prepare_fill_value(fill_value, input_5d) - if fill_tensor is not None: - ones = torch.ones_like(input_5d) - mask = functional.grid_sample( - ones, - grid_b, - padding_mode="zeros", - align_corners=True, - ) - sampled = torch.where(mask > 0.5, sampled, fill_tensor) - - # (B, C, K, J, I) -> (B, C, I, J, K) - return rearrange(sampled, "b c k j i -> b c i j k").to(data.dtype) - - -def _sample_batch_interpol( - data: Tensor, - voxel_grid: Tensor, - *, - order: int, - fill_value: float | Tensor, -) -> Tensor: - """High-quality path: resample with interpol.grid_pull (orders 2+). - - `interpol.grid_pull` works in voxel coordinates natively and - supports B-spline orders up to 7. - """ - import interpol - - batch_size = data.shape[0] - # interpol expects: input (B, C, *spatial), grid (B, *spatial, D) - # Our grid is (I_out, J_out, K_out, 3). Add batch dim. - grid_b = voxel_grid.unsqueeze(0).expand(batch_size, -1, -1, -1, -1) - - sampled = interpol.grid_pull( - data.float(), - grid_b, - interpolation=order, - bound="dct2", - extrapolate=False, - prefilter=True, - ) - return sampled.to(data.dtype) - - -def _sample_batch_per_sample( - data: Tensor, - voxel_grid: Tensor, - *, - input_shape: TypeThreeInts, - interpolation: str, - fill_value: float | Tensor, -) -> Tensor: - """Resample a 5D batch using a per-sample sampling grid. - - Unlike [`_sample_batch`][], the grid carries its own batch dimension - so each element of the batch is resampled with independent - coordinates (per-instance augmentation). - - Args: - data: `(B, C, I, J, K)` image batch. - voxel_grid: `(B, I_out, J_out, K_out, 3)` per-sample sampling - grid in input voxel coordinates. - input_shape: Spatial shape of the input volume `(I, J, K)`. - interpolation: Interpolation mode name. - fill_value: Scalar or per-channel fill for out-of-bounds samples. - - Returns: - Resampled `(B, C, I_out, J_out, K_out)` tensor. - """ - order = _INTERPOLATION_TO_ORDER[interpolation] - if order <= 1: - return _sample_batch_grid_sample_per_sample( - data, - voxel_grid, - input_shape=input_shape, - mode=_TORCH_INTERPOLATION_MODE[interpolation], - fill_value=fill_value, - ) - return _sample_batch_interpol_per_sample( - data, - voxel_grid, - order=order, - fill_value=fill_value, - ) - - -def _voxel_coordinates_to_grid_batched( - coords: Tensor, - input_shape: TypeThreeInts, -) -> Tensor: - """Normalize `(B, I, J, K, 3)` voxel coords to the `[-1, 1]` grid. - - Like [`_voxel_coordinates_to_grid`][] but keeps a leading batch - dimension, producing `(B, K, J, I, 3)` for `F.grid_sample`. - """ - size_i = max(input_shape[0] - 1, 1) - size_j = max(input_shape[1] - 1, 1) - size_k = max(input_shape[2] - 1, 1) - sizes = torch.tensor( - [size_i, size_j, size_k], - dtype=torch.float32, - device=coords.device, - ) - grid = 2.0 * coords / sizes - 1.0 - return rearrange(grid, "b i j k d -> b k j i d") - - -def _sample_batch_grid_sample_per_sample( - data: Tensor, - voxel_grid: Tensor, - *, - input_shape: TypeThreeInts, - mode: str, - fill_value: float | Tensor, -) -> Tensor: - """Per-sample fast path: resample with F.grid_sample (orders 0-1).""" - grid = _voxel_coordinates_to_grid_batched(voxel_grid, input_shape) - input_5d = rearrange(data, "b c i j k -> b c k j i").float() - sampled = functional.grid_sample( - input_5d, - grid, - mode=mode, - padding_mode="zeros", - align_corners=True, - ) - - fill_tensor = _prepare_fill_value(fill_value, input_5d) - if fill_tensor is not None: - ones = torch.ones_like(input_5d) - mask = functional.grid_sample( - ones, - grid, - padding_mode="zeros", - align_corners=True, - ) - sampled = torch.where(mask > 0.5, sampled, fill_tensor) - - return rearrange(sampled, "b c k j i -> b c i j k").to(data.dtype) - - -def _sample_batch_interpol_per_sample( - data: Tensor, - voxel_grid: Tensor, - *, - order: int, - fill_value: float | Tensor, -) -> Tensor: - """Per-sample high-quality path: resample with interpol.grid_pull (orders 2+).""" - import interpol - - sampled = interpol.grid_pull( - data.float(), - voxel_grid, - interpolation=order, - bound="dct2", - extrapolate=False, - prefilter=True, - ) - return sampled.to(data.dtype) - - -def _build_per_sample_grids( - *, - input_shape: TypeThreeInts, - input_affine: AffineMatrix, - output_shape: TypeThreeInts, - output_affine: AffineMatrix, - affine_matrices: list[np.ndarray | None], - control_points_list: list[Tensor | None], - max_displacements: list[tuple[float, float, float] | None], - affine_first: bool, - device: torch.device, -) -> Tensor: - """Build one sampling grid per batch element and stack them. - - Each element reuses the shared single-grid builder with its own - affine matrix and elastic field, so the per-instance path inherits - the exact geometry of the batch-shared path. A `None` affine and - `None` field yield an identity grid (used for gated-out elements). - - Returns: - A `(B, I_out, J_out, K_out, 3)` tensor of input voxel - coordinates. - """ - grids = [ - _build_sampling_grid( - input_shape=input_shape, - input_affine=input_affine, - output_shape=output_shape, - output_affine=output_affine, - affine_matrix=affine_matrices[index], - control_points=control_points_list[index], - max_displacement=max_displacements[index], - affine_first=affine_first, - device=device, - ) - for index in range(len(affine_matrices)) - ] - return torch.stack(grids, dim=0) - - -def _antialias_batch( - data: Tensor, - input_affine: AffineMatrix, - output_affine: AffineMatrix, -) -> Tensor: - """Apply Gaussian smoothing before downsampling. - - Only axes whose output spacing is larger than the input spacing - (i.e., axes being downsampled) are smoothed. The standard deviations - follow Cardoso et al., `Scale factor point spread function matching - `_, - MICCAI 2015. - - Args: - data: `(B, C, I, J, K)` image batch. - input_affine: Affine of the input grid. - output_affine: Affine of the output grid. - - Returns: - Smoothed `(B, C, I, J, K)` tensor. - """ - input_spacing = np.asarray(input_affine.spacing, dtype=np.float64) - output_spacing = np.asarray(output_affine.spacing, dtype=np.float64) - factors = output_spacing / input_spacing - sigmas = _antialias_sigmas(factors, input_spacing) - if np.all(sigmas == 0): - return data - return _gaussian_smooth_batch(data, sigmas) - - -def _antialias_sigmas( - factors: np.ndarray, - spacing: np.ndarray, -) -> np.ndarray: - """Compute per-axis Gaussian sigma for antialiasing. - - From Cardoso et al., MICCAI 2015, Eq. at top of p. 678: - `variance = (k^2 - 1) / (2 * sqrt(2 * ln(2)))^2` - - Args: - factors: Per-axis downsampling factor (output / input spacing). - spacing: Input voxel spacing in mm. - - Returns: - Per-axis sigma in voxels. Zero for axes not being downsampled. - """ - sigmas = np.zeros(3, dtype=np.float64) - for axis in range(3): - k = factors[axis] - if k <= 1.0: - continue - # Cardoso et al. formula: sigma in mm - variance = (k**2 - 1) * (2 * np.sqrt(2 * np.log(2))) ** (-2) - sigma_mm = spacing[axis] * np.sqrt(variance) - # Convert to voxels for the separable convolution - sigmas[axis] = sigma_mm / spacing[axis] - return sigmas - - -def _gaussian_smooth_batch(data: Tensor, sigmas: np.ndarray) -> Tensor: - """Apply separable Gaussian smoothing along spatial axes of a 5D tensor. - - Args: - data: `(B, C, I, J, K)` tensor. - sigmas: Per-axis sigma in voxels. Zero means skip that axis. - - Returns: - Smoothed `(B, C, I, J, K)` tensor. - """ - result = data.float() - b, c = result.shape[:2] - for axis_idx in range(3): - sigma = float(sigmas[axis_idx]) - if sigma <= 0: - continue - # Kernel radius: 3 sigma, at least 1 - radius = max(int(np.ceil(3 * sigma)), 1) - kernel_size = 2 * radius + 1 - x = ( - torch.arange( - kernel_size, - dtype=torch.float32, - device=data.device, - ) - - radius - ) - kernel_1d = torch.exp(-0.5 * (x / sigma) ** 2) - kernel_1d = kernel_1d / kernel_1d.sum() - - # Build a 5D depthwise kernel for F.conv3d with groups=C. - # Shape: (C, 1, kI, kJ, kK) with the kernel along the target axis. - k_shape = [1, 1, 1] - k_shape[axis_idx] = kernel_size - kernel_3d = kernel_1d.reshape(1, 1, *k_shape).expand(c, 1, -1, -1, -1) - - # Replicate-pad along the target spatial axis. - # F.pad order: (K_before, K_after, J_before, J_after, I_before, I_after) - pad = [0] * 6 - pad_idx = 2 * (2 - axis_idx) - pad[pad_idx] = radius - pad[pad_idx + 1] = radius - - padded = functional.pad(result, pad, mode="replicate") - # Convolve each (B, C, ...) slice with groups=C so channels are independent. - result = functional.conv3d( - rearrange(padded, "b c i j k -> (b c) 1 i j k"), - kernel_3d[:1], - padding=0, - ) - result = rearrange(result, "(b c) 1 i j k -> b c i j k", b=b, c=c) - return result.to(data.dtype) - - -def _batch_fill_value( - img_batch: ImagesBatch, - *, - default_pad_value: TypePadValue | float, - default_pad_label: float, -) -> float | Tensor: - """Compute a single fill value for the whole image batch.""" - if issubclass(img_batch._image_class, LabelMap): - return float(default_pad_label) - - if isinstance(default_pad_value, Number): - return float(default_pad_value) - - if not isinstance(default_pad_value, str): - msg = ( - "default_pad_value must be a string or number, got" - f" {type(default_pad_value)}" - ) - raise TypeError(msg) - - # Use the first sample to compute the channel-wise fill value - first_tensor = img_batch.data[0] - values = [ - _compute_channel_pad_value(channel, default_pad_value) - for channel in first_tensor - ] - return torch.as_tensor(values, dtype=torch.float32) - - -def _prepare_fill_value( - fill_value: float | Tensor, - reference: Tensor, -) -> Tensor | None: - """Convert a fill value to a broadcast-compatible tensor. - - Returns `None` when the fill is zero (so the caller can skip the - masking step, since `grid_sample` already pads with zeros). - """ - if isinstance(fill_value, Tensor): - fill_tensor = fill_value.to(device=reference.device, dtype=reference.dtype) - else: - if float(fill_value) == 0.0: - return None - fill_tensor = torch.as_tensor( - fill_value, - device=reference.device, - dtype=reference.dtype, - ) - if fill_tensor.ndim == 0: - return fill_tensor - if fill_tensor.ndim == 1: - return rearrange(fill_tensor, "c -> 1 c 1 1 1") - return fill_tensor - - -def _compute_channel_pad_value( - tensor: Tensor, - default_pad_value: TypePadValue, -) -> float: - """Compute a scalar fill value for a single 3D channel tensor.""" - if default_pad_value == "minimum": - return float(tensor.min().item()) - if default_pad_value == "mean": - return _border_mean(tensor, filter_otsu=False) - if default_pad_value == "otsu": - return _border_mean(tensor, filter_otsu=True) - msg = f'Unknown default_pad_value "{default_pad_value}"' - raise ValueError(msg) - - -def _border_mean( - tensor: Tensor, - *, - filter_otsu: bool, -) -> float: - """Mean intensity of the six boundary faces of a 3D tensor. - - When *filter_otsu* is `True`, only voxels below the Otsu threshold - are averaged, giving a background-aware fill value. - """ - borders = torch.cat( - [ - tensor[0, :, :].ravel(), - tensor[-1, :, :].ravel(), - tensor[:, 0, :].ravel(), - tensor[:, -1, :].ravel(), - tensor[:, :, 0].ravel(), - tensor[:, :, -1].ravel(), - ] - ).float() - if not filter_otsu: - return float(borders.mean().item()) - threshold = _otsu_threshold(borders) - values = borders[borders < threshold] - if values.numel() > 0: - return float(values.mean().item()) - return float(borders.mean().item()) - - -def _otsu_threshold(values: Tensor) -> float: - """Compute the Otsu threshold for a 1D tensor of values. - - Sweeps over sorted values, maximizing the between-class variance - to find the threshold that best separates foreground from background. - """ - sorted_values, _ = values.sort() - num_values = sorted_values.numel() - if num_values == 0: - return 0.0 - - total_sum = float(sorted_values.sum().item()) - best_threshold = float(sorted_values[0].item()) - best_variance = 0.0 - background_sum = 0.0 - background_count = 0 - - for background_count, item in enumerate(sorted_values[:-1], start=1): - value = float(item.item()) - foreground_count = num_values - background_count - background_sum += value - - mean_background = background_sum / background_count - mean_foreground = (total_sum - background_sum) / foreground_count - weight_background = background_count / num_values - weight_foreground = foreground_count / num_values - # Between-class variance: maximize this to find the best split. - between_variance = ( - weight_background - * weight_foreground - * (mean_background - mean_foreground) ** 2 - ) - if between_variance > best_variance: - best_variance = between_variance - best_threshold = value - return best_threshold - - -def _upsample_displacement_field( - control_points: Tensor, - output_shape: TypeThreeInts, -) -> Tensor: - """Trilinearly upsample a coarse `(n_i, n_j, n_k, 3)` field. - - The result has shape `(*output_shape, 3)` and approximates - cubic B-spline interpolation for smooth deformations. - """ - # (I, J, K, 3) -> (1, 3, I, J, K) for interpolate - field = rearrange(control_points, "i j k d -> 1 d i j k").float() - dense = functional.interpolate( - field, - size=list(output_shape), - mode="trilinear", - align_corners=True, - ) - # (1, 3, I, J, K) -> (I, J, K, 3) - return rearrange(dense, "1 d i j k -> i j k d") - - -def _check_folding( - control_points: np.ndarray, - max_displacement: tuple[float, float, float], - shape: TypeThreeInts, - spacing: np.ndarray, -) -> None: - """Warn if the displacement magnitude may cause grid folding. - - Folding occurs when a control point moves past its neighbor, - inverting the local Jacobian. The heuristic checks whether the - maximum displacement exceeds half the coarse-grid spacing. - """ - num_control_points = np.array(control_points.shape[:-1], dtype=np.float64) - image_bounds = np.array(shape, dtype=np.float64) * spacing - mesh_shape = num_control_points - _SPLINE_ORDER - grid_spacing = image_bounds / mesh_shape - conflicts = np.array(max_displacement, dtype=np.float64) > grid_spacing / 2 - if np.any(conflicts): - (where,) = np.where(conflicts) - warnings.warn( - "The maximum displacement is larger than half the coarse-grid" - f" spacing for dimensions {where.tolist()}, so folding may occur", - RuntimeWarning, - stacklevel=3, - ) - - -def _resolve_control_points( - control_points: Tensor | None, - num_control_points: TypeThreeInts, - max_displacement: _ParameterRange, - locked_borders: int, -) -> tuple[Tensor | None, tuple[float, float, float] | None]: - """Return a concrete control-point field and its max displacement. - - If *control_points* is already provided, clone it. Otherwise sample - a random field from the displacement range. Returns `(None, None)` - when the sampled displacement is zero everywhere. - """ - if control_points is not None: - return control_points.clone(), _max_abs_displacement(control_points) - - sampled = max_displacement.sample() - if all(value == 0.0 for value in sampled): - return None, None - field = _sample_control_points(num_control_points, sampled, locked_borders) - return field, sampled - - -def _sample_control_points( - grid_shape: TypeThreeInts, - max_displacement: tuple[float, float, float], - locked_borders: int, -) -> Tensor: - """Sample a random control-point displacement field. - - Each component is drawn uniformly from `[-max, +max]` along each - axis, then the outermost *locked_borders* layers are zeroed to - prevent boundary artifacts. - """ - field = torch.rand(*grid_shape, 3, dtype=torch.float32) - field -= 0.5 - field *= 2 - for axis in range(3): - field[..., axis] *= max_displacement[axis] - - # Zero out outermost control-point layers to avoid boundary artifacts. - for border in range(locked_borders): - field[border, :] = 0 - field[-1 - border, :] = 0 - field[:, border] = 0 - field[:, -1 - border] = 0 - field[:, :, border] = 0 - field[:, :, -1 - border] = 0 - return field - - -def _build_forward_affine( - *, - scales: tuple[float, float, float], - degrees: tuple[float, float, float], - translation: tuple[float, float, float], - center: TypeCenter, - shape: TypeThreeInts, - affine: AffineMatrix, -) -> np.ndarray: - """Build a 4x4 world-space affine from scale, rotation, translation. - - When *center* is `"image"`, the rotation and scaling pivot around - the image center. For 2D slices (last axis size 1), out-of-plane - components are suppressed. - """ - scaling = np.asarray(scales, dtype=np.float64) - rotation = np.asarray(degrees, dtype=np.float64) - shift = np.asarray(translation, dtype=np.float64) - - # Suppress out-of-plane components for 2D (single-slice) images. - if shape[-1] == 1: - scaling[2] = 1.0 - rotation[0] = 0.0 - rotation[1] = 0.0 - shift[2] = 0.0 - - center_world = _image_center_world(shape, affine) if center == "image" else None - return _physical_affine_matrix( - scales=scaling, - degrees=rotation, - translation=shift, - center_world=center_world, - ) - - -def _physical_affine_matrix( - *, - scales: np.ndarray, - degrees: np.ndarray, - translation: np.ndarray, - center_world: np.ndarray | None, -) -> np.ndarray: - """Compose rotation, scaling, and translation into a 4x4 matrix. - - If *center_world* is given the transform pivots around that point: - `T = R @ S` with `t = center - R @ S @ center + translation`. - """ - rotation = _euler_to_rotation_matrix(degrees) - scale = np.diag(scales) - transform = np.eye(4, dtype=np.float64) - rotation_scale = rotation @ scale - transform[:3, :3] = rotation_scale - # Pivot: translate so center is at origin, apply R@S, translate back. - if center_world is not None: - transform[:3, 3] = center_world - rotation_scale @ center_world - transform[:3, 3] += translation - return transform - - -def _euler_to_rotation_matrix(degrees: np.ndarray) -> np.ndarray: - """Convert XYZ Euler angles in degrees to a 3x3 rotation matrix. - - Uses the ZYX extrinsic (= XYZ intrinsic) convention: - `R = Rz @ Ry @ Rx`. - """ - radians = np.radians(degrees) - rx, ry, rz = radians - - cos_x, sin_x = np.cos(rx), np.sin(rx) - cos_y, sin_y = np.cos(ry), np.sin(ry) - cos_z, sin_z = np.cos(rz), np.sin(rz) - - rotation_x = np.array( - [ - [1, 0, 0], - [0, cos_x, -sin_x], - [0, sin_x, cos_x], - ], - dtype=np.float64, - ) - rotation_y = np.array( - [ - [cos_y, 0, sin_y], - [0, 1, 0], - [-sin_y, 0, cos_y], - ], - dtype=np.float64, - ) - rotation_z = np.array( - [ - [cos_z, -sin_z, 0], - [sin_z, cos_z, 0], - [0, 0, 1], - ], - dtype=np.float64, - ) - return rotation_z @ rotation_y @ rotation_x - - -def _image_center_world( - shape: TypeThreeInts, - affine: AffineMatrix, -) -> np.ndarray: - """Return the world-space coordinates of the image center.""" - center_index = (np.asarray(shape, dtype=np.float64) - 1) / 2 - matrix = affine.numpy() - return matrix[:3, 3] + matrix[:3, :3] @ center_index - - -def _check_shared_space( - images: dict[str, ImagesBatch], - reference_shape: TypeThreeInts, - reference_affine: AffineMatrix, -) -> None: - """Raise if any image has a different shape or affine than the first.""" - reference_matrix = reference_affine.data - for name, img_batch in images.items(): - current_shape = _get_spatial_shape(img_batch) - if current_shape != reference_shape: - msg = ( - f'Image "{name}" has shape {current_shape}, expected {reference_shape}' - ) - raise RuntimeError(msg) - for affine in img_batch.affines: - if not torch.allclose( - affine.data, - reference_matrix, - rtol=1e-6, - atol=1e-6, - ): - msg = ( - "Spatial transforms with affine or elastic components require" - " selected images to share the same affine" - ) - raise RuntimeError(msg) - - -def _get_spatial_shape(img_batch: ImagesBatch) -> TypeThreeInts: - """Extract the `(I, J, K)` spatial shape from a batched image.""" - return ( - int(img_batch.data.shape[-3]), - int(img_batch.data.shape[-2]), - int(img_batch.data.shape[-1]), - ) - - -def _interpolation_for_batch( - img_batch: ImagesBatch, - *, - image_interpolation: TypeImageInterpolation, - label_interpolation: TypeLabelInterpolation, -) -> str: - """Choose the interpolation mode based on the image class.""" - if issubclass(img_batch._image_class, LabelMap): - return label_interpolation - return image_interpolation - - -def _serialize_space(space: TypeTargetSpace | None) -> dict[str, Any] | None: - """Convert a `(shape, AffineMatrix)` pair to a JSON-safe dict.""" - if space is None: - return None - shape, affine = space - return { - "shape": list(shape), - "affine": affine.numpy().tolist(), - } - - -def _deserialize_space(data: dict[str, Any] | None) -> TypeTargetSpace | None: - """Reconstruct a `(shape, AffineMatrix)` pair from a dict.""" - if data is None: - return None - shape = ( - int(data["shape"][0]), - int(data["shape"][1]), - int(data["shape"][2]), - ) - return shape, AffineMatrix(np.asarray(data["affine"], dtype=np.float64)) - - -def _serialize_matrix(matrix: np.ndarray | None) -> list[list[float]] | None: - """Convert a numpy matrix to a nested list for JSON serialization.""" - if matrix is None: - return None - return matrix.tolist() - - -def _deserialize_matrix(data: list[list[float]] | None) -> np.ndarray | None: - """Reconstruct a numpy matrix from a nested list.""" - if data is None: - return None - return np.asarray(data, dtype=np.float64) - - -def _serialize_control_points(control_points: Tensor | None) -> list | None: - """Convert a control-point tensor to a nested list.""" - if control_points is None: - return None - return control_points.cpu().tolist() - - -def _deserialize_control_points(data: list | None) -> Tensor | None: - """Reconstruct a control-point tensor from a nested list.""" - if data is None: - return None - return torch.as_tensor(data, dtype=torch.float32) - - -def _deserialize_max_displacement( - values: list[float] | None, -) -> tuple[float, float, float] | None: - """Reconstruct a max-displacement 3-tuple from a list.""" - if values is None: - return None - return (float(values[0]), float(values[1]), float(values[2])) - - -def _max_abs_displacement(control_points: Tensor) -> tuple[float, float, float]: - """Return the per-axis maximum absolute displacement.""" - absolute = control_points.abs() - return ( - float(absolute[..., 0].max().item()), - float(absolute[..., 1].max().item()), - float(absolute[..., 2].max().item()), - ) - - -def _sample_scales( - scales: _ParameterRange, - isotropic: bool, -) -> tuple[float, float, float]: - """Sample a 3-tuple of scale factors, optionally isotropic.""" - if isotropic: - value = scales.sample_1d() - return (value, value, value) - return scales.sample() - - -def _has_affine_component( - scales: tuple[float, float, float], - degrees: tuple[float, float, float], - translation: tuple[float, float, float], -) -> bool: - """Return `True` if any sampled affine parameter is non-identity.""" - return not ( - np.allclose(scales, (1.0, 1.0, 1.0)) - and np.allclose(degrees, (0.0, 0.0, 0.0)) - and np.allclose(translation, (0.0, 0.0, 0.0)) - ) - - -def _parse_target_space_tuple( - shape: Sequence[int], - affine: AffineMatrix | Tensor | npt.ArrayLike, -) -> TypeTargetSpace: - """Validate and convert a `(shape, affine)` target pair.""" - if len(shape) != 3: - msg = f"Target shape must have length 3, got {len(shape)}" - raise ValueError(msg) - target_shape = (int(shape[0]), int(shape[1]), int(shape[2])) - return target_shape, AffineMatrix(affine) - - -def _is_spacing_sequence(target: Sequence[Any]) -> bool: - """Return `True` if *target* looks like a 3-element numeric spacing.""" - return len(target) == 3 and all(isinstance(value, Number) for value in target) - - -def _is_spacing_tuple( - target: object, -) -> TypeGuard[tuple[int | float, int | float, int | float]]: - """Type-guard: *target* is a 3-number tuple.""" - return isinstance(target, tuple) and _is_spacing_sequence(target) - - -def _is_spacing_list(target: object) -> TypeGuard[list[int | float]]: - """Type-guard: *target* is a 3-number list.""" - return isinstance(target, list) and _is_spacing_sequence(target) - - -def _is_target_space_tuple( - target: object, -) -> TypeGuard[tuple[Sequence[int], AffineMatrix | Tensor | npt.ArrayLike]]: - """Type-guard: *target* is a 2-element `(shape, affine)` tuple. - - The first element of a target-space tuple is a shape sequence, never a - plain number, so a 2-tuple of numbers like `(2, 4)` is treated as a - spacing range rather than a `(shape, affine)` pair. - """ - return ( - isinstance(target, tuple) - and len(target) == 2 - and not isinstance(target[0], Number) - ) - - -def _parse_spacing( - value: TypeSpacing | Sequence[float] | np.ndarray | float | int, -) -> TypeSpacing: - """Normalize a spacing specification to a strictly-positive 3-tuple.""" - if isinstance(value, (int, float)): - spacing = (float(value), float(value), float(value)) - elif isinstance(value, np.ndarray): - if value.size != 3: - msg = f"Spacing array must have 3 values, got {value.size}" - raise ValueError(msg) - flat = [float(v) for v in value.flat] - spacing = (flat[0], flat[1], flat[2]) - else: - if len(value) != 3: - msg = f"Spacing must have 3 values, got {len(value)}" - raise ValueError(msg) - spacing = (float(value[0]), float(value[1]), float(value[2])) - if any(v <= 0 for v in spacing): - msg = f"Spacing must be strictly positive, got {spacing}" - raise ValueError(msg) - return spacing - - -def _parse_interpolation( - interpolation: TypeInterpolation | int, -) -> TypeInterpolation: - """Validate an interpolation mode (string or integer order). - - Integer orders (0-7) are converted to their string equivalents. - """ - if isinstance(interpolation, int): - order_to_name = {v: k for k, v in _INTERPOLATION_TO_ORDER.items()} - if interpolation not in order_to_name: - msg = f"Interpolation order {interpolation} is not supported. Must be 0-7." - raise ValueError(msg) - return cast(TypeInterpolation, order_to_name[interpolation]) - if not isinstance(interpolation, str): - msg = f"Interpolation must be a string or int, got {type(interpolation)}" - raise TypeError(msg) - lowered = interpolation.lower() - if lowered not in _SUPPORTED_INTERPOLATIONS: - msg = ( - f'Interpolation "{lowered}" is not supported. Supported values are' - f" {_SUPPORTED_INTERPOLATIONS}" - ) - raise ValueError(msg) - return lowered - - -def _parse_one_hot_label_interpolation( - interpolation: TypeImageInterpolation | int, -) -> TypeImageInterpolation: - """Validate the per-channel interpolation used by the `"label"` mode. - - Accepts the same modes as image interpolation (`"nearest"`, `"linear"`, - or B-spline orders `"quadratic"`-`"seventh"` / integer orders 0-7) but - rejects `"label"`, which would recurse. - - Args: - interpolation: Interpolation mode (string or integer order). - - Returns: - The validated interpolation mode as a lowercase string. - - Raises: - ValueError: If `"label"` is passed. - """ - parsed = _parse_interpolation(interpolation) - if parsed == LABEL_INTERPOLATION: - msg = ( - f'one_hot_label_interpolation cannot be "{LABEL_INTERPOLATION}"; choose' - ' an interpolation for the one-hot channels (e.g. "linear")' - ) - raise ValueError(msg) - return parsed - - -def _parse_default_pad_value(value: TypePadValue | float) -> TypePadValue | float: - """Validate a pad-value specification (string keyword or number).""" - if isinstance(value, Number): - return float(value) - if value in _SUPPORTED_PAD_VALUES: - return value - msg = 'default_pad_value must be "minimum", "mean", "otsu", or a numeric value' - raise ValueError(msg) - - -def _parse_center(center: TypeCenter) -> TypeCenter: - """Validate the *center* argument.""" - if center not in ("image", "origin"): - msg = f'center must be "image" or "origin", got "{center}"' - raise ValueError(msg) - return center - - -def _to_positive_range( - value: TypeParameterValue, -) -> _ParameterRange: - """Convert to a `_ParameterRange`, rejecting non-positive scales.""" - result = _to_parameter_range(value) - if result._distribution is None: - for low, high in result._ranges: - if low <= 0 or high <= 0: - msg = f"Scale factors must be strictly positive, got {value}" - raise ValueError(msg) - return result - - -def _validate_isotropic( - value: TypeParameterValue, - isotropic: bool, -) -> None: - """Raise if *isotropic* is `True` but per-axis values were given.""" - if not isotropic or isinstance(value, Distribution): - return - if isinstance(value, tuple) and len(value) in (3, 6): - msg = "If isotropic=True, scales must be a single value or a 2-value range" - raise ValueError(msg) - - -def _parse_num_control_points( - value: int | TypeThreeInts, -) -> TypeThreeInts: - """Normalize to a 3-tuple and validate each axis has >= 4 points.""" - parsed = (value, value, value) if isinstance(value, int) else value - for axis, number in enumerate(parsed): - if not isinstance(number, int) or number < 4: - msg = ( - "Each num_control_points value must be an integer greater than 3;" - f" axis {axis} got {number}" - ) - raise ValueError(msg) - return parsed - - -def _parse_locked_borders(value: int) -> int: - """Validate that *value* is 0, 1, or 2.""" - if value not in (0, 1, 2): - msg = f"locked_borders must be 0, 1, or 2, got {value}" - raise ValueError(msg) - return value - - -def _parse_control_points(control_points: TypeControlPoints) -> Tensor: - """Validate and convert a control-point field to a contiguous float32 tensor.""" - tensor = ( - control_points.clone().detach().to(torch.float32) - if isinstance(control_points, Tensor) - else torch.as_tensor(np.asarray(control_points), dtype=torch.float32) - ) - if tensor.ndim != 4 or tensor.shape[-1] != 3: - msg = ( - "control_points must have shape (n_i, n_j, n_k, 3), got" - f" {tuple(tensor.shape)}" - ) - raise ValueError(msg) - for axis, size in enumerate(tensor.shape[:-1]): - if size < 4: - msg = ( - "Each control-point axis must have at least 4 elements;" - f" axis {axis} got {size}" - ) - raise ValueError(msg) - return tensor.contiguous() - - -def _normalize_parameter_value( - value: TypeParameterValue, -) -> float | tuple | Distribution | Choice: - """Cast ints to floats so `_ParameterRange` always receives floats.""" - if isinstance(value, (Distribution, Choice)): - return value - if isinstance(value, (int, float)): - return float(value) - # Tuple: may contain mixed specs (Choice, Distribution, sub-tuples) - # so pass through if any element is non-numeric. - if isinstance(value, tuple): - if all(isinstance(v, (int, float)) for v in value): - return tuple(float(v) for v in value) - return value - return value - - -def _to_parameter_range(value: TypeParameterValue) -> _ParameterRange: - """Convert a `TypeParameterValue` to a `_ParameterRange`.""" - return _ParameterRange(_normalize_parameter_value(value)) - - -def _to_nonnegative_parameter_range(value: TypeParameterValue) -> _ParameterRange: - """Like `_to_parameter_range`, but rejects negative values.""" - result = _to_parameter_range(value) - if result._distribution is None: - for low, high in result._ranges: - if low < 0 or high < 0: - msg = f"Value must be non-negative, got {value}" - raise ValueError(msg) - return result diff --git a/src/torchio/transforms/spatial/to_reference_space.py b/src/torchio/transforms/spatial/to_reference_space.py deleted file mode 100644 index 30510d41b..000000000 --- a/src/torchio/transforms/spatial/to_reference_space.py +++ /dev/null @@ -1,132 +0,0 @@ -"""ToReferenceSpace: set spatial metadata to match a reference image.""" - -from __future__ import annotations - -from typing import Any - -import numpy as np -from torch import Tensor - -from ...data.affine import AffineMatrix -from ...data.batch import SubjectsBatch -from ...data.image import Image -from ...types import TypeThreeInts -from ..transform import SpatialTransform - - -class ToReferenceSpace(SpatialTransform): - r"""Set the spatial metadata of an image to match a reference space. - - This is useful for assigning meaningful spatial metadata to a - tensor that has lost it, such as a neural network embedding or a - downsampled feature map. The data is left unchanged; only the - affine is updated so that the (possibly lower-resolution) grid - covers the same field of view, orientation, and physical center - as the *reference* image. - - A typical use case is visualizing or resampling the output of a - network whose spatial resolution differs from its input: - - Args: - reference: Full-resolution reference image whose field of view - and orientation will be matched. - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torch - >>> import torchio as tio - >>> reference = tio.ScalarImage(tensor=torch.rand(1, 64, 64, 64)) - >>> # A network embedding loses spatial metadata: - >>> embedding = torch.rand(8, 16, 16, 16) - >>> image = tio.ToReferenceSpace.from_tensor(embedding, reference) - >>> image.spatial_shape - (16, 16, 16) - """ - - def __init__(self, reference: Image, **kwargs: Any) -> None: - super().__init__(**kwargs) - if not isinstance(reference, Image): - msg = f"reference must be a TorchIO Image, got {type(reference).__name__}" - raise TypeError(msg) - self.reference = reference - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Replace each image's affine with the reference-space affine.""" - for _name, img_batch in self._get_images(batch).items(): - output_shape = ( - int(img_batch.data.shape[2]), - int(img_batch.data.shape[3]), - int(img_batch.data.shape[4]), - ) - new_affine = _reference_space_affine(self.reference, output_shape) - img_batch.affines[:] = [new_affine.clone() for _ in img_batch.affines] - return batch - - @staticmethod - def from_tensor(tensor: Tensor, reference: Image) -> Image: - """Build a TorchIO image from a tensor and a reference image. - - Args: - tensor: A `(C, I, J, K)` tensor (e.g., a network - embedding) whose spatial metadata should match the - reference space. - reference: Reference image whose field of view and - orientation will be matched. - - Returns: - A new image with *tensor* as data and a reference-space - affine. The image class matches that of *reference*. - """ - output_shape = ( - int(tensor.shape[-3]), - int(tensor.shape[-2]), - int(tensor.shape[-1]), - ) - new_affine = _reference_space_affine(reference, output_shape) - cls = type(reference) - return cls(tensor, affine=new_affine) - - -def _reference_space_affine( - reference: Image, - output_shape: TypeThreeInts, -) -> AffineMatrix: - """Compute an affine placing a grid in the reference field of view. - - The output grid shares the reference's physical center and - orientation; the voxel spacing is scaled so the grid covers the - same field of view, regardless of its resolution. - - Args: - reference: The reference image. - output_shape: Spatial shape `(I, J, K)` of the target grid. - - Returns: - The reference-space affine for the target grid. - """ - ref_affine = reference.affine - rotation = ref_affine.direction.cpu().numpy().astype(np.float64) - ref_spacing = np.asarray(ref_affine.spacing, dtype=np.float64) - ref_origin = np.asarray(ref_affine.origin, dtype=np.float64) - ref_shape = np.asarray(reference.spatial_shape, dtype=np.float64) - new_shape = np.asarray(output_shape, dtype=np.float64) - - downsampling = ref_shape / new_shape - new_spacing = ref_spacing * downsampling - - # Keep the physical center fixed so the grid covers the same FOV. - center = ref_origin + rotation @ (((ref_shape - 1) / 2) * ref_spacing) - new_origin = center - rotation @ (((new_shape - 1) / 2) * new_spacing) - - matrix = np.eye(4, dtype=np.float64) - matrix[:3, :3] = rotation * new_spacing - matrix[:3, 3] = new_origin - return AffineMatrix(matrix) diff --git a/src/torchio/transforms/spatial/transpose.py b/src/torchio/transforms/spatial/transpose.py deleted file mode 100644 index 40c5a12b0..000000000 --- a/src/torchio/transforms/spatial/transpose.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Transpose: swap the first and last spatial axes.""" - -from __future__ import annotations - -from typing import Any - -from ...data.batch import SubjectsBatch -from ..transform import SpatialTransform - - -class Transpose(SpatialTransform): - r"""Swap the first and last spatial dimensions. - - Transforms an image of shape $(C, I, J, K)$ into $(C, K, J, I)$. - The affine matrix is updated to reflect the reordering so that - world coordinates remain consistent. - - This is the v2 equivalent of v1's `Transpose`, which reversed - the orientation string. The transform is its own inverse. - - Args: - **kwargs: See [`Transform`][torchio.Transform]. - - Examples: - >>> import torchio as tio - >>> transform = tio.Transpose() - """ - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """No random parameters.""" - return {} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Swap first and last spatial axes for all images.""" - for _name, img_batch in batch.images.items(): - # Swap axes 2 (I) and 4 (K) in the (B, C, I, J, K) tensor. - img_batch.data = img_batch.data.permute(0, 1, 4, 3, 2).contiguous() - # Update affines: swap columns 0 and 2 (I↔K). - for affine in img_batch.affines: - m = affine._matrix.clone() - affine._matrix[:, 0] = m[:, 2] - affine._matrix[:, 2] = m[:, 0] - return batch - - @property - def invertible(self) -> bool: - """Transpose is its own inverse.""" - return True - - def inverse(self, params: dict[str, Any]) -> Transpose: - """Transposing twice is identity.""" - return Transpose(copy=False) diff --git a/src/torchio/transforms/spatial_transform.py b/src/torchio/transforms/spatial_transform.py new file mode 100644 index 000000000..616c192a0 --- /dev/null +++ b/src/torchio/transforms/spatial_transform.py @@ -0,0 +1,15 @@ +from ..data import Image +from ..data.subject import Subject +from .transform import Transform + + +class SpatialTransform(Transform): + """Transform that modifies image bounds or voxels positions.""" + + def get_images(self, subject: Subject) -> list[Image]: + images = subject.get_images( + intensity_only=False, + include=self.include, + exclude=self.exclude, + ) + return images diff --git a/src/torchio/transforms/to.py b/src/torchio/transforms/to.py deleted file mode 100644 index f69a5e48a..000000000 --- a/src/torchio/transforms/to.py +++ /dev/null @@ -1,49 +0,0 @@ -"""To transform: move data to a device and/or cast dtype.""" - -from __future__ import annotations - -from typing import Any - -from ..data.batch import SubjectsBatch -from .transform import Transform - - -class To(Transform): - """Move all data to a device and/or cast to a dtype. - - Wraps the `to()` method as a transform so it can be used inside - [`Compose`][torchio.Compose] pipelines. - - Args: - *to_args: Positional arguments forwarded to - [`torch.Tensor.to()`](https://pytorch.org/docs/stable/generated/torch.Tensor.to.html). - Typically a device string (`"cpu"`, `"cuda"`, - `"mps"`) or a `torch.dtype` (`torch.float16`). - **to_kwargs: Keyword arguments forwarded to - `torch.Tensor.to()`. - - Examples: - >>> import torchio as tio - >>> transform = tio.To(torch.float16) - >>> transform = tio.To("cuda") - >>> pipeline = tio.Compose([ - ... tio.To("cuda"), - ... tio.Noise(std=0.1), - ... ]) - """ - - def __init__(self, *to_args: Any, **to_kwargs: Any) -> None: - super().__init__() - self.to_args = to_args - self.to_kwargs = to_kwargs - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - return {"to_args": self.to_args, "to_kwargs": self.to_kwargs} - - def apply_transform( - self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - batch.to(*params["to_args"], **params["to_kwargs"]) - return batch diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py index d7d9f3bc0..056e1c4fa 100644 --- a/src/torchio/transforms/transform.py +++ b/src/torchio/transforms/transform.py @@ -1,693 +1,758 @@ -"""Transform base classes.""" - from __future__ import annotations -import contextlib -import copy as _copy -import inspect +import copy import warnings -from dataclasses import dataclass -from dataclasses import field -from typing import Any +from abc import ABC +from abc import abstractmethod +from collections.abc import MutableMapping +from collections.abc import Sequence +from contextlib import contextmanager +from typing import TYPE_CHECKING +from typing import Literal +from typing import TypeGuard +from typing import TypeVar +from typing import cast from typing import overload import nibabel as nib import numpy as np import SimpleITK as sitk import torch -from einops import rearrange -from torch import Tensor -from torch import nn -from ..data.batch import ImagesBatch -from ..data.batch import SubjectsBatch from ..data.image import Image -from ..data.image import ScalarImage +from ..data.image import LabelMap +from ..data.io import nib_to_sitk +from ..data.io import sitk_to_nib from ..data.subject import Subject - - -@dataclass -class AppliedTransform: - """Record of a transform application, stored in Subject history. - - Attributes: - name: Class name of the transform. - params: Sampled parameters (JSON-serializable). - include: Original include scope of the applied transform. - exclude: Original exclude scope of the applied transform. - """ - - name: str - params: dict[str, Any] = field(default_factory=dict) - include: list[str] | None = None - exclude: list[str] | None = None - - -#: Registry mapping transform class names to classes, for inverse lookup. -_TRANSFORM_REGISTRY: dict[str, type[Transform]] = {} - - -def _all_elements_gated_out(params: dict[str, Any]) -> bool: - """Whether per-element gating masked out every batch element. - - Args: - params: The parameter dict produced by `make_params`, possibly - carrying a `_keep` mask added by `_tag_batched`. - - Returns: - `True` only when a `_keep` mask is present and none of its - elements are kept, i.e. the transform was an exact no-op. - """ - keep = params.get("_keep") - return keep is not None and not any(keep) - - -def _copy_optional_list(value: list[str] | None) -> list[str] | None: - return None if value is None else list(value) - - -class Transform(nn.Module): +from ..types import TypeCallable +from ..types import TypeData +from ..types import TypeKeys +from ..types import TypeNumber +from ..types import TypeTripletInt +from ..utils import is_iterable +from .data_parser import DataParser +from .data_parser import TypeTransformInput +from .interpolation import Interpolation +from .interpolation import get_sitk_interpolator + +TypeSixBounds = tuple[int, int, int, int, int, int] +TypeBounds = int | Sequence[int] | None +TypeMaskingMethod = str | TypeCallable | TypeBounds | None +ANATOMICAL_AXES = ( + 'Left', + 'Right', + 'Posterior', + 'Anterior', + 'Inferior', + 'Superior', +) + +ArgumentsDictT = TypeVar('ArgumentsDictT', bound=MutableMapping[str, object]) +ImageT = TypeVar('ImageT', bound=Image) + +__all__ = ['Transform', 'TypeBounds', 'TypeMaskingMethod', 'TypeTripletInt'] + + +class Transform(ABC): """Abstract class for all TorchIO transforms. When called, the input can be an instance of - [`Subject`][torchio.Subject], - [`Image`][torchio.Image], - [`torch.Tensor`][torch.Tensor], + [`torchio.Subject`][torchio.Subject], + [`torchio.Image`][torchio.Image], [`numpy.ndarray`][numpy.ndarray], + [`torch.Tensor`][torch.Tensor], [`SimpleITK.Image`](https://simpleitk.org/doxygen/latest/html/classitk_1_1simple_1_1Image.html), - [`nibabel.Nifti1Image`](https://nipy.org/nibabel/reference/nibabel.nifti1.html), - [`dict`][dict] containing 4D tensors as values, - [`ImagesBatch`][torchio.ImagesBatch], or - [`SubjectsBatch`][torchio.SubjectsBatch]. - The output type always matches the input type. + or [`dict`][dict] containing 4D tensors as values. - All subclasses must override - [`apply_transform()`][torchio.Transform.apply_transform], - which receives a [`SubjectsBatch`][torchio.SubjectsBatch] and - returns the transformed batch. + All subclasses must overwrite + [`apply_transform()`][torchio.transforms.Transform.apply_transform], + which takes an instance of [`Subject`][torchio.Subject], + modifies it and returns the result. Args: - p: Probability that this transform will be applied. When - per-instance probability is active (see `per_instance`), - this is instead the per-element probability and each batch - element is gated independently. - copy: Make a deep copy of the input before applying the - transform. When transforms are composed with - [`Compose`][torchio.Compose], the outer `Compose` - copies once and sets `copy=False` on inner transforms - to avoid redundant copies. - per_instance: If `True` (default), transforms that support it - sample independent parameters for each element of a batch - (and gate each element independently with `p`). If - `False`, a single parameter set is sampled and applied - identically to every element, reproducing the legacy - batch-shared behavior. Single-element inputs (including a - single [`Subject`][torchio.Subject]) are unaffected by this - flag. - include: Sequence of strings with the names of the only images - to which the transform will be applied. - exclude: Sequence of strings with the names of the images to - which the transform will *not* be applied. + p: Probability that this transform will be applied. + copy: Make a deep copy of the input before applying the transform. + include: Sequence of strings with the names of the only images to which + the transform will be applied. + Mandatory if the input is a [`dict`][dict]. + exclude: Sequence of strings with the names of the images to which the + the transform will not be applied, apart from the ones that are + excluded because of the transform type. + For example, if a subject includes an MRI, a CT and a label map, + and the CT is added to the list of exclusions of an intensity + transform such as [`RandomBlur`][torchio.transforms.RandomBlur], + the transform will be only applied to the MRI, as the label map is + excluded by default by spatial transforms. + keep: Dictionary with the names of the input images that will be kept + in the output and their new names. For example: + `{'t1': 't1_original'}`. This might be useful for autoencoders + or registration tasks. + parse_input: If `True`, the input will be converted to an instance of + [`Subject`][torchio.Subject]. This is used internally by some special + transforms like + [`Compose`][torchio.transforms.augmentation.composition.Compose]. + label_keys: If the input is a dictionary, names of images that + correspond to label maps. """ + if TYPE_CHECKING: + invert_transform: bool + def __init__( self, - *, - p: float = 1.0, + p: float = 1, copy: bool = True, - per_instance: bool = True, - include: list[str] | None = None, - exclude: list[str] | None = None, - ) -> None: - super().__init__() - if not 0 <= p <= 1: - msg = f"Probability must be in [0, 1], got {p}" - raise ValueError(msg) - self.p = p + include: TypeKeys = None, + exclude: TypeKeys = None, + keys: TypeKeys = None, + keep: dict[str, str] | None = None, + parse_input: bool = True, + label_keys: TypeKeys = None, + ): + self.probability = self.parse_probability(p) self.copy = copy - self.per_instance = per_instance - self.include = include - self.exclude = exclude - - def __init_subclass__(cls, **kwargs: Any) -> None: - super().__init_subclass__(**kwargs) - _TRANSFORM_REGISTRY[cls.__name__] = cls - - def _warn_if_noop(self, *, is_noop: bool, hint: str) -> None: - """Warn that the transform leaves the data unchanged. - - Augmentation transforms whose parameters are sampled from a - range default to an identity (no-op) when constructed with no - arguments, so that randomness must be requested explicitly. This - warns the user when that happens (or whenever the given - parameters produce a no-op). - - Args: - is_noop: Whether the configured transform is an identity. - hint: Example argument to suggest in the warning message. - """ - if is_noop: - warnings.warn( - f"{type(self).__name__} is a no-op with the given parameters" - " and will not change the data. Pass arguments to apply an" - f" effect (e.g. {hint}), or a range like (a, b) for random" - " augmentation.", - stacklevel=3, + if keys is not None: + message = ( + 'The "keys" argument is deprecated and will be removed in the' + ' future. Use "include" instead' ) + warnings.warn(message, FutureWarning, stacklevel=2) + include = keys + self.include, self.exclude = self.parse_include_and_exclude_keys( + include, + exclude, + label_keys, + ) + self.keep = keep + self.parse_input = parse_input + self.label_keys = label_keys + # args_names is the sequence of parameters from self that need to be + # passed to a non-random version of a random transform. They are also + # used to invert invertible transforms + self.args_names: list[str] = [] - def __repr__(self) -> str: - """Show only non-default fields for a compact repr.""" - from .parameter_range import _ParameterRange - - parts = [] - for name, default in _collect_init_params(type(self)).items(): - value = getattr(self, name, default) - if isinstance(value, _ParameterRange): - if value._original == default: - continue - elif value == default: - continue - parts.append(f"{name}={value!r}") - return f"{type(self).__name__}({', '.join(parts)})" - - def __add__(self, other: object) -> Transform: - """Compose two transforms: `t1 + t2` → `Compose([t1, t2])`.""" - if not isinstance(other, Transform): - return NotImplemented - from .compose import Compose - - left = self.transforms if isinstance(self, Compose) else [self] - right = other.transforms if isinstance(other, Compose) else [other] - return Compose([*left, *right]) - - def __or__(self, other: object) -> Transform: - """Random choice: `t1 | t2` → `OneOf([t1, t2])`.""" - if not isinstance(other, Transform): - return NotImplemented - from .compose import OneOf - - left = self.transforms if isinstance(self, OneOf) else [self] - right = other.transforms if isinstance(other, OneOf) else [other] - return OneOf([*left, *right]) - - @overload - def forward(self, data: Subject) -> Subject: ... @overload - def forward(self, data: Image) -> Image: ... - @overload - def forward(self, data: Tensor) -> Tensor: ... - @overload - def forward(self, data: np.ndarray) -> np.ndarray: ... + def __call__(self, data: Subject) -> Subject: ... + @overload - def forward(self, data: sitk.Image) -> sitk.Image: ... + def __call__(self, data: ImageT) -> ImageT: ... + @overload - def forward(self, data: nib.Nifti1Image) -> nib.Nifti1Image: ... + def __call__(self, data: torch.Tensor) -> torch.Tensor: ... + @overload - def forward(self, data: dict) -> dict: ... + def __call__(self, data: np.ndarray) -> np.ndarray: ... + @overload - def forward(self, data: ImagesBatch) -> ImagesBatch: ... + def __call__(self, data: sitk.Image) -> sitk.Image: ... + @overload - def forward(self, data: SubjectsBatch) -> SubjectsBatch: ... + def __call__(self, data: dict[str, object]) -> dict[str, object]: ... - def forward(self, data: Any) -> Any: - """Apply the transform. + @overload + def __call__(self, data: nib.Nifti1Image) -> nib.Nifti1Image: ... - The output type always matches the input type. + def __call__(self, data: TypeTransformInput) -> TypeTransformInput: + """Transform data and return a result of the same type. Args: - data: Input data to transform. + data: Instance of [`torchio.Subject`][torchio.Subject], 4D + [`torch.Tensor`][torch.Tensor] or [`numpy.ndarray`][numpy.ndarray] with dimensions + $(C, W, H, D)$, where $C$ is the number of channels + and $W, H, D$ are the spatial dimensions. If the input is + a tensor, the affine matrix will be set to identity. Other + valid input types are a SimpleITK image, a + [`torchio.Image`][torchio.Image], a NiBabel Nifti1 image or a + [`dict`][dict]. The output type is the same as the input type. """ - if self.copy: - data = _copy.deepcopy(data) - batch, unwrap = self._wrap(data) - # When per-element gating is active, the transform handles the - # probability itself (masked-out elements get identity params), - # so skip the batch-wide coin flip here. Apply iff rand < p, so - # p=0 is always a no-op and p=1 always applies. - if not self._per_instance_p_active(batch) and torch.rand(1).item() >= self.p: - return unwrap(batch) - params = self.make_params(batch) - batch = self.apply_transform(batch, params) - # Record history on the batch, unless every element was gated out by - # per-element probability: that is an exact no-op, and recording it - # would let history replay (e.g. an invertible spatial transform) - # trigger an unnecessary identity resample. - if not _all_elements_gated_out(params): - trace = AppliedTransform( - name=type(self).__name__, - params=params, - include=_copy_optional_list(self.include), - exclude=_copy_optional_list(self.exclude), + if torch.rand(1).item() > self.probability: + return data + + # Some transforms such as Compose should not modify the input data + if self.parse_input: + data_parser = DataParser( + data, + keys=self.include, + label_keys=self.label_keys, ) - if not hasattr(batch, "applied_transforms"): - batch.applied_transforms = [] - batch.applied_transforms.append(trace) - result = unwrap(batch) - # Propagate history to outputs that can carry it - if ( - hasattr(batch, "applied_transforms") - and not isinstance(result, (SubjectsBatch, Tensor, np.ndarray)) - and not isinstance(result, dict) - ): - with contextlib.suppress(AttributeError): - result.applied_transforms = list(batch.applied_transforms) - return result - - @property - def supports_per_instance_params(self) -> bool: - """Whether this transform can sample parameters per batch element. - - Defaults to `False`. Transforms that implement per-instance - parameter sampling override this to return `True`. When `False`, - the transform always uses batch-shared parameters regardless of - the `per_instance` flag, preserving the legacy behavior. - """ - return False - - @property - def supports_per_instance_p(self) -> bool: - """Whether this transform can gate each batch element independently. - - Defaults to `False`. Shape-preserving transforms that implement - per-element probability override this to return `True`. - Shape-changing transforms must leave it `False` because masked - and unmasked elements would have incompatible shapes. - """ - return False - - def _per_instance_active(self, batch: SubjectsBatch) -> bool: - """Whether per-instance parameter sampling applies to *batch*. + subject = data_parser.get_subject() + else: + subject = cast(Subject, data) - Per-instance sampling only kicks in for genuine batches - (`batch_size > 1`); single-element inputs always use the legacy - scalar path. - """ - return ( - self.per_instance - and self.supports_per_instance_params - and batch.batch_size > 1 - ) + if self.keep is not None: + images_to_keep: dict[str, Image] = {} + for name, new_name in self.keep.items(): + images_to_keep[new_name] = copy.deepcopy(subject.get_image(name)) + if self.copy: + subject = copy.deepcopy(subject) + with np.errstate(all='raise', under='ignore'): + transformed = self.apply_transform(subject) + if self.keep is not None: + for name, image in images_to_keep.items(): + transformed.add_image(image, name) + + if self.parse_input: + self._add_transform_to_subject_history(transformed) + for image in transformed.get_images(intensity_only=False): + ndim = image.data.ndim + assert ndim == 4, f'Output of {self.name} is {ndim}D' + output = data_parser.get_output(transformed) + else: + output = transformed + + return output + + def __repr__(self): + if hasattr(self, 'args_names'): + named_args = self._get_named_arguments() + args_strings = [f'{arg}={value}' for arg, value in named_args.items()] + if hasattr(self, 'invert_transform') and self.invert_transform: + args_strings.append('invert=True') + args_string = ', '.join(args_strings) + return f'{self.name}({args_string})' + else: + return super().__repr__() - def _per_instance_p_active(self, batch: SubjectsBatch) -> bool: - """Whether per-element probability gating applies to *batch*.""" - return ( - self.per_instance - and self.supports_per_instance_p - and batch.batch_size > 1 - and 0.0 < self.p < 1.0 - ) + def _get_base_args(self) -> dict[str, object]: + r"""Provides easy access to the arguments used to instantiate the base class + ([`Transform`][torchio.transforms.transform.Transform]) of any transform. - def _resolve_n(self, batch: SubjectsBatch) -> int | None: - """Return the number of parameter sets to sample. + This method is particularly useful when a new transform can be represented as a variant + of an existing transform (e.g. all random transforms), allowing for seamless instantiation + of the existing transform with the same arguments as the new transform during `apply_transform`. - Returns: - The batch size when per-instance sampling is active, - otherwise `None` (the legacy single-sample path). + Note: + The `p` argument (probability of applying the transform) is excluded to avoid + multiplying the probability of both existing and new transform. """ - return batch.batch_size if self._per_instance_active(batch) else None + return { + 'copy': self.copy, + 'include': self.include, + 'exclude': self.exclude, + 'keep': self.keep, + 'parse_input': self.parse_input, + 'label_keys': self.label_keys, + } - def _keep_mask( + def _add_base_args( self, - batch: SubjectsBatch, - n: int | None, - ) -> Tensor | None: - """Sample a per-element keep mask for per-instance probability. - - Args: - batch: The batch being transformed. - n: The resolved number of parameter sets (from - `_resolve_n`). + arguments: ArgumentsDictT, + overwrite_on_existing: bool = False, + ) -> ArgumentsDictT: + """Add the init args to existing arguments""" + for key, value in self._get_base_args().items(): + if key in arguments and not overwrite_on_existing: + continue + arguments[key] = value + return arguments - Returns: - A boolean tensor of shape `(n,)` where `True` marks - elements that receive the transform, or `None` when - per-element gating is not active (all elements are kept). - """ - if n is None or not self._per_instance_p_active(batch): - return None - return torch.rand(n) < self.p + @property + def name(self): + return self.__class__.__name__ - @staticmethod - def _mask_identity( - value: Tensor | float, - keep: Tensor | None, - *, - identity: float, - ) -> Tensor | float: - """Replace masked-out elements of *value* with an identity value. + @abstractmethod + def apply_transform(self, subject: Subject) -> Subject: + """Apply the transform to a parsed subject. Args: - value: Sampled parameter, either a scalar (legacy path) or a - `(B,)` tensor (per-instance path). - keep: Per-element keep mask, or `None` for no masking. - identity: The value that makes the transform a no-op for an - element (for example `0.0` for additive or log-space - parameters). + subject: Subject to be modified by the transform. Returns: - The masked parameter. + The transformed subject. """ - if keep is None or not isinstance(value, Tensor): - return value - return torch.where(keep, value, torch.full_like(value, identity)) + raise NotImplementedError - @staticmethod - def _serialize_param(value: Tensor | Any) -> Any: - """Convert a possibly-tensor parameter to a JSON-serializable form.""" - if isinstance(value, Tensor): - return value.tolist() - return value + def _add_transform_to_subject_history(self, subject): + from . import Compose + from . import CropOrPad + from . import EnsureShapeMultiple + from . import OneOf + from .augmentation import RandomTransform + from .preprocessing import Resize + from .preprocessing import SequentialLabels + + call_others = ( + RandomTransform, + Compose, + OneOf, + CropOrPad, + EnsureShapeMultiple, + SequentialLabels, + Resize, + ) + if not isinstance(self, call_others): + subject.add_transform(self, self._get_reproducing_arguments()) @staticmethod - def _is_per_instance_params(params: dict[str, Any]) -> bool: - """Whether *params* holds per-element (batched) values.""" - return "_batched_keys" in params + def to_range(n: TypeNumber, around: float | None) -> tuple[float, float]: + if around is None: + return 0.0, float(n) + else: + return float(around - n), float(around + n) - def _tag_batched( + @overload + def parse_params( self, - params: dict[str, Any], - batch: SubjectsBatch, - n: int | None, - keep: Tensor | None, - batched_keys: list[str], - ) -> None: - """Annotate *params* with per-instance bookkeeping for history. - - Adds the batch size, the names of the per-element keys, and the - keep mask so that [`SubjectsBatch.unbatch`][torchio.SubjectsBatch.unbatch] - can split the history per subject. + params: TypeNumber | Sequence[TypeNumber], + around: float | None, + name: str, + make_ranges: Literal[True] = True, + min_constraint: TypeNumber | None = None, + max_constraint: TypeNumber | None = None, + type_constraint: type[int] | type[float] | None = None, + ) -> tuple[float, float, float, float, float, float]: ... - Args: - params: The parameter dict to annotate in place. - batch: The batch being transformed. - n: The resolved number of parameter sets. - keep: The per-element keep mask, or `None`. - batched_keys: Names of the params that hold one value per - element. - """ - if n is None: - return - params["_batch_size"] = batch.batch_size - params["_batched_keys"] = list(batched_keys) - if keep is not None: - params["_keep"] = keep.tolist() - - def make_params(self, batch: SubjectsBatch) -> dict[str, Any]: - """Sample random parameters for this transform. - - Override in subclasses that have random behavior. - - Args: - batch: A `SubjectsBatch`. - - Returns: - Dict of sampled parameters. - """ - return {} - - def apply_transform( + @overload + def parse_params( self, - batch: SubjectsBatch, - params: dict[str, Any], - ) -> SubjectsBatch: - """Apply the transform with the given parameters. - - Must be overridden by subclasses. Receives a `SubjectsBatch` - whose `ImagesBatch` entries contain 5D tensors - `(B, C, I, J, K)`. Use negative indexing (`-3`, `-2`, - `-1`) for spatial dims. - - Args: - batch: A `SubjectsBatch` to transform. - params: Parameters from `make_params`. - - Returns: - Transformed `SubjectsBatch`. - """ - raise NotImplementedError + params: TypeNumber | Sequence[TypeNumber], + around: float | None, + name: str, + make_ranges: Literal[False], + min_constraint: TypeNumber | None = None, + max_constraint: TypeNumber | None = None, + type_constraint: type[int] | type[float] | None = None, + ) -> tuple[float, ...]: ... + + def parse_params( + self, + params: TypeNumber | Sequence[TypeNumber], + around: float | None, + name: str, + make_ranges: bool = True, + min_constraint: TypeNumber | None = None, + max_constraint: TypeNumber | None = None, + type_constraint: type[int] | type[float] | None = None, + ) -> tuple[float, float, float, float, float, float] | tuple[float, ...]: + if isinstance(params, torch.Tensor): + params_tuple = tuple(float(param) for param in params.reshape(-1)) + elif isinstance(params, np.ndarray): + params_tuple = tuple(float(param) for param in np.ravel(params)) + elif isinstance(params, Sequence): + params_sequence = cast(Sequence[TypeNumber], params) + params_tuple = tuple(float(param) for param in params_sequence) + else: + params_tuple = (float(params),) + # d or (a, b) + if len(params_tuple) == 1 or (len(params_tuple) == 2 and make_ranges): + params_tuple *= 3 # (d, d, d) or (a, b, a, b, a, b) + if len(params_tuple) == 3 and make_ranges: # (a, b, c) + items = [self.to_range(n, around) for n in params_tuple] + # (-a, a, -b, b, -c, c) or (1-a, 1+a, 1-b, 1+b, 1-c, 1+c) + params_tuple = tuple(n for prange in items for n in prange) + if make_ranges: + if len(params_tuple) != 6: + message = ( + f'If "{name}" is a sequence, it must have length 2, 3 or' + f' 6, not {len(params_tuple)}' + ) + raise ValueError(message) + for param_range in zip( + params_tuple[::2], + params_tuple[1::2], + strict=True, + ): + self._parse_range( + cast(tuple[float, float], param_range), + name, + min_constraint=min_constraint, + max_constraint=max_constraint, + type_constraint=type_constraint, + ) + a, b, c, d, e, f = params_tuple + return a, b, c, d, e, f + return params_tuple - @property - def invertible(self) -> bool: - """Whether this transform can be inverted.""" - return False + @overload + @staticmethod + def _parse_range( + nums_range: int | tuple[int, int], + name: str, + min_constraint: int | None = None, + max_constraint: int | None = None, + type_constraint: type[int] = int, + ) -> tuple[int, int]: ... - def inverse(self, params: dict[str, Any]) -> Transform: - """Return a transform that undoes this one. + @overload + @staticmethod + def _parse_range( + nums_range: float | tuple[float, float], + name: str, + min_constraint: float | None = None, + max_constraint: float | None = None, + type_constraint: type[float] | None = None, + ) -> tuple[float, float]: ... - Override in invertible subclasses. The returned transform, - when applied, reverses the effect of the forward pass with - the given parameters. + @staticmethod + def _parse_range( + nums_range: TypeNumber | tuple[TypeNumber, TypeNumber], + name: str, + min_constraint: TypeNumber | None = None, + max_constraint: TypeNumber | None = None, + type_constraint: type[int] | type[float] | None = None, + ) -> tuple[TypeNumber, TypeNumber]: + r"""Adapted from [torchvision.transforms.RandomRotation][torchvision.transforms.RandomRotation]. Args: - params: The parameters recorded in the forward pass. + nums_range: Tuple of two numbers $(n_{min}, n_{max})$, + where $n_{min} \leq n_{max}$. + If a single positive number $n$ is provided, + $n_{min} = -n$ and $n_{max} = n$. + name: Name of the parameter, so that an informative error message + can be printed. + min_constraint: Minimal value that $n_{min}$ can take, + default is None, i.e. there is no minimal value. + max_constraint: Maximal value that $n_{max}$ can take, + default is None, i.e. there is no maximal value. + type_constraint: Precise type that $n_{max}$ and + $n_{min}$ must take. Returns: - A new `Transform` instance that inverts this one. + A tuple of two numbers $(n_{min}, n_{max})$. + + Raises: + ValueError: if `nums_range` is negative + ValueError: if $n_{max}$ or $n_{min}$ is not a number + ValueError: if $n_{max} \lt n_{min}$ + ValueError: if `min_constraint` is not None and + $n_{min}$ is smaller than `min_constraint` + ValueError: if `max_constraint` is not None and + $n_{max}$ is greater than `max_constraint` + ValueError: if `type_constraint` is not None and + $n_{max}$ and $n_{max}$ are not of type + `type_constraint`. """ - msg = f"{type(self).__name__} is not invertible" - raise NotImplementedError(msg) - - def _get_images(self, batch: SubjectsBatch) -> dict[str, ImagesBatch]: - """Get image batches filtered by include/exclude.""" - images = batch.images - if self.include is not None: - images = {k: v for k, v in images.items() if k in self.include} - if self.exclude is not None: - images = {k: v for k, v in images.items() if k not in self.exclude} - return images - - def to_hydra(self) -> dict[str, Any]: - """Export as a Hydra-compatible config dict. - - Returns a dict with `_target_` set to the fully qualified - class name and only non-default field values included. - - Returns: - Dict suitable for `hydra.utils.instantiate()`. - """ - from .parameter_range import _ParameterRange - - cls = type(self) - target = f"torchio.{cls.__qualname__}" - cfg: dict[str, Any] = {"_target_": target} - - for name, default in _collect_init_params(cls).items(): - value = getattr(self, name, default) - if isinstance(value, _ParameterRange): - if value._original == default: - continue - value = _hydra_value(value._original) - elif value == default: - continue - else: - value = _hydra_value(value) - cfg[name] = value - return cfg - - @staticmethod - def _wrap( - data: Any, - ) -> tuple[Any, Any]: - """Wrap any input into a SubjectsBatch; return (batch, unwrap_fn).""" - from ..data.batch import ImagesBatch - from ..data.batch import SubjectsBatch - - match data: - case SubjectsBatch(): - return data, _unwrap_subjects_batch - case ImagesBatch(): - sb = SubjectsBatch({"tio_default_image": data}) - return sb, _unwrap_images_batch - case Subject(): - sb = SubjectsBatch.from_subjects([data]) - return sb, _unwrap_subject - case dict(): - return _wrap_dict(data) - case _: - return _wrap_scalar_input(data) - - -def _wrap_single_image(img: Image, unwrap_fn: Any) -> tuple[Any, Any]: - """Wrap a single Image into a SubjectsBatch.""" - from ..data.batch import SubjectsBatch - - sub = Subject(tio_default_image=img) - sb = SubjectsBatch.from_subjects([sub]) - return sb, unwrap_fn - - -def _wrap_scalar_input(data: Any) -> tuple[Any, Any]: - """Wrap a scalar input (Tensor, ndarray, SimpleITK, NIfTI) into a batch.""" - match data: - case Image(): - return _wrap_single_image(data, _unwrap_image) - case Tensor(): - return _wrap_single_image(ScalarImage(data), _unwrap_tensor) - case np.ndarray(): - tensor = torch.as_tensor(data.copy(), dtype=torch.float32) - if tensor.ndim == 3: - tensor = rearrange(tensor, "i j k -> 1 i j k") - return _wrap_single_image( - ScalarImage(tensor), - _unwrap_ndarray, + if isinstance(nums_range, (int, float)): # single number given + if nums_range < 0: + raise ValueError( + f'If {name} is a single number,' + f' it must be positive, not {nums_range}', + ) + if min_constraint is not None and nums_range < min_constraint: + raise ValueError( + f'If {name} is a single number, it must be greater' + f' than {min_constraint}, not {nums_range}', + ) + if max_constraint is not None and nums_range > max_constraint: + raise ValueError( + f'If {name} is a single number, it must be smaller' + f' than {max_constraint}, not {nums_range}', + ) + if type_constraint is not None: + if not isinstance(nums_range, type_constraint): + raise ValueError( + f'If {name} is a single number, it must be of' + f' type {type_constraint}, not {nums_range}', + ) + min_range = -nums_range if min_constraint is None else min_constraint + return (min_range, nums_range) + + try: + values = tuple(nums_range) + except TypeError as err: + message = ( + f'If {name} is not a single number, it must be' + f' a sequence of len 2, not {nums_range}' ) - case sitk.Image(): - return _wrap_single_image(ScalarImage(data), _unwrap_sitk) - case nib.Nifti1Image(): - return _wrap_single_image(ScalarImage(data), _unwrap_nifti) - case _: - msg = ( - "Expected Subject, Image, Tensor, ndarray, dict," - f" SimpleITK Image, NIfTI, ImagesBatch, or SubjectsBatch," - f" got {type(data).__name__}" + raise ValueError(message) from err + if len(values) != 2: + message = ( + f'If {name} is not a single number, it must be' + f' a sequence of len 2, not {nums_range}' + ) + raise ValueError(message) + min_value, max_value = values + + min_is_number = isinstance(min_value, (int, float)) + max_is_number = isinstance(max_value, (int, float)) + if not min_is_number or not max_is_number: + message = f'{name} values must be numbers, not {nums_range}' + raise ValueError(message) + + if min_value > max_value: + raise ValueError( + f'If {name} is a sequence, the second value must be' + f' equal or greater than the first, but it is {nums_range}', ) - raise TypeError(msg) - - -def _wrap_dict(data: dict) -> tuple[Any, Any]: - """Wrap a MONAI-style dict into a SubjectsBatch.""" - from ..data.batch import SubjectsBatch - - kwargs: dict[str, Any] = {} - for k, v in data.items(): - match v: - case Image(): - kwargs[k] = v - case Tensor(): - kwargs[k] = ScalarImage(v) - case _: - kwargs[k] = v - sub = Subject(**kwargs) - keys: list[str] = [str(k) for k in data] - sb = SubjectsBatch.from_subjects([sub]) - return sb, lambda b: _unwrap_dict(b, keys) - - -def _collect_init_params(cls: type) -> dict[str, Any]: - """Collect all __init__ params with defaults from the full MRO. - - Walks from the leaf class up through parent classes, collecting - named parameters (skipping `self`, `*args`, `**kwargs`). - Returns an ordered dict of `{name: default}`. - """ - params: dict[str, Any] = {} - for klass in cls.__mro__: - if klass is object or klass is nn.Module: - break - init = klass.__dict__.get("__init__") - if init is None: - continue - sig = inspect.signature(init) - for name, param in sig.parameters.items(): - if name == "self": - continue - if param.kind in ( - inspect.Parameter.VAR_POSITIONAL, - inspect.Parameter.VAR_KEYWORD, - ): - continue - if name not in params: - params[name] = param.default - return params - - -def _hydra_value(value: Any) -> Any: - """Convert a value to a plain Python type for Hydra/YAML.""" - if isinstance(value, tuple): - return list(value) - if isinstance(value, Tensor): - return value.tolist() - if isinstance(value, np.ndarray): - return value.tolist() - return value - - -def _unwrap_subjects_batch(batch: SubjectsBatch) -> SubjectsBatch: - return batch - -def _unwrap_images_batch(batch: SubjectsBatch) -> ImagesBatch: - return batch.images["tio_default_image"] + if min_constraint is not None and min_value < min_constraint: + raise ValueError( + f'If {name} is a sequence, the first value must be greater' + f' than {min_constraint}, but it is {min_value}', + ) + if max_constraint is not None and max_value > max_constraint: + raise ValueError( + f'If {name} is a sequence, the second value must be' + f' smaller than {max_constraint}, but it is {max_value}', + ) -def _unwrap_subject(batch: SubjectsBatch) -> Subject: - return batch.unbatch()[0] + if type_constraint is not None: + min_type_ok = isinstance(min_value, type_constraint) + max_type_ok = isinstance(max_value, type_constraint) + if not min_type_ok or not max_type_ok: + raise ValueError( + f'If "{name}" is a sequence, its values must be of' + f' type "{type_constraint}", not "{type(nums_range)}"', + ) + return min_value, max_value + @staticmethod + def parse_interpolation(interpolation: str) -> str: + if not isinstance(interpolation, str): + itype = type(interpolation) + raise TypeError(f'Interpolation must be a string, not {itype}') + interpolation = interpolation.lower() + is_string = isinstance(interpolation, str) + supported_values = [key.name.lower() for key in Interpolation] + is_supported = interpolation.lower() in supported_values + if is_string and is_supported: + return interpolation + message = ( + f'Interpolation "{interpolation}" of type {type(interpolation)}' + f' must be a string among the supported values: {supported_values}' + ) + raise ValueError(message) -def _unwrap_image(batch: SubjectsBatch) -> Image: - sub = batch.unbatch()[0] - return sub.tio_default_image + @staticmethod + def parse_probability(probability: float) -> float: + is_number = isinstance(probability, (int, float)) + if not (is_number and 0 <= probability <= 1): + message = f'Probability must be a number in [0, 1], not {probability}' + raise ValueError(message) + return probability + @staticmethod + def parse_include_and_exclude_keys( + include: TypeKeys, + exclude: TypeKeys, + label_keys: TypeKeys, + ) -> tuple[TypeKeys, TypeKeys]: + if include is not None and exclude is not None: + raise ValueError('Include and exclude cannot both be specified') + Transform._validate_keys_sequence(include, 'include') + Transform._validate_keys_sequence(exclude, 'exclude') + Transform._validate_keys_sequence(label_keys, 'label_keys') + return include, exclude -def _unwrap_tensor(batch: SubjectsBatch) -> Tensor: - sub = batch.unbatch()[0] - return sub.tio_default_image.data + @staticmethod + def _validate_keys_sequence(keys: TypeKeys, name: str) -> None: + """Ensure that the input is not a string but a sequence of strings.""" + if keys is None: + return + if isinstance(keys, str): + message = f'"{name}" must be a sequence of strings, not a string "{keys}"' + raise ValueError(message) + if not is_iterable(keys): + message = f'"{name}" must be a sequence of strings, not {type(keys)}' + raise ValueError(message) + @staticmethod + def nib_to_sitk(data: TypeData, affine: TypeData) -> sitk.Image: + return nib_to_sitk(data, affine) -def _unwrap_ndarray(batch: SubjectsBatch) -> np.ndarray: - sub = batch.unbatch()[0] - return sub.tio_default_image.data.cpu().numpy() + @staticmethod + def sitk_to_nib(image: sitk.Image) -> tuple[np.ndarray, np.ndarray]: + return sitk_to_nib(image) + + def _get_reproducing_arguments(self): + """Return a dictionary with the arguments that would be necessary to + reproduce the transform exactly.""" + reproducing_arguments = { + 'include': self.include, + 'exclude': self.exclude, + 'copy': self.copy, + } + reproducing_arguments.update(self._get_named_arguments()) + return reproducing_arguments + def _get_named_arguments(self) -> dict[str, object]: + return {name: getattr(self, name) for name in self.args_names} -def _unwrap_sitk(batch: SubjectsBatch) -> sitk.Image: - sub = batch.unbatch()[0] - image = sub.tio_default_image - data = image.data - affine = image.affine - array = data.cpu().numpy() - array = array[0] if data.shape[0] == 1 else np.moveaxis(array, 0, -1) - sitk_image = sitk.GetImageFromArray(array) - if data.shape[0] > 1: - sitk_image = sitk.GetImageFromArray(array, isVector=True) - sitk_image.SetSpacing(affine.spacing) - sitk_image.SetOrigin(affine.origin) - sitk_image.SetDirection(rearrange(affine.direction, "i j -> (i j)").tolist()) - return sitk_image + def is_invertible(self): + return hasattr(self, 'invert_transform') + def inverse(self): + if not self.is_invertible(): + raise RuntimeError(f'{self.name} is not invertible') + new = copy.deepcopy(self) + new.invert_transform = not self.invert_transform + return new -def _unwrap_nifti(batch: SubjectsBatch) -> nib.Nifti1Image: - sub = batch.unbatch()[0] - image = sub.tio_default_image - array = image.data.cpu().numpy() - array = array[0] if array.shape[0] == 1 else np.moveaxis(array, 0, -1) - return nib.Nifti1Image(array, image.affine.numpy()) + @staticmethod + @contextmanager + def _use_seed(seed): + """Perform an operation using a specific seed for the PyTorch RNG.""" + torch_rng_state = torch.random.get_rng_state() + torch.manual_seed(seed) + yield + torch.random.set_rng_state(torch_rng_state) + @staticmethod + def get_sitk_interpolator(interpolation: str) -> int: + return get_sitk_interpolator(interpolation) -def _unwrap_dict(batch: SubjectsBatch, keys: list[str]) -> dict[str, Any]: - sub = batch.unbatch()[0] - result: dict[str, Any] = {} - for k in keys: - entry = getattr(sub, k, None) - if isinstance(entry, Image): - result[k] = entry.data + @staticmethod + def parse_bounds(bounds_parameters: TypeBounds) -> TypeSixBounds | None: + if bounds_parameters is None: + return None + if isinstance(bounds_parameters, int): + values: tuple[int, ...] = (bounds_parameters,) else: - result[k] = entry - return result - + values = tuple(bounds_parameters) + + # Check that numbers are integers + for number in values: + if not isinstance(number, (int, np.integer)) or number < 0: + message = ( + 'Bounds values must be integers greater or equal to zero,' + f' not "{bounds_parameters}" of type {type(number)}' + ) + raise ValueError(message) + bounds_parameters_tuple = tuple(int(n) for n in values) + bounds_parameters_length = len(bounds_parameters_tuple) + if bounds_parameters_length == 6: + i0, i1, j0, j1, k0, k1 = bounds_parameters_tuple + return i0, i1, j0, j1, k0, k1 + if bounds_parameters_length == 1: + (value,) = bounds_parameters_tuple + return value, value, value, value, value, value + if bounds_parameters_length == 3: + i, j, k = bounds_parameters_tuple + return i, i, j, j, k, k + message = ( + 'Bounds parameter must be an integer or a tuple of' + f' 3 or 6 integers, not {bounds_parameters_tuple}' + ) + raise ValueError(message) -class SpatialTransform(Transform): - """Base for transforms that modify spatial geometry. + @staticmethod + def _is_mask_callable( + masking_method: object, + ) -> TypeGuard[TypeCallable]: + return callable(masking_method) - Spatial transforms apply to all images (ScalarImage and LabelMap), - and also transform any Points and BoundingBoxes attached to the - Subject. - """ + @staticmethod + def ones(tensor: torch.Tensor) -> torch.Tensor: + return torch.ones_like(tensor, dtype=torch.bool) + @staticmethod + def mean(tensor: torch.Tensor) -> torch.Tensor: + mask = tensor > tensor.float().mean() + return mask -class IntensityTransform(Transform): - """Base for transforms that modify voxel intensities. + def get_mask_from_masking_method( + self, + masking_method: TypeMaskingMethod, + subject: Subject, + tensor: torch.Tensor, + labels: Sequence[int] | None = None, + ) -> torch.Tensor: + if masking_method is None: + return self.ones(tensor) + elif self._is_mask_callable(masking_method): + return masking_method(tensor) + elif type(masking_method) is str: + in_subject = masking_method in subject + if in_subject: + label_map = subject[masking_method] + else: + label_map = None + if isinstance(label_map, LabelMap): + if labels is None: + return label_map.data.bool() + else: + mask_data = label_map.data + volumes = [mask_data == label for label in labels] + return torch.stack(volumes).sum(0).bool() + possible_axis = masking_method.capitalize() + if possible_axis in ANATOMICAL_AXES: + return self.get_mask_from_anatomical_label( + possible_axis, + tensor, + ) + elif isinstance(masking_method, int): + return self.get_mask_from_bounds(masking_method, tensor) + elif isinstance(masking_method, (tuple, list)): + if all(isinstance(number, (int, np.integer)) for number in masking_method): + bounds_list: list[int] = [] + for number in masking_method: + assert isinstance(number, (int, np.integer)) + bounds_list.append(int(number)) + return self.get_mask_from_bounds(bounds_list, tensor) + first_anat_axes = tuple(s[0] for s in ANATOMICAL_AXES) + message = ( + 'Masking method must be one of:\n' + ' 1) A callable object, such as a function\n' + ' 2) The name of a label map in the subject' + f' ({subject.get_images_names()})\n' + f' 3) An anatomical label {ANATOMICAL_AXES + first_anat_axes}\n' + ' 4) A bounds parameter' + ' (int, tuple of 3 ints, or tuple of 6 ints)\n' + f' The passed value, "{masking_method}",' + f' of type "{type(masking_method)}", is not valid' + ) + raise ValueError(message) - Intensity transforms apply only to `ScalarImage` instances, - leaving `LabelMap` and annotations unchanged. - """ + @staticmethod + def get_mask_from_anatomical_label( + anatomical_label: str, + tensor: torch.Tensor, + ) -> torch.Tensor: + # Assume the image is in RAS orientation + anatomical_label = anatomical_label.capitalize() + if anatomical_label not in ANATOMICAL_AXES: + message = ( + f'Anatomical label must be one of {ANATOMICAL_AXES}' + f' not {anatomical_label}' + ) + raise ValueError(message) + mask = torch.zeros_like(tensor, dtype=torch.bool) + _, width, height, depth = tensor.shape + if anatomical_label == 'Right': + mask[:, width // 2 :] = True + elif anatomical_label == 'Left': + mask[:, : width // 2] = True + elif anatomical_label == 'Anterior': + mask[:, :, height // 2 :] = True + elif anatomical_label == 'Posterior': + mask[:, :, : height // 2] = True + elif anatomical_label == 'Superior': + mask[:, :, :, depth // 2 :] = True + elif anatomical_label == 'Inferior': + mask[:, :, :, : depth // 2] = True + return mask + + def get_mask_from_bounds( + self, + bounds_parameters: TypeBounds, + tensor: torch.Tensor, + ) -> torch.Tensor: + bounds_parameters = self.parse_bounds(bounds_parameters) + assert bounds_parameters is not None + low = bounds_parameters[::2] + high = bounds_parameters[1::2] + i0, j0, k0 = low + i1, j1, k1 = np.array(tensor.shape[1:]) - high + mask = torch.zeros_like(tensor, dtype=torch.bool) + mask[:, i0:i1, j0:j1, k0:k1] = True + return mask + + def _get_name_with_module(self) -> str: + """Return the name of the transform including its module.""" + return f'{self.__class__.__module__}.{self.__class__.__name__}' - def _get_images(self, batch: SubjectsBatch) -> dict[str, ImagesBatch]: - """Filter to ScalarImage batches only, then apply include/exclude.""" - images = { - k: v for k, v in batch.images.items() if v._image_class is ScalarImage - } - if self.include is not None: - images = {k: v for k, v in images.items() if k in self.include} - if self.exclude is not None: - images = {k: v for k, v in images.items() if k not in self.exclude} - return images + @staticmethod + def _tuples_to_lists(obj): + if isinstance(obj, (tuple, list)): + return [Transform._tuples_to_lists(x) for x in obj] + if isinstance(obj, dict): + return {k: Transform._tuples_to_lists(v) for k, v in obj.items()} + return obj + + def to_hydra_config(self) -> dict: + """Return a dictionary representation of the transform for Hydra instantiation.""" + target = self._get_name_with_module() + transform_dict = {'_target_': target} + transform_dict.update(self._get_reproducing_arguments()) + return self._tuples_to_lists(transform_dict) diff --git a/src/torchio/types.py b/src/torchio/types.py old mode 100644 new mode 100755 index b8d0a0d3a..7cefcd334 --- a/src/torchio/types.py +++ b/src/torchio/types.py @@ -1,32 +1,53 @@ -"""Type aliases for TorchIO.""" - from __future__ import annotations -import os +from collections.abc import Callable +from collections.abc import Sequence +from pathlib import Path from typing import TypeAlias -from jaxtyping import Float -from torch import Tensor +import numpy as np +import torch +from jaxtyping import Shaped -# Path type for user-facing APIs -TypePath: TypeAlias = str | os.PathLike[str] +# For typing hints +TypePath: TypeAlias = str | Path +TypeNumber: TypeAlias = int | float +TypeKeys: TypeAlias = Sequence[str] | None +TypeData: TypeAlias = torch.Tensor | np.ndarray +TypeImageTensor: TypeAlias = Shaped[torch.Tensor, 'channels width height depth'] # noqa: F722 +TypeImageArray: TypeAlias = Shaped[np.ndarray, 'channels width height depth'] # noqa: F722 +TypeImageData: TypeAlias = TypeImageTensor | TypeImageArray +TypeAffineMatrix: TypeAlias = Shaped[np.ndarray, '4 4'] # noqa: F722 +TypeDataAffine: TypeAlias = tuple[torch.Tensor, np.ndarray] +TypeImageDataAffine: TypeAlias = tuple[TypeImageTensor, TypeAffineMatrix] +TypeSlice: TypeAlias = int | slice -# Jaxtyping tensor types (document shape in annotations) -TypeImageData = Float[Tensor, "channels size_i size_j size_k"] -TypeAffineMatrix = Float[Tensor, "4 4"] -TypeDirection = Float[Tensor, "3 3"] -TypeWorldPoints = Float[Tensor, "num_points 3"] +TypeDoubletInt: TypeAlias = tuple[int, int] +TypeTripletInt: TypeAlias = tuple[int, int, int] +TypeQuartetInt: TypeAlias = tuple[int, int, int, int] +TypeSextetInt: TypeAlias = tuple[int, int, int, int, int, int] -# Tuple types -TypeThreeInts: TypeAlias = tuple[int, int, int] -TypeFourInts: TypeAlias = tuple[int, int, int, int] -TypeSixInts: TypeAlias = tuple[int, int, int, int, int, int] -TypeThreeFloats: TypeAlias = tuple[float, float, float] -TypeSpatialShape = TypeThreeInts -TypeTensorShape = TypeFourInts -TypeSpacing = TypeThreeFloats -TypeOrigin = TypeThreeFloats -TypeOrientationCodes: TypeAlias = tuple[str, str, str] +TypeDoubleFloat: TypeAlias = tuple[float, float] +TypeTripletFloat: TypeAlias = tuple[float, float, float] +TypeQuartetFloat: TypeAlias = tuple[float, float, float, float] +TypeSextetFloat: TypeAlias = tuple[float, float, float, float, float, float] -#: Index type accepted by backend `__getitem__`. -SliceIndex: TypeAlias = int | slice | tuple[int | slice, ...] +TypeTuple: TypeAlias = int | TypeTripletInt +TypeRangeInt: TypeAlias = int | TypeDoubletInt +TypeSpacing: TypeAlias = float | TypeTripletFloat +TypeSpatialShape: TypeAlias = int | TypeTripletInt +TypeRangeFloat: TypeAlias = float | TypeDoubleFloat +TypeCallable: TypeAlias = Callable[[torch.Tensor], torch.Tensor] +TypeDirection2D: TypeAlias = TypeQuartetFloat +TypeDirection3D: TypeAlias = tuple[ + float, + float, + float, + float, + float, + float, + float, + float, + float, +] +TypeDirection: TypeAlias = TypeDirection2D | TypeDirection3D diff --git a/src/torchio/utils.py b/src/torchio/utils.py new file mode 100644 index 000000000..0727ec9d6 --- /dev/null +++ b/src/torchio/utils.py @@ -0,0 +1,429 @@ +from __future__ import annotations + +import ast +import gzip +import os +import shutil +import sys +import tempfile +from collections.abc import Iterable +from collections.abc import Sequence +from pathlib import Path +from typing import Any +from typing import TypeVar +from typing import cast +from typing import overload + +import numpy as np +import SimpleITK as sitk +import torch +from nibabel.nifti1 import Nifti1Image +from torch.utils.data import DataLoader +from torch.utils.data._utils.collate import default_collate +from tqdm.auto import trange + +from . import constants +from .types import TypePath + +T = TypeVar('T') + + +@overload +def to_tuple( + value: Iterable[T], + length: int = 1, +) -> tuple[T, ...]: ... + + +@overload +def to_tuple( + value: T, + length: int = 1, +) -> tuple[T, ...]: ... + + +def to_tuple( + value: T | Iterable[T], + length: int = 1, +) -> tuple[T, ...]: + """Convert variable to tuple of length n. + + Examples: + >>> from torchio.utils import to_tuple + >>> to_tuple(1, length=1) + (1,) + >>> to_tuple(1, length=3) + (1, 1, 1) + + If value is an iterable, n is ignored and tuple(value) is returned + + Examples: + >>> to_tuple((1,), length=1) + (1,) + >>> to_tuple((1, 2), length=1) + (1, 2) + >>> to_tuple([1, 2], length=3) + (1, 2) + """ + if isinstance(value, Iterable) and not isinstance(value, (str, bytes)): + iterable_value = cast(Iterable[T], value) + return tuple(iterable_value) + scalar_value = cast(T, value) + return length * (scalar_value,) + + +def get_stem( + path: TypePath | Sequence[TypePath], +) -> str | list[str]: + """Get stem of path or paths. + + Examples: + >>> from torchio.utils import get_stem + >>> get_stem('/home/user/my_image.nii.gz') + 'my_image' + """ + + def _get_stem(path_string: TypePath) -> str: + return Path(path_string).name.split('.')[0] + + if isinstance(path, Sequence) and not isinstance(path, (str, Path)): + return [_get_stem(p) for p in path] + assert isinstance(path, (str, Path)) + return _get_stem(path) + + +def create_dummy_dataset( + num_images: int, + size_range: tuple[int, int], + directory: TypePath | None = None, + suffix: str = '.nii.gz', + force: bool = False, + verbose: bool = False, +): + from .data import LabelMap + from .data import ScalarImage + from .data import Subject + + output_dir = tempfile.gettempdir() if directory is None else directory + output_dir = Path(output_dir) + images_dir = output_dir / 'dummy_images' + labels_dir = output_dir / 'dummy_labels' + + if force: + shutil.rmtree(images_dir) + shutil.rmtree(labels_dir) + + subjects: list[Subject] = [] + if images_dir.is_dir(): + for i in trange(num_images): + image_path = images_dir / f'image_{i}{suffix}' + label_path = labels_dir / f'label_{i}{suffix}' + subject = Subject( + one_modality=ScalarImage(image_path), + segmentation=LabelMap(label_path), + ) + subjects.append(subject) + else: + images_dir.mkdir(exist_ok=True, parents=True) + labels_dir.mkdir(exist_ok=True, parents=True) + iterable: Iterable[int] + if verbose: + print('Creating dummy dataset...') # noqa: T201 + iterable = trange(num_images) + else: + iterable = range(num_images) + for i in iterable: + shape = np.random.randint(*size_range, size=3) + affine = np.eye(4) + image = np.random.rand(*shape) + label = np.ones_like(image) + label[image < 0.33] = 0 + label[image > 0.66] = 2 + image *= 255 + + image_path = images_dir / f'image_{i}{suffix}' + nii = Nifti1Image(image.astype(np.uint8), affine) + nii.to_filename(str(image_path)) + + label_path = labels_dir / f'label_{i}{suffix}' + nii = Nifti1Image(label.astype(np.uint8), affine) + nii.to_filename(str(label_path)) + + subject = Subject( + one_modality=ScalarImage(image_path), + segmentation=LabelMap(label_path), + ) + subjects.append(subject) + return subjects + + +def apply_transform_to_file( + input_path: TypePath, + transform, # : Transform seems to create a circular import + output_path: TypePath, + class_: str = 'ScalarImage', + verbose: bool = False, +): + from . import data + + image = getattr(data, class_)(input_path) + subject = data.Subject(image=image) + transformed = transform(subject) + transformed.image.save(output_path) + if verbose and transformed.history: + print('Applied transform:', transformed.history[0]) # noqa: T201 + + +def guess_type(string: str) -> Any: + # Adapted from + # https://www.reddit.com/r/learnpython/comments/4599hl/module_to_guess_type_from_a_string/czw3f5s + string = string.replace(' ', '') + result_type: Any + try: + value = ast.literal_eval(string) + except ValueError: + result_type = str + else: + result_type = type(value) + if result_type in (list, tuple): + string = string[1:-1] # remove brackets + split = string.split(',') + list_result = [guess_type(n) for n in split] + value = tuple(list_result) if result_type is tuple else list_result + return value + try: + value = result_type(string) + except TypeError: + value = None + return value + + +def get_torchio_cache_dir() -> Path: + return Path('~/.cache/torchio').expanduser() + + +def compress( + input_path: TypePath, + output_path: TypePath | None = None, +) -> Path: + if output_path is None: + output_path = Path(input_path).with_suffix('.nii.gz') + with open(input_path, 'rb') as f_in: + with gzip.open(output_path, 'wb') as f_out: + shutil.copyfileobj(f_in, f_out) + return Path(output_path) + + +def check_sequence(sequence: Sequence, name: str) -> None: + try: + iter(sequence) + except TypeError as err: + message = f'"{name}" must be a sequence, not {type(name)}' + raise TypeError(message) from err + + +def get_major_sitk_version() -> int: + # This attribute was added in version 2 + # https://github.com/SimpleITK/SimpleITK/pull/1171 + version = getattr(sitk, '__version__', None) + major_version = 1 if version is None else 2 + return major_version + + +def history_collate(batch: Sequence, collate_transforms=True) -> dict: + attr = constants.HISTORY if collate_transforms else 'applied_transforms' + # Adapted from + # https://github.com/romainVala/torchQC/blob/master/segmentation/collate_functions.py + from .data import Subject + + first_element = batch[0] + if isinstance(first_element, Subject): + dictionary = { + key: default_collate([d[key] for d in batch]) for key in first_element + } + if hasattr(first_element, attr): + dictionary.update({attr: [getattr(d, attr) for d in batch]}) + else: + dictionary = {} + return dictionary + + +def get_subclasses(target_class: type) -> list[type]: + subclasses = target_class.__subclasses__() + subclasses += sum((get_subclasses(cls) for cls in subclasses), []) + return subclasses + + +def get_first_item(data_loader: DataLoader): + return next(iter(data_loader)) + + +def get_batch_images_and_size(batch: dict) -> tuple[list[str], int]: + """Get number of images and images names in a batch. + + Args: + batch: Dictionary generated by a [`SubjectsLoader`][torchio.SubjectsLoader] + extracting data from a [`torchio.SubjectsDataset`][torchio.SubjectsDataset]. + + Raises: + RuntimeError: If the batch does not seem to contain any dictionaries + that seem to represent a [`torchio.Image`][torchio.Image]. + """ + names = [] + for key, value in batch.items(): + if isinstance(value, dict) and constants.DATA in value: + size = len(value[constants.DATA]) + names.append(key) + if not names: + raise RuntimeError('The batch does not seem to contain any images') + return names, size + + +def get_subjects_from_batch(batch: dict) -> list: + """Get list of subjects from collated batch. + + Args: + batch: Dictionary generated by a [`SubjectsLoader`][torchio.SubjectsLoader] + extracting data from a [`torchio.SubjectsDataset`][torchio.SubjectsDataset]. + """ + from .data import LabelMap + from .data import ScalarImage + from .data import Subject + + subjects = [] + image_names, batch_size = get_batch_images_and_size(batch) + + for i in range(batch_size): + subject_dict = {} + + for key, value in batch.items(): + if key in image_names: + image_name = key + image_dict = value + data = image_dict[constants.DATA][i] + affine = image_dict[constants.AFFINE][i] + path = Path(image_dict[constants.PATH][i]) + is_label = image_dict[constants.TYPE][i] == constants.LABEL + klass = LabelMap if is_label else ScalarImage + image = klass(tensor=data, affine=affine, filename=path.name) + subject_dict[image_name] = image + else: + instance_value = value[i] + subject_dict[key] = instance_value + + subject = Subject(subject_dict) + + if constants.HISTORY in batch: + applied_transforms = batch[constants.HISTORY][i] + for transform in applied_transforms: + transform._add_transform_to_subject_history(subject) + + subjects.append(subject) + return subjects + + +def add_images_from_batch( + subjects: list, + tensor: torch.Tensor, + class_=None, + name='prediction', +) -> None: + """Add images to subjects in a list, typically from a network prediction. + + The spatial metadata (affine matrices) will be extracted from one of the + images of each subject. + + Args: + subjects: List of instances of [`torchio.Subject`][torchio.Subject] to which images + will be added. + tensor: PyTorch tensor of shape $(B, C, W, H, D)$, where + $B$ is the batch size. + class_: Class used to instantiate the images, + e.g., [`torchio.LabelMap`][torchio.LabelMap]. + If `None`, [`torchio.ScalarImage`][torchio.ScalarImage] will be used. + name: Name of the images added to the subjects. + """ + if class_ is None: + from . import ScalarImage + + class_ = ScalarImage + for subject, data in zip(subjects, tensor, strict=True): + one_image = subject.get_first_image() + kwargs = {'tensor': data, 'affine': one_image.affine} + if 'filename' in one_image: + kwargs['filename'] = one_image['filename'] + image = class_(**kwargs) + subject.add_image(image, name) + + +def guess_external_viewer() -> Path | None: + """Guess the path to an executable that could be used to visualize images. + + Currently, it looks for 1) ITK-SNAP and 2) 3D Slicer. Implemented + for macOS and Windows. + """ + if 'SITK_SHOW_COMMAND' in os.environ: + return Path(os.environ['SITK_SHOW_COMMAND']) + platform = sys.platform + itk = 'ITK-SNAP' + slicer = 'Slicer' + if platform == 'darwin': + app_path = '/Applications/{}.app/Contents/MacOS/{}' + itk_snap_path = Path(app_path.format(2 * (itk,))) + if itk_snap_path.is_file(): + return itk_snap_path + slicer_path = Path(app_path.format(2 * (slicer,))) + if slicer_path.is_file(): + return slicer_path + elif platform == 'win32': + program_files_dir = Path(os.environ['ProgramW6432']) + itk_snap_dirs = list(program_files_dir.glob('ITK-SNAP*')) + if itk_snap_dirs: + itk_snap_dir = itk_snap_dirs[-1] + itk_snap_path = itk_snap_dir / 'bin/itk-snap.exe' + if itk_snap_path.is_file(): + return itk_snap_path + slicer_dirs = list(program_files_dir.glob('Slicer*')) + if slicer_dirs: + slicer_dir = slicer_dirs[-1] + slicer_path = slicer_dir / 'slicer.exe' + if slicer_path.is_file(): + return slicer_path + elif 'linux' in platform: + itk_snap_which = shutil.which('itksnap') + if itk_snap_which is not None: + return Path(itk_snap_which) + slicer_which = shutil.which('Slicer') + if slicer_which is not None: + return Path(slicer_which) + return None # for mypy + + +def parse_spatial_shape(shape): + result = to_tuple(shape, length=3) + for n in result: + if n < 1 or n % 1: + message = ( + 'All elements in a spatial shape must be positive integers,' + f' but the following shape was passed: {shape}' + ) + raise ValueError(message) + if len(result) != 3: + message = ( + 'Spatial shapes must have 3 elements, but the following shape' + f' was passed: {shape}' + ) + raise ValueError(message) + return result + + +def normalize_path(path: TypePath): + return Path(path).expanduser().resolve() + + +def is_iterable(object: Any) -> bool: + try: + iter(object) + return True + except TypeError: + return False diff --git a/src/torchio/visualization.py b/src/torchio/visualization.py index 4998b0253..e7906fa97 100644 --- a/src/torchio/visualization.py +++ b/src/torchio/visualization.py @@ -1,998 +1,357 @@ -"""Visualization utilities for TorchIO images. - -Requires the `[plot]` extras: `pip install torchio[plot]`. -""" - from __future__ import annotations import warnings -from collections.abc import Sequence -from importlib import import_module +from itertools import cycle from pathlib import Path from typing import TYPE_CHECKING from typing import Any -from typing import Literal -from typing import cast -from typing import overload import numpy as np import torch +from einops import rearrange +from .data.image import Image from .data.image import LabelMap -from .external.imports import get_colorcet +from .data.image import ScalarImage +from .data.subject import Subject from .external.imports import get_ffmpeg -from .external.imports import get_pillow -from .transforms import Normalize -from .transforms import Reorient +from .transforms.preprocessing.intensity.rescale import RescaleIntensity +from .transforms.preprocessing.intensity.to import To +from .transforms.preprocessing.spatial.ensure_shape_multiple import EnsureShapeMultiple +from .transforms.preprocessing.spatial.resample import Resample +from .transforms.preprocessing.spatial.to_canonical import ToCanonical +from .transforms.preprocessing.spatial.to_orientation import ToOrientation +from .types import TypePath if TYPE_CHECKING: + from collections.abc import Sequence + from matplotlib.axes import Axes - from matplotlib.colors import Colormap + from matplotlib.colors import BoundaryNorm + from matplotlib.colors import ListedColormap from matplotlib.figure import Figure - from .data.image import Image - from .data.subject import Subject - from .types import TypePath - -# Opposite anatomical direction for each orientation code. -_OPPOSITE: dict[str, str] = { - "R": "L", - "L": "R", - "A": "P", - "P": "A", - "S": "I", - "I": "S", -} - -_FULL_NAME: dict[str, str] = { - "R": "Right", - "L": "Left", - "A": "Anterior", - "P": "Posterior", - "S": "Superior", - "I": "Inferior", -} - -# Each view is defined by: -# (name, slice_pair, x_pair, y_pair, x_positive_on_left, y_positive_on_top) -# "pair" means the L/R, A/P, or S/I axis pair. -# x_positive_on_left / y_positive_on_top: the code that should appear on the -# left side (x) or top (y) of the display. -_VIEWS: list[tuple[str, str, str, str, str, str]] = [ - ("Sagittal", "LR", "AP", "SI", "A", "S"), - ("Coronal", "AP", "LR", "SI", "R", "S"), - ("Axial", "SI", "LR", "AP", "R", "A"), -] - -# Intersection line colors (from 3D Slicer, via v1). -# Each color identifies the slice position being shown. -_COLOR_SAGITTAL = "#42A5F5" # blue -_COLOR_CORONAL = "#8FE561" # green -_COLOR_AXIAL = "#FF8372" # red -# Map view name to its intersection color. -_VIEW_COLOR: dict[str, str] = { - "Sagittal": _COLOR_SAGITTAL, - "Coronal": _COLOR_CORONAL, - "Axial": _COLOR_AXIAL, -} -_CODE_TO_PAIR: dict[str, str] = { - "R": "LR", - "L": "LR", - "A": "AP", - "P": "AP", - "S": "SI", - "I": "SI", -} - - -def _get_mpl(): - """Lazy-import matplotlib, raising a helpful error if missing.""" +def import_mpl_plt(): try: - import matplotlib + import matplotlib as mpl import matplotlib.pyplot as plt - except ImportError: - msg = ( - "matplotlib is required for plotting. " - "Install it with: pip install torchio[plot]" - ) - raise ImportError(msg) from None - return matplotlib, plt + except ImportError as e: + raise ImportError('Install matplotlib for plotting support') from e + return mpl, plt -def _display_figure(fig: Figure) -> None: - """Display a figure inline (notebooks) or interactively (scripts).""" - try: - from IPython.display import display +def _figure_to_html(fig: Figure) -> str: + """Convert a matplotlib Figure to an HTML img tag with base64-encoded PNG.""" + import base64 + import io - display(fig) - except ImportError: - import matplotlib.pyplot as plt + from matplotlib import pyplot as plt - plt.show() - plt.close(fig) + buffer = io.BytesIO() + fig.savefig(buffer, format='png', bbox_inches='tight') + plt.close(fig) + buffer.seek(0) + img_str = base64.b64encode(buffer.read()).decode('utf-8') + return f'' -def _get_categorical_cmap( - slices_2d: list[np.ndarray], - cmap_name: str = "glasbey_category10", -) -> Any: - """Build a categorical ListedColormap for label maps. +def rotate(image: np.ndarray, *, radiological: bool = True, n: int = -1) -> np.ndarray: + # Rotate for visualization purposes + image = np.rot90(image, n, axes=(0, 1)) + if radiological: + image = np.fliplr(image) + return image - Uses colorcet if available, otherwise falls back to matplotlib's - `tab10` colors. - """ - from itertools import cycle - mpl, _ = _get_mpl() - num_classes = max(int(s.max()) for s in slices_2d) - colors: list[tuple[float, ...]] = [ - (0.0, 0.0, 0.0), # black for background - (1.0, 1.0, 1.0), # white for class 1 +def _create_categorical_colormap( + data: torch.Tensor, + cmap_name: str = 'glasbey_category10', +) -> tuple[ListedColormap, BoundaryNorm]: + num_classes = int(data.max()) + mpl, _ = import_mpl_plt() + + colors = [ + (0, 0, 0), # black for background + (1, 1, 1), # white for class 1 ] if num_classes > 1: - if (cc := get_colorcet()) is not None: - cc_cmap = getattr(cc.cm, cmap_name) - color_cycle = cycle(cc_cmap.colors) - else: - tab10 = mpl.colormaps["tab10"] - color_cycle = cycle(tab10.colors) - colors.extend(next(color_cycle) for _ in range(num_classes - 1)) + from .external.imports import get_colorcet + + colorcet = get_colorcet() + cmap = getattr(colorcet.cm, cmap_name) + color_cycle = cycle(cmap.colors) + distinct_colors = [next(color_cycle) for _ in range(num_classes - 1)] + colors.extend(distinct_colors) boundaries = np.arange(-0.5, num_classes + 1.5, 1) colormap = mpl.colors.ListedColormap(colors) - norm = mpl.colors.BoundaryNorm(boundaries, ncolors=colormap.N) - return colormap, norm + boundary_norm = mpl.colors.BoundaryNorm(boundaries, ncolors=colormap.N) + return colormap, boundary_norm -def _find_axis(orientation: tuple[str, str, str], pair: str) -> int: - """Find which tensor axis (0, 1, 2) corresponds to an anatomical pair.""" - for i, code in enumerate(orientation): - if _CODE_TO_PAIR[code] == pair: - return i - msg = f"No axis found for pair {pair!r} in orientation {orientation}" - raise ValueError(msg) - - -def _extract_slices( +def plot_volume( image: Image, - channel: int, - resolved: tuple[int, ...], - axis_for: dict[str, int], -) -> list[np.ndarray]: - """Extract oriented 2D slices for each anatomical view.""" - orientation = image.orientation - slices_2d: list[np.ndarray] = [] - for _view_name, slice_pair, x_pair, y_pair, x_left, y_top in _VIEWS: - slice_axis = axis_for[slice_pair] - x_axis = axis_for[x_pair] - y_axis = axis_for[y_pair] - - sl: list[slice | int] = [slice(None), slice(None), slice(None), slice(None)] - sl[0] = channel - sl[slice_axis + 1] = resolved[slice_axis] - plane = image[tuple(sl)] - data_2d = plane.data.squeeze().cpu().numpy() - - if x_axis < y_axis: - data_2d = data_2d.T + radiological=True, + channel=None, + axes: Sequence[Axes] | None = None, + cmap=None, + output_path=None, + show=True, + xlabels=True, + percentiles: tuple[float, float] = (0, 100), + figsize=None, + title=None, + reorient=True, + indices=None, + rgb=True, + savefig_kwargs: dict[str, Any] | None = None, + **imshow_kwargs, +) -> Figure | None: + _, plt = import_mpl_plt() + fig: Figure | None = None + if axes is None: + fig, generated_axes = plt.subplots(1, 3, figsize=figsize) + axes = tuple(np.ravel(generated_axes).tolist()) - x_code = orientation[x_axis] - y_code = orientation[y_axis] - if x_code == x_left: - data_2d = np.flip(data_2d, axis=1) - if y_code != y_top: - data_2d = np.flip(data_2d, axis=0) + if reorient: + image = ToCanonical()(image) - slices_2d.append(np.ascontiguousarray(data_2d)) - return slices_2d + is_label = isinstance(image, LabelMap) + if is_label: # probabilistic label map + data = image.data[np.newaxis, -1] + elif rgb and image.num_channels == 3: + data = image.data # keep image as it is + elif channel is None: + if image.num_channels > 1: + message = ( + 'Multiple channels found in the image. ' + 'Plotting the first channel (0). ' + 'To plot a different channel, please specify the channel ' + 'index using the "channel" argument.' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + data = image.data[0:1] # just use the first channel + else: + data = image.data[np.newaxis, channel] + data = rearrange(data, 'c x y z -> x y z c') + data_numpy: np.ndarray = data.cpu().numpy() + if indices is None: + indices = np.array(data_numpy.shape[:3]) // 2 + i, j, k = indices + slice_x = rotate(data_numpy[i, :, :], radiological=radiological) + slice_y = rotate(data_numpy[:, j, :], radiological=radiological) + slice_z = rotate(data_numpy[:, :, k], radiological=radiological) -def _resolve_label_colors( - image: Image, - cmap: str | Colormap | dict[int, tuple[int, int, int]] | None, - slices_2d: list[np.ndarray], - kw: dict[str, Any], -) -> tuple[list[np.ndarray], bool]: - """If cmap is a color dict (or image carries one), colorize slices.""" - color_map: dict[int, tuple[int, int, int]] | None = None if isinstance(cmap, dict): - color_map = cast("dict[int, tuple[int, int, int]]", cmap) - elif cmap is None and hasattr(image, "color_map"): - meta = image.color_map - if isinstance(meta, dict): - color_map = cast("dict[int, tuple[int, int, int]]", meta) - - if color_map is not None: - slices_2d = _colorize_labels(slices_2d, color_map) - kw["origin"] = "lower" - kw.setdefault("interpolation", "none") - return slices_2d, True - return slices_2d, False - - -def _build_imshow_kwargs( - image: Image, - slices_2d: list[np.ndarray], - cmap: str | Colormap | dict[int, tuple[int, int, int]] | None, - percentiles: tuple[float, float], - imshow_kwargs: dict[str, Any], -) -> tuple[dict[str, Any], list[np.ndarray]]: - """Prepare the keyword arguments for `ax.imshow()`. - - Returns the kwargs dict and (possibly RGB-converted) slices. - """ - kw = dict(imshow_kwargs) - is_label = isinstance(image, LabelMap) - - slices_2d, colorized = _resolve_label_colors(image, cmap, slices_2d, kw) - if colorized: - return kw, slices_2d + slices = slice_x, slice_y, slice_z + slice_x, slice_y, slice_z = color_labels(slices, cmap) + else: + boundary_norm = None + if cmap is None: + if is_label: + cmap, boundary_norm = _create_categorical_colormap(data) + else: + cmap = 'gray' + imshow_kwargs['cmap'] = cmap + imshow_kwargs['norm'] = boundary_norm - if cmap is None: - if is_label: - cmap, norm = _get_categorical_cmap(slices_2d) - kw.setdefault("norm", norm) - else: - cmap = "gray" - kw.setdefault("cmap", cmap) - kw["origin"] = "lower" if is_label: - kw.setdefault("interpolation", "none") + imshow_kwargs['interpolation'] = 'none' else: - kw.setdefault("interpolation", "bilinear") - - if not is_label: - all_values = np.concatenate([s.ravel() for s in slices_2d]) - vmin, vmax = np.percentile(all_values, percentiles) - kw.setdefault("vmin", vmin) - kw.setdefault("vmax", vmax) - return kw, slices_2d - - -def _colorize_labels( - slices_2d: list[np.ndarray], - color_map: dict[int, tuple[int, int, int]], -) -> list[np.ndarray]: - """Convert label slices to RGB using a label-to-color mapping.""" - result: list[np.ndarray] = [] - for label_slice in slices_2d: - h, w = label_slice.shape[:2] - rgb = np.zeros((h, w, 3), dtype=np.uint8) - for label, color in color_map.items(): - rgb[label_slice == label] = color - result.append(rgb) - return result - - -def _plot_image_on_axes( - image: Image, - plot_axes: Sequence[Axes], - *, - channel: int = 0, - resolved: tuple[int, ...], - cmap: str | Colormap | dict[int, tuple[int, int, int]] | None = None, - percentiles: tuple[float, float] = (0.5, 99.5), - voxels: bool = False, - intersections: bool = True, - show_titles: bool = True, - **imshow_kwargs: Any, -) -> None: - """Plot 3 orthogonal views of a single image onto pre-created axes.""" - spatial_shape = image.spatial_shape - spacing = image.spacing - orientation = image.orientation - origin = image.origin - - axis_for: dict[str, int] = {} - for pair in ("LR", "AP", "SI"): - axis_for[pair] = _find_axis(orientation, pair) - - slices_2d = _extract_slices(image, channel, resolved, axis_for) - kw, slices_2d = _build_imshow_kwargs( - image, - slices_2d, - cmap, - percentiles, - imshow_kwargs, - ) - - for view_idx, (view_name, slice_pair, x_pair, y_pair, x_left, y_top) in enumerate( - _VIEWS, - ): - ax = plot_axes[view_idx] - - slice_axis = axis_for[slice_pair] - x_axis = axis_for[x_pair] - y_axis = axis_for[y_pair] + if 'interpolation' not in imshow_kwargs: + imshow_kwargs['interpolation'] = 'bicubic' - aspect = spacing[y_axis] / spacing[x_axis] - ax.imshow(slices_2d[view_idx], aspect=aspect, **kw) + imshow_kwargs['origin'] = 'lower' - if voxels: - x_label = f"{_axis_name(x_axis)} ({x_left} ↔ {_OPPOSITE[x_left]})" - y_label = f"{_axis_name(y_axis)} ({_OPPOSITE[y_top]} ↔ {y_top})" - else: - x_label = f"{_FULL_NAME[x_left]} [mm] ({_axis_name(x_axis)})" - y_label = f"{_FULL_NAME[y_top]} [mm] ({_axis_name(y_axis)})" - ax.set_xlabel(x_label) - ax.set_ylabel(y_label) - - _set_ticks( - ax, - x_axis=x_axis, - y_axis=y_axis, - x_code=orientation[x_axis], - y_code=orientation[y_axis], - x_left=x_left, - y_top=y_top, - spacing=spacing, - origin_mm=origin, - spatial_shape=spatial_shape, - voxels=voxels, - ) - - if show_titles: - ax.set_title(f"{view_name} [{resolved[slice_axis]}]") - - if intersections: - _draw_intersections( - plot_axes, - axis_for=axis_for, - orientation=orientation, - spatial_shape=spatial_shape, - resolved=resolved, + if not is_label: + displayed_data = np.concatenate( + [ + slice_x.flatten(), + slice_y.flatten(), + slice_z.flatten(), + ] ) - - -@overload -def plot_image( - image: Image, - *, - show: Literal[False], - channel: int = ..., - indices: tuple[int | None, int | None, int | None] | None = ..., - coordinates: tuple[float | None, float | None, float | None] | None = ..., - axes: Sequence[Axes] | None = ..., - cmap: str | Colormap | dict[int, tuple[int, int, int]] | None = ..., - percentiles: tuple[float, float] = ..., - figsize: tuple[float, float] | None = ..., - title: str | None = ..., - output_path: TypePath | None = ..., - savefig_kwargs: dict[str, Any] | None = ..., - voxels: bool = ..., - figsize_multiplier: float = ..., - intersections: bool = ..., - **imshow_kwargs: Any, -) -> Figure: ... - - -@overload -def plot_image( - image: Image, - *, - show: Literal[True] = ..., - channel: int = ..., - indices: tuple[int | None, int | None, int | None] | None = ..., - coordinates: tuple[float | None, float | None, float | None] | None = ..., - axes: Sequence[Axes] | None = ..., - cmap: str | Colormap | dict[int, tuple[int, int, int]] | None = ..., - percentiles: tuple[float, float] = ..., - figsize: tuple[float, float] | None = ..., - title: str | None = ..., - output_path: TypePath | None = ..., - savefig_kwargs: dict[str, Any] | None = ..., - voxels: bool = ..., - figsize_multiplier: float = ..., - intersections: bool = ..., - **imshow_kwargs: Any, -) -> None: ... - - -def plot_image( - image: Image, - *, - channel: int = 0, - indices: tuple[int | None, int | None, int | None] | None = None, - coordinates: tuple[float | None, float | None, float | None] | None = None, - axes: Sequence[Axes] | None = None, - cmap: str | Colormap | dict[int, tuple[int, int, int]] | None = None, - percentiles: tuple[float, float] = (0.5, 99.5), - figsize: tuple[float, float] | None = None, - title: str | None = None, - output_path: TypePath | None = None, - show: bool = True, - savefig_kwargs: dict[str, Any] | None = None, - voxels: bool = False, - figsize_multiplier: float = 2.0, - intersections: bool = True, - **imshow_kwargs: Any, -) -> Figure | None: - """Plot 3 orthogonal slices of a 3D image. - - Always displays Sagittal, Coronal, Axial with fixed anatomical - positions regardless of image orientation. Data is flipped and - transposed as needed. Uses lazy `Image.__getitem__` so only - the 3 requested planes are read from disk. - - Args: - image: The image to plot. - channel: Which channel to display. - indices: Slice index for each spatial axis. `None` entries - default to the mid-slice. Pass `None` for all mid-slices. - Mutually exclusive with `coordinates`. - coordinates: World coordinates in mm for each slice. - `None` entries default to the mid-slice. Converted to - the nearest voxel index via the inverse affine. Mutually - exclusive with `indices`. - axes: Pre-created sequence of 3 matplotlib `Axes`. If - `None`, a new figure with correct proportions is created. - cmap: Colormap. Defaults to `'gray'` for intensity images. - percentiles: Intensity percentile range for display windowing. - Ignored for label maps. - figsize: Figure size in inches `(width, height)`. - title: Figure super-title. - output_path: Save figure to this path. - show: Call `plt.show()` after plotting. - savefig_kwargs: Extra keyword arguments for `fig.savefig()`. - voxels: Show voxel indices on ticks instead of world - coordinates in mm. - figsize_multiplier: Scale factor applied to the default - `rcParams["figure.figsize"]` when `figsize` is `None`. - intersections: Draw coloured cross-hair lines showing where - the other two slices intersect each view. - **imshow_kwargs: Forwarded to `ax.imshow()`. - - Returns: - The matplotlib `Figure`, or `None` when `show=True` - (the figure is displayed and closed to prevent duplicate - rendering in notebooks). - """ - mpl, plt = _get_mpl() - - resolved = _resolve_indices(image, indices, coordinates) - - # Read spatial metadata from headers (no data load) - spatial_shape = image.spatial_shape - spacing = image.spacing - orientation = image.orientation - - # Find tensor axis for each anatomical pair - axis_for: dict[str, int] = {} - for pair in ("LR", "AP", "SI"): - axis_for[pair] = _find_axis(orientation, pair) - - # Compute physical extents for proportional subplot sizing - lr_mm = spatial_shape[axis_for["LR"]] * spacing[axis_for["LR"]] - ap_mm = spatial_shape[axis_for["AP"]] * spacing[axis_for["AP"]] - width_ratios = [ap_mm, lr_mm, lr_mm] - - # Create figure if needed - fig: Figure - if axes is None: - if figsize is None: - default_w, default_h = plt.rcParams["figure.figsize"] - figsize = ( - default_w * figsize_multiplier, - default_h * figsize_multiplier, - ) - gs = mpl.gridspec.GridSpec(1, 3, width_ratios=width_ratios) - fig = plt.figure(figsize=figsize) - plt.close(fig) - plot_axes: Sequence[Axes] = [fig.add_subplot(gs[0, i]) for i in range(3)] - else: - if len(axes) < 3: - msg = f"Expected 3 axes, got {len(axes)}" - raise ValueError(msg) - plot_axes = axes - fig = cast("Figure", plot_axes[0].get_figure()) - - _plot_image_on_axes( - image=image, - plot_axes=plot_axes, - channel=channel, - resolved=resolved, - cmap=cmap, - percentiles=percentiles, - voxels=voxels, - intersections=intersections, - **imshow_kwargs, + p1, p2 = np.percentile(displayed_data, percentiles) + if 'vmin' not in imshow_kwargs: + imshow_kwargs['vmin'] = p1 + if 'vmax' not in imshow_kwargs: + imshow_kwargs['vmax'] = p2 + + spacing_r, spacing_a, spacing_s = image.spacing + sag_axis, cor_axis, axi_axis = axes + slices_data = ( + ('Sagittal', sag_axis, slice_x, spacing_s / spacing_a, 'A', 'S'), + ('Coronal', cor_axis, slice_y, spacing_s / spacing_r, 'R', 'S'), + ('Axial', axi_axis, slice_z, spacing_a / spacing_r, 'R', 'A'), ) + for axis_title, axis, axis_slice, aspect, xlabel, ylabel in slices_data: + axis.imshow(axis_slice, aspect=aspect, **imshow_kwargs) + if xlabels: + axis.set_xlabel(xlabel) + axis.set_ylabel(ylabel) + axis.invert_xaxis() + axis.set_title(axis_title) + + plt.tight_layout() if title is not None: - fig.suptitle(title) - fig.tight_layout() + plt.suptitle(title) - if output_path is not None: - fig.savefig(output_path, **(savefig_kwargs or {})) + if output_path is not None and fig is not None: + if savefig_kwargs is None: + savefig_kwargs = {} + fig.savefig(output_path, **savefig_kwargs) if show: - _display_figure(fig) - return None - + plt.show() return fig -def _coordinates_to_indices( - image: Image, - coordinates: tuple[float | None, float | None, float | None], -) -> tuple[int | None, int | None, int | None]: - """Convert world coordinates (mm) to voxel indices.""" - inv_affine = image.affine.inverse() - voxel_coords = inv_affine.apply( - torch.tensor( - [[c if c is not None else float("nan") for c in coordinates]], - dtype=torch.float64, - ), - )[0] - c0, c1, c2 = coordinates - return ( - None if c0 is None else round(float(voxel_coords[0])), - None if c1 is None else round(float(voxel_coords[1])), - None if c2 is None else round(float(voxel_coords[2])), - ) - - -def _resolve_indices( - image: Image, - indices: tuple[int | None, int | None, int | None] | None, - coordinates: tuple[float | None, float | None, float | None] | None, -) -> tuple[int, ...]: - """Resolve indices/coordinates to concrete voxel indices.""" - if indices is not None and coordinates is not None: - msg = "indices and coordinates are mutually exclusive" - raise ValueError(msg) - - if coordinates is not None: - indices = _coordinates_to_indices(image, coordinates) - - if indices is None: - indices = (None, None, None) - return tuple( - s // 2 if idx is None else idx - for idx, s in zip(indices, image.spatial_shape, strict=True) - ) - - -@overload def plot_subject( subject: Subject, - *, - show: Literal[False], - channel: int = ..., - indices: tuple[int | None, int | None, int | None] | None = ..., - coordinates: tuple[float | None, float | None, float | None] | None = ..., - cmap_dict: dict[str, Any] | None = ..., - percentiles: tuple[float, float] = ..., - figsize: tuple[float, float] | None = ..., - title: str | None = ..., - output_path: TypePath | None = ..., - savefig_kwargs: dict[str, Any] | None = ..., - voxels: bool = ..., - figsize_multiplier: float = ..., - intersections: bool = ..., - **imshow_kwargs: Any, -) -> Figure: ... - - -@overload -def plot_subject( - subject: Subject, - *, - show: Literal[True] = ..., - channel: int = ..., - indices: tuple[int | None, int | None, int | None] | None = ..., - coordinates: tuple[float | None, float | None, float | None] | None = ..., - cmap_dict: dict[str, Any] | None = ..., - percentiles: tuple[float, float] = ..., - figsize: tuple[float, float] | None = ..., - title: str | None = ..., - output_path: TypePath | None = ..., - savefig_kwargs: dict[str, Any] | None = ..., - voxels: bool = ..., - figsize_multiplier: float = ..., - intersections: bool = ..., - **imshow_kwargs: Any, -) -> None: ... - - -def plot_subject( - subject: Subject, - *, - channel: int = 0, - indices: tuple[int | None, int | None, int | None] | None = None, - coordinates: tuple[float | None, float | None, float | None] | None = None, - cmap_dict: dict[str, Any] | None = None, - percentiles: tuple[float, float] = (0.5, 99.5), - figsize: tuple[float, float] | None = None, - title: str | None = None, - output_path: TypePath | None = None, - show: bool = True, + cmap_dict=None, + show=True, + output_path=None, + figsize=None, + clear_axes=True, savefig_kwargs: dict[str, Any] | None = None, - voxels: bool = False, - figsize_multiplier: float = 2.0, - intersections: bool = True, - **imshow_kwargs: Any, -) -> Figure | None: - """Plot all images in a subject as a grid. - - Each image gets a row (or column if >3 images) of Sagittal, - Coronal, Axial views. LabelMaps are automatically detected and - use categorical colormaps. - - Args: - subject: The subject to plot. - channel: Which channel to display. - indices: Voxel indices for each slice. Mutually exclusive - with `coordinates`. - coordinates: World coordinates in mm. Mutually exclusive - with `indices`. - cmap_dict: Per-image colormap overrides, keyed by image name. - percentiles: Intensity percentile range for windowing. - figsize: Figure size in inches. - title: Figure super-title. - output_path: Save figure to this path. - show: Call `plt.show()` after plotting. - savefig_kwargs: Extra keyword arguments for `fig.savefig()`. - voxels: Show voxel ticks instead of world coordinates. - figsize_multiplier: Scale factor for default figure size. - intersections: Draw slice intersection cross-hairs. - **imshow_kwargs: Forwarded to `ax.imshow()`. - - Returns: - The `Figure`, or `None` when `show=True`. - """ - mpl, plt = _get_mpl() - - images = subject.images - num_images = len(images) - if num_images == 0: - msg = "Subject has no images to plot" - raise ValueError(msg) - - first_image = next(iter(images.values())) - _resolve_indices(first_image, indices, coordinates) - - many = num_images > 3 - fig, all_axes = _create_subject_grid( - first_image, - num_images, - many, - figsize, - figsize_multiplier, - mpl, - plt, - ) - - _populate_subject_grid( - images, - all_axes, - many, - indices, - coordinates, - channel=channel, - cmap_dict=cmap_dict, - percentiles=percentiles, - voxels=voxels, - intersections=intersections, - **imshow_kwargs, - ) - - if title is not None: - fig.suptitle(title) - fig.tight_layout() - + **plot_volume_kwargs, +) -> Figure: + _, plt = import_mpl_plt() + num_images = len(subject) + many_images = num_images > 2 + subplots_kwargs = {'figsize': figsize} + try: + if clear_axes: + subject.check_consistent_spatial_shape() + subplots_kwargs['sharex'] = 'row' if many_images else 'col' + subplots_kwargs['sharey'] = 'row' if many_images else 'col' + except RuntimeError: # different shapes in subject + pass + args = (3, num_images) if many_images else (num_images, 3) + fig, axes = plt.subplots(*args, **subplots_kwargs) + # The array of axes must be 2D so that it can be indexed correctly within + # the plot_volume() function + axes = axes.T if many_images else axes.reshape(-1, 3) + iterable = enumerate(subject.get_images_dict(intensity_only=False).items()) + axes_names = 'sagittal', 'coronal', 'axial' + for image_index, (name, image) in iterable: + image_axes = axes[image_index] + cmap = None + if cmap_dict is not None and name in cmap_dict: + cmap = cmap_dict[name] + last_row = image_index == len(axes) - 1 + plot_volume( + image, + axes=image_axes, + show=False, + cmap=cmap, + xlabels=last_row, + **plot_volume_kwargs, + ) + for axis, axis_name in zip(image_axes, axes_names, strict=True): + axis.set_title(f'{name} ({axis_name})') + plt.tight_layout() if output_path is not None: - fig.savefig(output_path, **(savefig_kwargs or {})) + if savefig_kwargs is None: + savefig_kwargs = {} + fig.savefig(output_path, **savefig_kwargs) if show: - _display_figure(fig) - return None - + plt.show() return fig -def _create_subject_grid( - first_image: Image, - num_images: int, - many: bool, - figsize: tuple[float, float] | None, - figsize_multiplier: float, - mpl: Any, - plt: Any, -) -> tuple[Any, list[list[Any]]]: - """Create the figure and axes grid for `plot_subject`.""" - orientation = first_image.orientation - spacing = first_image.spacing - spatial_shape = first_image.spatial_shape - axis_for: dict[str, int] = {} - for pair in ("LR", "AP", "SI"): - axis_for[pair] = _find_axis(orientation, pair) - lr_mm = spatial_shape[axis_for["LR"]] * spacing[axis_for["LR"]] - ap_mm = spatial_shape[axis_for["AP"]] * spacing[axis_for["AP"]] - width_ratios = [ap_mm, lr_mm, lr_mm] - - if figsize is None: - default_w, default_h = plt.rcParams["figure.figsize"] - figsize = (default_w * figsize_multiplier, default_h * figsize_multiplier) - - n_views = 3 - if many: - nrows, ncols = n_views, num_images - gs = mpl.gridspec.GridSpec(nrows, ncols) - else: - nrows, ncols = num_images, n_views - gs = mpl.gridspec.GridSpec(nrows, ncols, width_ratios=width_ratios) - - fig = plt.figure(figsize=figsize) - plt.close(fig) - all_axes = [[fig.add_subplot(gs[r, c]) for c in range(ncols)] for r in range(nrows)] - return fig, all_axes - - -def _populate_subject_grid( - images: dict[str, Image], - all_axes: list[list[Any]], - many: bool, - indices: tuple[int | None, int | None, int | None] | None, - coordinates: tuple[float | None, float | None, float | None] | None, - *, - channel: int, - cmap_dict: dict[str, Any] | None, - percentiles: tuple[float, float], - voxels: bool, - intersections: bool, - **imshow_kwargs: Any, -) -> None: - """Plot each image into its row/column of the subject grid.""" - n_views = 3 - - for img_idx, (name, image) in enumerate(images.items()): - cmap = cmap_dict.get(name) if cmap_dict else None - img_resolved = _resolve_indices(image, indices, coordinates) - img_axes = _get_image_axes(all_axes, img_idx, n_views, many) - - _plot_image_on_axes( - image=image, - plot_axes=img_axes, - channel=channel, - resolved=img_resolved, - cmap=cmap, - percentiles=percentiles, - voxels=voxels, - intersections=intersections, - show_titles=False, - **imshow_kwargs, - ) - - _label_image_header(img_axes, name, many) - - -def _get_image_axes( - all_axes: list[list[Any]], - img_idx: int, - n_views: int, - many: bool, -) -> list[Any]: - """Get the 3 axes for a given image in the grid.""" - if many: - return [all_axes[v][img_idx] for v in range(n_views)] - return all_axes[img_idx] - +def get_num_bins(x: np.ndarray) -> int: + """Get the optimal number of bins for a histogram. -def _label_image_header( - img_axes: list[Any], - name: str, - many: bool, -) -> None: - """Add image name as a row/column header.""" - if many: - img_axes[0].set_title(name) - else: - # Prepend image name to the existing ylabel (orientation label) - existing = img_axes[0].get_ylabel() - img_axes[0].set_ylabel(f"{name}\n{existing}", fontsize=10) - - -def _axis_name(axis: int) -> str: - """Return the tensor axis name: i, j, or k.""" - return ("i", "j", "k")[axis] - - -def _draw_intersections( - plot_axes: Sequence[Axes], - *, - axis_for: dict[str, int], - orientation: tuple[str, str, str], - spatial_shape: tuple[int, int, int], - resolved: tuple[int, ...], -) -> None: - """Draw coloured cross-hair lines showing slice positions.""" - for view_idx, (view_name, _slice_pair, x_pair, y_pair, x_left, y_top) in enumerate( - _VIEWS, - ): - ax = plot_axes[view_idx] - x_axis = axis_for[x_pair] - y_axis = axis_for[y_pair] - x_size = spatial_shape[x_axis] - y_size = spatial_shape[y_axis] - x_code = orientation[x_axis] - y_code = orientation[y_axis] + This method uses the Freedman–Diaconis rule to compute the histogram that + minimizes "the integral of the squared difference between the histogram + (i.e., relative frequency density) and the density of the theoretical + probability distribution" ([Wikipedia ](https://en.wikipedia.org/wiki/Freedman%E2%80%93Diaconis_rule)). - for other_name, other_slice_pair, _, _, _, _ in _VIEWS: - if other_name == view_name: - continue - other_axis = axis_for[other_slice_pair] - other_pos = resolved[other_axis] - color = _VIEW_COLOR[other_name] - - if other_axis == x_axis: - display_x = _display_pos(other_pos, x_size, x_code == x_left) - ax.axvline(display_x, color=color, linewidth=0.8, alpha=0.8) - elif other_axis == y_axis: - display_y = _display_pos(other_pos, y_size, y_code != y_top) - ax.axhline(display_y, color=color, linewidth=0.8, alpha=0.8) - - -def _display_pos(voxel: int, size: int, flipped: bool) -> float: - """Convert a voxel index to display position, accounting for flips.""" - return float(size - 1 - voxel) if flipped else float(voxel) - - -def _set_ticks( - ax: Axes, - *, - x_axis: int, - y_axis: int, - x_code: str, - y_code: str, - x_left: str, - y_top: str, - spacing: tuple[float, float, float], - origin_mm: tuple[float, float, float], - spatial_shape: tuple[int, int, int], - voxels: bool, -) -> None: - """Set tick labels for a subplot.""" - x_size = spatial_shape[x_axis] - y_size = spatial_shape[y_axis] - x_flipped = x_code == x_left - y_flipped = y_code != y_top - - x_ticks = np.linspace(0, x_size - 1, min(5, x_size)) - y_ticks = np.linspace(0, y_size - 1, min(5, y_size)) - ax.set_xticks(x_ticks) - ax.set_yticks(y_ticks) - - if voxels: - ax.set_xticklabels(_voxel_tick_labels(x_ticks, x_size, x_flipped)) - ax.set_yticklabels(_voxel_tick_labels(y_ticks, y_size, y_flipped)) - else: - x_sp = spacing[x_axis] - y_sp = spacing[y_axis] - x_sign = -1.0 if x_code in ("L", "P", "I") else 1.0 - y_sign = -1.0 if y_code in ("L", "P", "I") else 1.0 - x_origin = origin_mm[_world_dim(x_code)] - y_origin = origin_mm[_world_dim(y_code)] - ax.set_xticklabels( - _mm_tick_labels(x_ticks, x_size, x_flipped, x_origin, x_sp, x_sign), - ) - ax.set_yticklabels( - _mm_tick_labels(y_ticks, y_size, y_flipped, y_origin, y_sp, y_sign), - ) - - -def _voxel_tick_labels( - ticks: np.ndarray, - size: int, - flipped: bool, -) -> list[str]: - """Generate voxel-index tick labels.""" - if flipped: - return [str(int(size - 1 - v)) for v in ticks] - return [str(int(v)) for v in ticks] - - -def _mm_tick_labels( - ticks: np.ndarray, - size: int, - flipped: bool, - origin: float, - sp: float, - sign: float, -) -> list[str]: - """Generate world-coordinate (mm) tick labels.""" - labels: list[str] = [] - for v in ticks: - voxel = (size - 1 - v) if flipped else v - mm = origin + voxel * sp * sign - labels.append(f"{mm:.0f}") - return labels - - -def _world_dim(code: str) -> int: - """Map an orientation code to the world coordinate dimension (0=x, 1=y, 2=z).""" - match code: - case "R" | "L": - return 0 - case "A" | "P": - return 1 - case _: - return 2 + Args: + x: Input values. + """ + # Freedman–Diaconis number of bins + q25, q75 = np.percentile(x, [25, 75]) + bin_width = 2 * (q75 - q25) * len(x) ** (-1 / 3) + bins = round((x.max() - x.min()) / bin_width) + return bins + + +def plot_histogram(x: np.ndarray, show=True, **kwargs) -> None: + _, plt = import_mpl_plt() + plt.hist(x, bins=get_num_bins(x), **kwargs) + plt.xlabel('Intensity') + density = kwargs.pop('density', False) + ylabel = 'Density' if density else 'Frequency' + plt.ylabel(ylabel) + if show: + plt.show() -# ── GIF and video export ───────────────────────────────────────────── +def color_labels(arrays, cmap_dict): + results = [] + for slice_array in arrays: + si, sj, _ = slice_array.shape + rgb = np.zeros((si, sj, 3), dtype=np.uint8) + for label, color in cmap_dict.items(): + if isinstance(color, str): + mpl, _ = import_mpl_plt() + color = mpl.colors.to_rgb(color) + color = [255 * n for n in color] + rgb[slice_array[..., 0] == label] = color + results.append(rgb) + return results def make_gif( - image: Image, + tensor: torch.Tensor, + axis: int, + duration: float, # of full gif output_path: TypePath, - *, - seconds: float = 5.0, - direction: str = "I", loop: int = 0, optimize: bool = True, rescale: bool = True, reverse: bool = False, ) -> None: - """Save an animated GIF sweeping through slices of a 3D image. - - The image is reoriented so slices appear in the expected anatomical - view for the given direction, matching `make_video` behavior. - - Args: - image: A [`Image`][torchio.Image] instance. - output_path: Path to the output `.gif` file. - seconds: Duration of the full animation in seconds. - direction: Anatomical sweep direction: one of - `"I"`, `"S"`, `"A"`, `"P"`, `"R"`, `"L"`. - loop: Number of loops (0 = infinite). - optimize: Attempt to compress the GIF palette. - rescale: Rescale intensities to `[0, 255]` before encoding. - reverse: Reverse the temporal order of frames. - """ - get_pillow() # raises ImportError with install hint if missing - pil_image = import_module("PIL.Image") - - # Reorient so the sweep direction is the first spatial axis and the - # remaining two axes produce an anatomically correct 2D view. - target = _video_orientation(direction) - image = Reorient(orientation=target)(image) - if rescale: - image = Normalize(out_min=0, out_max=255, copy=False)(image) - - single_channel = image.num_channels == 1 - mode = "P" if single_channel else "RGB" - - # Tensor is (C, sweep, H, W). Iterate over the sweep axis. - frames = image.data.cpu().byte().numpy() - images = [] - for i in range(frames.shape[1]): - # Single channel: (H, W); multi-channel: (C, H, W) -> (H, W, C) - frame_2d = frames[0, i] if single_channel else np.moveaxis(frames[:, i], 0, -1) - images.append(pil_image.fromarray(frame_2d).convert(mode)) - - if reverse: - images = list(reversed(images)) - + try: + from PIL import Image as ImagePIL + except ModuleNotFoundError as e: + message = 'Please install Pillow to use Image.to_gif(): pip install Pillow' + raise RuntimeError(message) from e + transform = RescaleIntensity((0, 255)) + tensor = transform(tensor) if rescale else tensor + single_channel = len(tensor) == 1 + + # Move channels dimension to the end and bring selected axis to 0 + axes = np.roll(range(1, 4), -axis) + tensor = tensor.permute(*axes, 0) + + if single_channel: + mode = 'P' + tensor = tensor[..., 0] + else: + mode = 'RGB' + array = tensor.byte().numpy() + n = 2 if axis == 1 else 1 + images = [ImagePIL.fromarray(rotate(i, n=n)).convert(mode) for i in array] num_images = len(images) - # GIF frame delay is stored in centiseconds (10ms units). - # Most browsers/viewers silently clamp delays ≤ 20ms to ~100ms, - # so we enforce a 20ms floor for reliable playback timing. - min_frame_ms = 20 - frame_duration_ms = round(seconds / num_images * 1000 / 10) * 10 - frame_duration_ms = max(frame_duration_ms, min_frame_ms) - actual_seconds = frame_duration_ms * num_images / 1000 - if abs(actual_seconds - seconds) > 0.5 * seconds / num_images: - warnings.warn( - f"GIF frame delay is quantized to 10ms steps (minimum" - f" {min_frame_ms}ms for browser compatibility). Actual" - f" duration will be {actual_seconds:.2f}s instead of" - f" {seconds:.2f}s. Consider reducing the number of slices" - f" or increasing the requested duration.", - RuntimeWarning, - stacklevel=2, + images = list(reversed(images)) if reverse else images + frame_duration_ms = duration / num_images * 1000 + if frame_duration_ms < 10: + fps = round(1000 / frame_duration_ms) + frame_duration_ms = 10 + new_duration = frame_duration_ms * num_images / 1000 + message = ( + 'The computed frame rate from the given duration is too high' + f' ({fps} fps). The highest possible frame rate in the GIF' + ' file format specification is 100 fps. The duration has been set' + f' to {new_duration:.1f} seconds, instead of {duration:.1f}' ) - + warnings.warn(message, RuntimeWarning, stacklevel=2) images[0].save( - Path(output_path), + output_path, save_all=True, append_images=images[1:], optimize=optimize, @@ -1002,126 +361,119 @@ def make_gif( def make_video( - image: Image, + image: ScalarImage, output_path: TypePath, - *, - seconds: float = 5.0, - direction: str = "I", - verbosity: str = "error", + seconds: float | None = None, + frame_rate: float | None = None, + direction: str = 'I', + verbosity: str = 'error', ) -> None: - """Create an MP4 video sweeping through slices of a 3D image. - - The image is reoriented so slices are shown in the expected - anatomical view for the given direction. Requires `ffmpeg-python`. - - Args: - image: A single-channel [`ScalarImage`][torchio.ScalarImage]. - output_path: Path to the output `.mp4` file. - seconds: Duration of the full video in seconds. - direction: Anatomical sweep direction: one of - `"I"`, `"S"`, `"A"`, `"P"`, `"R"`, `"L"`. - verbosity: ffmpeg log level. - """ ffmpeg = get_ffmpeg() + if seconds is None and frame_rate is None: + message = 'Either seconds or frame_rate must be provided.' + raise ValueError(message) + if seconds is not None and frame_rate is not None: + message = 'Provide either seconds or frame_rate, not both.' + raise ValueError(message) if image.num_channels > 1: - msg = "Only single-channel images are supported for video export." - raise ValueError(msg) - - # Reorient to the target sweep direction. - target_orientation = _video_orientation(direction) - reoriented = Reorient(orientation=target_orientation)(image) - tensor = reoriented.data - - # Rescale to [0, 255] uint8 if needed. - if tensor.min() < 0 or tensor.max() > 255: - warnings.warn( - "Tensor values outside [0, 256). Rescaling to [0, 255].", - RuntimeWarning, - stacklevel=2, + message = 'Only single-channel tensors are supported for video output for now.' + raise ValueError(message) + tmin, tmax = image.data.min(), image.data.max() + if tmin < 0 or tmax > 255: + message = ( + 'The tensor must be in the range [0, 256) for video output.' + ' The image data will be rescaled to this range.' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + image = RescaleIntensity((0, 255))(image) + if image.data.dtype != torch.uint8: + message = ( + 'Only uint8 tensors are supported for video output. The image data' + ' will be cast to uint8.' ) - tensor = _rescale_to_uint8(tensor) - if tensor.dtype != torch.uint8: - tensor = tensor.byte() + warnings.warn(message, RuntimeWarning, stacklevel=2) + image = To(torch.uint8)(image) - # Crop to even dimensions (required by H.265). - num_frames, height, width = ( - tensor.shape[-3], - tensor.shape[-2], - tensor.shape[-1], - ) - if height % 2 != 0: - tensor = tensor[:, :, : height - 1, :] - height -= 1 - if width % 2 != 0: - tensor = tensor[:, :, :, : width - 1] - width -= 1 + # Reorient so the output looks like in typical visualization software + direction = direction.upper() + if direction == 'I': # axial top to bottom + target = 'IPL' + elif direction == 'S': # axial bottom to top + target = 'SPL' + elif direction == 'A': # coronal back to front + target = 'AIL' + elif direction == 'P': # coronal front to back + target = 'PIL' + elif direction == 'R': # sagittal left to right + target = 'RIP' + elif direction == 'L': # sagittal right to left + target = 'LIP' + else: + message = ( + 'Direction must be one of "I", "S", "P", "A", "R" or "L".' + f' Got {direction!r}.' + ) + raise ValueError(message) + image = ToOrientation(target)(image) + + # Check isotropy + spacing_f, spacing_h, spacing_w = image.spacing + if spacing_h != spacing_w: + message = ( + 'The height and width spacings should be the same video output.' + f' Got {spacing_h:.2f} and {spacing_w:.2f}.' + f' Resampling both to {spacing_f:.2f}.' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + spacing_iso = min(spacing_h, spacing_w) + target_spacing = spacing_f, spacing_iso, spacing_iso + image = Resample(target_spacing)(image) + + # Check that height and width are multiples of 2 for H.265 encoding + num_frames, height, width = image.spatial_shape + if height % 2 != 0 or width % 2 != 0: + message = ( + f'The height ({height}) and width ({width}) must be even.' + ' The image will be cropped to the nearest even number.' + ) + warnings.warn(message, RuntimeWarning, stacklevel=2) + image = EnsureShapeMultiple((1, 2, 2), method='crop')(image) - frame_rate = num_frames / seconds + if seconds is not None: + frame_rate = num_frames / seconds - out = Path(output_path) - if out.suffix.lower() != ".mp4": - msg = "Only .mp4 output is supported." - raise NotImplementedError(msg) + output_path = Path(output_path) + if output_path.suffix.lower() != '.mp4': + message = 'Only .mp4 files are supported for video output.' + raise NotImplementedError(message) - frames = tensor[0].cpu().numpy() + frames = image.numpy()[0] + first = frames[0] + height, width = first.shape process = ( ffmpeg.input( - "pipe:", - format="rawvideo", - pix_fmt="gray", - s=f"{width}x{height}", + 'pipe:', + format='rawvideo', + pix_fmt='gray', + s=f'{width}x{height}', framerate=frame_rate, ) .output( - str(out), - vcodec="libx264", - pix_fmt="yuv420p", - movflags="+faststart", - # Baseline profile for maximum browser/Jupyter compatibility. - profile="baseline", - level="3.0", + str(output_path), + vcodec='libx265', + pix_fmt='yuv420p', loglevel=verbosity, + **{'x265-params': f'log-level={verbosity}'}, ) .overwrite_output() .run_async(pipe_stdin=True) ) - for frame in frames: - process.stdin.write(frame.tobytes()) + for array in frames: + buffer = array.tobytes() + process.stdin.write(buffer) process.stdin.close() process.wait() - - -def _rescale_to_uint8(tensor: Any) -> Any: - """Rescale a tensor to `[0, 255]` uint8.""" - - t = tensor.float() - tmin = t.min() - tmax = t.max() - if tmax - tmin > 0: - t = (t - tmin) / (tmax - tmin) * 255 - return t.byte() - - -_VIDEO_ORIENTATIONS: dict[str, str] = { - "I": "IPL", - "S": "SPL", - "A": "AIL", - "P": "PIL", - "R": "RIP", - "L": "LIP", -} - - -def _video_orientation(direction: str) -> str: - """Map a sweep direction letter to a 3-character orientation string.""" - direction = direction.upper() - if direction not in _VIDEO_ORIENTATIONS: - msg = ( - f"Direction must be one of {list(_VIDEO_ORIENTATIONS)}, got {direction!r}." - ) - raise ValueError(msg) - return _VIDEO_ORIENTATIONS[direction] diff --git a/tests/__init__.py b/tests/__init__.py index e69de29bb..3aaf2afa0 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Unit test package for torchio.""" diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index 95cd5456d..000000000 --- a/tests/conftest.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Shared pytest fixtures for the test suite.""" - -from __future__ import annotations - -import copy -from collections.abc import Callable -from typing import Any - -import pytest -import torch - -from torchio.data.batch import SubjectsBatch -from torchio.data.batch import _slice_params - - -@pytest.fixture -def assert_vectorized() -> Callable[..., None]: - """Return an assertion that a per-instance transform is vectorized faithfully. - - The returned helper applies *transform* to a batch (per-instance) and then, - for every element, re-applies the same per-element parameters to that - element alone. The two must match, which proves the vectorized whole-batch - computation is equivalent to processing each element independently (no - cross-element contamination or broadcasting mistakes). - - This is only valid for transforms whose `apply_transform` is deterministic - given the recorded parameters (the randomness lives in `make_params`). - Transforms that sample inside `apply_transform` (e.g. Noise, LabelsToImage) - need a different check. - """ - - def _assert( - transform: Any, - batch: SubjectsBatch, - *, - rtol: float = 1e-5, - atol: float = 1e-6, - ) -> None: - original = copy.deepcopy(batch) - result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params, "per-instance path was not active" - batched_keys = params["_batched_keys"] - keep = params.get("_keep") - image_names = list(transform._get_images(result).keys()) - result_images = transform._get_images(result) - original_subjects = original.unbatch() - for index in range(original.batch_size): - single = SubjectsBatch.from_subjects([original_subjects[index]]) - single_input = { - name: image.data.clone() - for name, image in transform._get_images(single).items() - } - element_params = _slice_params(params, index, batched_keys) - single = transform.apply_transform(single, element_params) - single_images = transform._get_images(single) - gated_out = keep is not None and not keep[index] - for name in image_names: - result_row = result_images[name].data[index : index + 1] - torch.testing.assert_close( - result_row, - single_images[name].data, - rtol=rtol, - atol=atol, - ) - if gated_out: - # A gated-out element must be a bit-for-bit no-op. - torch.testing.assert_close( - result_row, - single_input[name], - rtol=0, - atol=0, - ) - - return _assert diff --git a/tests/data/__init__.py b/tests/data/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/inference/__init__.py b/tests/data/inference/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/inference/test_aggregator.py b/tests/data/inference/test_aggregator.py new file mode 100644 index 000000000..78cc47bc7 --- /dev/null +++ b/tests/data/inference/test_aggregator.py @@ -0,0 +1,189 @@ +from typing import cast + +import pytest +import torch +from torch.utils.data import Dataset + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestAggregator(TorchioTestCase): + """Tests for `aggregator` module.""" + + def aggregate(self, mode, fixture): + image_shape = 1, 1, 4, 4 + tensor = torch.ones(image_shape) + fixture = torch.as_tensor(fixture).reshape(image_shape) + image_name = 'img' + subject = tio.Subject({image_name: tio.ScalarImage(tensor=tensor)}) + patch_size = 1, 3, 3 + patch_overlap = 0, 2, 2 + sampler = tio.data.GridSampler(subject, patch_size, patch_overlap) + aggregator = tio.data.GridAggregator(sampler, overlap_mode=mode) + sampler_dataset = cast(Dataset[tio.Subject], sampler) + loader = tio.SubjectsLoader(sampler_dataset, batch_size=3) + values_dict = { + (0, 0): 0, + (0, 1): 2, + (1, 0): 4, + (1, 1): 6, + } + for batch in loader: + location = batch[tio.LOCATION] + data = batch[image_name][tio.DATA] + locations_and_channels = zip(location, data, strict=True) + for location, data in locations_and_channels: + coords_2d = tuple(location[1:3].tolist()) + data *= values_dict[coords_2d] + batch_data = batch[image_name][tio.DATA] + aggregator.add_batch(batch_data, batch[tio.LOCATION]) + output = aggregator.get_output_tensor() + self.assert_tensor_equal(output, fixture) + + def test_overlap_crop(self): + fixture = ( + (0, 0, 2, 2), + (0, 0, 2, 2), + (4, 4, 6, 6), + (4, 4, 6, 6), + ) + self.aggregate('crop', fixture) + + def test_overlap_average(self): + fixture = ( + (0, 1, 1, 2), + (2, 3, 3, 4), + (2, 3, 3, 4), + (4, 5, 5, 6), + ) + self.aggregate('average', fixture) + + def test_overlap_hann(self): + fixture = ( + (0 / 3, 2 / 3, 4 / 3, 6 / 3), # noqa: E201, E241 + (4 / 3, 6 / 3, 8 / 3, 10 / 3), # noqa: E201, E241 + (8 / 3, 10 / 3, 12 / 3, 14 / 3), # noqa: E201, E241 + (12 / 3, 14 / 3, 16 / 3, 18 / 3), + ) + self.aggregate('hann', fixture) + + def run_sampler_aggregator(self, overlap_mode='crop'): + patch_size = 10 + patch_overlap = 2 + grid_sampler = tio.inference.GridSampler( + self.sample_subject, + patch_size, + patch_overlap, + ) + sampler_dataset = cast(Dataset[tio.Subject], grid_sampler) + patch_loader = tio.SubjectsLoader(sampler_dataset) + aggregator = tio.inference.GridAggregator( + grid_sampler, + overlap_mode=overlap_mode, + ) + for batch in patch_loader: + data = batch['t1'][tio.DATA].long() + aggregator.add_batch(data, batch[tio.LOCATION]) + return aggregator + + def test_warning_int64(self): + aggregator = self.run_sampler_aggregator() + with pytest.warns(RuntimeWarning): + aggregator.get_output_tensor() + + def run_patch_crop_issue(self, *, padding_mode): + # https://github.com/TorchIO-project/torchio/issues/813 + pao, pas, ims, bb1, bb2 = 4, 102, 320, 100, 120 + + patch_overlap = pao, 0, 0 + patch_size = pas, 1, 1 + img = torch.zeros((1, ims, 1, 1)) + bbox = [bb1, bb2] + + img[:, bbox[0] : bbox[1]] = 1 + image = tio.LabelMap(tensor=img) + subject = tio.Subject(image=image) + grid_sampler = tio.inference.GridSampler( + subject, + patch_size, + patch_overlap, + ) + sampler_dataset = cast(Dataset[tio.Subject], grid_sampler) + patch_loader = tio.SubjectsLoader(sampler_dataset) + aggregator = tio.inference.GridAggregator(grid_sampler) + for patches_batch in patch_loader: + input_tensor = patches_batch['image'][tio.DATA] + locations = patches_batch[tio.LOCATION] + aggregator.add_batch(input_tensor, locations) + output_tensor = aggregator.get_output_tensor() + self.assert_tensor_equal(image.tensor, output_tensor) + + def test_patch_crop_issue_no_padding(self): + self.run_patch_crop_issue(padding_mode=None) + + def test_patch_crop_issue_padding(self): + self.run_patch_crop_issue(padding_mode='constant') + + def test_bad_aggregator_shape(self): + # https://github.com/microsoft/InnerEye-DeepLearning/pull/677/checks?check_run_id=5395915817 + tensor = torch.ones(1, 40, 40, 40) + image_name = 'img' + subject = tio.Subject({image_name: tio.ScalarImage(tensor=tensor)}) + patch_size = 40 + patch_overlap = 30 + sampler = tio.data.GridSampler( + subject, + patch_size, + patch_overlap, + padding_mode='edge', + ) + aggregator = tio.data.GridAggregator(sampler) + sampler_dataset = cast(Dataset[tio.Subject], sampler) + loader = tio.SubjectsLoader(sampler_dataset, batch_size=3) + for batch in loader: + input_batch = batch[image_name][tio.DATA] + crop = tio.CropOrPad(12) + patches = [crop(patch) for patch in input_batch] + inference_batch = torch.stack(patches) + with pytest.raises(RuntimeError): + aggregator.add_batch(inference_batch, batch[tio.LOCATION]) + + def test_downsampling_model(self): + # This might be useful to compute image embeddings using a sliding window + downsampling_factor = 4 # e.g. patch size in a ViT + embedding_dim = 5 + net_input_size = 20 + image_size = 40 + + def network(x): + down = x[ + ..., + ::downsampling_factor, + ::downsampling_factor, + ::downsampling_factor, + ] + embeddings = torch.cat(embedding_dim * [down], dim=1) + return embeddings + + tensor = torch.ones(1, image_size, image_size, image_size) + image_name = 'img' + subject = tio.Subject({image_name: tio.ScalarImage(tensor=tensor)}) + sampler = tio.data.GridSampler( + subject, + patch_size=net_input_size, + ) + aggregator = tio.data.GridAggregator( + sampler, + downsampling_factor=downsampling_factor, + ) + sampler_dataset = cast(Dataset[tio.Subject], sampler) + loader = tio.SubjectsLoader(sampler_dataset, batch_size=3) + for batch in loader: + input_batch = batch[image_name][tio.DATA] + embeddings = network(input_batch) + aggregator.add_batch(embeddings, batch[tio.LOCATION]) + output = aggregator.get_output_tensor() + expected_shape = (embedding_dim,) + (image_size // downsampling_factor,) * 3 + self.assertEqual(output.shape, expected_shape) diff --git a/tests/data/inference/test_grid_sampler.py b/tests/data/inference/test_grid_sampler.py new file mode 100644 index 000000000..6e7d1529a --- /dev/null +++ b/tests/data/inference/test_grid_sampler.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python +from copy import copy + +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestGridSampler(TorchioTestCase): + """Tests for `GridSampler`.""" + + def test_locations(self): + patch_size = 5, 20, 20 + patch_overlap = 2, 4, 6 + sampler = tio.GridSampler( + subject=self.sample_subject, + patch_size=patch_size, + patch_overlap=patch_overlap, + ) + fixture = [ + [0, 0, 0, 5, 20, 20], + [0, 0, 10, 5, 20, 30], + [3, 0, 0, 8, 20, 20], + [3, 0, 10, 8, 20, 30], + [5, 0, 0, 10, 20, 20], + [5, 0, 10, 10, 20, 30], + ] + locations = sampler.locations.tolist() + assert locations == fixture + + def test_generate_patches(self): + patch_size = 5, 15, 15 + sampler = tio.GridSampler(self.sample_subject, patch_size) + for patch in sampler(): + assert patch.spatial_shape == patch_size + + def test_large_patch(self): + with pytest.raises(ValueError): + tio.GridSampler(self.sample_subject, (5, 21, 5), (0, 2, 0)) + + def test_large_overlap(self): + with pytest.raises(ValueError): + tio.GridSampler(self.sample_subject, (5, 20, 5), (2, 4, 6)) + + def test_odd_overlap(self): + with pytest.raises(ValueError): + tio.GridSampler(self.sample_subject, (5, 20, 5), (2, 4, 3)) + + def test_single_location(self): + sampler = tio.GridSampler(self.sample_subject, (10, 20, 30), 0) + fixture = [[0, 0, 0, 10, 20, 30]] + assert sampler.locations.tolist() == fixture + + def test_subject_shape(self): + patch_size = 5, 20, 20 + patch_overlap = 2, 4, 6 + initial_shape = copy(self.sample_subject.shape) + tio.GridSampler( + self.sample_subject, + patch_size, + patch_overlap, + padding_mode='reflect', + ) + final_shape = self.sample_subject.shape + assert initial_shape == final_shape diff --git a/tests/data/inference/test_inference.py b/tests/data/inference/test_inference.py new file mode 100644 index 000000000..f6f289a3c --- /dev/null +++ b/tests/data/inference/test_inference.py @@ -0,0 +1,55 @@ +from typing import cast + +from torch.utils.data import Dataset + +import torchio as tio +from torchio import DATA +from torchio import LOCATION +from torchio.data.inference import GridAggregator +from torchio.data.inference import GridSampler + +from ...utils import TorchioTestCase + + +class TestInference(TorchioTestCase): + """Tests for `inference` module.""" + + def test_inference_no_padding(self): + self.try_inference(None) + + def test_inference_padding(self): + self.try_inference(3) + + def try_inference(self, padding_mode): + for n in 17, 27: + patch_size = 10, 15, n + patch_overlap = 4, 6, 8 + batch_size = 6 + + grid_sampler = GridSampler( + self.sample_subject, + patch_size, + patch_overlap, + padding_mode=padding_mode, + ) + aggregator = GridAggregator(grid_sampler) + sampler_dataset = cast(Dataset[tio.Subject], grid_sampler) + patch_loader = tio.SubjectsLoader( + sampler_dataset, + batch_size=batch_size, + ) + for patches_batch in patch_loader: + input_tensor = patches_batch['t1'][DATA] + locations = patches_batch[LOCATION] + logits = model(input_tensor) # some model + outputs = logits + aggregator.add_batch(outputs, locations) + + output = aggregator.get_output_tensor() + assert (output == -5).all() + assert output.shape == self.sample_subject.t1.shape + + +def model(tensor): + tensor[:] = -5 + return tensor diff --git a/tests/data/sampler/__init__.py b/tests/data/sampler/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/data/sampler/test_label_sampler.py b/tests/data/sampler/test_label_sampler.py new file mode 100644 index 000000000..03e76a2a9 --- /dev/null +++ b/tests/data/sampler/test_label_sampler.py @@ -0,0 +1,81 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestLabelSampler(TorchioTestCase): + """Tests for `LabelSampler` class.""" + + def test_label_sampler(self): + sampler = tio.LabelSampler(5) + for patch in sampler(self.sample_subject, num_patches=10): + patch_center = patch.get_label_map('label').data[0, 2, 2, 2] + assert patch_center == 1 + + def test_label_probabilities(self): + labels = torch.Tensor((0, 0, 1, 1, 2, 1, 0)).reshape(1, 1, 1, -1) + subject = tio.Subject( + label=tio.Image(tensor=labels, type=tio.LABEL), + ) + subject = tio.SubjectsDataset([subject])[0] + probs_dict: dict[int, float] = {0: 0.0, 1: 50.0, 2: 25.0, 3: 25.0} + patch_size = (1, 1, 5) + sampler = tio.LabelSampler(patch_size, label_probabilities=probs_dict) + probabilities = sampler.get_probability_map(subject) + fixture = torch.Tensor((0, 0, 1 / 4, 1 / 4, 1 / 4, 0, 0)) + assert torch.all(probabilities.squeeze().eq(fixture)) + + def test_inconsistent_shape(self): + # https://github.com/TorchIO-project/torchio/issues/234#issuecomment-675029767 + subject = tio.Subject( + im1=tio.ScalarImage(tensor=torch.rand(2, 4, 5, 6)), + im2=tio.LabelMap(tensor=torch.rand(1, 4, 5, 6)), + ) + patch_size = 2 + sampler = tio.LabelSampler(patch_size, 'im2') + next(sampler(subject)) + + def test_multichannel_label_sampler(self): + subject = tio.Subject( + label=tio.LabelMap( + tensor=torch.tensor( + [ + [[[1, 1]]], + [[[0, 1]]], + ], + ), + ), + ) + patch_size = 1 + sampler = tio.LabelSampler( + patch_size, + 'label', + label_probabilities={0: 1, 1: 1}, + ) + # There are 2 voxels in the image, channels have same probabilities, + # 1st voxel has probability 0.5 * 0.5 + 0 * 0.5 of being chosen while + # 2nd voxel has probability 0.5 * 0.5 + 1 * 0.5 of being chosen. + probabilities = sampler.get_probability_map(subject) + fixture = torch.Tensor((1 / 4, 3 / 4)) + assert torch.all(probabilities.squeeze().eq(fixture)) + + def test_no_labelmap(self): + im = tio.ScalarImage(tensor=torch.rand(1, 1, 1, 1)) + subject = tio.Subject(image=im, no_label=im) + sampler = tio.LabelSampler(1) + with pytest.raises(RuntimeError): + next(sampler(subject)) + + def test_empty_map(self): + # https://github.com/TorchIO-project/torchio/issues/392 + im = tio.ScalarImage(tensor=torch.rand(1, 6, 6, 6)) + label = torch.zeros(1, 6, 6, 6) + label[..., 0] = 1 # voxels far from center + label_im = tio.LabelMap(tensor=label) + subject = tio.Subject(image=im, label=label_im) + sampler = tio.LabelSampler(4) + with pytest.raises(RuntimeError): + next(sampler(subject)) diff --git a/tests/data/sampler/test_patch_sampler.py b/tests/data/sampler/test_patch_sampler.py new file mode 100644 index 000000000..47246f240 --- /dev/null +++ b/tests/data/sampler/test_patch_sampler.py @@ -0,0 +1,23 @@ +from typing import Any +from typing import cast + +import pytest + +from torchio.data import PatchSampler + +from ...utils import TorchioTestCase + + +class TestPatchSampler(TorchioTestCase): + """Tests for `PatchSampler` class.""" + + def test_bad_patch_size(self): + with pytest.raises(ValueError): + PatchSampler(0) + with pytest.raises(ValueError): + PatchSampler(-1) + with pytest.raises(ValueError): + PatchSampler(cast(Any, 1.5)) + + def test_extract_patch(self): + PatchSampler(1).extract_patch(self.sample_subject, (3, 4, 5)) diff --git a/tests/data/sampler/test_random_sampler.py b/tests/data/sampler/test_random_sampler.py new file mode 100644 index 000000000..3c9e0fb57 --- /dev/null +++ b/tests/data/sampler/test_random_sampler.py @@ -0,0 +1,14 @@ +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRandomSampler(TorchioTestCase): + def test_not_implemented(self): + sampler = tio.data.sampler.RandomSampler(1) + with pytest.raises(NotImplementedError): + sampler(self.sample_subject, 5) + with pytest.raises(NotImplementedError): + sampler.get_probability_map(self.sample_subject) diff --git a/tests/data/sampler/test_uniform_sampler.py b/tests/data/sampler/test_uniform_sampler.py new file mode 100644 index 000000000..37d6838db --- /dev/null +++ b/tests/data/sampler/test_uniform_sampler.py @@ -0,0 +1,26 @@ +import torch + +import torchio +from torchio.data import UniformSampler + +from ...utils import TorchioTestCase + + +class TestUniformSampler(TorchioTestCase): + """Tests for `UniformSampler` class.""" + + def test_uniform_probabilities(self): + sampler = UniformSampler(5) + probabilities = sampler.get_probability_map(self.sample_subject) + fixtures = torch.ones_like(probabilities) + assert torch.all(probabilities.eq(fixtures)) + + def test_incosistent_shape(self): + # https://github.com/TorchIO-project/torchio/issues/234#issuecomment-675029767 + subject = torchio.Subject( + im1=torchio.ScalarImage(tensor=torch.rand(1, 4, 5, 6)), + im2=torchio.ScalarImage(tensor=torch.rand(2, 4, 5, 6)), + ) + patch_size = 2 + sampler = UniformSampler(patch_size) + next(sampler(subject)) diff --git a/tests/data/sampler/test_weighted_sampler.py b/tests/data/sampler/test_weighted_sampler.py new file mode 100644 index 000000000..e9e891d6f --- /dev/null +++ b/tests/data/sampler/test_weighted_sampler.py @@ -0,0 +1,39 @@ +import torch + +import torchio as tio +from torchio.data import WeightedSampler + +from ...utils import TorchioTestCase + + +class TestWeightedSampler(TorchioTestCase): + """Tests for `WeightedSampler` class.""" + + def test_weighted_sampler(self): + subject = self.get_sample((1, 7, 7, 7)) + sampler = WeightedSampler(5, 'prob') + patch = next(sampler(subject)) + location = patch[tio.LOCATION] + assert isinstance(location, torch.Tensor) + assert tuple(location[:3].tolist()) == (1, 1, 1) + + def get_sample(self, image_shape): + t1 = torch.rand(*image_shape) + prob = torch.zeros_like(t1) + prob[0, 3, 3, 3] = 1 + subject = tio.Subject( + t1=tio.ScalarImage(tensor=t1), + prob=tio.ScalarImage(tensor=prob), + ) + subject = tio.SubjectsDataset([subject])[0] + return subject + + def test_inconsistent_shape(self): + # https://github.com/TorchIO-project/torchio/issues/234#issuecomment-675029767 + subject = tio.Subject( + im1=tio.ScalarImage(tensor=torch.rand(1, 4, 5, 6)), + im2=tio.ScalarImage(tensor=torch.rand(2, 4, 5, 6)), + ) + patch_size = 2 + sampler = tio.data.WeightedSampler(patch_size, 'im1') + next(sampler(subject)) diff --git a/tests/data/test_image.py b/tests/data/test_image.py new file mode 100644 index 000000000..97287a9f9 --- /dev/null +++ b/tests/data/test_image.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python +"""Tests for Image.""" + +import copy +import sys +import tempfile +from pathlib import Path + +import nibabel as nib +import numpy as np +import pytest +import torch + +import torchio as tio + +from ..utils import TorchioTestCase + + +class TestImage(TorchioTestCase): + """Tests for `Image`.""" + + def test_image_not_found(self): + with pytest.raises(FileNotFoundError): + tio.ScalarImage('nopath') + + @pytest.mark.skipif(sys.platform == 'win32', reason='Path not valid') + def test_wrong_path_value(self): + with pytest.raises(RuntimeError): + tio.ScalarImage('~&./@#"!?X7=+') + + def test_wrong_path_type(self): + with pytest.raises(TypeError): + tio.ScalarImage(5) + + def test_wrong_affine(self): + with pytest.raises(TypeError): + tio.ScalarImage(5, affine=1) + + def test_tensor_flip(self): + sample_input = torch.ones((4, 30, 30, 30)) + tio.RandomFlip()(sample_input) + + def test_tensor_affine(self): + sample_input = torch.ones((4, 10, 10, 10)) + tio.RandomAffine()(sample_input) + + def test_wrong_scalar_image_type(self): + data = torch.ones((1, 10, 10, 10)) + with pytest.raises(ValueError): + tio.ScalarImage(tensor=data, type=tio.LABEL) + + def test_wrong_label_map_type(self): + data = torch.ones((1, 10, 10, 10)) + with pytest.raises(ValueError): + tio.LabelMap(tensor=data, type=tio.INTENSITY) + + def test_no_input(self): + with pytest.raises(ValueError): + tio.ScalarImage() + + def test_bad_key(self): + with pytest.raises(ValueError): + tio.ScalarImage(path='', data=5) + + def test_repr(self): + subject = tio.Subject( + t1=tio.ScalarImage(self.get_image_path('repr_test')), + ) + assert 'memory' not in repr(subject['t1']) + subject.load() + assert 'memory' in repr(subject['t1']) + + def test_data_tensor(self): + subject = copy.deepcopy(self.sample_subject) + subject.load() + assert subject.t1.data is subject.t1.tensor + + def test_bad_affine(self): + with pytest.raises(ValueError): + tio.ScalarImage(tensor=torch.rand(1, 2, 3, 4), affine=np.eye(3)) + + def test_nans_tensor(self): + tensor = np.random.rand(1, 2, 3, 4) + tensor[0, 0, 0, 0] = np.nan + with pytest.warns(RuntimeWarning): + image = tio.ScalarImage(tensor=tensor, check_nans=True) + image.set_check_nans(False) + + def test_get_center(self): + tensor = torch.rand(1, 3, 3, 3) + image = tio.ScalarImage(tensor=tensor) + ras = image.get_center() + lps = image.get_center(lps=True) + assert ras == (1, 1, 1) + assert lps == (-1, -1, 1) + + def test_with_list_of_missing_files(self): + with pytest.raises(FileNotFoundError): + tio.ScalarImage(path=['nopath', 'error']) + + def test_with_sequences_of_paths(self): + shape = (5, 5, 5) + path1 = self.get_image_path('path1', shape=shape) + path2 = self.get_image_path('path2', shape=shape) + paths_tuple = path1, path2 + paths_list = list(paths_tuple) + for sequence in (paths_tuple, paths_list): + image = tio.ScalarImage(path=sequence) + assert image.shape == (2, 5, 5, 5) + assert image[tio.STEM] == ['path1', 'path2'] + + def test_with_a_list_of_images_with_different_shapes(self): + path1 = self.get_image_path('path1', shape=(5, 5, 5)) + path2 = self.get_image_path('path2', shape=(7, 5, 5)) + image = tio.ScalarImage(path=[path1, path2]) + with pytest.raises(RuntimeError): + image.load() + + def test_with_a_list_of_images_with_different_affines(self): + path1 = self.get_image_path('path1', spacing=(1, 1, 1)) + path2 = self.get_image_path('path2', spacing=(1, 2, 1)) + image = tio.ScalarImage(path=[path1, path2]) + with pytest.warns(RuntimeWarning): + image.load() + + def test_with_a_list_of_2d_paths(self): + shape = (5, 6) + path1 = self.get_image_path('path1', shape=shape, suffix='.nii') + path2 = self.get_image_path('path2', shape=shape, suffix='.img') + path3 = self.get_image_path('path3', shape=shape, suffix='.hdr') + image = tio.ScalarImage(path=[path1, path2, path3]) + assert image.shape == (3, 5, 6, 1) + assert image[tio.STEM] == ['path1', 'path2', 'path3'] + + def test_axis_name_2d(self): + path = self.get_image_path('im2d', shape=(5, 6)) + image = tio.ScalarImage(path) + height_idx = image.axis_name_to_index('t') + width_idx = image.axis_name_to_index('l') + assert image.height == image.shape[height_idx] + assert image.width == image.shape[width_idx] + + def test_different_shape(self): + path_1 = self.get_image_path('im_shape1', shape=(5, 5, 5)) + path_2 = self.get_image_path('im_shape2', shape=(7, 5, 5)) + + image = tio.ScalarImage([path_1, path_2]) + with pytest.raises(RuntimeError): + image.load() + + @pytest.mark.slow + @pytest.mark.skipif(sys.platform == 'win32', reason='Unstable on Windows') + def test_plot(self): + image = self.sample_subject.t1 + image.plot(show=False, output_path=self.dir / 'image.png') + + def test_data_type_uint16_array(self): + tensor = np.random.rand(1, 3, 3, 3).astype(np.uint16) + image = tio.ScalarImage(tensor=tensor) + assert image.data.dtype == torch.int32 + + def test_data_type_uint32_array(self): + tensor = np.random.rand(1, 3, 3, 3).astype(np.uint32) + image = tio.ScalarImage(tensor=tensor) + assert image.data.dtype == torch.int64 + + def test_save_image_with_data_type_boolean(self): + tensor = np.random.rand(1, 3, 3, 3).astype(bool) + image = tio.ScalarImage(tensor=tensor) + image.save(self.dir / 'image.nii') + + def test_load_uint(self): + affine = np.eye(4) + for dtype in np.uint16, np.uint32: + data = np.ones((3, 3, 3), dtype=dtype) + img = nib.Nifti1Image(data, affine) + with tempfile.NamedTemporaryFile(suffix='.nii', delete=False) as f: + nib.save(img, f.name) + tio.ScalarImage(f.name).load() + + def test_pil_3d(self): + with pytest.raises(RuntimeError): + tio.ScalarImage(tensor=torch.rand(1, 2, 3, 4)).as_pil() + + def test_pil_1(self): + tio.ScalarImage(tensor=torch.rand(1, 2, 3, 1)).as_pil() + + def test_pil_2(self): + with pytest.raises(RuntimeError): + tio.ScalarImage(tensor=torch.rand(2, 2, 3, 1)).as_pil() + + def test_pil_3(self): + tio.ScalarImage(tensor=torch.rand(3, 2, 3, 1)).as_pil() + + def test_set_data(self): + im = self.sample_subject.t1 + with pytest.deprecated_call(): + im.data = im.data + + def test_no_type(self): + with pytest.warns(FutureWarning): + tio.Image(tensor=torch.rand(1, 2, 3, 4)) + + def test_custom_reader(self): + path = self.dir / 'im.npy' + + def numpy_reader(path): + return np.load(path), np.eye(4) + + def assert_shape(shape_in, shape_out): + np.save(path, np.random.rand(*shape_in)) + image = tio.ScalarImage(path, reader=numpy_reader) + assert image.shape == shape_out + + assert_shape((5, 5), (1, 5, 5, 1)) + assert_shape((5, 5, 3), (3, 5, 5, 1)) + assert_shape((3, 5, 5), (3, 5, 5, 1)) + assert_shape((5, 5, 5), (1, 5, 5, 5)) + assert_shape((1, 5, 5, 5), (1, 5, 5, 5)) + assert_shape((4, 5, 5, 5), (4, 5, 5, 5)) + + def test_fast_gif(self): + with pytest.warns(RuntimeWarning): + with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as f: + self.sample_subject.t1.to_gif(0, 0.0001, f.name) + + def test_gif_rgb(self): + with tempfile.NamedTemporaryFile(suffix='.gif', delete=False) as f: + tio.ScalarImage(tensor=torch.rand(3, 4, 5, 6)).to_gif(0, 1, f.name) + + @pytest.mark.slow + def test_hist(self): + self.sample_subject.t1.hist(density=False, show=False) + self.sample_subject.t1.hist(density=True, show=False) + + def test_count(self): + image = self.sample_subject.label + max_n = image.data.numel() + nonzero = image.count_nonzero() + assert 0 <= nonzero <= max_n + counts = image.count_labels() + assert tuple(counts) == (0, 1) + assert 0 <= counts[0] <= max_n + assert 0 <= counts[1] <= max_n + + def test_affine_multipath(self): + # https://github.com/TorchIO-project/torchio/issues/762 + path1 = self.get_image_path('multi1') + path2 = self.get_image_path('multi2') + paths = path1, path2 + image = tio.ScalarImage(paths) + self.assert_tensor_equal(image.affine, np.eye(4)) + + def test_bad_numpy_type_reader(self): + # https://github.com/TorchIO-project/torchio/issues/764 + def numpy_reader(path): + return np.load(path), np.eye(4) + + tensor = np.random.rand(1, 2, 3, 4).astype(np.uint16) + test_path = self.dir / 'test_image.npy' + np.save(test_path, tensor) + image = tio.ScalarImage(test_path, reader=numpy_reader) + image.load() + + def test_load_unload(self): + path = self.get_image_path('unload') + image = tio.ScalarImage(path) + with self.assertRaises(RuntimeError): + image.unload() + image.load() + assert image._loaded + image.unload() + assert not image._loaded + assert image[tio.DATA] is None + assert image[tio.AFFINE] is None + assert not image._loaded + + def test_unload_no_path(self): + tensor = torch.rand(1, 2, 3, 4) + image = tio.ScalarImage(tensor=tensor) + with self.assertRaises(RuntimeError): + image.unload() + + def test_copy_no_data(self): + # https://github.com/TorchIO-project/torchio/issues/974 + path = self.get_image_path('im_copy') + my_image = tio.LabelMap(path) + assert not my_image._loaded + new_image = copy.copy(my_image) + assert not my_image._loaded + assert not new_image._loaded + + my_image.load() + new_image = copy.copy(my_image) + assert my_image._loaded + assert new_image._loaded + + def test_slicing(self): + path = self.get_image_path('im_slicing') + image = tio.ScalarImage(path) + + assert image.shape == (1, 10, 20, 30) + + cropped = image[0] + assert cropped.shape == (1, 1, 20, 30) + + cropped = image[:, 2:-3] + assert cropped.shape == (1, 10, 15, 30) + + cropped = image[-5:, 5:] + assert cropped.shape == (1, 5, 15, 30) + + with pytest.raises(NotImplementedError): + image[..., 5] + + with pytest.raises(ValueError): + image[0:8:-1] + + with pytest.raises(ValueError): + image[3::-1] + + def test_verify_path(self): + path = Path(self.get_image_path('im_verify')) + + image = tio.ScalarImage(path, verify_path=False) + assert image.path == path + + image = tio.ScalarImage(path, verify_path=True) + assert image.path == path + + fake_path = Path('fake_path.nii') + + image = tio.ScalarImage(fake_path, verify_path=False) + assert image.path == fake_path + + with pytest.raises(FileNotFoundError): + tio.ScalarImage(fake_path, verify_path=True) diff --git a/tests/data/test_io.py b/tests/data/test_io.py new file mode 100644 index 000000000..0128bc798 --- /dev/null +++ b/tests/data/test_io.py @@ -0,0 +1,216 @@ +import tempfile +from pathlib import Path + +import numpy as np +import pytest +import SimpleITK as sitk +import torch + +from torchio.data import ScalarImage +from torchio.data import io + +from ..utils import TorchioTestCase + + +class TestIO(TorchioTestCase): + """Tests for `io` module.""" + + def setUp(self): + super().setUp() + self.nii_path = self.get_image_path('read_image') + self.dicom_dir = self.get_tests_data_dir() / 'dicom' + self.dicom_path = self.dicom_dir / 'IMG0001.dcm' + string = ( + '1.5 0.18088 -0.124887 0.65072 ' + '-0.20025 0.965639 -0.165653 -11.6452 ' + '0.0906326 0.18661 0.978245 11.4002 ' + '0 0 0 1 ' + ) + tensor = torch.as_tensor(np.fromstring(string, sep=' ').reshape(4, 4)) + self.matrix = tensor + + def test_read_image(self): + # I need to find something readable by nib but not sitk + io.read_image(self.nii_path) + + def test_save_rgb(self): + im = ScalarImage(tensor=torch.rand(1, 4, 5, 1)) + with pytest.warns(RuntimeWarning): + im.save(self.dir / 'test.jpg') + + def test_read_dicom_file(self): + tensor, _ = io.read_image(self.dicom_path) + assert tuple(tensor.shape) == (1, 88, 128, 1) + + def test_read_dicom_dir(self): + tensor, _ = io.read_image(self.dicom_dir) + assert tuple(tensor.shape) == (1, 88, 128, 17) + + def test_dicom_dir_missing(self): + with pytest.raises(FileNotFoundError): + io._read_dicom('missing') + + def test_dicom_dir_no_files(self): + empty = self.dir / 'empty' + empty.mkdir() + sitk.ProcessObject_SetGlobalWarningDisplay(False) + with pytest.raises(FileNotFoundError): + io._read_dicom(empty) + sitk.ProcessObject_SetGlobalWarningDisplay(True) + + def write_read_matrix(self, suffix): + out_path = self.dir / f'matrix{suffix}' + io.write_matrix(self.matrix, out_path) + matrix = io.read_matrix(out_path) + assert torch.allclose(matrix, self.matrix) + + def test_matrix_itk(self): + self.write_read_matrix('.tfm') + self.write_read_matrix('.h5') + + def test_matrix_txt(self): + self.write_read_matrix('.txt') + + def test_ensure_4d_5d(self): + tensor = torch.rand(3, 4, 5, 1, 2) + assert io.ensure_4d(tensor).shape == (2, 3, 4, 5) + + def test_ensure_4d_5d_t_gt_1(self): + tensor = torch.rand(3, 4, 5, 2, 2) + with pytest.raises(ValueError): + io.ensure_4d(tensor) + + def test_ensure_4d_2d(self): + tensor = torch.rand(4, 5) + assert io.ensure_4d(tensor).shape == (1, 4, 5, 1) + + def test_ensure_4d_2d_3dims_rgb_first(self): + tensor = torch.rand(3, 4, 5) + assert io.ensure_4d(tensor).shape == (3, 4, 5, 1) + + def test_ensure_4d_2d_3dims_rgb_last(self): + tensor = torch.rand(4, 5, 3) + assert io.ensure_4d(tensor).shape == (3, 4, 5, 1) + + def test_ensure_4d_3d(self): + tensor = torch.rand(4, 5, 6) + assert io.ensure_4d(tensor).shape == (1, 4, 5, 6) + + def test_ensure_4d_2_spatial_dims(self): + tensor = torch.rand(4, 5, 6) + assert io.ensure_4d(tensor, num_spatial_dims=2).shape == (4, 5, 6, 1) + + def test_ensure_4d_3_spatial_dims(self): + tensor = torch.rand(4, 5, 6) + assert io.ensure_4d(tensor, num_spatial_dims=3).shape == (1, 4, 5, 6) + + def test_ensure_4d_nd_not_supported(self): + tensor = torch.rand(1, 2, 3, 4, 5) + with pytest.raises(ValueError): + io.ensure_4d(tensor) + + def test_sitk_to_nib(self): + data = np.random.rand(10, 12) + image = sitk.GetImageFromArray(data) + tensor, _ = io.sitk_to_nib(image) + assert data.sum() == pytest.approx(tensor.sum()) + + def test_sitk_to_affine(self): + spacing = 1, 2, 3 + direction_lps = -1, 0, 0, 0, -1, 0, 0, 0, 1 + origin_lps = left, posterior, superior = -10, -20, 30 + image = sitk.GetImageFromArray(np.random.rand(10, 20, 30)) + image.SetDirection(direction_lps) + image.SetSpacing(spacing) + image.SetOrigin(origin_lps) + origin_ras = -left, -posterior, superior + fixture = np.diag((*spacing, 1)) + fixture[:3, 3] = origin_ras + affine = io.get_ras_affine_from_sitk(image) + self.assert_tensor_almost_equal(fixture, affine) + + +# This doesn't work as a method of the class +libs = 'sitk', 'nibabel' +parameters = [] +for save_lib in libs: + for load_lib in libs: + for dims in 2, 3, 4: + parameters.append((save_lib, load_lib, dims)) + + +@pytest.mark.parametrize(('save_lib', 'load_lib', 'dims'), parameters) +def test_write_nd_with_a_read_it_with_b(save_lib, load_lib, dims): + shape = [1, 4, 5, 6] + if dims == 2: + shape[-1] = 1 + elif dims == 4: + shape[0] = 2 + tensor = torch.randn(*shape) + affine = np.eye(4) + tempdir = Path(tempfile.gettempdir()) / '.torchio_tests' + tempdir.mkdir(exist_ok=True) + path = tempdir / 'test_io.nii' + save_function = getattr(io, f'_write_{save_lib}') + load_function = getattr(io, f'_read_{save_lib}') + save_function(tensor, affine, path) + loaded_tensor, loaded_affine = load_function(path) + TorchioTestCase.assert_tensor_equal( + tensor.squeeze(), + loaded_tensor.squeeze(), + msg=f'Save lib: {save_lib}; load lib: {load_lib}; dims: {dims}', + check_stride=False, + ) + TorchioTestCase.assert_tensor_equal(affine, loaded_affine) + + +class TestNibabelToSimpleITK(TorchioTestCase): + def setUp(self): + super().setUp() + self.affine = np.eye(4) + + def test_wrong_num_dims(self): + with pytest.raises(ValueError): + io.nib_to_sitk(np.random.rand(10, 10), self.affine) + + def test_2d_single(self): + data = np.random.rand(1, 10, 12, 1) + image = io.nib_to_sitk(data, self.affine) + assert image.GetDimension() == 2 + assert image.GetSize() == (10, 12) + assert image.GetNumberOfComponentsPerPixel() == 1 + + def test_2d_multi(self): + data = np.random.rand(5, 10, 12, 1) + image = io.nib_to_sitk(data, self.affine) + assert image.GetDimension() == 2 + assert image.GetSize() == (10, 12) + assert image.GetNumberOfComponentsPerPixel() == 5 + + def test_2d_3d_single(self): + data = np.random.rand(1, 10, 12, 1) + image = io.nib_to_sitk(data, self.affine, force_3d=True) + assert image.GetDimension() == 3 + assert image.GetSize() == (10, 12, 1) + assert image.GetNumberOfComponentsPerPixel() == 1 + + def test_2d_3d_multi(self): + data = np.random.rand(5, 10, 12, 1) + image = io.nib_to_sitk(data, self.affine, force_3d=True) + assert image.GetDimension() == 3 + assert image.GetSize() == (10, 12, 1) + assert image.GetNumberOfComponentsPerPixel() == 5 + + def test_3d_single(self): + data = np.random.rand(1, 8, 10, 12) + image = io.nib_to_sitk(data, self.affine) + assert image.GetDimension() == 3 + assert image.GetSize() == (8, 10, 12) + assert image.GetNumberOfComponentsPerPixel() == 1 + + def test_3d_multi(self): + data = np.random.rand(5, 8, 10, 12) + image = io.nib_to_sitk(data, self.affine) + assert image.GetDimension() == 3 + assert image.GetSize() == (8, 10, 12) + assert image.GetNumberOfComponentsPerPixel() == 5 diff --git a/tests/data/test_queue.py b/tests/data/test_queue.py new file mode 100644 index 000000000..1b8b98e72 --- /dev/null +++ b/tests/data/test_queue.py @@ -0,0 +1,81 @@ +import sys + +import pytest +import torch +from parameterized import parameterized + +import torchio as tio +from torchio.data import UniformSampler +from torchio.utils import create_dummy_dataset + +from ..utils import TorchioTestCase + + +class TestQueue(TorchioTestCase): + """Tests for `queue` module.""" + + def setUp(self): + super().setUp() + self.subjects_list = create_dummy_dataset( + num_images=10, + size_range=(10, 20), + directory=self.dir, + suffix='.nii', + force=False, + ) + + def run_queue(self, num_workers=0, **kwargs): + subjects_dataset = tio.SubjectsDataset(self.subjects_list) + patch_size = 10 + sampler = UniformSampler(patch_size) + queue_dataset = tio.Queue( + subjects_dataset, + max_length=6, + samples_per_volume=2, + sampler=sampler, + num_workers=num_workers, + **kwargs, + ) + _ = str(queue_dataset) + batch_loader = tio.SubjectsLoader(queue_dataset, batch_size=4) + for batch in batch_loader: + _ = batch['one_modality'][tio.DATA] + _ = batch['segmentation'][tio.DATA] + return queue_dataset + + def test_queue(self): + self.run_queue(num_workers=0) + + @pytest.mark.skipif(sys.platform == 'darwin', reason='Takes too long on macOS') + def test_queue_multiprocessing(self): + self.run_queue(num_workers=2) + + def test_queue_no_start_background(self): + self.run_queue(num_workers=0, start_background=False) + + @parameterized.expand([(11,), (12,)]) + def test_different_samples_per_volume(self, max_length): + image2 = tio.ScalarImage(tensor=2 * torch.ones(1, 1, 1, 1)) + image10 = tio.ScalarImage(tensor=10 * torch.ones(1, 1, 1, 1)) + subject2 = tio.Subject(im=image2, num_samples=2) + subject10 = tio.Subject(im=image10, num_samples=10) + dataset = tio.SubjectsDataset([subject2, subject10]) + patch_size = 1 + sampler = UniformSampler(patch_size) + queue_dataset = tio.Queue( + dataset, + max_length=max_length, + samples_per_volume=3, # should be ignored + sampler=sampler, + shuffle_patches=False, + ) + batch_loader = tio.SubjectsLoader(queue_dataset, batch_size=6) + tensors = [batch['im'][tio.DATA] for batch in batch_loader] + all_numbers = torch.stack(tensors).flatten().tolist() + assert all_numbers.count(10) == 10 + assert all_numbers.count(2) == 2 + + def test_get_memory_string(self): + queue = self.run_queue() + memory_string = queue.get_max_memory_pretty() + assert isinstance(memory_string, str) diff --git a/tests/data/test_subject.py b/tests/data/test_subject.py new file mode 100644 index 000000000..74ce3196e --- /dev/null +++ b/tests/data/test_subject.py @@ -0,0 +1,218 @@ +import copy +import sys +import tempfile +from typing import cast + +import numpy as np +import pytest +import torch + +import torchio as tio + +from ..utils import TorchioTestCase + + +class TestSubject(TorchioTestCase): + """Tests for `Subject`.""" + + def test_positional_args(self): + with pytest.raises(ValueError): + tio.Subject(cast(dict[str, object], 0)) + + def test_input_dict(self): + with tempfile.NamedTemporaryFile(delete=False) as f: + input_dict = {'image': tio.ScalarImage(f.name)} + tio.Subject(input_dict) + tio.Subject(**input_dict) + + def test_no_sample(self): + with tempfile.NamedTemporaryFile(delete=False) as f: + input_dict = {'image': tio.ScalarImage(f.name)} + subject = tio.Subject(input_dict) + with pytest.raises(RuntimeError): + with pytest.warns(UserWarning): + tio.RandomFlip()(subject) + + def test_history(self): + transformed = tio.RandomGamma()(self.sample_subject) + assert len(transformed.history) == 1 + + def test_inconsistent_shape(self): + subject = tio.Subject( + a=tio.ScalarImage(tensor=torch.rand(1, 2, 3, 4)), + b=tio.ScalarImage(tensor=torch.rand(2, 2, 3, 4)), + ) + _ = subject.spatial_shape + with pytest.raises(RuntimeError): + _ = subject.shape + + def test_inconsistent_spatial_shape(self): + subject = tio.Subject( + a=tio.ScalarImage(tensor=torch.rand(1, 3, 3, 4)), + b=tio.ScalarImage(tensor=torch.rand(2, 2, 3, 4)), + ) + with pytest.raises(RuntimeError): + _ = subject.spatial_shape + + @pytest.mark.slow + @pytest.mark.skipif(sys.platform == 'win32', reason='Unstable on Windows') + def test_plot(self): + self.sample_subject.plot( + show=False, + output_path=self.dir / 'figure.png', + cmap_dict={ + 't2': 'viridis', + 'label': {0: 'yellow', 1: 'blue'}, + }, + ) + + @pytest.mark.slow + @pytest.mark.skipif(sys.platform == 'win32', reason='Unstable on Windows') + def test_plot_one_image(self): + path = self.get_image_path('t1_plot') + subject = tio.Subject(t1=tio.ScalarImage(path)) + subject.plot(show=False) + + def test_same_space(self): + # https://github.com/TorchIO-project/torchio/issues/381 + affine1 = np.array( + [ + [ + 4.27109375e-14, + -8.71264808e-03, + 9.99876633e-01, + -3.39850907e01, + ], + [ + -5.54687500e-01, + -2.71630469e-12, + 8.75148028e-17, + 1.62282930e02, + ], + [ + 2.71575000e-12, + -5.54619070e-01, + -1.57073092e-02, + 2.28515784e02, + ], + [0.00000000e00, 0.00000000e00, 0.00000000e00, 1.00000000e00], + ] + ) + affine2 = np.array( + [ + [ + 3.67499773e-08, + -8.71257665e-03, + 9.99876635e-01, + -3.39850922e01, + ], + [ + -5.54687500e-01, + 3.67499771e-08, + 6.73024385e-08, + 1.62282928e02, + ], + [ + -3.73318194e-08, + -5.54619071e-01, + -1.57071802e-02, + 2.28515778e02, + ], + [0.00000000e00, 0.00000000e00, 0.00000000e00, 1.00000000e00], + ] + ) + t = torch.rand(1, 2, 3, 4) + subject = tio.Subject( + im1=tio.ScalarImage(tensor=t, affine=affine1), + im2=tio.ScalarImage(tensor=t, affine=affine2), + ) + subject.check_consistent_space() + + def test_delete_image(self): + subject = copy.deepcopy(self.sample_subject) + subject.remove_image('t1') + with pytest.raises(KeyError): + subject['t1'] + with pytest.raises(AttributeError): + _ = subject.t1 + + def test_2d(self): + subject = self.make_2d(self.sample_subject) + assert subject.is_2d() + + def test_different_non_numeric(self): + with pytest.raises(RuntimeError): + self.sample_subject.check_consistent_attribute('path') + + def test_bad_arg(self): + with pytest.raises(ValueError): + tio.Subject(cast(dict[str, object], 0)) + + def test_no_images(self): + with pytest.raises(TypeError): + tio.Subject(a=0) + + def test_copy_subject(self): + sub_copy = copy.copy(self.sample_subject) + assert isinstance(sub_copy, tio.data.Subject) + sub_deep_copy = copy.deepcopy(self.sample_subject) + assert isinstance(sub_deep_copy, tio.data.Subject) + + def test_copy_subclass(self): + class DummySubjectSubClass(tio.data.Subject): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + dummy_sub = DummySubjectSubClass( + attr_1='abcd', + attr_2=tio.ScalarImage(tensor=torch.zeros(1, 1, 1, 1)), + ) + sub_copy = copy.copy(dummy_sub) + assert isinstance(sub_copy, tio.data.Subject) + assert isinstance(sub_copy, DummySubjectSubClass) + sub_deep_copy = copy.deepcopy(dummy_sub) + assert isinstance(sub_deep_copy, tio.data.Subject) + assert isinstance(sub_deep_copy, DummySubjectSubClass) + + def test_load_unload(self): + self.sample_subject.load() + for image in self.sample_subject.get_images(intensity_only=False): + assert image._loaded + self.sample_subject.unload() + for image in self.sample_subject.get_images(intensity_only=False): + assert not image._loaded + + def test_subjects_batch(self): + subjects = tio.SubjectsDataset(10 * [self.sample_subject]) + loader = tio.SubjectsLoader(subjects, batch_size=4) + batch = next(iter(loader)) + assert batch.__class__ is dict + + def test_deep_copy_subject(self): + sub_copy = copy.deepcopy(self.sample_subject) + assert isinstance(sub_copy, tio.data.Subject) + + sub_copy_t1 = sub_copy.get_scalar_image('t1') + sample_t1 = self.sample_subject.get_scalar_image('t1') + new_tensor = torch.ones_like(sub_copy_t1.data) + sub_copy_t1.set_data(new_tensor) + # The data of the original subject should not be modified + assert not torch.allclose(sub_copy_t1.data, sample_t1.data) + + def test_shallow_copy_subject(self): + # We are creating a deep copy of the original subject first to not modify the original subject + copy_original_subj = copy.deepcopy(self.sample_subject) + sub_copy = copy.copy(copy_original_subj) + assert isinstance(sub_copy, tio.data.Subject) + + sub_copy_t1 = sub_copy.get_scalar_image('t1') + copy_original_t1 = copy_original_subj.get_scalar_image('t1') + sample_t1 = self.sample_subject.get_scalar_image('t1') + new_tensor = torch.ones_like(sub_copy_t1.data) + sub_copy_t1.set_data(new_tensor) + + # The data of both copies needs to be the same as we are using a shallow copy + assert torch.allclose(sub_copy_t1.data, copy_original_t1.data) + # The data of the original subject should not be modified + assert not torch.allclose(sub_copy_t1.data, sample_t1.data) + assert not torch.allclose(copy_original_t1.data, sample_t1.data) diff --git a/tests/data/test_subjects_dataset.py b/tests/data/test_subjects_dataset.py new file mode 100644 index 000000000..681591c13 --- /dev/null +++ b/tests/data/test_subjects_dataset.py @@ -0,0 +1,71 @@ +from collections.abc import Callable +from typing import cast + +import pytest +import torch + +import torchio as tio + +from ..utils import TorchioTestCase + + +class TestSubjectsDataset(TorchioTestCase): + def test_indexing_nonint(self): + dset = tio.SubjectsDataset(self.subjects_list) + dset[cast(int, torch.tensor(0))] + + def test_images(self): + self.iterate_dataset(self.subjects_list) + + def test_empty_subjects_list(self): + with pytest.raises(ValueError): + self.iterate_dataset([]) + + def test_empty_subjects_tuple(self): + with pytest.raises(ValueError): + self.iterate_dataset(()) + + def test_wrong_subjects_type(self): + with pytest.raises(TypeError): + self.iterate_dataset(0) + + def test_wrong_subject_type_int(self): + with pytest.raises(TypeError): + self.iterate_dataset([0]) + + def test_wrong_subject_type_dict(self): + with pytest.raises(TypeError): + self.iterate_dataset([{}]) + + def test_wrong_index(self): + with pytest.raises(ValueError): + self.dataset[cast(int, slice(None, 3))] + + def test_wrong_transform_init(self): + with pytest.raises(ValueError): + invalid_transform = cast(Callable[[tio.Subject], tio.Subject], {}) + tio.SubjectsDataset( + self.subjects_list, + transform=invalid_transform, + ) + + def test_wrong_transform_arg(self): + with pytest.raises(ValueError): + invalid_transform = cast(Callable[[tio.Subject], tio.Subject], 1) + self.dataset.set_transform(invalid_transform) + + @staticmethod + def iterate_dataset(subjects_list): + dataset = tio.SubjectsDataset(subjects_list) + for _ in dataset: + pass + + def test_from_batch(self): + dataset = tio.SubjectsDataset([self.sample_subject]) + loader = tio.SubjectsLoader(dataset) + batch = tio.utils.get_first_item(loader) + new_dataset = tio.SubjectsDataset.from_batch(batch) + self.assert_tensor_equal( + dataset[0].t1.data, + new_dataset[0].t1.data, + ) diff --git a/tests/datasets/__init__.py b/tests/datasets/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/datasets/test_ixi.py b/tests/datasets/test_ixi.py new file mode 100644 index 000000000..e4b8e9922 --- /dev/null +++ b/tests/datasets/test_ixi.py @@ -0,0 +1,17 @@ +import pytest + +import torchio as tio + +from ..utils import TorchioTestCase + + +class TestIXI(TorchioTestCase): + """Tests for `ixi` module.""" + + def test_not_downloaded(self): + with pytest.raises(RuntimeError): + tio.datasets.IXI('testing123', download=False) + + def test_tiny_not_downloaded(self): + with pytest.raises(RuntimeError): + tio.datasets.IXITiny('testing123', download=False) diff --git a/tests/datasets/test_medmnist.py b/tests/datasets/test_medmnist.py new file mode 100644 index 000000000..c9872baee --- /dev/null +++ b/tests/datasets/test_medmnist.py @@ -0,0 +1,34 @@ +import os + +import pytest + +import torchio as tio +from torchio.datasets.medmnist import AdrenalMNIST3D +from torchio.datasets.medmnist import FractureMNIST3D +from torchio.datasets.medmnist import NoduleMNIST3D +from torchio.datasets.medmnist import OrganMNIST3D +from torchio.datasets.medmnist import SynapseMNIST3D +from torchio.datasets.medmnist import VesselMNIST3D + +classes = ( + OrganMNIST3D, + NoduleMNIST3D, + AdrenalMNIST3D, + FractureMNIST3D, + VesselMNIST3D, + SynapseMNIST3D, +) + + +@pytest.mark.slow +@pytest.mark.skipif('CI' in os.environ, reason='Unstable on GitHub Actions') +@pytest.mark.parametrize('class_', classes) +@pytest.mark.parametrize('split', ('train', 'val', 'test')) +def test_load_all(class_, split): + dataset = class_(split) + loader = tio.SubjectsLoader( + dataset, + batch_size=256, + ) + for _ in loader: + pass diff --git a/tests/test_affine.py b/tests/test_affine.py deleted file mode 100644 index 1dd5cc396..000000000 --- a/tests/test_affine.py +++ /dev/null @@ -1,290 +0,0 @@ -"""Tests for AffineMatrix.""" - -from __future__ import annotations - -import numpy as np -import pytest - -from torchio import AffineMatrix - - -class TestAffineCreation: - def test_identity(self): - affine = AffineMatrix() - np.testing.assert_array_equal(affine.numpy(), np.eye(4)) - - def test_from_numpy(self): - matrix = np.diag([2.0, 3.0, 4.0, 1.0]) - affine = AffineMatrix(matrix) - np.testing.assert_array_equal(affine.numpy(), matrix) - - def test_from_list(self): - matrix = np.eye(4).tolist() - affine = AffineMatrix(matrix) - np.testing.assert_array_equal(affine.numpy(), np.eye(4)) - - def test_must_be_4x4(self): - with pytest.raises(ValueError, match=r"4.*4"): - AffineMatrix(np.eye(3)) - - def test_must_be_2d(self): - with pytest.raises(ValueError, match=r"4.*4"): - AffineMatrix(np.ones((4, 4, 4))) - - def test_always_float64(self): - matrix = np.eye(4, dtype=np.float32) - affine = AffineMatrix(matrix) - assert affine.numpy().dtype == np.float64 - - def test_copies_input(self): - matrix = np.eye(4) - affine = AffineMatrix(matrix) - matrix[0, 0] = 999 - assert affine.numpy()[0, 0] == 1.0 - - -class TestAffineFromSpacingOrigin: - def test_isotropic(self): - affine = AffineMatrix.from_spacing(spacing=(2.0, 2.0, 2.0)) - np.testing.assert_allclose(affine.spacing, (2.0, 2.0, 2.0)) - np.testing.assert_allclose(affine.origin, (0.0, 0.0, 0.0)) - - def test_anisotropic(self): - affine = AffineMatrix.from_spacing(spacing=(0.5, 0.8, 1.2)) - np.testing.assert_allclose(affine.spacing, (0.5, 0.8, 1.2)) - - def test_with_origin(self): - affine = AffineMatrix.from_spacing( - spacing=(1.0, 1.0, 1.0), - origin=(100.0, 200.0, 300.0), - ) - np.testing.assert_allclose(affine.origin, (100.0, 200.0, 300.0)) - - def test_with_direction(self): - # 90-degree rotation around z-axis - direction = np.array( - [ - [0, -1, 0], - [1, 0, 0], - [0, 0, 1], - ], - dtype=np.float64, - ) - affine = AffineMatrix.from_spacing( - spacing=(2.0, 2.0, 2.0), - direction=direction, - ) - np.testing.assert_allclose(affine.spacing, (2.0, 2.0, 2.0)) - np.testing.assert_allclose(affine.direction, direction, atol=1e-10) - - -class TestAffineProperties: - def test_spacing_identity(self): - affine = AffineMatrix() - np.testing.assert_allclose(affine.spacing, (1.0, 1.0, 1.0)) - - def test_spacing_scaled(self): - affine = AffineMatrix(np.diag([0.5, 0.8, 1.2, 1.0])) - np.testing.assert_allclose(affine.spacing, (0.5, 0.8, 1.2)) - - def test_origin_identity(self): - affine = AffineMatrix() - np.testing.assert_allclose(affine.origin, (0.0, 0.0, 0.0)) - - def test_origin_translated(self): - matrix = np.eye(4) - matrix[:3, 3] = [10, 20, 30] - affine = AffineMatrix(matrix) - np.testing.assert_allclose(affine.origin, (10.0, 20.0, 30.0)) - - def test_direction_identity(self): - affine = AffineMatrix() - np.testing.assert_allclose(affine.direction, np.eye(3)) - - def test_direction_with_rotation(self): - # 90-degree rotation around z - matrix = np.eye(4) - matrix[:3, :3] = [[0, -1, 0], [1, 0, 0], [0, 0, 1]] - affine = AffineMatrix(matrix) - expected = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=np.float64) - np.testing.assert_allclose(affine.direction, expected, atol=1e-10) - - def test_orientation_ras(self): - affine = AffineMatrix() - assert affine.orientation == ("R", "A", "S") - - def test_orientation_las(self): - matrix = np.diag([-1.0, 1.0, 1.0, 1.0]) - affine = AffineMatrix(matrix) - assert affine.orientation == ("L", "A", "S") - - -class TestAffineInverse: - def test_inverse_identity(self): - affine = AffineMatrix() - inv = affine.inverse() - np.testing.assert_allclose(inv.numpy(), np.eye(4)) - - def test_inverse_scaling(self): - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - inv = affine.inverse() - np.testing.assert_allclose(inv.spacing, (0.5, 1 / 3, 0.25)) - - def test_inverse_roundtrip(self): - matrix = np.eye(4) - matrix[:3, :3] = [[0, -2, 0], [3, 0, 0], [0, 0, 4]] - matrix[:3, 3] = [10, 20, 30] - affine = AffineMatrix(matrix) - roundtrip = affine.inverse().inverse() - np.testing.assert_allclose(roundtrip.numpy(), affine.numpy(), atol=1e-10) - - -class TestAffineCompose: - def test_compose_identity(self): - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - result = affine.compose(AffineMatrix()) - np.testing.assert_allclose(result.numpy(), affine.numpy()) - - def test_compose_translations(self): - a = AffineMatrix.from_spacing(spacing=(1, 1, 1), origin=(10, 0, 0)) - b = AffineMatrix.from_spacing(spacing=(1, 1, 1), origin=(0, 20, 0)) - result = a.compose(b) - np.testing.assert_allclose(result.origin, (10, 20, 0)) - - def test_compose_matmul(self): - m1 = np.eye(4) - m1[:3, 3] = [1, 2, 3] - m2 = np.diag([2.0, 2.0, 2.0, 1.0]) - result = AffineMatrix(m1).compose(AffineMatrix(m2)) - np.testing.assert_allclose(result.numpy(), m1 @ m2) - - -class TestAffineMatmul: - def test_matmul_operator(self): - m1 = np.eye(4) - m1[:3, 3] = [1, 2, 3] - m2 = np.diag([2.0, 2.0, 2.0, 1.0]) - result = AffineMatrix(m1) @ AffineMatrix(m2) - np.testing.assert_allclose(result.numpy(), m1 @ m2) - - def test_matmul_identity(self): - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - result = affine @ AffineMatrix() - np.testing.assert_allclose(result.numpy(), affine.numpy()) - - def test_matmul_inverse(self): - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - result = affine @ affine.inverse() - np.testing.assert_allclose(result.numpy(), np.eye(4), atol=1e-10) - - def test_matmul_returns_affine(self): - result = AffineMatrix() @ AffineMatrix() - assert isinstance(result, AffineMatrix) - - def test_matmul_not_implemented_for_other_types(self): - with pytest.raises(TypeError): - AffineMatrix() @ "not an affine" - - -class TestAffineApply: - def test_apply_identity(self): - affine = AffineMatrix() - points = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - result = affine.apply(points) - np.testing.assert_allclose(result, points) - - def test_apply_translation(self): - matrix = np.eye(4) - matrix[:3, 3] = [10, 20, 30] - affine = AffineMatrix(matrix) - points = np.array([[0.0, 0.0, 0.0]]) - result = affine.apply(points) - np.testing.assert_allclose(result, [[10.0, 20.0, 30.0]]) - - def test_apply_scaling(self): - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - points = np.array([[1.0, 1.0, 1.0]]) - result = affine.apply(points) - np.testing.assert_allclose(result, [[2.0, 3.0, 4.0]]) - - def test_apply_single_point(self): - affine = AffineMatrix(np.diag([2.0, 2.0, 2.0, 1.0])) - point = np.array([[5.0, 5.0, 5.0]]) - result = affine.apply(point) - np.testing.assert_allclose(result, [[10.0, 10.0, 10.0]]) - - -class TestAffineNumpyInterop: - def test_array_protocol(self): - matrix = np.diag([2.0, 3.0, 4.0, 1.0]) - affine = AffineMatrix(matrix) - result = np.asarray(affine) - np.testing.assert_array_equal(result, matrix) - - def test_array_with_dtype(self): - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - result = np.array(affine, dtype=np.float32) - assert result.dtype == np.float32 - np.testing.assert_allclose(result, np.diag([2.0, 3.0, 4.0, 1.0])) - - def test_array_with_copy(self): - affine = AffineMatrix() - result = np.array(affine, copy=True) - result[0, 0] = 999 - assert affine.numpy()[0, 0] == 1.0 - - def test_matmul_with_numpy(self): - affine = AffineMatrix(np.diag([2.0, 2.0, 2.0, 1.0])) - vec = np.array([1.0, 1.0, 1.0, 1.0]) - result = np.asarray(affine) @ vec - np.testing.assert_allclose(result, [2.0, 2.0, 2.0, 1.0]) - - -class TestAffineRepr: - def test_repr_identity(self): - affine = AffineMatrix() - r = repr(affine) - assert "AffineMatrix" in r - assert "1.00" in r - - def test_repr_scaled(self): - affine = AffineMatrix(np.diag([0.5, 0.8, 1.2, 1.0])) - r = repr(affine) - assert "0.50" in r - assert "0.80" in r - assert "1.20" in r - - -class TestAffineCopy: - def test_copy(self): - import copy - - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - copied = copy.copy(affine) - np.testing.assert_array_equal(copied.numpy(), affine.numpy()) - # Verify independence - copied.numpy()[0, 0] = 999 - assert affine.numpy()[0, 0] == 2.0 - - def test_deepcopy(self): - import copy - - affine = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - copied = copy.deepcopy(affine) - np.testing.assert_array_equal(copied.numpy(), affine.numpy()) - - -class TestAffineEquality: - def test_equal(self): - a = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - b = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - assert a == b - - def test_not_equal(self): - a = AffineMatrix(np.diag([2.0, 3.0, 4.0, 1.0])) - b = AffineMatrix() - assert a != b - - def test_not_equal_to_other_type(self): - a = AffineMatrix() - assert a != "not an affine" diff --git a/tests/test_anisotropy.py b/tests/test_anisotropy.py deleted file mode 100644 index 5692b28cd..000000000 --- a/tests/test_anisotropy.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Tests for Anisotropy transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestAnisotropy: - def test_changes_data(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Anisotropy(downsampling=3.0)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_preserves_shape(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Anisotropy(downsampling=2.0)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_specific_axis(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Anisotropy(axes=(0,), downsampling=3.0)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_labels_use_nearest(self) -> None: - subject = _make_subject() - result = tio.Anisotropy(downsampling=2.0)(subject) - unique = result.seg.data.unique().tolist() - for v in unique: - assert v == int(v) - - def test_factor_one_is_identity(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Anisotropy(downsampling=1.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - -class TestAnisotropyPerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - data = torch.rand(1, 12, 12, 12) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Anisotropy(downsampling=(2.0, 5.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["factor"]) == batch.batch_size - assert not torch.allclose(result.t1.data[0], result.t1.data[1]) - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Anisotropy(downsampling=(2.0, 5.0), per_instance=False)(batch) - torch.testing.assert_close(result.t1.data[0], result.t1.data[1]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 12, 12, 12))) - result = tio.Anisotropy(downsampling=(2.0, 5.0))(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params - - -class TestAnisotropyAxisValidation: - def test_out_of_range_axis_raises(self) -> None: - # An active per-element axis outside {0, 1, 2} must raise, matching the - # scalar path, rather than silently becoming a no-op. - from torchio.transforms.spatial.anisotropy import ( - _simulate_anisotropy_per_instance, - ) - - with pytest.raises(ValueError, match="axis must be in"): - _simulate_anisotropy_per_instance( - torch.rand(2, 1, 8, 8, 8), - axes=[0, 3], - factors=[2.0, 2.0], - mode="linear", - ) diff --git a/tests/test_axes.py b/tests/test_axes.py deleted file mode 100644 index b10616807..000000000 --- a/tests/test_axes.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Tests for axis validation and conversion utilities.""" - -from __future__ import annotations - -import itertools - -import pytest - -from torchio.data.axes import ANATOMICAL_PAIRS -from torchio.data.axes import AxesType -from torchio.data.axes import axes_type -from torchio.data.axes import get_axis_mapping -from torchio.data.axes import validate_axes - - -class TestValidateAxes: - """Test axis string validation.""" - - # --- Valid voxel axes --- - - @pytest.mark.parametrize( - "axes", - ["".join(p) for p in itertools.permutations("IJK")], - ) - def test_all_voxel_permutations_valid(self, axes: str): - assert validate_axes(axes) == axes - - # --- Valid anatomical axes --- - - def test_ras_valid(self): - assert validate_axes("RAS") == "RAS" - - def test_lpi_valid(self): - assert validate_axes("LPI") == "LPI" - - def test_air_valid(self): - """One from each pair, unusual order.""" - assert validate_axes("AIR") == "AIR" - - def test_all_anatomical_combinations_valid(self): - """There are 48 valid anatomical axis strings (8 sign combos x 6 orders).""" - count = 0 - for choices in itertools.product(*ANATOMICAL_PAIRS): - for perm in itertools.permutations(choices): - validate_axes("".join(perm)) - count += 1 - assert count == 48 - - # --- Invalid axes --- - - def test_xyz_invalid(self): - with pytest.raises(ValueError, match="Invalid"): - validate_axes("XYZ") - - def test_too_short(self): - with pytest.raises(ValueError, match="3 characters"): - validate_axes("IJ") - - def test_too_long(self): - with pytest.raises(ValueError, match="3 characters"): - validate_axes("IJKL") - - def test_duplicate_voxel(self): - with pytest.raises(ValueError, match="Invalid"): - validate_axes("IIJ") - - def test_duplicate_anatomical(self): - with pytest.raises(ValueError, match="Invalid"): - validate_axes("RRS") - - def test_same_pair_twice(self): - """R and L are from the same pair (invalid).""" - with pytest.raises(ValueError, match="Invalid"): - validate_axes("RLS") - - def test_mixed_voxel_anatomical(self): - """I and J are voxel, R is anatomical: neither system matches.""" - with pytest.raises(ValueError, match="Invalid"): - validate_axes("IJR") - - def test_lowercase_invalid(self): - with pytest.raises(ValueError, match=r"3 characters|Invalid"): - validate_axes("ijk") - - def test_empty_string(self): - with pytest.raises(ValueError, match="3 characters"): - validate_axes("") - - -class TestAxesType: - """Test axis type detection.""" - - def test_ijk_is_voxel(self): - assert axes_type("IJK") == AxesType.VOXEL - - def test_kji_is_voxel(self): - assert axes_type("KJI") == AxesType.VOXEL - - def test_ras_is_anatomical(self): - assert axes_type("RAS") == AxesType.ANATOMICAL - - def test_lpi_is_anatomical(self): - assert axes_type("LPI") == AxesType.ANATOMICAL - - -class TestAxisMapping: - """Test computing permutation and flip between axis systems.""" - - def test_ijk_to_ijk_identity(self): - perm, flips = get_axis_mapping("IJK", "IJK") - assert perm == (0, 1, 2) - assert flips == (False, False, False) - - def test_ijk_to_kji(self): - perm, flips = get_axis_mapping("IJK", "KJI") - assert perm == (2, 1, 0) - assert flips == (False, False, False) - - def test_ijk_to_jki(self): - perm, flips = get_axis_mapping("IJK", "JKI") - assert perm == (1, 2, 0) - assert flips == (False, False, False) - - def test_ras_to_ras_identity(self): - perm, flips = get_axis_mapping("RAS", "RAS") - assert perm == (0, 1, 2) - assert flips == (False, False, False) - - def test_ras_to_lpi(self): - """R→L (flip), A→P (flip), S→I (flip), same order.""" - perm, flips = get_axis_mapping("RAS", "LPI") - assert perm == (0, 1, 2) - assert flips == (True, True, True) - - def test_ras_to_asr(self): - """R→R (col 0→2), A→A (col 1→0), S→S (col 2→1). No flips.""" - perm, flips = get_axis_mapping("RAS", "ASR") - assert perm == (1, 2, 0) - assert flips == (False, False, False) - - def test_ras_to_lai(self): - """R→L (flip, col 0→0), A→A (col 1→1), S→I (flip, col 2→2).""" - perm, flips = get_axis_mapping("RAS", "LAI") - assert perm == (0, 1, 2) - assert flips == (True, False, True) - - def test_ras_to_psl(self): - """A→P (flip, col 1→0), S→S (no flip, col 2→1), R→L (flip, col 0→2).""" - perm, flips = get_axis_mapping("RAS", "PSL") - assert perm == (1, 2, 0) - assert flips == (True, False, True) - - def test_cross_type_raises(self): - """Cannot map between voxel and anatomical directly.""" - with pytest.raises(ValueError, match="same type"): - get_axis_mapping("IJK", "RAS") diff --git a/tests/test_backends.py b/tests/test_backends.py deleted file mode 100644 index 3352d1aa5..000000000 --- a/tests/test_backends.py +++ /dev/null @@ -1,679 +0,0 @@ -"""Tests for lazy image data backends.""" - -from __future__ import annotations - -from pathlib import Path - -import nibabel as nib -import numpy as np -import pytest -import torch - -from torchio import ScalarImage -from torchio.data.backends import NibabelBackend -from torchio.data.backends import TensorBackend - - -class TestNibabelBackend: - @pytest.fixture - def nifti_path(self, tmp_path: Path) -> Path: - data = np.random.randn(10, 12, 14).astype(np.float32) - nii = nib.Nifti1Image(data, np.diag([2.0, 2.0, 2.0, 1.0])) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - return path - - @pytest.fixture - def multichannel_nifti_path(self, tmp_path: Path) -> Path: - data = np.random.randn(10, 12, 14, 3).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "multi.nii.gz" - nib.save(nii, path) - return path - - def test_shape_3d(self, nifti_path: Path): - nii = nib.load(nifti_path) - backend = NibabelBackend(nii) - assert backend.shape == (1, 10, 12, 14) - - def test_shape_4d(self, multichannel_nifti_path: Path): - nii = nib.load(multichannel_nifti_path) - backend = NibabelBackend(nii) - assert backend.shape == (3, 10, 12, 14) - - def test_affine(self, nifti_path: Path): - nii = nib.load(nifti_path) - backend = NibabelBackend(nii) - np.testing.assert_array_equal( - backend.affine, - np.diag([2.0, 2.0, 2.0, 1.0]), - ) - - def test_to_tensor_3d(self, nifti_path: Path): - nii = nib.load(nifti_path) - backend = NibabelBackend(nii) - tensor = backend.to_tensor() - assert tensor.shape == (1, 10, 12, 14) - assert tensor.dtype == torch.float32 - - def test_to_tensor_4d(self, multichannel_nifti_path: Path): - nii = nib.load(multichannel_nifti_path) - backend = NibabelBackend(nii) - tensor = backend.to_tensor() - assert tensor.shape == (3, 10, 12, 14) - - @pytest.mark.parametrize( - ("np_dtype", "torch_dtype"), - [ - (np.int16, torch.int16), - (np.int32, torch.int32), - (np.uint8, torch.uint8), - (np.uint16, torch.int32), # upcast: torch has no uint16 - (np.float32, torch.float32), - (np.float64, torch.float64), - ], - ) - def test_to_tensor_preserves_dtype( - self, - tmp_path: Path, - np_dtype: type, - torch_dtype: torch.dtype, - ) -> None: - data = np.zeros((10, 12, 14), dtype=np_dtype) - path = tmp_path / f"dtype_{np.dtype(np_dtype).name}.nii.gz" - nib.save(nib.Nifti1Image(data, np.eye(4)), path) - backend = NibabelBackend(nib.load(path)) - tensor = backend.to_tensor() - assert tensor.dtype == torch_dtype - - def test_getitem_3d(self, nifti_path: Path): - nii = nib.load(nifti_path) - backend = NibabelBackend(nii) - # Spatial slice in (I, J, K) space - sliced = backend[:, 2:5, 3:7, 4:8] - assert sliced.shape == (1, 3, 4, 4) - - def test_does_not_load_full_data_for_shape(self, nifti_path: Path): - nii = nib.load(nifti_path) - backend = NibabelBackend(nii) - _ = backend.shape - # dataobj should still be a proxy, not a loaded array - assert not isinstance(nii.dataobj, np.ndarray) - - def test_invalid_ndim_raises(self, tmp_path: Path): - data = np.random.randn(10, 12, 14, 3, 2).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "bad.nii.gz" - nib.save(nii, path) - nii = nib.load(path) - with pytest.raises(ValueError, match="3D or 4D"): - NibabelBackend(nii) - - -class TestImageWithBackends: - def test_from_tensor_uses_tensor_backend(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - assert isinstance(image._backend, TensorBackend) - - def test_nifti_uses_nibabel_backend(self, tmp_path: Path): - data = np.random.randn(10, 10, 10).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - assert not image.is_loaded - _ = image.data # trigger load - assert isinstance(image._backend, NibabelBackend) - - def test_shape_without_loading_uses_backend(self, tmp_path: Path): - data = np.random.randn(10, 12, 14).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - assert image.shape == (1, 10, 12, 14) - # Backend should be created but data not materialized - assert image._backend is not None - assert image._data is None - - def test_dataobj_returns_backend(self, tmp_path: Path): - data = np.random.randn(10, 10, 10).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - dataobj = image.dataobj - assert isinstance(dataobj, NibabelBackend) - - def test_dataobj_from_tensor(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - dataobj = image.dataobj - assert isinstance(dataobj, TensorBackend) - - def test_data_caches_tensor(self, tmp_path: Path): - data = np.random.randn(10, 10, 10).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - tensor1 = image.data - tensor2 = image.data - assert tensor1 is tensor2 # same object, cached - - def test_lazy_slice_via_dataobj(self, tmp_path: Path): - data = np.random.randn(10, 12, 14).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - backend = image.dataobj - sliced = backend[:, 2:5, 3:7, 4:8] - assert sliced.shape == (1, 3, 4, 4) - assert image._data is None # full tensor never materialized - - -class TestZarrBackend: - """Tests for ZarrBackend using nifti-zarr.""" - - @pytest.fixture(scope="class") - def zarr_path(self, tmp_path_factory: pytest.TempPathFactory) -> Path: - try: - import niizarr - except ImportError: - pytest.skip("nifti-zarr not installed") - tmp_path = tmp_path_factory.mktemp("zarr") - data = np.random.rand(16, 16, 16).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - nii_path = tmp_path / "test.nii" - nib.save(nii, nii_path) - zarr_path = tmp_path / "test.nii.zarr" - niizarr.nii2zarr(str(nii_path), str(zarr_path)) - return zarr_path - - def test_zarr_image_shape(self, zarr_path: Path): - image = ScalarImage(zarr_path) - assert image.shape == (1, 16, 16, 16) - - def test_zarr_lazy_load(self, zarr_path: Path): - from torchio.data.backends import ZarrBackend - - image = ScalarImage(zarr_path) - backend = image.dataobj - assert isinstance(backend, (NibabelBackend, ZarrBackend)) - assert image._data is None - - def test_zarr_slice(self, zarr_path: Path): - image = ScalarImage(zarr_path) - backend = image.dataobj - sliced = backend[:, 2:12, 2:12, 2:12] - assert sliced.shape == (1, 10, 10, 10) - assert image._data is None - - def test_zarr_materialize(self, zarr_path: Path): - image = ScalarImage(zarr_path) - tensor = image.data - assert isinstance(tensor, torch.Tensor) - assert tensor.shape == (1, 16, 16, 16) - - def test_read_nifti_zarr(self, zarr_path: Path): - from torchio.data.io import read_nifti_zarr - - tensor, affine = read_nifti_zarr(zarr_path) - assert tensor.shape == (1, 16, 16, 16) - assert affine.shape == (4, 4) - - -class TestBackendCoherence: - """Regression tests: `_backend` must stay coherent with `data`. - - See assessment flaws on stale backend state after `set_data()` and - `to()`, and strategy section "Fix correctness issues". - """ - - @pytest.fixture - def nifti_path(self, tmp_path: Path) -> Path: - data = np.random.randn(10, 12, 14).astype(np.float32) - nii = nib.Nifti1Image(data, np.diag([2.0, 2.0, 2.0, 1.0])) - path = tmp_path / "coherence.nii.gz" - nib.save(nii, path) - return path - - def test_set_data_refreshes_backend_shape(self) -> None: - image = ScalarImage(torch.randn(1, 4, 4, 4)) - image.set_data(torch.randn(2, 6, 6, 6)) - assert image.dataobj.shape == (2, 6, 6, 6) - assert tuple(image.dataobj.to_tensor().shape) == (2, 6, 6, 6) - - def test_set_data_refreshes_backend_values_tensor_source(self) -> None: - image = ScalarImage(torch.zeros(1, 4, 4, 4)) - new = torch.ones(1, 4, 4, 4) - image.set_data(new) - assert torch.equal(image.dataobj.to_tensor(), new) - - def test_set_data_refreshes_backend_path_source(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path) - new = torch.full((1, 10, 12, 14), 7.0) - image.set_data(new) - # dataobj must reflect the in-memory tensor, not the on-disk data. - assert torch.equal(image.dataobj.to_tensor(), new) - - def test_set_data_on_empty_image_defaults_to_identity(self) -> None: - # Created empty then filled: no affine source, so default to identity - # without crashing, and keep dataobj consistent. - image = ScalarImage() - image.set_data(torch.zeros(1, 2, 3, 4)) - assert image.dataobj.shape == (1, 2, 3, 4) - np.testing.assert_allclose(image.affine.numpy(), np.eye(4)) - np.testing.assert_allclose( - image.affine.numpy(), np.asarray(image.dataobj.affine) - ) - - def test_set_data_preserves_disk_affine(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path) - image.set_data(torch.full((1, 10, 12, 14), 7.0)) - disk = np.diag([2.0, 2.0, 2.0, 1.0]) - np.testing.assert_allclose(image.affine.numpy(), disk) - np.testing.assert_allclose(np.asarray(image.dataobj.affine), disk) - - def test_to_refreshes_backend_dtype(self) -> None: - image = ScalarImage(torch.randn(1, 4, 4, 4).float()) - image.to(torch.float64) - assert image.dataobj.dtype == np.dtype("float64") - assert image.dataobj.to_tensor().dtype == torch.float64 - - def test_to_refreshes_backend_path_source(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path) - image.to(torch.float64) - assert image.dataobj.dtype == np.dtype("float64") - assert image.dataobj.to_tensor().dtype == torch.float64 - - -class TestAffineOverride: - """Regression tests: an overridden affine must be reported consistently. - - See assessment flaw that `image.affine` and `image.dataobj.affine` can - disagree when the affine is overridden in the constructor. - """ - - @pytest.fixture - def nifti_path(self, tmp_path: Path) -> Path: - data = np.random.randn(10, 12, 14).astype(np.float32) - nii = nib.Nifti1Image(data, np.diag([2.0, 2.0, 2.0, 1.0])) - path = tmp_path / "override.nii.gz" - nib.save(nii, path) - return path - - def test_override_matches_dataobj_nifti(self, nifti_path: Path) -> None: - custom = np.diag([3.0, 4.0, 5.0, 1.0]) - image = ScalarImage(nifti_path, affine=custom) - np.testing.assert_allclose(image.affine.numpy(), custom) - np.testing.assert_allclose(np.asarray(image.dataobj.affine), custom) - - def test_no_override_uses_disk_affine(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path) - disk = np.diag([2.0, 2.0, 2.0, 1.0]) - np.testing.assert_allclose(image.affine.numpy(), disk) - np.testing.assert_allclose(np.asarray(image.dataobj.affine), disk) - - def test_override_matches_dataobj_tensor(self) -> None: - custom = np.diag([3.0, 4.0, 5.0, 1.0]) - image = ScalarImage(torch.randn(1, 4, 4, 4), affine=custom) - np.testing.assert_allclose(image.affine.numpy(), custom) - np.testing.assert_allclose(np.asarray(image.dataobj.affine), custom) - - -class TestVectorNifti5D: - """Lazy slicing of 5D vector NIfTI as written by SimpleITK: (I, J, K, 1, C).""" - - @pytest.fixture - def path_5d(self, tmp_path: Path) -> Path: - data = np.random.randn(8, 9, 10, 1, 3).astype(np.float32) - path = tmp_path / "vector.nii.gz" - nib.save(nib.Nifti1Image(data, np.eye(4)), path) - return path - - def test_shape(self, path_5d: Path) -> None: - image = ScalarImage(path_5d) - assert image.shape == (3, 8, 9, 10) - assert image._data is None - - def test_lazy_spatial_slice(self, path_5d: Path) -> None: - image = ScalarImage(path_5d) - ref = ScalarImage(path_5d).data - sliced = image.dataobj[:, 1:4, 2:5, 0:6] - assert image._data is None # full tensor never materialized - assert isinstance(sliced, torch.Tensor) - assert tuple(sliced.shape) == (3, 3, 3, 6) - assert torch.allclose(sliced.float(), ref[:, 1:4, 2:5, 0:6].float()) - - def test_lazy_channel_slice_preserves_dim(self, path_5d: Path) -> None: - image = ScalarImage(path_5d) - ref = ScalarImage(path_5d).data - sliced = image.dataobj[1] - assert tuple(sliced.shape) == (1, 8, 9, 10) - assert torch.allclose(sliced.float(), ref[1:2].float()) - - -class TestBackendSlicingContract: - """The backend slicing contract: always 4D `(C, I, J, K)` tensors. - - Direct `dataobj[...]` access must match `ref[normalize_index(index)]` for - every backend, where `ref` is the materialized `(C, I, J, K)` tensor. - """ - - @pytest.fixture(params=["tensor", "nifti_3d", "nifti_4d", "nifti_5d"]) - def image_and_ref( - self, - request: pytest.FixtureRequest, - tmp_path: Path, - ) -> tuple[ScalarImage, torch.Tensor]: - kind = request.param - if kind == "tensor": - data = torch.randn(3, 8, 9, 10) - return ScalarImage(data), data.clone() - if kind == "nifti_3d": - arr = np.random.randn(8, 9, 10).astype(np.float32) - elif kind == "nifti_4d": - arr = np.random.randn(8, 9, 10, 3).astype(np.float32) - else: # nifti_5d - arr = np.random.randn(8, 9, 10, 1, 3).astype(np.float32) - path = tmp_path / f"{kind}.nii.gz" - nib.save(nib.Nifti1Image(arr, np.eye(4)), path) - return ScalarImage(path), ScalarImage(path).data - - @pytest.mark.parametrize( - "index", - [ - slice(None), - 0, - -1, - slice(0, 1), - (slice(None), slice(1, 4)), - (slice(None), slice(1, 4), slice(2, 5), slice(0, 6)), - ..., - (..., slice(0, 5)), - (slice(None), slice(-4, None)), - (0, 1, 2, 3), - ], - ) - def test_matches_reference( - self, - image_and_ref: tuple[ScalarImage, torch.Tensor], - index, - ) -> None: - from torchio.data.backends import normalize_index - - image, ref = image_and_ref - result = image.dataobj[index] - expected = ref[normalize_index(index)] - assert isinstance(result, torch.Tensor) - assert result.ndim == 4 - assert tuple(result.shape) == tuple(expected.shape) - assert torch.allclose(result.float(), expected.float()) - - def test_multichannel_selection(self, tmp_path: Path) -> None: - arr = np.random.randn(8, 9, 10, 3).astype(np.float32) - path = tmp_path / "multi.nii.gz" - nib.save(nib.Nifti1Image(arr, np.eye(4)), path) - image = ScalarImage(path) - ref = ScalarImage(path).data - sliced = image.dataobj[0:2, 1:4] - assert tuple(sliced.shape) == (2, 3, 9, 10) - assert torch.allclose(sliced.float(), ref[0:2, 1:4].float()) - - -class TestSlicingReturnTypes: - """`__getitem__` returns tensors, preserving device/dtype where possible.""" - - def test_tensor_backend_preserves_dtype(self) -> None: - data = torch.randn(2, 4, 5, 6, dtype=torch.float64) - image = ScalarImage(data) - sliced = image.dataobj[:, 1:3] - assert isinstance(sliced, torch.Tensor) - assert sliced.dtype == torch.float64 - assert sliced.device == data.device - - def test_nibabel_backend_returns_tensor_lazily(self, tmp_path: Path) -> None: - arr = np.random.randn(8, 9, 10).astype(np.float32) - path = tmp_path / "lazy.nii.gz" - nib.save(nib.Nifti1Image(arr, np.eye(4)), path) - image = ScalarImage(path) - sliced = image.dataobj[:, 2:5] - assert isinstance(sliced, torch.Tensor) - assert image._data is None # still lazy - - def test_too_many_indices_raises(self) -> None: - image = ScalarImage(torch.randn(1, 4, 4, 4)) - with pytest.raises(IndexError, match="Too many indices"): - _ = image.dataobj[0, 0, 0, 0, 0] - - -class TestBackendResolver: - """Backend selection is delegated to the registry-based resolver.""" - - def test_resolve_nifti_path(self, tmp_path: Path) -> None: - from torchio.data.backends import BackendRequest - from torchio.data.backends import resolve_backend - - arr = np.random.randn(8, 9, 10).astype(np.float32) - path = tmp_path / "r.nii.gz" - nib.save(nib.Nifti1Image(arr, np.eye(4)), path) - backend = resolve_backend(BackendRequest(path=path)) - assert isinstance(backend, NibabelBackend) - - def test_resolve_non_nifti_returns_none(self, tmp_path: Path) -> None: - from torchio.data.backends import BackendRequest - from torchio.data.backends import resolve_backend - - path = tmp_path / "image.png" - backend = resolve_backend(BackendRequest(path=path)) - assert backend is None - - def test_resolve_applies_affine_override(self, tmp_path: Path) -> None: - from torchio.data.backends import BackendRequest - from torchio.data.backends import resolve_backend - - arr = np.random.randn(8, 9, 10).astype(np.float32) - path = tmp_path / "r.nii.gz" - nib.save(nib.Nifti1Image(arr, np.diag([2.0, 2.0, 2.0, 1.0])), path) - custom = torch.diag(torch.tensor([3.0, 4.0, 5.0, 1.0], dtype=torch.float64)) - backend = resolve_backend(BackendRequest(path=path, affine=custom)) - assert backend is not None - np.testing.assert_allclose(np.asarray(backend.affine), custom.numpy()) - - def test_image_delegates_to_resolver(self, tmp_path: Path, monkeypatch) -> None: - import torchio.data.image as image_module - - arr = np.random.randn(8, 9, 10).astype(np.float32) - path = tmp_path / "r.nii.gz" - nib.save(nib.Nifti1Image(arr, np.eye(4)), path) - - calls: list = [] - real_resolve = image_module.resolve_backend - - def spy(request): - calls.append(request) - return real_resolve(request) - - monkeypatch.setattr(image_module, "resolve_backend", spy) - image = ScalarImage(path) - _ = image.shape # triggers backend resolution - assert len(calls) == 1 - assert calls[0].path == path - - -class TestBackendRegistration: - """Controlled extensibility: register custom backends and lazy readers.""" - - def test_register_and_resolve_custom_backend(self, tmp_path: Path) -> None: - from torchio.data import register_backend - from torchio.data import resolve_backend - from torchio.data import unregister_backend - from torchio.data.backends import BackendRequest - - marker = tmp_path / "custom.fake" - marker.touch() - - def matcher(request: BackendRequest) -> bool: - return request.path is not None and request.path.suffix == ".fake" - - def factory(request: BackendRequest): - return TensorBackend(torch.zeros(1, 2, 3, 4)) - - register_backend("fake", matcher, factory) - try: - backend = resolve_backend(BackendRequest(path=marker)) - assert isinstance(backend, TensorBackend) - assert backend.shape == (1, 2, 3, 4) - finally: - unregister_backend("fake") - assert resolve_backend(BackendRequest(path=marker)) is None - - def test_custom_backend_takes_priority_over_builtin(self, tmp_path: Path) -> None: - from torchio.data import register_backend - from torchio.data import resolve_backend - from torchio.data import unregister_backend - from torchio.data.backends import BackendRequest - - arr = np.random.randn(8, 9, 10).astype(np.float32) - path = tmp_path / "p.nii.gz" - nib.save(nib.Nifti1Image(arr, np.eye(4)), path) - - sentinel = TensorBackend(torch.ones(1, 1, 1, 1)) - - register_backend( - "override-nifti", - lambda req: req.path is not None and req.path.name.endswith(".nii.gz"), - lambda req: sentinel, - ) - try: - backend = resolve_backend(BackendRequest(path=path)) - assert backend is sentinel - finally: - unregister_backend("override-nifti") - - def test_image_uses_registered_backend(self, tmp_path: Path) -> None: - from torchio.data import register_backend - from torchio.data import unregister_backend - - marker = tmp_path / "image.fake" - marker.touch() - data = torch.arange(2 * 4 * 5 * 6, dtype=torch.float32).reshape(2, 4, 5, 6) - - register_backend( - "fake-image", - lambda req: req.path is not None and req.path.suffix == ".fake", - lambda req: TensorBackend(data, affine=torch.eye(4, dtype=torch.float64)), - ) - try: - image = ScalarImage(marker) - assert image.shape == (2, 4, 5, 6) - assert image._data is None # shape read lazily via the backend - assert torch.equal(image.data, data) - finally: - unregister_backend("fake-image") - - -class _LazyNiftiReader: - """A custom reader that is also a LazyReader (builds a NibabelBackend).""" - - def __call__(self, path: Path, **kwargs): # simple-reader fallback - backend = self.create_backend(_request_for(path, kwargs)) - return backend.to_tensor(), backend.affine.numpy() - - def create_backend(self, request): - nii = nib.load(request.path, **dict(request.reader_kwargs)) - return NibabelBackend(nii, affine=request.affine) - - -def _request_for(path, kwargs): - from torchio.data.backends import BackendRequest - - return BackendRequest(path=path, reader_kwargs=kwargs) - - -class TestLazyCustomReader: - """A custom reader exposing `create_backend` enables lazy access.""" - - @pytest.fixture - def nifti_path(self, tmp_path: Path) -> Path: - arr = np.random.randn(8, 9, 10).astype(np.float32) - path = tmp_path / "lazy_reader.nii.gz" - nib.save(nib.Nifti1Image(arr, np.diag([2.0, 3.0, 4.0, 1.0])), path) - return path - - def test_shape_is_lazy(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path, reader=_LazyNiftiReader()) - assert image.shape == (1, 8, 9, 10) - assert image._data is None # not materialized - - def test_dtype_and_affine_lazy(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path, reader=_LazyNiftiReader()) - assert image.dtype == np.dtype("float32") - np.testing.assert_allclose(image.affine.numpy(), np.diag([2.0, 3.0, 4.0, 1.0])) - assert image._data is None - - def test_lazy_slice(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path, reader=_LazyNiftiReader()) - sliced = image.dataobj[:, 1:4, 2:5, 0:6] - assert tuple(sliced.shape) == (1, 3, 3, 6) - assert image._data is None - - def test_materialization_still_works(self, nifti_path: Path) -> None: - image = ScalarImage(nifti_path, reader=_LazyNiftiReader()) - assert tuple(image.data.shape) == (1, 8, 9, 10) - - -class TestSimpleCustomReaderUnchanged: - """Simple `(tensor, affine)` readers keep loading eagerly, as before.""" - - def test_simple_reader_loads_eagerly(self, tmp_path: Path) -> None: - calls: list = [] - - def reader(path: Path, **kwargs): - calls.append(path) - return torch.ones(1, 3, 4, 5), np.eye(4) - - path = tmp_path / "anything.mha" - image = ScalarImage(path, reader=reader) - assert image.shape == (1, 3, 4, 5) - assert calls # reader was invoked (full load), no lazy backend - assert image._data is not None - - -class TestNormalizeIndex: - """Error handling of the shared index-normalization helper.""" - - def test_invalid_top_level_type(self) -> None: - from torchio.data.backends import normalize_index - - with pytest.raises(TypeError, match="not understood"): - normalize_index("foo") # type: ignore[arg-type] - - def test_invalid_element_in_tuple(self) -> None: - from torchio.data.backends import normalize_index - - with pytest.raises(TypeError, match="not understood"): - normalize_index((0, "bad")) # type: ignore[arg-type] - - def test_multiple_ellipsis(self) -> None: - from torchio.data.backends import normalize_index - - with pytest.raises(IndexError, match="one ellipsis"): - normalize_index((..., 0, ...)) - - def test_too_many_indices(self) -> None: - from torchio.data.backends import normalize_index - - with pytest.raises(IndexError, match="Too many indices"): - normalize_index((0, 0, 0, 0, 0)) - - def test_negative_one_keeps_last(self) -> None: - from torchio.data.backends import normalize_index - - assert normalize_index(-1)[0] == slice(-1, None) diff --git a/tests/test_batch.py b/tests/test_batch.py deleted file mode 100644 index d1a3b3f2f..000000000 --- a/tests/test_batch.py +++ /dev/null @@ -1,336 +0,0 @@ -"""Tests for ImagesBatch and SubjectsBatch.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio -from torchio.data.batch import ImagesBatch -from torchio.data.batch import SubjectsBatch - - -class TestImagesBatch: - def test_from_images(self) -> None: - images = [tio.ScalarImage(torch.rand(1, 8, 8, 8)) for _ in range(4)] - batch = ImagesBatch.from_images(images) - assert batch.data.shape == (4, 1, 8, 8, 8) - - def test_batch_size(self) -> None: - batch = ImagesBatch( - data=torch.rand(4, 1, 8, 8, 8), - affines=[tio.AffineMatrix() for _ in range(4)], - ) - assert batch.batch_size == 4 - - def test_to_device(self) -> None: - batch = ImagesBatch( - data=torch.rand(2, 1, 4, 4, 4), - affines=[tio.AffineMatrix() for _ in range(2)], - ) - result = batch.to(torch.float64) - assert result.data.dtype == torch.float64 - - def test_unbatch(self) -> None: - images = [tio.ScalarImage(torch.rand(1, 8, 8, 8)) for _ in range(3)] - batch = ImagesBatch.from_images(images) - restored = batch.unbatch() - assert len(restored) == 3 - for img in restored: - assert isinstance(img, tio.ScalarImage) - assert img.shape == (1, 8, 8, 8) - - def test_getitem_int(self) -> None: - batch = ImagesBatch( - data=torch.rand(4, 1, 8, 8, 8), - affines=[tio.AffineMatrix() for _ in range(4)], - ) - img = batch[0] - assert isinstance(img, tio.ScalarImage) - assert img.shape == (1, 8, 8, 8) - - def test_per_sample_affines(self) -> None: - affine_a = tio.AffineMatrix.from_spacing((1.0, 1.0, 1.0)) - affine_b = tio.AffineMatrix.from_spacing((2.0, 2.0, 2.0)) - batch = ImagesBatch( - data=torch.rand(2, 1, 8, 8, 8), - affines=[affine_a, affine_b], - ) - assert batch[0].affine.spacing == (1.0, 1.0, 1.0) - assert batch[1].affine.spacing == (2.0, 2.0, 2.0) - - def test_repr(self) -> None: - batch = ImagesBatch( - data=torch.rand(4, 1, 8, 8, 8), - affines=[tio.AffineMatrix() for _ in range(4)], - ) - assert "4" in repr(batch) - assert "8" in repr(batch) - - -class TestSubjectsBatch: - def test_from_subjects(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 8, 8, 8))), - age=42 + i, - ) - for i in range(4) - ] - batch = SubjectsBatch.from_subjects(subjects) - assert batch["t1"].data.shape == (4, 1, 8, 8, 8) - assert batch["seg"].data.shape == (4, 1, 8, 8, 8) - - def test_attribute_access(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - ) - for _ in range(2) - ] - batch = SubjectsBatch.from_subjects(subjects) - assert batch.t1.data.shape == (2, 1, 8, 8, 8) - - def test_batch_size(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - ) - for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - assert batch.batch_size == 3 - - def test_unbatch(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - age=42 + i, - ) - for i in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - restored = batch.unbatch() - assert len(restored) == 3 - for i, sub in enumerate(restored): - assert isinstance(sub, tio.Subject) - assert sub.t1.shape == (1, 8, 8, 8) - assert sub.age == 42 + i - - def test_to_device(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - ) - for _ in range(2) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = batch.to(torch.float64) - assert result.t1.data.dtype == torch.float64 - - def test_metadata_preserved(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - age=42 + i, - name=f"sub_{i}", - ) - for i in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - assert batch.metadata["age"] == [42, 43, 44] - assert batch.metadata["name"] == ["sub_0", "sub_1", "sub_2"] - - -class TestBatchTransforms: - def test_flip_images_batch(self) -> None: - images = [ - tio.ScalarImage(torch.arange(8).reshape(1, 2, 2, 2).float()) - for _ in range(3) - ] - batch = ImagesBatch.from_images(images) - original = batch.data.clone() - result = tio.Flip(axes=(0,))(batch) - assert isinstance(result, ImagesBatch) - assert result.data.shape == (3, 1, 2, 2, 2) - # Check data was actually flipped - assert not torch.equal(result.data, original) - - def test_flip_subjects_batch(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - ) - for _ in range(4) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.Flip(axes=(0,))(batch) - assert isinstance(result, SubjectsBatch) - assert result.t1.data.shape == (4, 1, 8, 8, 8) - - def test_noise_images_batch(self) -> None: - images = [tio.ScalarImage(torch.zeros(1, 4, 4, 4)) for _ in range(3)] - batch = ImagesBatch.from_images(images) - result = tio.Noise(std=1.0)(batch) - # Noise should have been added - assert result.data.abs().sum() > 0 - - def test_batch_preserves_affines(self) -> None: - affine_a = tio.AffineMatrix.from_spacing((1.0, 1.0, 1.0)) - affine_b = tio.AffineMatrix.from_spacing((2.0, 2.0, 2.0)) - images = [ - tio.ScalarImage(torch.rand(1, 8, 8, 8), affine=affine_a), - tio.ScalarImage(torch.rand(1, 8, 8, 8), affine=affine_b), - ] - batch = ImagesBatch.from_images(images) - result = tio.Flip(axes=(0,))(batch) - assert result.affines[0].spacing == (1.0, 1.0, 1.0) - assert result.affines[1].spacing == (2.0, 2.0, 2.0) - - def test_batch_copy_preserves_original(self) -> None: - images = [tio.ScalarImage(torch.zeros(1, 4, 4, 4)) for _ in range(2)] - batch = ImagesBatch.from_images(images) - original = batch.data.clone() - tio.Noise(std=1.0)(batch) - # Original should be unchanged (copy=True default) - torch.testing.assert_close(batch.data, original) - - -# ── Coverage gap tests ─────────────────────────────────────────────── - - -class TestImagesBatchValidation: - def test_non_5d_raises(self) -> None: - from torchio.data.batch import ImagesBatch - - with pytest.raises(ValueError, match="5"): - ImagesBatch( - data=torch.rand(1, 10, 10), - affines=[tio.AffineMatrix()], - image_class=tio.ScalarImage, - ) - - def test_affine_count_mismatch_raises(self) -> None: - from torchio.data.batch import ImagesBatch - - with pytest.raises(ValueError, match="affines"): - ImagesBatch( - data=torch.rand(2, 1, 5, 5, 5), - affines=[tio.AffineMatrix()], # only 1 for batch of 2 - image_class=tio.ScalarImage, - ) - - def test_from_images_empty_raises(self) -> None: - from torchio.data.batch import ImagesBatch - - with pytest.raises(ValueError, match="empty"): - ImagesBatch.from_images([]) - - def test_data_setter_non_5d_raises(self) -> None: - from torchio.data.batch import ImagesBatch - - batch = ImagesBatch( - data=torch.rand(1, 1, 5, 5, 5), - affines=[tio.AffineMatrix()], - image_class=tio.ScalarImage, - ) - with pytest.raises(ValueError, match="5"): - batch.data = torch.rand(1, 5, 5) - - def test_device_property(self) -> None: - from torchio.data.batch import ImagesBatch - - batch = ImagesBatch( - data=torch.rand(1, 1, 5, 5, 5), - affines=[tio.AffineMatrix()], - image_class=tio.ScalarImage, - ) - assert batch.device.type == "cpu" - - def test_len(self) -> None: - from torchio.data.batch import ImagesBatch - - batch = ImagesBatch( - data=torch.rand(3, 1, 5, 5, 5), - affines=[tio.AffineMatrix() for _ in range(3)], - image_class=tio.ScalarImage, - ) - assert len(batch) == 3 - - -class TestSubjectsBatchEdgeCases: - def test_from_subjects_empty_raises(self) -> None: - from torchio.data.batch import SubjectsBatch - - with pytest.raises(ValueError, match="empty"): - SubjectsBatch.from_subjects([]) - - def test_device_property(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - from torchio.data.batch import SubjectsBatch - - batch = SubjectsBatch.from_subjects([subject]) - assert batch.device.type == "cpu" - - def test_getattr_invalid_raises(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - from torchio.data.batch import SubjectsBatch - - batch = SubjectsBatch.from_subjects([subject]) - with pytest.raises(AttributeError): - _ = batch.nonexistent_image - - def test_len(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - from torchio.data.batch import SubjectsBatch - - batch = SubjectsBatch.from_subjects([subject]) - assert len(batch) == 1 - - def test_repr(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - from torchio.data.batch import SubjectsBatch - - batch = SubjectsBatch.from_subjects([subject]) - r = repr(batch) - assert "SubjectsBatch" in r - assert "t1" in r - - -class TestPerElementHistory: - def _batch(self, batch_size: int = 4) -> SubjectsBatch: - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.rand(1, 6, 6, 6))) - for _ in range(batch_size) - ] - return SubjectsBatch.from_subjects(subjects) - - def test_adopt_history_preserves_per_element(self) -> None: - # Simulate the adapter pattern: a per-element batch is unbatched, - # processed, and re-stacked; history must survive. - torch.manual_seed(0) - batch = self._batch() - branched = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))])(batch) - subjects = branched.unbatch() - rebuilt = SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(branched, subjects) - for original, restored in zip( - branched.unbatch(), - rebuilt.unbatch(), - strict=True, - ): - assert [t.name for t in restored.applied_transforms] == [ - t.name for t in original.applied_transforms - ] - - def test_adopt_history_shared_case(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transformed = tio.Gamma(log_gamma=0.3, per_instance=False)(batch) - subjects = transformed.unbatch() - rebuilt = SubjectsBatch.from_subjects(subjects) - rebuilt.adopt_history(transformed, subjects) - assert rebuilt._per_element_history is None - for subject in rebuilt.unbatch(): - assert [t.name for t in subject.applied_transforms] == ["Gamma"] diff --git a/tests/test_bboxes.py b/tests/test_bboxes.py deleted file mode 100644 index d7afc99ba..000000000 --- a/tests/test_bboxes.py +++ /dev/null @@ -1,473 +0,0 @@ -"""Tests for BoundingBoxes.""" - -from __future__ import annotations - -import copy - -import numpy as np -import pytest -import torch - -from torchio.data.bboxes import BoundingBoxes -from torchio.data.bboxes import BoundingBoxFormat -from torchio.data.bboxes import Representation - - -class TestBoundingBoxFormat: - """Test the format class.""" - - def test_ijk_corners(self): - fmt = BoundingBoxFormat("IJK", Representation.CORNERS) - assert fmt.axes == "IJK" - assert fmt.representation == Representation.CORNERS - - def test_ras_center_size(self): - fmt = BoundingBoxFormat("RAS", Representation.CENTER_SIZE) - assert fmt.axes == "RAS" - assert fmt.representation == Representation.CENTER_SIZE - - def test_string_representation(self): - fmt = BoundingBoxFormat("IJK", "corners") - assert fmt.representation == Representation.CORNERS - - def test_string_center_size(self): - fmt = BoundingBoxFormat("IJK", "center_size") - assert fmt.representation == Representation.CENTER_SIZE - - def test_invalid_axes_raises(self): - with pytest.raises(ValueError, match="Invalid"): - BoundingBoxFormat("XYZ", Representation.CORNERS) - - def test_equality(self): - a = BoundingBoxFormat("IJK", Representation.CORNERS) - b = BoundingBoxFormat("IJK", Representation.CORNERS) - assert a == b - - def test_inequality_axes(self): - a = BoundingBoxFormat("IJK", Representation.CORNERS) - b = BoundingBoxFormat("KJI", Representation.CORNERS) - assert a != b - - def test_inequality_representation(self): - a = BoundingBoxFormat("IJK", Representation.CORNERS) - b = BoundingBoxFormat("IJK", Representation.CENTER_SIZE) - assert a != b - - def test_hashable(self): - a = BoundingBoxFormat("IJK", Representation.CORNERS) - b = BoundingBoxFormat("IJK", Representation.CORNERS) - assert hash(a) == hash(b) - assert len({a, b}) == 1 - - def test_repr(self): - fmt = BoundingBoxFormat("RAS", Representation.CORNERS) - r = repr(fmt) - assert "RAS" in r - assert "corners" in r - - def test_predefined_ijkijk(self): - assert ( - BoundingBoxFormat( - "IJK", - Representation.CORNERS, - ) - == BoundingBoxFormat.IJKIJK - ) - - def test_predefined_ijkwhd(self): - assert ( - BoundingBoxFormat( - "IJK", - Representation.CENTER_SIZE, - ) - == BoundingBoxFormat.IJKWHD - ) - - -class TestBoundingBoxesCreation: - def test_from_tensor(self): - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - assert bboxes.data.shape == (1, 6) - - def test_from_numpy(self): - data = np.array([[10, 20, 30, 50, 60, 70]], dtype=np.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - assert isinstance(bboxes.data, torch.Tensor) - - def test_multiple_boxes(self): - data = torch.tensor( - [ - [10, 20, 30, 50, 60, 70], - [0, 0, 0, 10, 10, 10], - ], - dtype=torch.float32, - ) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - assert len(bboxes) == 2 - - def test_empty_boxes(self): - bboxes = BoundingBoxes( - torch.zeros(0, 6), - format=BoundingBoxFormat.IJKIJK, - ) - assert len(bboxes) == 0 - - def test_with_labels(self): - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - labels = torch.tensor([3]) - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - labels=labels, - ) - assert bboxes.labels is not None - assert bboxes.labels[0] == 3 - - def test_with_affine(self): - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - affine=affine, - ) - np.testing.assert_array_equal(bboxes.affine.numpy(), affine) - - def test_with_metadata(self): - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - metadata={"source": "manual"}, - ) - assert bboxes.metadata["source"] == "manual" - - def test_wrong_shape_raises(self): - with pytest.raises(ValueError, match="N, 6"): - BoundingBoxes( - torch.tensor([[1, 2, 3]]), - format=BoundingBoxFormat.IJKIJK, - ) - - def test_wrong_ndim_raises(self): - with pytest.raises(ValueError, match="N, 6"): - BoundingBoxes( - torch.tensor([1, 2, 3, 4, 5, 6]), - format=BoundingBoxFormat.IJKIJK, - ) - - def test_labels_length_mismatch_raises(self): - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - labels = torch.tensor([1, 2]) # 2 labels for 1 box - with pytest.raises(ValueError, match="labels"): - BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - labels=labels, - ) - - -class TestBoundingBoxesProperties: - def test_len(self): - data = torch.randn(5, 6) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - assert len(bboxes) == 5 - - def test_num_boxes(self): - data = torch.randn(3, 6) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - assert bboxes.num_boxes == 3 - - -class TestRepresentationConversion: - """Test corners ↔ center_size conversion (same axes).""" - - def test_corners_to_center_size(self): - # Box from (10, 20, 30) to (50, 60, 70) - # Center: (30, 40, 50), Size: (40, 40, 40) - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - converted = bboxes.to_format(BoundingBoxFormat.IJKWHD) - expected = torch.tensor([[30, 40, 50, 40, 40, 40]], dtype=torch.float32) - torch.testing.assert_close(converted.data, expected) - assert converted.format == BoundingBoxFormat.IJKWHD - - def test_center_size_to_corners(self): - data = torch.tensor([[30, 40, 50, 40, 40, 40]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKWHD) - converted = bboxes.to_format(BoundingBoxFormat.IJKIJK) - expected = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - torch.testing.assert_close(converted.data, expected) - - def test_same_format_noop(self): - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - converted = bboxes.to_format(BoundingBoxFormat.IJKIJK) - torch.testing.assert_close(converted.data, data) - - def test_roundtrip(self): - data = torch.tensor( - [ - [10, 20, 30, 50, 60, 70], - [0, 0, 0, 100, 100, 100], - ], - dtype=torch.float32, - ) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - roundtrip = bboxes.to_format(BoundingBoxFormat.IJKWHD).to_format( - BoundingBoxFormat.IJKIJK, - ) - torch.testing.assert_close(roundtrip.data, data) - - def test_preserves_labels(self): - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - labels = torch.tensor([5]) - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - labels=labels, - ) - converted = bboxes.to_format(BoundingBoxFormat.IJKWHD) - assert converted.labels is not None - assert converted.labels[0] == 5 - - -class TestVoxelAxisPermutation: - """Test conversion between voxel axis orderings.""" - - def test_ijk_to_kji_corners(self): - # Box: i=[10,50], j=[20,60], k=[30,70] - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - target = BoundingBoxFormat("KJI", Representation.CORNERS) - converted = bboxes.to_format(target) - # KJI corners: k1, j1, i1, k2, j2, i2 - expected = torch.tensor([[30, 20, 10, 70, 60, 50]], dtype=torch.float32) - torch.testing.assert_close(converted.data, expected) - - def test_ijk_to_jki_corners(self): - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - target = BoundingBoxFormat("JKI", Representation.CORNERS) - converted = bboxes.to_format(target) - # JKI corners: j1, k1, i1, j2, k2, i2 - expected = torch.tensor([[20, 30, 10, 60, 70, 50]], dtype=torch.float32) - torch.testing.assert_close(converted.data, expected) - - def test_ijk_to_kji_center_size(self): - # center=(30,40,50), size_i=40, size_j=40, size_k=40 - data = torch.tensor([[30, 40, 50, 40, 40, 40]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKWHD) - target = BoundingBoxFormat("KJI", Representation.CENTER_SIZE) - converted = bboxes.to_format(target) - # KJI center_size: kc, jc, ic, size_k, size_j, size_i - expected = torch.tensor([[50, 40, 30, 40, 40, 40]], dtype=torch.float32) - torch.testing.assert_close(converted.data, expected) - - def test_roundtrip_ijk_kji(self): - data = torch.tensor( - [[10, 20, 30, 50, 60, 70]], - dtype=torch.float32, - ) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - kji = BoundingBoxFormat("KJI", Representation.CORNERS) - roundtrip = bboxes.to_format(kji).to_format(BoundingBoxFormat.IJKIJK) - torch.testing.assert_close(roundtrip.data, data) - - -class TestAnatomicalAxisConversion: - """Test conversion between anatomical axis systems.""" - - def test_ras_to_lpi_corners(self): - # Box in RAS: r=[10,50], a=[20,60], s=[30,70] - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - src = BoundingBoxFormat("RAS", Representation.CORNERS) - tgt = BoundingBoxFormat("LPI", Representation.CORNERS) - bboxes = BoundingBoxes(data, format=src) - converted = bboxes.to_format(tgt) - # L=-R, P=-A, I=-S. Negation swaps min/max for corners. - # L: [-50, -10], P: [-60, -20], I: [-70, -30] - expected = torch.tensor( - [[-50, -60, -70, -10, -20, -30]], - dtype=torch.float32, - ) - torch.testing.assert_close(converted.data, expected) - - def test_ras_to_asr_corners(self): - """Pure reorder, no flips.""" - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - src = BoundingBoxFormat("RAS", Representation.CORNERS) - tgt = BoundingBoxFormat("ASR", Representation.CORNERS) - bboxes = BoundingBoxes(data, format=src) - converted = bboxes.to_format(tgt) - # ASR: a1, s1, r1, a2, s2, r2 - expected = torch.tensor([[20, 30, 10, 60, 70, 50]], dtype=torch.float32) - torch.testing.assert_close(converted.data, expected) - - def test_ras_to_lpi_center_size(self): - # Center in RAS: (30, 40, 50), sizes: (40, 40, 40) - data = torch.tensor([[30, 40, 50, 40, 40, 40]], dtype=torch.float32) - src = BoundingBoxFormat("RAS", Representation.CENTER_SIZE) - tgt = BoundingBoxFormat("LPI", Representation.CENTER_SIZE) - bboxes = BoundingBoxes(data, format=src) - converted = bboxes.to_format(tgt) - # Center negated: (-30, -40, -50). Sizes unchanged: (40, 40, 40). - expected = torch.tensor( - [[-30, -40, -50, 40, 40, 40]], - dtype=torch.float32, - ) - torch.testing.assert_close(converted.data, expected) - - def test_roundtrip_ras_lpi(self): - data = torch.tensor( - [[10, 20, 30, 50, 60, 70]], - dtype=torch.float32, - ) - src = BoundingBoxFormat("RAS", Representation.CORNERS) - tgt = BoundingBoxFormat("LPI", Representation.CORNERS) - bboxes = BoundingBoxes(data, format=src) - roundtrip = bboxes.to_format(tgt).to_format(src) - torch.testing.assert_close(roundtrip.data, data) - - -class TestVoxelAnatomicalConversion: - """Test conversion between voxel and anatomical using the affine.""" - - def test_ijk_to_ras_identity_affine(self): - """With identity affine, IJK == RAS numerically.""" - data = torch.tensor([[10, 20, 30, 50, 60, 70]], dtype=torch.float32) - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - ) # identity affine - tgt = BoundingBoxFormat("RAS", Representation.CORNERS) - converted = bboxes.to_format(tgt) - torch.testing.assert_close(converted.data, data) - - def test_ijk_to_ras_with_spacing(self): - """Spacing scales coordinates.""" - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - affine=affine, - ) - tgt = BoundingBoxFormat("RAS", Representation.CORNERS) - converted = bboxes.to_format(tgt) - expected = torch.tensor([[0, 0, 0, 20, 30, 40]], dtype=torch.float32) - torch.testing.assert_close(converted.data, expected) - - def test_ijk_to_ras_with_origin(self): - """Origin shifts coordinates.""" - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - affine = np.eye(4) - affine[:3, 3] = [100, 200, 300] - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - affine=affine, - ) - tgt = BoundingBoxFormat("RAS", Representation.CORNERS) - converted = bboxes.to_format(tgt) - expected = torch.tensor( - [[100, 200, 300, 110, 210, 310]], - dtype=torch.float32, - ) - torch.testing.assert_close(converted.data, expected) - - def test_ras_to_ijk_roundtrip(self): - data = torch.tensor([[5, 10, 15, 25, 30, 35]], dtype=torch.float32) - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - affine[:3, 3] = [10, 20, 30] - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - affine=affine, - ) - ras_fmt = BoundingBoxFormat("RAS", Representation.CORNERS) - roundtrip = bboxes.to_format(ras_fmt).to_format(BoundingBoxFormat.IJKIJK) - torch.testing.assert_close(roundtrip.data, data, atol=1e-5, rtol=1e-5) - - def test_ijk_to_lpi(self): - """Voxel to non-RAS anatomical (combines affine + flip).""" - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - # Identity affine → world is RAS - bboxes = BoundingBoxes( - data, - format=BoundingBoxFormat.IJKIJK, - ) - tgt = BoundingBoxFormat("LPI", Representation.CORNERS) - converted = bboxes.to_format(tgt) - # In RAS: [0,0,0,10,10,10]. LPI = negate all, swap corners. - expected = torch.tensor( - [[-10, -10, -10, 0, 0, 0]], - dtype=torch.float32, - ) - torch.testing.assert_close(converted.data, expected) - - def test_no_affine_cross_type_raises(self): - """Voxel↔anatomical conversion requires an affine, but identity - is the default so it should work. This test documents that the - affine is used implicitly.""" - data = torch.tensor([[0, 0, 0, 10, 10, 10]], dtype=torch.float32) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - tgt = BoundingBoxFormat("RAS", Representation.CORNERS) - # Should work (uses default identity affine) - converted = bboxes.to_format(tgt) - assert converted.format == tgt - - -class TestNewLike: - def test_new_like_preserves_format(self): - bboxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKWHD, - ) - new = bboxes.new_like(data=torch.randn(2, 6)) - assert new.format == BoundingBoxFormat.IJKWHD - - def test_new_like_preserves_affine(self): - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - bboxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - affine=affine, - ) - new = bboxes.new_like(data=torch.randn(1, 6)) - np.testing.assert_array_equal(new.affine.numpy(), affine) - - def test_new_like_preserves_metadata(self): - bboxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - metadata={"organ": "liver"}, - ) - new = bboxes.new_like(data=torch.randn(1, 6)) - assert new.metadata["organ"] == "liver" - - -class TestRepr: - def test_repr_contains_info(self): - data = torch.randn(3, 6) - bboxes = BoundingBoxes(data, format=BoundingBoxFormat.IJKIJK) - r = repr(bboxes) - assert "BoundingBoxes" in r - assert "3" in r - assert "IJK" in r - - -class TestCopy: - def test_deepcopy(self): - bboxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - labels=torch.tensor([1, 2, 3]), - metadata={"a": 1}, - ) - copied = copy.deepcopy(bboxes) - assert torch.equal(copied.data, bboxes.data) - assert torch.equal(copied.labels, bboxes.labels) - copied.metadata["a"] = 2 - assert bboxes.metadata["a"] == 1 diff --git a/tests/test_bias_field.py b/tests/test_bias_field.py deleted file mode 100644 index 9167fd7f1..000000000 --- a/tests/test_bias_field.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Tests for BiasField transform.""" - -from __future__ import annotations - -import numpy as np -import pytest -import torch - -import torchio as tio - - -def _make_subject() -> tio.Subject: - data = torch.rand(1, 16, 16, 16) + 1.0 # positive intensities - return tio.Subject( - t1=tio.ScalarImage(data), - seg=tio.LabelMap(torch.zeros(1, 16, 16, 16)), - ) - - -class TestBasic: - def test_changes_data(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.BiasField()(subject) - assert not torch.allclose(result.t1.data, original) - - def test_field_is_multiplicative(self) -> None: - subject = _make_subject() - result = tio.BiasField(std=0.3)(subject) - # All values should remain positive (exp is always > 0) - assert result.t1.data.min() > 0 - - def test_zero_std_is_identity(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.BiasField(std=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_leaves_label_maps_unchanged(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.BiasField()(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_random_std(self) -> None: - subject = _make_subject() - transform = tio.BiasField(std=(0.0, 1.0)) - results = [transform(subject).t1.data.mean().item() for _ in range(5)] - assert len({f"{v:.4f}" for v in results}) > 1 - - def test_custom_scale(self) -> None: - subject = _make_subject() - result = tio.BiasField(std=0.5, scale=0.1)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - -class TestValidation: - def test_negative_scale_raises(self) -> None: - with pytest.raises(ValueError, match="scale"): - tio.BiasField(scale=-0.1) - - def test_scale_above_one_raises(self) -> None: - with pytest.raises(ValueError, match="scale"): - tio.BiasField(scale=1.5) - - -class TestInverse: - def test_inverse_restores_values(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - transformed = tio.BiasField(std=0.3)(subject) - restored = transformed.apply_inverse_transform() - np.testing.assert_allclose( - restored.t1.data.numpy(), - original.numpy(), - atol=1e-5, - ) - - -class TestExports: - def test_available_at_top_level(self) -> None: - assert hasattr(tio, "BiasField") - - -class TestBiasFieldPerInstance: - def _batch(self, batch_size: int = 5) -> tio.SubjectsBatch: - data = torch.rand(1, 12, 12, 12) + 0.5 - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.BiasField(std=(0.3, 0.6))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["std"]) == batch.batch_size - assert not torch.allclose(result.t1.data[0], result.t1.data[1]) - - def test_per_instance_false_shares_std(self) -> None: - # The bias field is sampled at full batch shape, so the spatial - # realization still varies per element; only the std is shared. - torch.manual_seed(0) - batch = self._batch() - result = tio.BiasField(std=(0.3, 0.6), per_instance=False)(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["std"], float) - assert "_batched_keys" not in params - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 12, 12, 12) + 0.5)) - result = tio.BiasField(std=(0.3, 0.6))(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params - - def test_per_instance_inverse_round_trip(self) -> None: - torch.manual_seed(0) - batch = self._batch() - original = batch.t1.data.clone() - result = tio.BiasField(std=(0.3, 0.6))(batch) - restored = result.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original, atol=1e-4, rtol=0) - - def test_per_instance_inverse_after_unbatch(self) -> None: - """Each unbatched subject must invert using its own field.""" - torch.manual_seed(0) - batch = self._batch() - original = batch.t1.data.clone() - result = tio.BiasField(std=(0.3, 0.6))(batch) - for index, subject in enumerate(result.unbatch()): - restored = subject.apply_inverse_transform() - torch.testing.assert_close( - restored.t1.data, - original[index], - atol=1e-4, - rtol=0, - ) diff --git a/tests/test_blur.py b/tests/test_blur.py deleted file mode 100644 index 101b84fbb..000000000 --- a/tests/test_blur.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Tests for Blur transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestBlur: - def test_changes_data(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Blur(std=2.0)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_zero_std_is_identity(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Blur(std=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_leaves_labels_unchanged(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.Blur(std=1.0)(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - -class TestBlurPerInstance: - def _batch(self, batch_size: int = 5) -> tio.SubjectsBatch: - data = torch.rand(1, 10, 10, 10) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Blur(std=(1.0, 4.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["std"]) == batch.batch_size - data = result.t1.data - assert not torch.allclose(data[0], data[1]) - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Blur(std=(1.0, 4.0), per_instance=False)(batch) - data = result.t1.data - torch.testing.assert_close(data[0], data[1]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8))) - result = tio.Blur(std=(1.0, 4.0))(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params - - def test_per_instance_p_gates_some_elements(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=64) - original = batch.t1.data.clone() - result = tio.Blur(std=(2.0, 4.0), p=0.5)(batch) - changed = [ - not torch.allclose(result.t1.data[i], original[i]) - for i in range(batch.batch_size) - ] - assert any(changed) - assert not all(changed) - - def test_per_instance_p_masked_float64_elements_unchanged(self) -> None: - """Zero-sigma elements must be exact, even for float64 data.""" - torch.manual_seed(0) - data = (torch.rand(1, 8, 8, 8) + 0.1).double() - subjects = [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(32)] - batch = tio.SubjectsBatch.from_subjects(subjects) - original = batch.t1.data.clone() - result = tio.Blur(std=(2.0, 4.0), p=0.5)(batch) - unchanged = [ - torch.equal(result.t1.data[i], original[i]) for i in range(batch.batch_size) - ] - assert any(unchanged) - assert not all(unchanged) diff --git a/tests/test_clamp.py b/tests/test_clamp.py deleted file mode 100644 index b0231264f..000000000 --- a/tests/test_clamp.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Tests for Clamp transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestClamp: - def test_clamps_values(self) -> None: - data = torch.tensor([-10.0, 0.0, 50.0, 200.0]).reshape(1, 1, 1, 4) - subject = tio.Subject(t1=tio.ScalarImage(data)) - result = tio.Clamp(out_min=0, out_max=100)(subject) - assert result.t1.data.min() >= 0 - assert result.t1.data.max() <= 100 - - def test_clamp_min_only(self) -> None: - data = torch.tensor([-5.0, 0.0, 5.0]).reshape(1, 1, 1, 3) - subject = tio.Subject(t1=tio.ScalarImage(data)) - result = tio.Clamp(out_min=0)(subject) - assert result.t1.data.min() >= 0 - - def test_invalid_range_raises(self) -> None: - with pytest.raises(ValueError, match="out_min"): - tio.Clamp(out_min=100, out_max=0) diff --git a/tests/test_cli.py b/tests/test_cli.py index 5410fc180..431dec365 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,161 +1,50 @@ -"""Tests for the torchio CLI. - -Commands are invoked in-process via their dataclass `.run()` method -rather than through `subprocess`, avoiding the ~1s per-test overhead -of spawning a new Python interpreter and re-importing PyTorch. -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -import nibabel as nib -import numpy as np -import pytest - -import torchio as tio -from torchio.cli import Animate -from torchio.cli import Cache -from torchio.cli import Convert -from torchio.cli import Dir -from torchio.cli import Info -from torchio.cli import Plot -from torchio.cli import Transform -from torchio.cli import main - - -@pytest.fixture -def nii_path(tmp_path: Path) -> Path: - path = tmp_path / "test.nii.gz" - nib.save(nib.Nifti1Image(np.zeros((10, 10, 10)), np.eye(4)), path) - return path - - -class TestInfo: - def test_prints_metadata( - self, - nii_path: Path, - capsys: pytest.CaptureFixture[str], - ) -> None: - Info(path=nii_path).run() - captured = capsys.readouterr() - assert "spatial:" in captured.out - assert "spacing:" in captured.out - assert "orientation:" in captured.out - - -class TestConvert: - def test_convert_nii_to_nii(self, nii_path: Path, tmp_path: Path) -> None: - output = tmp_path / "out.nii" - Convert(input=nii_path, output=output).run() - assert output.exists() - - def test_convert_nonexistent(self, tmp_path: Path) -> None: - with pytest.raises(FileNotFoundError): - Convert( - input=Path("nonexistent.nii"), - output=tmp_path / "out.nii", - ).run() - - def test_preserves_dtype(self, tmp_path: Path) -> None: - input_path = tmp_path / "in.nii.gz" - data = np.zeros((4, 5, 6), dtype=np.int16) - nib.save(nib.Nifti1Image(data, np.eye(4)), input_path) - output = tmp_path / "out.nii.gz" - Convert(input=input_path, output=output).run() - loaded = nib.load(output) - assert loaded.header.get_data_dtype() == np.int16 - - def test_no_stdout( - self, - nii_path: Path, - tmp_path: Path, - capsys: pytest.CaptureFixture[str], - ) -> None: - output = tmp_path / "out.nii.gz" - Convert(input=nii_path, output=output).run() - captured = capsys.readouterr() - assert captured.out == "" - - -class TestTransform: - def test_apply_noise(self, nii_path: Path, tmp_path: Path) -> None: - output = tmp_path / "noisy.nii.gz" - Transform( - input=nii_path, - output=output, - name="Noise", - args=["std=0.1"], - ).run() - assert output.exists() - - def test_unknown_transform(self, nii_path: Path, tmp_path: Path) -> None: - output = tmp_path / "out.nii.gz" - with pytest.raises(SystemExit): - Transform( - input=nii_path, - output=output, - name="FakeTransform", - args=[], - ).run() - - -class TestCacheDir: - def test_prints_path(self, capsys: pytest.CaptureFixture[str]) -> None: - Cache(command=Dir()).run() - captured = capsys.readouterr() - assert "torchio" in captured.out.strip() - - -class TestPlot: - def test_plot_to_file(self, nii_path: Path, tmp_path: Path) -> None: - output = tmp_path / "plot.png" - Plot(path=nii_path, output=output).run() - assert output.exists() - assert output.stat().st_size > 0 - - -class TestAnimate: - def test_animate_gif(self, nii_path: Path, tmp_path: Path) -> None: - output = tmp_path / "anim.gif" - Animate(path=nii_path, output=output, seconds=1.0, direction="I").run() - assert output.exists() - assert output.stat().st_size > 0 - - def test_animate_unsupported_format( - self, - nii_path: Path, - tmp_path: Path, - ) -> None: - output = tmp_path / "bad.avi" - with pytest.raises(SystemExit): - Animate(path=nii_path, output=output).run() - - -class TestVersion: - @pytest.mark.parametrize("flag", ["--version"]) - def test_version_flag_prints_version_and_exits( - self, - flag: str, - capsys: pytest.CaptureFixture[str], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(sys, "argv", ["torchio", flag]) - with pytest.raises(SystemExit) as exc_info: - main() - assert exc_info.value.code == 0 - captured = capsys.readouterr() - assert tio.__version__ in captured.out - - def test_version_flag_short_circuits_subcommand( - self, - capsys: pytest.CaptureFixture[str], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - monkeypatch.setattr(sys, "argv", ["torchio", "--version", "info"]) - with pytest.raises(SystemExit) as exc_info: - main() - assert exc_info.value.code == 0 - captured = capsys.readouterr() - assert tio.__version__ in captured.out +#!/usr/bin/env python +"""Tests for CLI tool package.""" + +from typer.testing import CliRunner + +from torchio.cli import apply_transform +from torchio.cli import print_info + +from .utils import TorchioTestCase + +runner = CliRunner() + + +class TestCLI(TorchioTestCase): + def test_cli_transform(self): + image = str(self.get_image_path('cli')) + args = [ + image, + 'RandomFlip', + '--seed', + '0', + '--kwargs', + 'axes=(0,1,2)', + '--hide-progress', + image, + ] + result = runner.invoke(apply_transform.app, args) + assert result.exit_code == 0 + assert result.output.strip() == '' + + def test_bad_transform(self): + image = str(self.get_image_path('cli')) + args = [image, 'RandomRandom', image] + result = runner.invoke(apply_transform.app, args) + assert result.exit_code == 1 + + def test_cli_hd(self): + image = str(self.get_image_path('cli')) + args = [image, '--load'] + result = runner.invoke(print_info.app, args) + assert result.exit_code == 0 + assert ( + result.output == 'ScalarImage(' + 'shape: (1, 10, 20, 30);' + ' spacing: (1.00, 1.00, 1.00);' + ' orientation: RAS+;' + ' dtype: torch.DoubleTensor;' + ' memory: 46.9 KiB' + ')\n' + ) diff --git a/tests/test_compose.py b/tests/test_compose.py deleted file mode 100644 index 04788d889..000000000 --- a/tests/test_compose.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Tests for Compose transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject() -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - seg=tio.LabelMap(torch.zeros(1, 10, 10, 10)), - ) - - -class TestCompose: - def test_identity_compose(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.Compose([])(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_single_transform(self) -> None: - subject = _make_subject() - result = tio.Compose([tio.Flip(axes=(0,))])(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_multiple_transforms(self) -> None: - subject = _make_subject() - pipeline = tio.Compose( - [ - tio.Flip(axes=(0,)), - tio.Gamma(log_gamma=0.0), - ] - ) - result = pipeline(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_nested_compose(self) -> None: - subject = _make_subject() - inner = tio.Compose([tio.Flip(axes=(0,))], copy=False) - outer = tio.Compose([inner]) - result = outer(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_copy_default(self) -> None: - """Compose deep-copies input by default.""" - subject = _make_subject() - original_data = subject.t1.data.clone() - tio.Compose([tio.Gamma(log_gamma=0.5)])(subject) - torch.testing.assert_close(subject.t1.data, original_data) - - def test_no_copy(self) -> None: - subject = _make_subject() - result = tio.Compose([tio.Gamma(log_gamma=0.0)], copy=False)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_history_recorded(self) -> None: - subject = _make_subject() - result = tio.Compose([tio.Flip(axes=(0,))])(subject) - assert len(result.applied_transforms) > 0 - - def test_dict_transforms(self) -> None: - subject = _make_subject() - pipeline = tio.Compose( - { - "flip": tio.Flip(axes=(0,)), - "gamma": tio.Gamma(log_gamma=0.0), - } - ) - assert len(pipeline.transforms) == 2 - assert isinstance(pipeline.transforms[0], tio.Flip) - assert isinstance(pipeline.transforms[1], tio.Gamma) - result = pipeline(subject) - assert result.t1.data.shape == subject.t1.data.shape diff --git a/tests/test_contour.py b/tests/test_contour.py deleted file mode 100644 index d5d53a55d..000000000 --- a/tests/test_contour.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Tests for Contour transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestContour: - def test_basic_contour(self) -> None: - subject = _make_subject() - result = tio.Contour()(subject) - unique = result.seg.data.unique().tolist() - assert set(unique) <= {0.0, 1.0} - - def test_solid_block_has_boundary(self) -> None: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 3:7, 3:7, 3:7] = 1 - subject = tio.Subject(seg=tio.LabelMap(seg)) - result = tio.Contour()(subject) - assert result.seg.data[0, 4, 5, 5] == 0 - assert result.seg.data[0, 3, 5, 5] == 1 - - def test_uniform_label_no_contour(self) -> None: - seg = torch.ones(1, 10, 10, 10, dtype=torch.float32) - subject = tio.Subject(seg=tio.LabelMap(seg)) - result = tio.Contour()(subject) - assert result.seg.data[0, 4, 4, 4] == 0 - - def test_leaves_scalar_unchanged(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.Contour()(subject) - torch.testing.assert_close(result.t1.data, original) diff --git a/tests/test_copy_affine.py b/tests/test_copy_affine.py deleted file mode 100644 index 9f57d957e..000000000 --- a/tests/test_copy_affine.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Tests for CopyAffine transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestCopyAffine: - def test_copies_affine(self) -> None: - t1 = tio.ScalarImage(torch.rand(1, 5, 5, 5)) - t2 = tio.ScalarImage(torch.rand(1, 5, 5, 5)) - t2.affine._matrix[0, 3] = 99.0 - subject = tio.Subject(t1=t1, t2=t2) - result = tio.CopyAffine(target="t1")(subject) - torch.testing.assert_close( - result.t2.affine._matrix, - result.t1.affine._matrix, - ) - - def test_missing_target_raises(self) -> None: - subject = _make_subject(with_label=False) - with pytest.raises(KeyError, match="not_here"): - tio.CopyAffine(target="not_here")(subject) - - def test_does_not_modify_target(self) -> None: - subject = _make_subject() - original = subject.t1.affine._matrix.clone() - tio.CopyAffine(target="t1")(subject) - torch.testing.assert_close(subject.t1.affine._matrix, original) diff --git a/tests/test_cornucopia_adapter.py b/tests/test_cornucopia_adapter.py deleted file mode 100644 index 6d54f581f..000000000 --- a/tests/test_cornucopia_adapter.py +++ /dev/null @@ -1,135 +0,0 @@ -"""Tests for CornucopiaAdapter transform.""" - -from __future__ import annotations - -import cornucopia as cc -import pytest -import torch - -import torchio as tio - - -def _make_subject() -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) * 100), - seg=tio.LabelMap(torch.zeros(1, 8, 8, 8)), - ) - - -# ── Adapter logic ──────────────────────────────────────────────────── - - -class TestCornucopiaAdapterLogic: - def test_not_callable_raises(self) -> None: - with pytest.raises(TypeError, match="callable"): - tio.CornucopiaAdapter(42) # type: ignore[arg-type] - - def test_p_zero_is_identity(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.CornucopiaAdapter(cc.GaussianNoiseTransform(), p=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_include_filter(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.CornucopiaAdapter(cc.GaussianNoiseTransform(), include=["t1"])( - subject - ) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_exclude_filter(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.CornucopiaAdapter(cc.GaussianNoiseTransform(), exclude=["seg"])( - subject - ) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_scalar_images_come_first(self) -> None: - """Scalars are passed before labels to the callable.""" - received: list[str] = [] - - def spy(*tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: - for t in tensors: - received.append("scalar" if t.sum() > 0 else "label") - return tensors - - subject = _make_subject() - tio.CornucopiaAdapter(spy)(subject) - assert received[0] == "scalar" - assert received[1] == "label" - - def test_not_invertible(self) -> None: - adapter = tio.CornucopiaAdapter(cc.FlipTransform()) - assert adapter.invertible is False - - def test_no_history_recorded(self) -> None: - subject = _make_subject() - result = tio.CornucopiaAdapter(cc.FlipTransform())(subject) - for at in result.applied_transforms: - assert not isinstance(at, tio.CornucopiaAdapter) - - def test_in_compose(self) -> None: - subject = _make_subject() - pipeline = tio.Compose( - [ - tio.CornucopiaAdapter(cc.FlipTransform()), - tio.Gamma(log_gamma=0.0), - ] - ) - result = pipeline(subject) - assert result.t1.data.shape == subject.t1.data.shape - - -# ── Real Cornucopia transforms ─────────────────────────────────────── - - -class TestCornucopiaAdapterTransforms: - def test_gaussian_noise(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - ) - original = subject.t1.data.clone() - result = tio.CornucopiaAdapter(cc.GaussianNoiseTransform())(subject) - assert not torch.allclose(result.t1.data, original) - - def test_flip(self) -> None: - subject = _make_subject() - result = tio.CornucopiaAdapter(cc.FlipTransform())(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_gamma(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8).clamp(0.01, 1)), - ) - original = subject.t1.data.clone() - result = tio.CornucopiaAdapter(cc.GammaTransform())(subject) - assert result.t1.data.shape == original.shape - - def test_elastic_shared(self) -> None: - """Elastic deformation is shared across image and label.""" - subject = _make_subject() - result = tio.CornucopiaAdapter(cc.ElasticTransform())(subject) - assert result.t1.data.shape == subject.t1.data.shape - assert result.seg.data.shape == subject.seg.data.shape - - def test_affine(self) -> None: - subject = _make_subject() - result = tio.CornucopiaAdapter(cc.AffineTransform())(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_sequential(self) -> None: - """Cornucopia SequentialTransform (compose via +).""" - seq = cc.GaussianNoiseTransform() + cc.FlipTransform() - subject = _make_subject() - result = tio.CornucopiaAdapter(seq)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - -# ── Export ─────────────────────────────────────────────────────────── - - -class TestCornucopiaExport: - def test_top_level(self) -> None: - assert hasattr(tio, "CornucopiaAdapter") diff --git a/tests/test_crop.py b/tests/test_crop.py deleted file mode 100644 index d3d798698..000000000 --- a/tests/test_crop.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Tests for the Crop transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -class TestCrop: - def test_crop_uniform(self) -> None: - image = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - subject = tio.Subject(t1=image) - result = tio.Crop(cropping=5)(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_crop_per_axis(self) -> None: - image = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - subject = tio.Subject(t1=image) - result = tio.Crop(cropping=(2, 4, 6))(subject) - assert result.t1.shape == (1, 16, 12, 8) - - def test_crop_six_values(self) -> None: - image = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - subject = tio.Subject(t1=image) - result = tio.Crop(cropping=(2, 3, 4, 5, 6, 7))(subject) - assert result.t1.shape == (1, 15, 11, 7) - - def test_crop_all_images(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 20, 20, 20)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 20, 20, 20))), - ) - result = tio.Crop(cropping=5)(subject) - assert result.t1.shape == (1, 10, 10, 10) - assert result.seg.shape == (1, 10, 10, 10) - - def test_crop_affine_updated(self) -> None: - image = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - subject = tio.Subject(t1=image) - original_origin = subject.t1.affine.origin - result = tio.Crop(cropping=(5, 0, 0, 0, 0, 0))(subject) - # Origin should shift along first axis - assert result.t1.affine.origin[0] != original_origin[0] - - def test_crop_lazy(self, tmp_path) -> None: - """Cropping an unloaded image should not load the full volume.""" - import nibabel as nib - import numpy as np - - path = tmp_path / "test.nii.gz" - nib.save( - nib.Nifti1Image(np.zeros((20, 20, 20)), np.eye(4)), - path, - ) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - assert not subject.t1.is_loaded - result = tio.Crop(cropping=5)(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_crop_preserves_laziness_of_original(self, tmp_path) -> None: - """Cropping should not load the original image's data.""" - import nibabel as nib - import numpy as np - - path = tmp_path / "test.nii.gz" - nib.save( - nib.Nifti1Image(np.zeros((20, 20, 20)), np.eye(4)), - path, - ) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - assert not subject.t1.is_loaded - tio.Crop(cropping=5)(subject) - assert not subject.t1.is_loaded - - def test_crop_history(self) -> None: - image = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - subject = tio.Subject(t1=image) - result = tio.Crop(cropping=5)(subject) - assert len(result.applied_transforms) == 1 - assert result.applied_transforms[0].name == "Crop" - - def test_crop_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - result = tio.Crop(cropping=5)(image) - assert isinstance(result, tio.Image) - assert result.shape == (1, 10, 10, 10) - - def test_crop_accepts_tensor(self) -> None: - tensor = torch.rand(1, 20, 20, 20) - result = tio.Crop(cropping=5)(tensor) - assert isinstance(result, torch.Tensor) - assert result.shape == (1, 10, 10, 10) - - def test_crop_batch(self) -> None: - from torchio.data.batch import SubjectsBatch - - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 20, 20, 20)), - ) - for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.Crop(cropping=5)(batch) - assert result.t1.data.shape == (3, 1, 10, 10, 10) - - def test_crop_inverse_on_image(self) -> None: - """apply_inverse_transform works directly on a cropped Image.""" - image = tio.ScalarImage(torch.rand(1, 200, 200, 200)) - cropped = tio.Crop(cropping=50)(image) - assert cropped.shape == (1, 100, 100, 100) - restored = tio.apply_inverse_transform(cropped) - assert restored.shape == (1, 200, 200, 200) diff --git a/tests/test_crop_or_pad.py b/tests/test_crop_or_pad.py deleted file mode 100644 index 5382f10f1..000000000 --- a/tests/test_crop_or_pad.py +++ /dev/null @@ -1,706 +0,0 @@ -"""Tests for the CropOrPad transform.""" - -from __future__ import annotations - -from pathlib import Path - -import nibabel as nib -import numpy as np -import pytest -import torch - -import torchio as tio -from torchio.data.affine import AffineMatrix -from torchio.data.batch import SubjectsBatch - -# --------------------------------------------------------------------------- -# Helper -# --------------------------------------------------------------------------- - - -def _make_subject( - shape: tuple[int, int, int] = (20, 20, 20), - spacing: tuple[float, float, float] = (1.0, 1.0, 1.0), - *, - with_label: bool = False, -) -> tio.Subject: - affine = AffineMatrix.from_spacing(spacing) - image = tio.ScalarImage( - torch.rand(1, *shape), - affine=affine, - ) - kwargs: dict = {"t1": image} - if with_label: - kwargs["seg"] = tio.LabelMap( - torch.randint(0, 3, (1, *shape)), - affine=affine, - ) - return tio.Subject(**kwargs) - - -# --------------------------------------------------------------------------- -# Basic crop / pad / no-op -# --------------------------------------------------------------------------- - - -class TestCropOrPadBasic: - def test_no_op_when_already_target_shape(self) -> None: - subject = _make_subject((10, 10, 10)) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_pad_when_smaller(self) -> None: - subject = _make_subject((8, 8, 8)) - result = tio.CropOrPad(target_shape=12)(subject) - assert result.t1.shape == (1, 12, 12, 12) - - def test_crop_when_larger(self) -> None: - subject = _make_subject((20, 20, 20)) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_mixed_crop_and_pad(self) -> None: - """Some axes need cropping, others padding.""" - subject = _make_subject((30, 10, 20)) - result = tio.CropOrPad(target_shape=(20, 20, 20))(subject) - assert result.t1.shape == (1, 20, 20, 20) - - def test_odd_difference_centering(self) -> None: - """When the difference is odd, ini gets ceil and fin gets floor.""" - subject = _make_subject((10, 10, 10)) - result = tio.CropOrPad(target_shape=13)(subject) - assert result.t1.shape == (1, 13, 13, 13) - - def test_crop_odd_difference_centering(self) -> None: - subject = _make_subject((13, 13, 13)) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.shape == (1, 10, 10, 10) - - -# --------------------------------------------------------------------------- -# target_shape specifications -# --------------------------------------------------------------------------- - - -class TestTargetShapeParam: - def test_single_int(self) -> None: - subject = _make_subject((20, 20, 20)) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_three_tuple(self) -> None: - subject = _make_subject((20, 20, 20)) - result = tio.CropOrPad(target_shape=(10, 15, 20))(subject) - assert result.t1.shape == (1, 10, 15, 20) - - def test_none_leaves_axis_unchanged(self) -> None: - subject = _make_subject((30, 20, 10)) - result = tio.CropOrPad(target_shape=(10, None, 20))(subject) - assert result.t1.shape == (1, 10, 20, 20) - - def test_all_none_is_no_op(self) -> None: - subject = _make_subject((30, 20, 10)) - result = tio.CropOrPad(target_shape=(None, None, None))(subject) - assert result.t1.shape == (1, 30, 20, 10) - - def test_none_with_units(self) -> None: - # 20 voxels at 2 mm = 40 mm, target None → keep 20 - subject = _make_subject((20, 20, 20), spacing=(2.0, 2.0, 2.0)) - result = tio.CropOrPad( - target_shape=(30.0, None, 30.0), - units="mm", - )(subject) - assert result.t1.shape == (1, 15, 20, 15) - - def test_invalid_tuple_length(self) -> None: - with pytest.raises(ValueError, match="1 or 3"): - tio.CropOrPad(target_shape=(1, 2)) # type: ignore[arg-type] - - def test_invalid_tuple_length_four(self) -> None: - with pytest.raises(ValueError, match="1 or 3"): - tio.CropOrPad(target_shape=(1, 2, 3, 4)) # type: ignore[arg-type] - - -# --------------------------------------------------------------------------- -# Units -# --------------------------------------------------------------------------- - - -class TestUnits: - def test_voxels_default(self) -> None: - subject = _make_subject((20, 20, 20), spacing=(2.0, 2.0, 2.0)) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_mm(self) -> None: - # 20 voxels at 2 mm spacing = 40 mm. Target 30 mm → 15 voxels. - subject = _make_subject((20, 20, 20), spacing=(2.0, 2.0, 2.0)) - result = tio.CropOrPad(target_shape=30.0, units="mm")(subject) - assert result.t1.shape == (1, 15, 15, 15) - - def test_cm(self) -> None: - # 20 voxels at 2 mm spacing = 40 mm. Target 3 cm = 30 mm → 15 voxels. - subject = _make_subject((20, 20, 20), spacing=(2.0, 2.0, 2.0)) - result = tio.CropOrPad(target_shape=3.0, units="cm")(subject) - assert result.t1.shape == (1, 15, 15, 15) - - def test_mm_per_axis(self) -> None: - # spacing (1, 2, 4) mm - # target (10, 20, 40) mm → (10, 10, 10) voxels - subject = _make_subject((20, 20, 20), spacing=(1.0, 2.0, 4.0)) - result = tio.CropOrPad( - target_shape=(10.0, 20.0, 40.0), - units="mm", - )(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_mm_rounds_to_nearest(self) -> None: - # spacing 3 mm, target 10 mm → 10/3 ≈ 3.33 → round → 3 voxels - subject = _make_subject((20, 20, 20), spacing=(3.0, 3.0, 3.0)) - result = tio.CropOrPad(target_shape=10.0, units="mm")(subject) - assert result.t1.shape == (1, 3, 3, 3) - - def test_mm_rounds_up_at_half(self) -> None: - # spacing 2 mm, target 5 mm → 5/2 = 2.5 → round → 2 voxels - # (Python round uses banker's rounding: 2.5 → 2) - subject = _make_subject((20, 20, 20), spacing=(2.0, 2.0, 2.0)) - result = tio.CropOrPad(target_shape=5.0, units="mm")(subject) - assert result.t1.shape == (1, 2, 2, 2) - - def test_invalid_units(self) -> None: - with pytest.raises(ValueError, match="units"): - tio.CropOrPad(target_shape=10, units="inches") # type: ignore[arg-type] - - -# --------------------------------------------------------------------------- -# only_crop / only_pad -# --------------------------------------------------------------------------- - - -class TestOnlyCropOnlyPad: - def test_only_crop_true_skips_padding(self) -> None: - subject = _make_subject((20, 10, 20)) - result = tio.CropOrPad( - target_shape=(15, 15, 15), - only_crop=True, - )(subject) - # Axis 0: 20→15 (crop), axis 1: 10→15 (skip), axis 2: 20→15 (crop) - assert result.t1.shape == (1, 15, 10, 15) - - def test_only_pad_true_skips_cropping(self) -> None: - subject = _make_subject((20, 10, 20)) - result = tio.CropOrPad( - target_shape=(15, 15, 15), - only_pad=True, - )(subject) - # Axis 0: 20→15 (skip), axis 1: 10→15 (pad), axis 2: 20→15 (skip) - assert result.t1.shape == (1, 20, 15, 20) - - def test_only_crop_no_op_when_all_smaller(self) -> None: - subject = _make_subject((5, 5, 5)) - result = tio.CropOrPad(target_shape=10, only_crop=True)(subject) - assert result.t1.shape == (1, 5, 5, 5) - - def test_only_pad_no_op_when_all_larger(self) -> None: - subject = _make_subject((20, 20, 20)) - result = tio.CropOrPad(target_shape=10, only_pad=True)(subject) - assert result.t1.shape == (1, 20, 20, 20) - - def test_both_raises(self) -> None: - with pytest.raises(ValueError, match="cannot both be True"): - tio.CropOrPad(target_shape=10, only_crop=True, only_pad=True) - - -# --------------------------------------------------------------------------- -# Padding mode / fill -# --------------------------------------------------------------------------- - - -class TestPaddingMode: - def test_constant_fill(self) -> None: - tensor = torch.ones(1, 4, 4, 4) - image = tio.ScalarImage(tensor) - subject = tio.Subject(t1=image) - result = tio.CropOrPad(target_shape=8, fill=-1)(subject) - # Padded corners should be -1 - assert result.t1.data[0, 0, 0, 0] == -1 - # Interior should be 1 - assert result.t1.data[0, 2, 2, 2] == 1 - - def test_reflect_mode(self) -> None: - subject = _make_subject((4, 4, 4)) - result = tio.CropOrPad( - target_shape=8, - padding_mode="reflect", - )(subject) - assert result.t1.shape == (1, 8, 8, 8) - - @pytest.mark.parametrize( - ("padding_mode", "expected"), - [ - ("mean", 3.5), - ("median", 3.5), - ("minimum", 0), - ], - ) - def test_statistic_mode_tensor_path( - self, - padding_mode: str, - expected: float, - ) -> None: - tensor = torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2) - result = tio.CropOrPad( - target_shape=4, - padding_mode=padding_mode, - )(tensor) - assert result[0, 0, 0, 0].item() == expected - - -# --------------------------------------------------------------------------- -# AffineMatrix correctness -# --------------------------------------------------------------------------- - - -class TestAffine: - def test_crop_shifts_origin_forward(self) -> None: - subject = _make_subject((20, 20, 20)) - orig = subject.t1.affine.origin - result = tio.CropOrPad(target_shape=10)(subject) - new = result.t1.affine.origin - # With identity direction and 1mm spacing, cropping 5 from start - # shifts origin by +5 on each axis - for o, n in zip(orig, new, strict=True): - assert n > o - - def test_pad_shifts_origin_backward(self) -> None: - subject = _make_subject((10, 10, 10)) - orig = subject.t1.affine.origin - result = tio.CropOrPad(target_shape=20)(subject) - new = result.t1.affine.origin - for o, n in zip(orig, new, strict=True): - assert n < o - - def test_affine_with_anisotropic_spacing(self) -> None: - spacing = (0.5, 1.0, 2.0) - subject = _make_subject((20, 20, 20), spacing=spacing) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.affine.spacing == pytest.approx(spacing) - - -# --------------------------------------------------------------------------- -# All images -# --------------------------------------------------------------------------- - - -class TestAllImages: - def test_crop_or_pad_all_images(self) -> None: - subject = _make_subject((20, 20, 20), with_label=True) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.shape == (1, 10, 10, 10) - assert result.seg.shape == (1, 10, 10, 10) - - -# --------------------------------------------------------------------------- -# Invertibility -# --------------------------------------------------------------------------- - - -class TestInvertibility: - def test_crop_then_inverse(self) -> None: - tensor = torch.rand(1, 20, 20, 20) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - transformed = tio.CropOrPad(target_shape=10)(subject) - assert transformed.t1.shape == (1, 10, 10, 10) - restored = transformed.apply_inverse_transform() - assert restored.t1.shape == (1, 20, 20, 20) - - def test_pad_then_inverse(self) -> None: - tensor = torch.rand(1, 10, 10, 10) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - transformed = tio.CropOrPad(target_shape=20)(subject) - assert transformed.t1.shape == (1, 20, 20, 20) - restored = transformed.apply_inverse_transform() - assert restored.t1.shape == (1, 10, 10, 10) - torch.testing.assert_close(restored.t1.data, tensor) - - def test_mixed_then_inverse(self) -> None: - tensor = torch.rand(1, 30, 10, 20) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - transformed = tio.CropOrPad(target_shape=20)(subject) - assert transformed.t1.shape == (1, 20, 20, 20) - restored = transformed.apply_inverse_transform() - assert restored.t1.shape == (1, 30, 10, 20) - - -# --------------------------------------------------------------------------- -# Input types -# --------------------------------------------------------------------------- - - -class TestInputTypes: - def test_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - result = tio.CropOrPad(target_shape=10)(image) - assert isinstance(result, tio.Image) - assert result.shape == (1, 10, 10, 10) - - def test_accepts_tensor(self) -> None: - tensor = torch.rand(1, 20, 20, 20) - result = tio.CropOrPad(target_shape=10)(tensor) - assert isinstance(result, torch.Tensor) - assert result.shape == (1, 10, 10, 10) - - def test_accepts_subject(self) -> None: - subject = _make_subject((20, 20, 20)) - result = tio.CropOrPad(target_shape=10)(subject) - assert isinstance(result, tio.Subject) - assert result.t1.shape == (1, 10, 10, 10) - - -# --------------------------------------------------------------------------- -# Batch mode -# --------------------------------------------------------------------------- - - -class TestBatch: - def test_batch_crop(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 20, 20, 20)), - ) - for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.CropOrPad(target_shape=10)(batch) - assert result.t1.data.shape == (3, 1, 10, 10, 10) - - def test_batch_pad(self) -> None: - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.CropOrPad(target_shape=20)(batch) - assert result.t1.data.shape == (3, 1, 20, 20, 20) - - -# --------------------------------------------------------------------------- -# Probability -# --------------------------------------------------------------------------- - - -class TestProbability: - def test_p_zero_is_no_op(self) -> None: - subject = _make_subject((20, 20, 20)) - result = tio.CropOrPad(target_shape=10, p=0)(subject) - assert result.t1.shape == (1, 20, 20, 20) - - -# --------------------------------------------------------------------------- -# Random location -# --------------------------------------------------------------------------- - - -class TestRandomLocation: - def test_random_crop_shape(self) -> None: - subject = _make_subject((30, 30, 30)) - result = tio.CropOrPad(target_shape=10, location="random")(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_random_crop_varies(self) -> None: - """Two random crops of the same subject should (usually) differ.""" - torch.manual_seed(0) - data = torch.arange(20 * 20 * 20, dtype=torch.float32).reshape(1, 20, 20, 20) - transform = tio.CropOrPad(target_shape=5, location="random") - r1 = transform(tio.ScalarImage(data.clone())) - r2 = transform(tio.ScalarImage(data.clone())) - assert not torch.equal(r1.data, r2.data) - - def test_random_pad_is_still_centered(self) -> None: - """Padding should be centered even with location='random'.""" - subject = _make_subject((10, 10, 10)) - result_center = tio.CropOrPad(target_shape=20, location="center")(subject) - result_random = tio.CropOrPad(target_shape=20, location="random")(subject) - # Pure padding: both should produce the same result - torch.testing.assert_close(result_center.t1.data, result_random.t1.data) - - def test_random_mixed_crop_and_pad(self) -> None: - subject = _make_subject((30, 5, 20)) - result = tio.CropOrPad(target_shape=10, location="random")(subject) - assert result.t1.shape == (1, 10, 10, 10) - - def test_random_with_none_axis(self) -> None: - subject = _make_subject((30, 20, 10)) - result = tio.CropOrPad( - target_shape=(10, None, 10), - location="random", - )(subject) - assert result.t1.shape == (1, 10, 20, 10) - - def test_random_batch(self) -> None: - from torchio.data.batch import SubjectsBatch - - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 20, 20, 20)), - ) - for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.CropOrPad(target_shape=10, location="random")(batch) - assert result.t1.data.shape == (3, 1, 10, 10, 10) - - def test_invalid_location(self) -> None: - with pytest.raises(ValueError, match="location"): - tio.CropOrPad(target_shape=10, location="top-left") # type: ignore[arg-type] - - -# --------------------------------------------------------------------------- -# Laziness preservation -# --------------------------------------------------------------------------- - - -class TestLaziness: - def test_preserves_laziness_of_original(self, tmp_path) -> None: - """CropOrPad should not load the original image's data.""" - import nibabel as nib - import numpy as np - - path = tmp_path / "test.nii.gz" - nib.save( - nib.Nifti1Image(np.zeros((20, 20, 20)), np.eye(4)), - path, - ) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - assert not subject.t1.is_loaded - tio.CropOrPad(target_shape=10)(subject) - assert not subject.t1.is_loaded - - -# --------------------------------------------------------------------------- -# Lazy backend coverage -# --------------------------------------------------------------------------- - - -class TestLazyBackends: - """Test _CroppedBackend and _PaddedBackend properties and data access.""" - - @staticmethod - def _make_nii(path, shape=(20, 20, 20)): - import nibabel as nib - import numpy as np - - data = np.random.rand(*shape).astype(np.float32) - nib.save(nib.Nifti1Image(data, np.eye(4)), path) - return data - - def test_crop_lazy_backend_shape(self, tmp_path) -> None: - path = tmp_path / "test.nii.gz" - self._make_nii(path) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - assert not subject.t1.is_loaded - result = tio.CropOrPad(target_shape=10)(subject) - # Should be cropped to 10x10x10. - assert result.t1.shape == (1, 10, 10, 10) - - def test_crop_lazy_backend_data(self, tmp_path) -> None: - path = tmp_path / "test.nii.gz" - self._make_nii(path) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - result = tio.CropOrPad(target_shape=10)(subject) - # Access the data (triggers lazy load). - data = result.t1.data - assert data.shape == (1, 10, 10, 10) - assert data.dtype == torch.float32 - - def test_crop_lazy_backend_affine(self, tmp_path) -> None: - path = tmp_path / "test.nii.gz" - self._make_nii(path) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - result = tio.CropOrPad(target_shape=10)(subject) - assert result.t1.affine is not None - - def test_pad_lazy_backend_shape(self, tmp_path) -> None: - path = tmp_path / "test.nii.gz" - self._make_nii(path, shape=(8, 8, 8)) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - assert not subject.t1.is_loaded - result = tio.CropOrPad(target_shape=12)(subject) - assert result.t1.shape == (1, 12, 12, 12) - - def test_pad_lazy_backend_data(self, tmp_path) -> None: - path = tmp_path / "test.nii.gz" - self._make_nii(path, shape=(8, 8, 8)) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - result = tio.CropOrPad(target_shape=12)(subject) - data = result.t1.data - assert data.shape == (1, 12, 12, 12) - - def test_pad_lazy_backend_affine(self, tmp_path) -> None: - path = tmp_path / "test.nii.gz" - self._make_nii(path, shape=(8, 8, 8)) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - result = tio.CropOrPad(target_shape=12)(subject) - assert result.t1.affine is not None - - @pytest.mark.parametrize( - ("padding_mode", "expected"), - [ - ("mean", 3.5), - ("median", 3.5), - ("minimum", 0), - ], - ) - def test_pad_lazy_backend_statistic_mode( - self, - tmp_path, - padding_mode: str, - expected: float, - ) -> None: - path = tmp_path / "test.nii.gz" - data = np.arange(8, dtype=np.float32).reshape(2, 2, 2) - nib.save(nib.Nifti1Image(data, np.eye(4)), path) - result = tio.CropOrPad( - target_shape=4, - padding_mode=padding_mode, - )(tio.Subject(t1=tio.ScalarImage(path))) - assert not result.t1.is_loaded - assert result.t1.data[0, 0, 0, 0].item() == expected - - def test_crop_and_pad_lazy_mixed(self, tmp_path) -> None: - """Anisotropic shape needing both crop and pad.""" - path = tmp_path / "test.nii.gz" - self._make_nii(path, shape=(20, 8, 15)) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - result = tio.CropOrPad(target_shape=12)(subject) - assert result.t1.shape == (1, 12, 12, 12) - data = result.t1.data - assert data.shape == (1, 12, 12, 12) - - def test_lazy_dtype(self, tmp_path) -> None: - path = tmp_path / "test.nii.gz" - self._make_nii(path) - image = tio.ScalarImage(path) - subject = tio.Subject(t1=image) - result = tio.CropOrPad(target_shape=10)(subject) - # Access dtype through the lazy backend. - assert result.t1.shape == (1, 10, 10, 10) - - def test_deepcopy_cropped_lazy_preserves_shape(self, tmp_path) -> None: - """Deep-copying a lazily cropped image must keep the cropped shape. - - Regression test: ``__deepcopy__`` used to rebuild the image from its - source path only, discarding the ``_CroppedBackend`` and reverting to - the full-resolution image. - """ - import copy - - path = tmp_path / "test.nii.gz" - self._make_nii(path, shape=(20, 20, 20)) - result = tio.CropOrPad(target_shape=10)(tio.Subject(t1=tio.ScalarImage(path))) - image = result.t1 - assert not image.is_loaded - - copied = copy.deepcopy(image) - assert copied.shape == (1, 10, 10, 10) - assert not copied.is_loaded - torch.testing.assert_close(copied.data, image.data) - - def test_deepcopy_padded_lazy_preserves_shape(self, tmp_path) -> None: - """Deep-copying a lazily padded image must keep the padded shape.""" - import copy - - path = tmp_path / "test.nii.gz" - self._make_nii(path, shape=(8, 8, 8)) - result = tio.CropOrPad(target_shape=12)(tio.Subject(t1=tio.ScalarImage(path))) - image = result.t1 - assert not image.is_loaded - - copied = copy.deepcopy(image) - assert copied.shape == (1, 12, 12, 12) - assert not copied.is_loaded - torch.testing.assert_close(copied.data, image.data) - - def test_transform_after_lazy_crop_uses_cropped_shape(self, tmp_path) -> None: - """A transform applied after a lazy crop must see the cropped shape. - - Transforms deep-copy their input in ``forward``; before the fix this - reverted the crop, so the second transform operated on the original - full-resolution image. - """ - path = tmp_path / "test.nii.gz" - self._make_nii(path, shape=(20, 20, 20)) - cropped = tio.CropOrPad(target_shape=10)(tio.Subject(t1=tio.ScalarImage(path))) - assert not cropped.t1.is_loaded - padded = tio.Pad(padding=2)(cropped) - assert padded.t1.shape == (1, 14, 14, 14) - - def test_deepcopy_nibabel_backed_lazy_image(self) -> None: - """Deep-copying a backend-only (nibabel) lazy image must work. - - Such images have ``_backend`` set but no ``_path`` or ``_data``; the - copy used to fall through to an empty image and fail on ``shape``. - """ - import copy - - nii = nib.Nifti1Image( - np.random.rand(6, 6, 6).astype(np.float32), - np.eye(4), - ) - image = tio.ScalarImage(nii) - assert not image.is_loaded - - copied = copy.deepcopy(image) - assert copied.shape == (1, 6, 6, 6) - assert not copied.is_loaded - torch.testing.assert_close(copied.data, image.data) - - -class TestLazyCropPadAffine: - """Lazy crop/pad must keep image.affine and image.dataobj.affine consistent. - - For unloaded path-based images, CropOrPad installs `_CroppedBackend` / - `_PaddedBackend`; their reported affine must match the cropped/padded - image's affine (shifted origin), not the original source affine. - """ - - def _path_subject( - self, - tmp_path: Path, - shape: tuple[int, int, int] = (10, 12, 14), - ) -> tio.Subject: - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - path = tmp_path / "t1.nii.gz" - nib.save(nib.Nifti1Image(np.zeros(shape, "float32"), affine), path) - return tio.Subject(t1=tio.ScalarImage(path)) - - def test_lazy_crop_affine_consistent(self, tmp_path: Path) -> None: - subject = self._path_subject(tmp_path) - out = tio.CropOrPad(target_shape=(6, 8, 10))(subject)["t1"] - assert not out.is_loaded # still lazy - np.testing.assert_allclose(out.affine.numpy(), np.asarray(out.dataobj.affine)) - - def test_lazy_pad_affine_consistent(self, tmp_path: Path) -> None: - subject = self._path_subject(tmp_path) - out = tio.CropOrPad(target_shape=(16, 18, 20))(subject)["t1"] - assert not out.is_loaded # still lazy - np.testing.assert_allclose(out.affine.numpy(), np.asarray(out.dataobj.affine)) - - def test_lazy_crop_origin_shifted(self, tmp_path: Path) -> None: - subject = self._path_subject(tmp_path) - out = tio.CropOrPad(target_shape=(6, 8, 10))(subject)["t1"] - # 2-voxel crop start on each axis * spacing (2, 3, 4) -> origin (4, 6, 8) - np.testing.assert_allclose(out.affine.numpy()[:3, 3], [4.0, 6.0, 8.0]) - np.testing.assert_allclose( - np.asarray(out.dataobj.affine)[:3, 3], [4.0, 6.0, 8.0] - ) diff --git a/tests/test_datasets.py b/tests/test_datasets.py deleted file mode 100644 index 2e0528008..000000000 --- a/tests/test_datasets.py +++ /dev/null @@ -1,146 +0,0 @@ -"""Tests for built-in datasets and new Image features.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio -from torchio.datasets.zone_plate import ZonePlate -from torchio.download import check_integrity -from torchio.download import compress -from torchio.download import get_torchio_cache_dir -from torchio.io import read_matrix -from torchio.io import write_matrix - -# --- ZonePlate --------------------------------------------------------------- - - -class TestZonePlate: - def test_default_size(self): - zp = ZonePlate(size=11) - assert zp.size == 11 - img = zp.image - assert img.data.shape == (1, 11, 11, 11) - - def test_custom_size(self): - zp = ZonePlate(size=11) - img = zp.image - assert img.spatial_shape == (11, 11, 11) - - def test_even_size(self): - zp = ZonePlate(size=10) - img = zp.image - assert img.spatial_shape == (10, 10, 10) - - def test_minimum_size(self): - zp = ZonePlate(size=3) - img = zp.image - assert img.spatial_shape == (3, 3, 3) - - def test_too_small(self): - with pytest.raises(ValueError, match="at least 3"): - ZonePlate(size=2) - - def test_is_subject(self): - zp = ZonePlate(size=5) - assert isinstance(zp, tio.Subject) - - def test_has_scalar_image(self): - zp = ZonePlate(size=5) - assert isinstance(zp.image, tio.ScalarImage) - - def test_affine_origin(self): - zp = ZonePlate(size=11) - img = zp.image - origin = img.affine.origin - assert origin == pytest.approx((-5.0, -5.0, -5.0)) - - -# --- Download utilities ------------------------------------------------------ - - -class TestDownloadUtils: - def test_cache_dir(self): - d = get_torchio_cache_dir() - # On Windows, platformdirs appends a "Cache" subdirectory, so the leaf - # name is not necessarily "torchio"; check it appears in the path parts. - assert "torchio" in d.parts - assert d.is_absolute() - - def test_compress(self, tmp_path): - inp = tmp_path / "test.nii" - inp.write_bytes(b"fake nifti content " * 100) - out = compress(inp) - assert out.suffix == ".gz" - assert out.exists() - assert out.stat().st_size < inp.stat().st_size - - def test_compress_custom_output(self, tmp_path): - inp = tmp_path / "test.nii" - inp.write_bytes(b"hello" * 50) - custom = tmp_path / "custom.nii.gz" - result = compress(inp, custom) - assert result == custom - assert custom.exists() - - def test_check_integrity_missing(self, tmp_path): - assert not check_integrity(tmp_path / "nonexistent.bin") - - def test_check_integrity_exists_no_md5(self, tmp_path): - f = tmp_path / "file.bin" - f.write_bytes(b"data") - assert check_integrity(f) - - -# --- read_matrix / write_matrix ---------------------------------------------- - - -class TestMatrixIO: - def test_roundtrip_tfm(self, tmp_path): - matrix = torch.eye(4, dtype=torch.float64) - matrix[0, 3] = 10.0 - matrix[1, 3] = -5.0 - path = tmp_path / "transform.tfm" - write_matrix(matrix, path) - loaded = read_matrix(path) - torch.testing.assert_close(loaded, matrix, atol=1e-6, rtol=1e-6) - - def test_roundtrip_txt(self, tmp_path): - matrix = torch.eye(4, dtype=torch.float64) - matrix[2, 3] = 7.0 - path = tmp_path / "transform.txt" - write_matrix(matrix, path) - loaded = read_matrix(path) - torch.testing.assert_close(loaded, matrix, atol=1e-6, rtol=1e-6) - - def test_unsupported_suffix(self, tmp_path): - with pytest.raises(ValueError, match="Unknown suffix"): - read_matrix(tmp_path / "bad.xyz") - - def test_write_unsupported_suffix(self, tmp_path): - with pytest.raises(ValueError, match="Unknown suffix"): - write_matrix(torch.eye(4), tmp_path / "bad.xyz") - - -# --- channels_last ----------------------------------------------------------- - - -class TestChannelsLast: - def test_from_tensor_channels_last(self): - data = torch.randn(10, 20, 30, 3) # (I, J, K, C) - img = tio.ScalarImage(data, channels_last=True) - assert img.data.shape == (3, 10, 20, 30) - - def test_from_tensor_channels_first(self): - data = torch.randn(3, 10, 20, 30) # (C, I, J, K) - img = tio.ScalarImage(data, channels_last=False) - assert img.data.shape == (3, 10, 20, 30) - - def test_channels_last_load(self, tmp_path): - # Create a standard 3D NIfTI image, then the reader returns (1,I,J,K) - # We test that channels_last permutes from (I,J,K,C) -> (C,I,J,K) - # by using from_tensor with explicit channels_last data - data = torch.randn(5, 6, 7, 3) # (I, J, K, C) - img = tio.LabelMap(data, channels_last=True) - assert img.data.shape == (3, 5, 6, 7) diff --git a/tests/test_device.py b/tests/test_device.py deleted file mode 100644 index cfa701d41..000000000 --- a/tests/test_device.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Tests for .to() device methods on Image and Subject, and To transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio -from torchio.data.bboxes import BoundingBoxes -from torchio.data.bboxes import BoundingBoxFormat -from torchio.data.points import Points - -HAS_CUDA = torch.cuda.is_available() -HAS_MPS = torch.backends.mps.is_available() - - -class TestImageTo: - def test_to_returns_self(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - result = image.to("cpu") - assert result is image - - def test_device_property(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - assert image.device == torch.device("cpu") - - def test_to_dtype(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - result = image.to(torch.float16) - assert result.data.dtype == torch.float16 - - @pytest.mark.skipif(not HAS_CUDA, reason="No CUDA") - def test_to_cuda(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - result = image.to("cuda") - assert result.device.type == "cuda" - assert result.data.is_cuda - - @pytest.mark.skipif(not HAS_MPS, reason="No MPS") - def test_to_mps(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - result = image.to("mps") - assert result.device.type == "mps" - - @pytest.mark.skipif(not HAS_MPS, reason="No MPS") - def test_mps_round_trip(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - original = image.data.clone() - image.to("mps").to("cpu") - torch.testing.assert_close(image.data, original) - - -class TestSubjectTo: - def test_to_returns_self(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - result = subject.to("cpu") - assert result is subject - - def test_moves_all_images(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 4, 4, 4))), - ) - result = subject.to(torch.float64) - assert result.t1.data.dtype == torch.float64 - assert result.seg.data.dtype == torch.float64 - - def test_moves_points(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - pts=Points(torch.rand(3, 3)), - ) - result = subject.to(torch.float64) - assert result.pts.data.dtype == torch.float64 - - def test_moves_bboxes(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - boxes=BoundingBoxes( - torch.rand(2, 6), - format=BoundingBoxFormat.IJKIJK, - ), - ) - result = subject.to(torch.float64) - assert result.boxes.data.dtype == torch.float64 - - @pytest.mark.skipif(not HAS_CUDA, reason="No CUDA") - def test_to_cuda(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - result = subject.to("cuda") - assert result.t1.data.is_cuda - - @pytest.mark.skipif(not HAS_MPS, reason="No MPS") - def test_to_mps(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - pts=Points(torch.rand(3, 3)), - ) - result = subject.to("mps") - assert result.t1.device.type == "mps" - assert result.pts.device.type == "mps" - assert result.device.type == "mps" - - -class TestToTransform: - def test_to_dtype(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - transform = tio.To(torch.float64) - result = transform(subject) - assert result.t1.data.dtype == torch.float64 - - def test_to_device_str(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - transform = tio.To("cpu") - result = transform(subject) - assert result.t1.device == torch.device("cpu") - - def test_history_recorded(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - transform = tio.To(torch.float64) - result = transform(subject) - assert len(result.applied_transforms) == 1 - assert result.applied_transforms[0].name == "To" - - def test_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - result = tio.To(torch.float64)(image) - assert isinstance(result, tio.Image) - assert result.data.dtype == torch.float64 - - def test_accepts_tensor(self) -> None: - tensor = torch.rand(1, 4, 4, 4) - result = tio.To(torch.float64)(tensor) - assert isinstance(result, torch.Tensor) - assert result.dtype == torch.float64 - - @pytest.mark.skipif(not HAS_MPS, reason="No MPS") - def test_to_mps_via_transform(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - result = tio.To("mps")(subject) - assert result.t1.device.type == "mps" diff --git a/tests/test_ensure_shape_multiple.py b/tests/test_ensure_shape_multiple.py deleted file mode 100644 index 6b024af57..000000000 --- a/tests/test_ensure_shape_multiple.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Tests for the EnsureShapeMultiple transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio -from torchio.data.affine import AffineMatrix -from torchio.data.batch import SubjectsBatch - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_subject( - shape: tuple[int, int, int] = (10, 10, 10), - spacing: tuple[float, float, float] = (1.0, 1.0, 1.0), - *, - with_label: bool = False, -) -> tio.Subject: - affine = AffineMatrix.from_spacing(spacing) - image = tio.ScalarImage(torch.rand(1, *shape), affine=affine) - kwargs: dict = {"t1": image} - if with_label: - kwargs["seg"] = tio.LabelMap( - torch.randint(0, 3, (1, *shape)), - affine=affine, - ) - return tio.Subject(**kwargs) - - -# --------------------------------------------------------------------------- -# Padding (default method) -# --------------------------------------------------------------------------- - - -class TestPad: - def test_pad_to_next_multiple(self) -> None: - subject = _make_subject((10, 10, 10)) - result = tio.EnsureShapeMultiple(8)(subject) - assert result.t1.spatial_shape == (16, 16, 16) - - def test_pad_asymmetric_shape(self) -> None: - subject = _make_subject((10, 17, 25)) - result = tio.EnsureShapeMultiple(8)(subject) - # 10→16, 17→24, 25→32 - assert result.t1.spatial_shape == (16, 24, 32) - - def test_pad_no_op_when_already_multiple(self) -> None: - subject = _make_subject((16, 24, 8)) - result = tio.EnsureShapeMultiple(8)(subject) - assert result.t1.spatial_shape == (16, 24, 8) - - def test_pad_per_axis_tuple(self) -> None: - subject = _make_subject((10, 10, 10)) - result = tio.EnsureShapeMultiple((4, 8, 16))(subject) - # 10→12, 10→16, 10→16 - assert result.t1.spatial_shape == (12, 16, 16) - - -# --------------------------------------------------------------------------- -# Cropping -# --------------------------------------------------------------------------- - - -class TestCrop: - def test_crop_to_previous_multiple(self) -> None: - subject = _make_subject((10, 10, 10)) - result = tio.EnsureShapeMultiple(8, method="crop")(subject) - assert result.t1.spatial_shape == (8, 8, 8) - - def test_crop_asymmetric_shape(self) -> None: - subject = _make_subject((10, 17, 25)) - result = tio.EnsureShapeMultiple(8, method="crop")(subject) - # 10→8, 17→16, 25→24 - assert result.t1.spatial_shape == (8, 16, 24) - - def test_crop_no_op_when_already_multiple(self) -> None: - subject = _make_subject((16, 24, 8)) - result = tio.EnsureShapeMultiple(8, method="crop")(subject) - assert result.t1.spatial_shape == (16, 24, 8) - - def test_crop_per_axis_tuple(self) -> None: - subject = _make_subject((10, 10, 10)) - result = tio.EnsureShapeMultiple((4, 6, 8), method="crop")(subject) - # 10→8, 10→6, 10→8 - assert result.t1.spatial_shape == (8, 6, 8) - - def test_crop_small_shape_clamps_to_one(self) -> None: - """When cropping would result in 0, clamp to at least 1.""" - subject = _make_subject((3, 3, 3)) - result = tio.EnsureShapeMultiple(8, method="crop")(subject) - # floor(3/8) = 0 → 0*8 = 0 → max(0, 1) = 1 - assert all(s >= 1 for s in result.t1.spatial_shape) - - -# --------------------------------------------------------------------------- -# Validation -# --------------------------------------------------------------------------- - - -class TestValidation: - def test_invalid_method_raises(self) -> None: - with pytest.raises(ValueError, match="method"): - tio.EnsureShapeMultiple(8, method="resize") - - def test_invalid_padding_mode_raises(self) -> None: - with pytest.raises(ValueError, match="padding_mode"): - tio.EnsureShapeMultiple(8, padding_mode="maximum") # type: ignore[arg-type] - - def test_method_must_be_crop_or_pad(self) -> None: - # Valid methods should not raise - tio.EnsureShapeMultiple(8, method="crop") - tio.EnsureShapeMultiple(8, method="pad") - - -# --------------------------------------------------------------------------- -# Input types -# --------------------------------------------------------------------------- - - -class TestInputTypes: - def test_accepts_subject(self) -> None: - subject = _make_subject((10, 10, 10)) - result = tio.EnsureShapeMultiple(8)(subject) - assert isinstance(result, tio.Subject) - assert result.t1.spatial_shape == (16, 16, 16) - - def test_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - result = tio.EnsureShapeMultiple(8)(image) - assert isinstance(result, tio.Image) - assert result.spatial_shape == (16, 16, 16) - - def test_accepts_tensor(self) -> None: - tensor = torch.rand(1, 10, 10, 10) - result = tio.EnsureShapeMultiple(8)(tensor) - assert isinstance(result, torch.Tensor) - assert result.shape == (1, 16, 16, 16) - - -# --------------------------------------------------------------------------- -# Batch -# --------------------------------------------------------------------------- - - -class TestBatch: - def test_batch_pad(self) -> None: - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.rand(1, 10, 10, 10))) for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.EnsureShapeMultiple(8)(batch) - assert result.t1.data.shape == (3, 1, 16, 16, 16) - - def test_batch_crop(self) -> None: - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.rand(1, 10, 10, 10))) for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.EnsureShapeMultiple(8, method="crop")(batch) - assert result.t1.data.shape == (3, 1, 8, 8, 8) - - -# --------------------------------------------------------------------------- -# Multiple images -# --------------------------------------------------------------------------- - - -class TestMultipleImages: - def test_all_images_transformed(self) -> None: - subject = _make_subject((10, 10, 10), with_label=True) - result = tio.EnsureShapeMultiple(8)(subject) - assert result.t1.spatial_shape == (16, 16, 16) - assert result.seg.spatial_shape == (16, 16, 16) - - -# --------------------------------------------------------------------------- -# Probability -# --------------------------------------------------------------------------- - - -class TestProbability: - def test_p_zero_is_no_op(self) -> None: - subject = _make_subject((10, 10, 10)) - result = tio.EnsureShapeMultiple(8, p=0)(subject) - assert result.t1.spatial_shape == (10, 10, 10) - - -# --------------------------------------------------------------------------- -# Power of 2 use case (common for U-Nets) -# --------------------------------------------------------------------------- - - -class TestPowerOfTwo: - def test_three_pooling_layers(self) -> None: - """Common U-Net use case: 3 pooling layers → multiple of 8.""" - subject = _make_subject((181, 217, 181)) - result = tio.EnsureShapeMultiple(2**3)(subject) - for s in result.t1.spatial_shape: - assert s % 8 == 0 - - def test_four_pooling_layers(self) -> None: - """4 pooling layers → multiple of 16.""" - subject = _make_subject((181, 217, 181)) - result = tio.EnsureShapeMultiple(2**4)(subject) - for s in result.t1.spatial_shape: - assert s % 16 == 0 - - -# ── Coverage gap tests ─────────────────────────────────────────────── - - -class TestEnsureShapeMultipleValidation: - def test_zero_multiple_raises(self) -> None: - with pytest.raises(ValueError, match=">= 1"): - tio.EnsureShapeMultiple(target_multiple=0) - - def test_wrong_tuple_length_raises(self) -> None: - with pytest.raises(ValueError, match="1 or 3"): - tio.EnsureShapeMultiple(target_multiple=(2, 4)) - - def test_negative_in_tuple_raises(self) -> None: - with pytest.raises(ValueError, match=">= 1"): - tio.EnsureShapeMultiple(target_multiple=(2, -1, 4)) diff --git a/tests/test_flip.py b/tests/test_flip.py deleted file mode 100644 index 60b766cdf..000000000 --- a/tests/test_flip.py +++ /dev/null @@ -1,315 +0,0 @@ -"""Tests for the Flip spatial transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - -HAS_MPS = torch.backends.mps.is_available() - - -class TestFlip: - def test_flip_axis_0(self) -> None: - tensor = torch.arange(8).reshape(1, 2, 2, 2).float() - image = tio.ScalarImage(tensor) - subject = tio.Subject(t1=image) - result = tio.Flip(axes=0)(subject) - expected = torch.flip(tensor, [1]) - torch.testing.assert_close(result.t1.data, expected) - - def test_flip_single_int_axis(self) -> None: - """axes=0 should work the same as axes=(0,).""" - tensor = torch.arange(8).reshape(1, 2, 2, 2).float() - s1 = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - s2 = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - r1 = tio.Flip(axes=0)(s1) - r2 = tio.Flip(axes=(0,))(s2) - torch.testing.assert_close(r1.t1.data, r2.t1.data) - - def test_flip_multiple_axes(self) -> None: - tensor = torch.arange(8).reshape(1, 2, 2, 2).float() - image = tio.ScalarImage(tensor) - subject = tio.Subject(t1=image) - result = tio.Flip(axes=(0, 1))(subject) - expected = torch.flip(tensor, [1, 2]) - torch.testing.assert_close(result.t1.data, expected) - - def test_flip_all_images(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 4, 4, 4))), - ) - original_t1 = subject.t1.data.clone() - original_seg = subject.seg.data.clone() - result = tio.Flip(axes=2)(subject) - assert not torch.equal(result.t1.data, original_t1) - assert not torch.equal(result.seg.data, original_seg) - - def test_flip_is_self_inverse(self) -> None: - tensor = torch.rand(1, 4, 5, 6) - image = tio.ScalarImage(tensor.clone()) - subject = tio.Subject(t1=image) - flip = tio.Flip(axes=(0, 1, 2)) - result = flip(flip(subject)) - torch.testing.assert_close(result.t1.data, tensor) - - def test_flip_with_probability_zero(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - original = subject.t1.data.clone() - result = tio.Flip(axes=0, p=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_flip_probability_per_axis(self) -> None: - """flip_probability=0 should not flip any axis.""" - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - original = subject.t1.data.clone() - result = tio.Flip(axes=(0, 1, 2), flip_probability=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_flip_probability_one(self) -> None: - """flip_probability=1 should always flip all specified axes.""" - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - original = subject.t1.data.clone() - result = tio.Flip(axes=(0, 1, 2), flip_probability=1.0)(subject) - expected = torch.flip(original, [1, 2, 3]) - torch.testing.assert_close(result.t1.data, expected) - - def test_flip_history_recorded(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - result = tio.Flip(axes=1)(subject) - assert len(result.applied_transforms) == 1 - assert result.applied_transforms[0].name == "Flip" - - def test_flip_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - result = tio.Flip(axes=0)(image) - assert isinstance(result, tio.Image) - - def test_flip_accepts_tensor(self) -> None: - tensor = torch.rand(1, 4, 4, 4) - result = tio.Flip(axes=0)(tensor) - assert isinstance(result, torch.Tensor) - - def test_flip_in_compose(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - pipeline = tio.Compose([tio.Flip(axes=0), tio.Flip(axes=1)]) - result = pipeline(subject) - assert len(result.applied_transforms) == 2 - - def test_flip_differentiable(self) -> None: - tensor = torch.rand(1, 4, 4, 4, requires_grad=True) - result = tio.Flip(axes=0, copy=False)(tensor) - loss = result.sum() - loss.backward() - assert tensor.grad is not None - - @pytest.mark.skipif(not HAS_MPS, reason="No MPS") - def test_flip_on_mps(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - subject.to("mps") - result = tio.Flip(axes=0)(subject) - assert result.t1.device.type == "mps" - - def test_invalid_axis(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - with pytest.raises(ValueError, match="0, 1, or 2"): - tio.Flip(axes=3)(subject) - - def test_string_axis(self) -> None: - """Anatomical label 'Left' should resolve to an axis.""" - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - original = subject.t1.data.clone() - # Default affine is RAS, so 'Left'/'Right' = axis 0 - result = tio.Flip(axes="Left")(subject) - expected = torch.flip(original, [1]) - torch.testing.assert_close(result.t1.data, expected) - - def test_string_axis_lr(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - result = tio.Flip(axes="LR")(subject) - assert result.t1.shape == (1, 4, 4, 4) - - def test_invalid_string_axis(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - with pytest.raises(ValueError, match="Unknown anatomical"): - tio.Flip(axes="X")(subject) - - def test_flip_invertible(self) -> None: - assert tio.Flip(axes=0).invertible - - def test_flip_inverse_round_trip(self) -> None: - """Subject.apply_inverse_transform round-trips a flip.""" - tensor = torch.rand(1, 4, 5, 6) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - flip = tio.Flip(axes=(0, 1, 2)) - flipped = flip(subject) - restored = flipped.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, tensor) - - def test_inverse_respects_include_scope(self) -> None: - a = torch.arange(8.0).reshape(1, 2, 2, 2) - b = torch.arange(100.0, 108.0).reshape(1, 2, 2, 2) - subject = tio.Subject( - a=tio.ScalarImage(a.clone()), - b=tio.ScalarImage(b.clone()), - ) - - transformed = tio.Flip(axes=(0,), include=["a"])(subject) - restored = transformed.apply_inverse_transform() - - torch.testing.assert_close(restored.a.data, a) - torch.testing.assert_close(restored.b.data, b) - - def test_compose_inverse(self) -> None: - """Compose inverse via Subject.apply_inverse_transform.""" - tensor = torch.rand(1, 4, 5, 6) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - pipeline = tio.Compose( - [ - tio.Flip(axes=0), - tio.Flip(axes=1), - ] - ) - transformed = pipeline(subject) - assert len(transformed.applied_transforms) == 2 - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, tensor) - - def test_inverse_on_image_via_subject(self) -> None: - """Inverse works by copying history to a renamed prediction subject.""" - tensor = torch.rand(1, 4, 5, 6) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - flipped = tio.Flip(axes=0)(subject) - # Create a prediction subject and copy history - pred_subject = tio.Subject( - pred=tio.ScalarImage(flipped.t1.data.clone()), - ) - pred_subject.applied_transforms = flipped.applied_transforms - restored = pred_subject.apply_inverse_transform() - torch.testing.assert_close(restored.pred.data, tensor) - - def test_inverse_skips_non_invertible(self) -> None: - """Non-invertible transforms are skipped with a warning.""" - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - pipeline = tio.Compose( - [ - tio.Flip(axes=0), - tio.Noise(std=0.1), # not invertible - ] - ) - transformed = pipeline(subject) - import warnings - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - restored = transformed.apply_inverse_transform() - assert any("not invertible" in str(x.message) for x in w) - assert restored.t1.shape == (1, 4, 4, 4) - - def test_ignore_intensity(self) -> None: - """ignore_intensity=True skips intensity transforms silently.""" - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - pipeline = tio.Compose( - [ - tio.Flip(axes=0), - tio.Noise(std=0.1), - ] - ) - transformed = pipeline(subject) - import warnings - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - transformed.apply_inverse_transform(ignore_intensity=True) - # No warning about Noise since intensity is ignored - assert not any("Noise" in str(x.message) for x in w) - - -class TestFlipPerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - data = torch.rand(1, 8, 8, 8) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_axes_differ_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Flip(axes=(0, 1, 2), flip_probability=0.5)(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["axes"]) == batch.batch_size - distinct = {tuple(a) for a in params["axes"]} - assert len(distinct) > 1 - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Flip(axes=(0, 1, 2), flip_probability=0.5, per_instance=False)( - batch - ) - torch.testing.assert_close(result.t1.data[0], result.t1.data[1]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8))) - result = tio.Flip(axes=(0, 1, 2), flip_probability=0.5)(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params - - def test_per_instance_inverse_round_trip(self) -> None: - torch.manual_seed(0) - batch = self._batch() - original = batch.t1.data.clone() - result = tio.Flip(axes=(0, 1, 2), flip_probability=0.5)(batch) - restored = result.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original) - - def test_per_instance_inverse_after_unbatch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - original = batch.t1.data.clone() - result = tio.Flip(axes=(0, 1, 2), flip_probability=0.5)(batch) - for index, subject in enumerate(result.unbatch()): - restored = subject.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original[index]) - - -class TestFlipInclude: - def test_no_selected_images_is_noop(self) -> None: - # make_params must resolve orientation from the selected images and - # return a no-op when include matches nothing. - torch.manual_seed(0) - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8))) for _ in range(3) - ] - batch = tio.SubjectsBatch.from_subjects(subjects) - original = batch.t1.data.clone() - result = tio.Flip(axes=0, include=["missing"])(batch) - assert torch.equal(result.t1.data, original) - assert result.applied_transforms[-1].params["axes"] == () diff --git a/tests/test_gamma.py b/tests/test_gamma.py deleted file mode 100644 index 8a82abb21..000000000 --- a/tests/test_gamma.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Tests for Gamma transform.""" - -from __future__ import annotations - -import numpy as np -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestGamma: - def test_changes_data(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Gamma(log_gamma=0.3)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_zero_log_gamma_is_identity(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Gamma(log_gamma=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_inverse(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - transformed = tio.Gamma(log_gamma=0.2)(subject) - restored = transformed.apply_inverse_transform() - np.testing.assert_allclose( - restored.t1.data.numpy(), - original.numpy(), - atol=1e-4, - ) - - def test_inverse_respects_include_scope(self) -> None: - a = torch.arange(8.0).reshape(1, 2, 2, 2) - b = torch.arange(100.0, 108.0).reshape(1, 2, 2, 2) - subject = tio.Subject( - a=tio.ScalarImage(a.clone()), - b=tio.ScalarImage(b.clone()), - ) - - transformed = tio.Gamma(log_gamma=0.5, include=["a"])(subject) - restored = transformed.apply_inverse_transform() - - torch.testing.assert_close(restored.a.data, a) - torch.testing.assert_close(restored.b.data, b) - - def test_inverse_respects_exclude_scope(self) -> None: - a = torch.arange(8.0).reshape(1, 2, 2, 2) - b = torch.arange(100.0, 108.0).reshape(1, 2, 2, 2) - subject = tio.Subject( - a=tio.ScalarImage(a.clone()), - b=tio.ScalarImage(b.clone()), - ) - - transformed = tio.Gamma(log_gamma=0.5, exclude=["b"])(subject) - restored = transformed.apply_inverse_transform() - - torch.testing.assert_close(restored.a.data, a) - torch.testing.assert_close(restored.b.data, b) - - def test_leaves_labels_unchanged(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.Gamma(log_gamma=0.3)(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - -class TestGammaPerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 0.1)) - for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_default_differs_across_batch(self) -> None: - """With a random range, each element gets its own gamma.""" - torch.manual_seed(0) - batch = self._batch() - transform = tio.Gamma(log_gamma=(0.2, 0.8)) - result = transform(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], list) - assert len(params["log_gamma"]) == batch.batch_size - assert len(set(params["log_gamma"])) > 1 - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.Gamma(log_gamma=(0.2, 0.8), per_instance=False) - result = transform(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], float) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 0.1)) - result = tio.Gamma(log_gamma=(0.2, 0.8))(subject) - assert isinstance(result.applied_transforms[-1].params["log_gamma"], float) - - def test_per_instance_values_applied(self) -> None: - """Element-wise gamma is actually applied to each element.""" - torch.manual_seed(0) - batch = self._batch(batch_size=4) - original = batch.t1.data.clone() - transform = tio.Gamma(log_gamma=(0.2, 0.8)) - result = transform(batch) - log_gammas = result.applied_transforms[-1].params["log_gamma"] - for i, log_gamma in enumerate(log_gammas): - gamma = torch.tensor(log_gamma).exp() - expected = original[i].sign() * original[i].abs().pow(gamma) - torch.testing.assert_close(result.t1.data[i], expected) - - def test_per_instance_p_gates_some_elements(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=64) - original = batch.t1.data.clone() - transform = tio.Gamma(log_gamma=(0.5, 1.0), p=0.5) - result = transform(batch) - changed = [ - not torch.allclose(result.t1.data[i], original[i]) - for i in range(batch.batch_size) - ] - assert any(changed) - assert not all(changed) - - def test_per_instance_p_masked_elements_have_no_history(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=32) - original = batch.t1.data.clone() - transform = tio.Gamma(log_gamma=(0.5, 1.0), p=0.5) - result = transform(batch) - subjects = result.unbatch() - for i, subject in enumerate(subjects): - changed = not torch.allclose(subject.t1.data, original[i]) - has_history = len(subject.applied_transforms) == 1 - assert changed == has_history - - def test_per_instance_inverse_round_trip(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=5) - original = batch.t1.data.clone() - transform = tio.Gamma(log_gamma=(0.2, 0.5)) - result = transform(batch) - restored = result.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original, atol=1e-4, rtol=0) diff --git a/tests/test_ghosting.py b/tests/test_ghosting.py deleted file mode 100644 index 6c97ef233..000000000 --- a/tests/test_ghosting.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Tests for Ghosting transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestGhosting: - def test_changes_data(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Ghosting(num_ghosts=5, intensity=0.8)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_zero_intensity_is_identity(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Ghosting(intensity=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_leaves_labels_unchanged(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.Ghosting(num_ghosts=5, intensity=0.8)(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_specific_axis(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Ghosting(axes=(1,), intensity=0.8)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_restore_fraction(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Ghosting(restore=0.2, intensity=0.8)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - -class TestGhostingPerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - data = torch.rand(1, 12, 12, 12) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Ghosting(intensity=(0.5, 1.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["intensity"]) == batch.batch_size - assert not torch.allclose(result.t1.data[0], result.t1.data[1]) - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Ghosting(intensity=(0.5, 1.0), per_instance=False)(batch) - torch.testing.assert_close(result.t1.data[0], result.t1.data[1]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 12, 12, 12))) - result = tio.Ghosting(intensity=(0.5, 1.0))(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params diff --git a/tests/test_histogram_standardization.py b/tests/test_histogram_standardization.py deleted file mode 100644 index ae0e73c79..000000000 --- a/tests/test_histogram_standardization.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Tests for HistogramStandardization transform.""" - -from __future__ import annotations - -from pathlib import Path - -import pytest -import torch - -import torchio as tio -from torchio.transforms.intensity.histogram_standardization import ( - compute_histogram_landmarks, -) - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestHistogramStandardization: - @staticmethod - def _make_images(n: int = 5) -> list[tio.ScalarImage]: - images = [] - for i in range(n): - data = torch.randn(1, 10, 10, 10) * (10 + i) + (50 + i * 5) - images.append(tio.ScalarImage(data)) - return images - - def test_compute_landmarks(self) -> None: - images = self._make_images() - landmarks = compute_histogram_landmarks(images) - assert landmarks.ndim == 1 - assert len(landmarks) > 2 - - def test_landmarks_monotonic(self) -> None: - images = self._make_images(n=10) - landmarks = compute_histogram_landmarks(images) - diffs = torch.diff(landmarks) - assert (diffs >= -1e-5).all() - - def test_apply_changes_data(self) -> None: - images = self._make_images() - landmarks = compute_histogram_landmarks(images) - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.HistogramStandardization(landmarks)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_apply_with_masking(self) -> None: - images = self._make_images() - mask_fn = lambda x: x > x.median() # noqa: E731 - landmarks = compute_histogram_landmarks( - images, - masking_method=mask_fn, - ) - assert landmarks.ndim == 1 - - def test_leaves_labels_unchanged(self) -> None: - images = self._make_images() - landmarks = compute_histogram_landmarks(images) - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.HistogramStandardization(landmarks)(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_landmark_count_mismatch_raises(self) -> None: - landmarks = torch.linspace(0, 100, 5) - subject = _make_subject(with_label=False) - with pytest.raises(ValueError, match="does not match"): - tio.HistogramStandardization(landmarks)(subject) - - def test_custom_quantiles(self) -> None: - images = self._make_images() - quantiles = (0.01, 0.25, 0.5, 0.75, 0.99) - landmarks = compute_histogram_landmarks( - images, - quantiles=quantiles, - ) - assert len(landmarks) == 5 - - def test_too_few_quantiles_raises(self) -> None: - images = self._make_images() - with pytest.raises(ValueError, match="at least 2"): - compute_histogram_landmarks(images, quantiles=(0.5,)) - - -class TestHistogramStandardizationEdgeCases: - def test_quantiles_out_of_range_raises(self) -> None: - images = [tio.ScalarImage(torch.randn(1, 5, 5, 5)) for _ in range(3)] - with pytest.raises(ValueError, match="\\[0, 1\\]"): - compute_histogram_landmarks(images, quantiles=(-0.1, 0.5, 1.1)) - - def test_cutoff_not_in_quantiles_raises(self) -> None: - images = [tio.ScalarImage(torch.randn(1, 5, 5, 5)) for _ in range(3)] - with pytest.raises(ValueError, match="Cutoff"): - compute_histogram_landmarks( - images, quantiles=(0.25, 0.5, 0.75), cutoff=(0.01, 0.99) - ) - - def test_load_landmarks_from_npy(self, tmp_path: Path) -> None: - import numpy as np - - arr = np.linspace(0, 100, 13).astype(np.float32) - npy_path = tmp_path / "landmarks.npy" - np.save(npy_path, arr) - subject = _make_subject(with_label=False) - result = tio.HistogramStandardization(npy_path)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_load_landmarks_from_pt(self, tmp_path: Path) -> None: - landmarks = torch.linspace(0, 100, 13) - pt_path = tmp_path / "landmarks.pt" - torch.save(landmarks, pt_path) - subject = _make_subject(with_label=False) - result = tio.HistogramStandardization(pt_path)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_unsupported_format_raises(self, tmp_path: Path) -> None: - bad_path = tmp_path / "landmarks.csv" - bad_path.write_text("1,2,3") - with pytest.raises(ValueError, match="Unsupported"): - tio.HistogramStandardization(bad_path) - - def test_pt_with_wrong_type_raises(self, tmp_path: Path) -> None: - pt_path = tmp_path / "landmarks.pt" - torch.save({"not": "a tensor"}, pt_path) - with pytest.raises(TypeError, match="Expected a Tensor"): - tio.HistogramStandardization(pt_path) - - def test_load_from_path_string(self) -> None: - images = [tio.ScalarImage(torch.randn(1, 5, 5, 5)) for _ in range(3)] - landmarks = compute_histogram_landmarks(images) - assert landmarks.ndim == 1 diff --git a/tests/test_hydra.py b/tests/test_hydra.py deleted file mode 100644 index 33e9b0c3b..000000000 --- a/tests/test_hydra.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Tests for Hydra YAML export.""" - -from __future__ import annotations - -import torchio as tio - - -class TestToHydra: - def test_noise_default(self) -> None: - n = tio.Noise() - cfg = n.to_hydra() - assert cfg["_target_"] == "torchio.Noise" - assert "std" not in cfg # default, omitted - - def test_noise_custom(self) -> None: - n = tio.Noise(std=(0.05, 0.2), p=0.5) - cfg = n.to_hydra() - assert cfg["_target_"] == "torchio.Noise" - assert cfg["std"] == [0.05, 0.2] - assert cfg["p"] == 0.5 - assert "mean" not in cfg # still default - - def test_flip(self) -> None: - f = tio.Flip(axes=(0, 1)) - cfg = f.to_hydra() - assert cfg["_target_"] == "torchio.Flip" - assert cfg["axes"] == [0, 1] - - def test_compose(self) -> None: - pipeline = tio.Compose( - [ - tio.Flip(axes=(0,), p=0.5), - tio.Noise(std=0.1), - ] - ) - cfg = pipeline.to_hydra() - assert cfg["_target_"] == "torchio.Compose" - assert len(cfg["transforms"]) == 2 - assert cfg["transforms"][0]["_target_"] == "torchio.Flip" - assert cfg["transforms"][1]["_target_"] == "torchio.Noise" - - def test_nested_compose(self) -> None: - pipeline = tio.Compose( - [ - tio.OneOf([tio.Noise(), tio.Flip()]), - ] - ) - cfg = pipeline.to_hydra() - inner = cfg["transforms"][0] - assert inner["_target_"] == "torchio.OneOf" - assert "transforms" in inner - - def test_round_trip_values(self) -> None: - """Hydra config values should be plain Python types.""" - n = tio.Noise(std=(0.05, 0.2), mean=0.5, p=0.8) - cfg = n.to_hydra() - # All values should be JSON-compatible types - for v in cfg.values(): - assert isinstance(v, (str, int, float, bool, list, type(None))) diff --git a/tests/test_identity_warning.py b/tests/test_identity_warning.py deleted file mode 100644 index 50459322e..000000000 --- a/tests/test_identity_warning.py +++ /dev/null @@ -1,88 +0,0 @@ -"""No-arg augmentation transforms are a deterministic no-op and warn. - -Transforms whose parameters are sampled from a range should, when -constructed with no arguments, default to an identity (no-op) and emit a -warning telling the user to pass arguments. Inherently-stochastic -transforms (which draw a random realisation rather than sampling a scalar -parameter) are exempt. -""" - -from __future__ import annotations - -import warnings - -import pytest -import torch - -import torchio as tio -from torchio.transforms.parameter_range import _ParameterRange - -# Transforms that must be a deterministic no-op (and warn) with no args. -NOOP_TRANSFORMS = ["Affine", "Anisotropy", "Blur", "Gamma", "Ghosting", "Spike"] - -# Args that activate each transform: should NOT warn and should change data. -# Ranges are kept away from the identity value to avoid flaky assertions. -ACTIVE_KWARGS: dict[str, dict] = { - "Affine": {"degrees": (10, 15)}, - "Anisotropy": {"downsampling": (2, 5)}, - "Blur": {"std": (1, 2)}, - "Gamma": {"log_gamma": (0.3, 0.5)}, - "Ghosting": {"intensity": (0.8, 1.0)}, - "Spike": {"intensity": (2, 3)}, -} - -# Inherently-stochastic transforms: no-arg construction must NOT warn. -EXEMPT_TRANSFORMS = ["Noise", "ElasticDeformation", "Swap", "BiasField"] - - -def _subject() -> tio.Subject: - torch.manual_seed(0) - return tio.Subject(t1=tio.ScalarImage(torch.rand(1, 12, 12, 12) * 100)) - - -@pytest.mark.parametrize("name", NOOP_TRANSFORMS) -def test_no_args_is_identity(name: str) -> None: - subject = _subject() - original = subject.t1.data.clone() - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - result = getattr(tio, name)()(subject) - torch.testing.assert_close(result.t1.data, original) - - -@pytest.mark.parametrize("name", NOOP_TRANSFORMS) -def test_no_args_warns(name: str) -> None: - with pytest.warns(UserWarning, match=name): - getattr(tio, name)() - - -@pytest.mark.parametrize("name", NOOP_TRANSFORMS) -def test_active_does_not_warn_and_changes(name: str) -> None: - subject = _subject() - original = subject.t1.data.clone() - with warnings.catch_warnings(): - warnings.simplefilter("error") # any warning becomes an error - transform = getattr(tio, name)(**ACTIVE_KWARGS[name]) - torch.manual_seed(0) - result = transform(subject) - assert not torch.allclose(result.t1.data, original) - - -@pytest.mark.parametrize("name", EXEMPT_TRANSFORMS) -def test_stochastic_no_args_does_not_warn(name: str) -> None: - with warnings.catch_warnings(): - warnings.simplefilter("error") - getattr(tio, name)() - - -class TestIsConstant: - def test_scalar(self) -> None: - assert _ParameterRange(0.0).is_constant(0.0) - assert _ParameterRange(1.0).is_constant(1.0) - assert not _ParameterRange(0.0).is_constant(1.0) - - def test_degenerate_range_is_constant(self) -> None: - assert _ParameterRange((0.0, 0.0)).is_constant(0.0) - - def test_real_range_is_not_constant(self) -> None: - assert not _ParameterRange((0.0, 2.0)).is_constant(0.0) diff --git a/tests/test_image.py b/tests/test_image.py deleted file mode 100644 index ea7f513dd..000000000 --- a/tests/test_image.py +++ /dev/null @@ -1,880 +0,0 @@ -"""Tests for Image, ScalarImage, and LabelMap.""" - -from __future__ import annotations - -from pathlib import Path - -import nibabel as nib -import numpy as np -import pytest -import SimpleITK as sitk -import torch -from einops import rearrange - -from torchio import Image -from torchio import LabelMap -from torchio import ScalarImage - - -class TestImageCreationFromPath: - def test_from_path_positional(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - array = rearrange(tensor.numpy(), "c i j k -> i j k c") - nii = nib.Nifti1Image(array, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - image = ScalarImage(path) - assert not image.is_loaded - _ = image.data - assert image.is_loaded - - def test_from_path_keyword(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - nii = nib.Nifti1Image(tensor.numpy()[0], np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - image = ScalarImage(source=path) - assert image.path == path - - def test_from_path_string(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - nii = nib.Nifti1Image(tensor.numpy()[0], np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - image = ScalarImage(str(path)) - assert image.path == path - - def test_path_with_affine(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - nii = nib.Nifti1Image(tensor.numpy()[0], np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - custom_affine = np.diag([2.0, 2.0, 2.0, 1.0]) - image = ScalarImage(path, affine=custom_affine) - assert not image.is_loaded - np.testing.assert_array_equal(image.affine, custom_affine) - - def test_path_property(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - nii = nib.Nifti1Image(tensor.numpy()[0], np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - image = ScalarImage(path) - assert image.path == path - - def test_no_path_creates_empty_image(self): - image = ScalarImage() - assert image.path is None - assert not image.is_loaded - - -class TestImageCreationFromTensor: - def test_from_tensor(self): - tensor = torch.randn(1, 10, 10, 10) - image = ScalarImage(tensor) - assert isinstance(image, ScalarImage) - assert torch.equal(image.data, tensor) - - def test_from_tensor_numpy(self): - array = np.random.randn(1, 10, 10, 10).astype(np.float32) - image = ScalarImage(array) - assert torch.equal(image.data, torch.from_numpy(array)) - - def test_from_tensor_default_affine(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - np.testing.assert_array_equal(image.affine, np.eye(4)) - - def test_from_tensor_custom_affine(self): - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - affine=affine, - ) - np.testing.assert_array_equal(image.affine, affine) - - def test_from_tensor_affine_object(self): - from torchio import AffineMatrix - - aff = AffineMatrix(np.diag([2.0, 2.0, 2.0, 1.0])) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - affine=aff, - ) - assert image.affine == aff - - def test_from_tensor_metadata(self): - image = ScalarImage( - torch.randn(1, 10, 10, 10), - scan_id="abc123", - ) - assert image.metadata == {"scan_id": "abc123"} - - def test_from_tensor_must_be_4d(self): - with pytest.raises(ValueError, match="4D"): - ScalarImage(torch.randn(10, 10, 10)) - - def test_from_tensor_affine_must_be_4x4(self): - with pytest.raises(ValueError, match=r"4.*4"): - ScalarImage( - torch.randn(1, 10, 10, 10), - affine=np.eye(3), - ) - - def test_from_tensor_path_is_none(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - assert image.path is None - - def test_from_tensor_preserves_subclass(self): - image = LabelMap(torch.randint(0, 5, (1, 10, 10, 10))) - assert isinstance(image, LabelMap) - - def test_from_tensor_is_loaded(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - assert image.is_loaded - - -class TestImageProperties: - @pytest.fixture - def image(self) -> ScalarImage: - return ScalarImage(torch.randn(2, 10, 20, 30)) - - def test_shape(self, image: ScalarImage): - assert image.shape == (2, 10, 20, 30) - - def test_spatial_shape(self, image: ScalarImage): - assert image.spatial_shape == (10, 20, 30) - - def test_num_channels(self, image: ScalarImage): - assert image.num_channels == 2 - - def test_spacing(self, image: ScalarImage): - assert image.spacing == (1.0, 1.0, 1.0) - - def test_spacing_with_custom_affine(self): - affine = np.diag([0.5, 0.8, 1.2, 1.0]) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - affine=affine, - ) - np.testing.assert_allclose(image.spacing, (0.5, 0.8, 1.2)) - - def test_origin(self, image: ScalarImage): - assert image.origin == (0.0, 0.0, 0.0) - - def test_is_loaded_from_tensor(self, image: ScalarImage): - assert image.is_loaded - - def test_memory(self, image: ScalarImage): - # 2 * 10 * 20 * 30 * 4 bytes (float32) - assert image.memory == 2 * 10 * 20 * 30 * 4 - - -class TestLabelMap: - def test_is_label_map(self): - label = LabelMap(torch.randint(0, 5, (1, 10, 10, 10))) - assert isinstance(label, LabelMap) - assert isinstance(label, Image) - - def test_is_not_scalar_image(self): - label = LabelMap(torch.randint(0, 5, (1, 10, 10, 10))) - assert not isinstance(label, ScalarImage) - - def test_is_label_subclass(self): - assert issubclass(LabelMap, Image) - - -class TestScalarImage: - def test_is_scalar_image(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - assert isinstance(image, ScalarImage) - assert isinstance(image, Image) - - def test_is_not_label_map(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - assert not isinstance(image, LabelMap) - - def test_is_image_subclass(self): - assert issubclass(ScalarImage, Image) - - -class TestNewLike: - def test_new_like_preserves_type(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - new = image.new_like(data=torch.randn(1, 5, 5, 5)) - assert isinstance(new, ScalarImage) - - def test_new_like_preserves_affine(self): - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - affine=affine, - ) - new = image.new_like(data=torch.randn(1, 5, 5, 5)) - np.testing.assert_array_equal(new.affine, affine) - - def test_new_like_with_new_affine(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - new_affine = np.diag([3.0, 3.0, 3.0, 1.0]) - new = image.new_like(data=torch.randn(1, 5, 5, 5), affine=new_affine) - np.testing.assert_array_equal(new.affine, new_affine) - - def test_new_like_preserves_metadata(self): - image = ScalarImage( - torch.randn(1, 10, 10, 10), - scan_id="abc123", - ) - new = image.new_like(data=torch.randn(1, 5, 5, 5)) - assert new.metadata == {"scan_id": "abc123"} - - def test_new_like_on_custom_subclass(self): - class MyImage(ScalarImage): - pass - - image = MyImage(torch.randn(1, 10, 10, 10)) - new = image.new_like(data=torch.randn(1, 5, 5, 5)) - assert isinstance(new, MyImage) - - def test_new_like_label_map(self): - label = LabelMap(torch.randint(0, 5, (1, 10, 10, 10))) - new = label.new_like(data=torch.randint(0, 5, (1, 5, 5, 5))) - assert isinstance(new, LabelMap) - assert not isinstance(new, ScalarImage) - - -class TestSetData: - def test_set_data(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - new_data = torch.randn(1, 5, 5, 5) - image.set_data(new_data) - assert torch.equal(image.data, new_data) - - def test_set_data_must_be_4d(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - with pytest.raises(ValueError, match="4D"): - image.set_data(torch.randn(10, 10, 10)) - - -class TestImageRepr: - def test_loaded_repr(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - r = repr(image) - assert "ScalarImage" in r - assert "10" in r - - def test_unloaded_repr(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - array = rearrange(tensor.numpy(), "c i j k -> i j k c") - nii = nib.Nifti1Image(array, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - image = ScalarImage(path) - r = repr(image) - assert "ScalarImage" in r - # Now shows shape/dtype/memory even without loading - assert "shape" in r or "spatial" in r - assert "dtype" in r - assert "memory" in r - - def test_repr_tensor_only_no_path(self): - image = Image(torch.randn(1, 4, 4, 4)) - image._data = None - r = repr(image) - # No data and no path: falls back to minimal repr - assert "Image" in r - - -class TestImageLoad: - def test_load_already_loaded_is_noop(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - original_data = image.data - image.load() - assert image.data is original_data - - def test_load_no_path_raises(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - image._data = None - image._path = None - image._backend = None - with pytest.raises(RuntimeError, match="no path or backend"): - image.load() - - def test_shape_no_data_no_path_raises(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - image._data = None - image._path = None - image._backend = None - with pytest.raises(RuntimeError, match="Cannot determine shape"): - image.shape - - -class TestImageCopy: - def test_copy(self): - import copy - - image = ScalarImage(torch.randn(1, 10, 10, 10)) - copied = copy.copy(image) - assert isinstance(copied, ScalarImage) - assert torch.equal(copied.data, image.data) - copied.set_data(torch.zeros(1, 10, 10, 10)) - assert not torch.equal(image.data, copied.data) - - def test_deepcopy_tensor_based(self): - import copy - - image = ScalarImage(torch.randn(1, 10, 10, 10)) - copied = copy.deepcopy(image) - assert isinstance(copied, ScalarImage) - assert torch.equal(copied.data, image.data) - copied.set_data(torch.zeros(1, 10, 10, 10)) - assert not torch.equal(image.data, copied.data) - - def test_deepcopy_path_based_unloaded(self, tmp_path: Path): - import copy - - tensor = torch.randn(1, 10, 10, 10) - array = rearrange(tensor.numpy(), "c i j k -> i j k c") - nii = nib.Nifti1Image(array, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - copied = copy.deepcopy(image) - assert isinstance(copied, ScalarImage) - assert not copied.is_loaded - assert copied.path == path - - def test_deepcopy_path_based_loaded(self, tmp_path: Path): - import copy - - tensor = torch.randn(1, 10, 10, 10) - array = rearrange(tensor.numpy(), "c i j k -> i j k c") - nii = nib.Nifti1Image(array, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - _ = image.data # trigger load - copied = copy.deepcopy(image) - assert isinstance(copied, ScalarImage) - assert copied.is_loaded - assert copied.path == path - - def test_deepcopy_degenerate_state(self): - import copy - - image = ScalarImage(torch.randn(1, 4, 4, 4)) - image._data = None - image._path = None - copied = copy.deepcopy(image) - assert isinstance(copied, ScalarImage) - assert not copied.is_loaded - assert copied.path is None - - -class TestNibabelReader: - def test_4d_nifti(self, tmp_path: Path): - data = np.random.randn(10, 10, 10, 3).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "multichannel.nii.gz" - nib.save(nii, path) - image = ScalarImage(path) - assert image.shape == (3, 10, 10, 10) - - def test_4d_nifti_shape_from_header(self, tmp_path: Path): - data = np.random.randn(10, 10, 10, 3).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "multichannel.nii.gz" - nib.save(nii, path) - image = ScalarImage(path) - assert not image.is_loaded - shape = image.shape - assert shape == (3, 10, 10, 10) - assert not image.is_loaded - - def test_invalid_ndim_raises(self, tmp_path: Path): - data = np.random.randn(10, 10, 10, 3, 2).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "bad.nii.gz" - nib.save(nii, path) - image = ScalarImage(path) - with pytest.raises(ValueError, match="3D or 4D"): - image.data - - def test_invalid_ndim_shape_from_header(self, tmp_path: Path): - data = np.random.randn(10, 10, 10, 3, 2).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "bad.nii.gz" - nib.save(nii, path) - image = ScalarImage(path) - with pytest.raises(ValueError, match="3D or 4D"): - image.shape - - -class TestSimpleITKReader: - def test_read_nrrd(self, tmp_path: Path): - data = np.random.randn(10, 12, 14).astype(np.float32) - sitk_image = sitk.GetImageFromArray( - rearrange(data, "i j k -> k j i"), - ) - sitk_image.SetSpacing((0.5, 0.8, 1.2)) - path = tmp_path / "test.nrrd" - sitk.WriteImage(sitk_image, str(path)) - - image = ScalarImage(path) - assert image.shape == (1, 10, 12, 14) - np.testing.assert_allclose(image.spacing, (0.5, 0.8, 1.2)) - - def test_shape_from_header_nrrd(self, tmp_path: Path): - data = np.random.randn(10, 12, 14).astype(np.float32) - sitk_image = sitk.GetImageFromArray( - rearrange(data, "i j k -> k j i"), - ) - path = tmp_path / "test.nrrd" - sitk.WriteImage(sitk_image, str(path)) - - image = ScalarImage(path) - assert not image.is_loaded - shape = image.shape - assert shape == (1, 10, 12, 14) - assert not image.is_loaded - - def test_read_multichannel_nrrd(self, tmp_path: Path): - data = np.random.randn(10, 12, 14, 3).astype(np.float32) - sitk_image = sitk.GetImageFromArray( - rearrange(data, "i j k c -> k j i c"), - isVector=True, - ) - path = tmp_path / "multi.nrrd" - sitk.WriteImage(sitk_image, str(path)) - - image = ScalarImage(path) - assert image.shape == (3, 10, 12, 14) - - -class TestReaderErrors: - def test_sitk_reader_invalid_ndim(self, tmp_path: Path): - """SimpleITK reader raises for unexpected dimensions.""" - from torchio.data.io import read_sitk - - data_2d = np.zeros((10, 10), dtype=np.float32) - sitk_image = sitk.GetImageFromArray(data_2d) - path = tmp_path / "flat.nrrd" - sitk.WriteImage(sitk_image, str(path)) - - with pytest.raises(ValueError, match="Expected 3D"): - read_sitk(path) - - def test_sitk_shape_reader_invalid_ndim(self, tmp_path: Path): - """SimpleITK shape reader raises for non-3D images.""" - data_2d = np.zeros((10, 10), dtype=np.float32) - sitk_image = sitk.GetImageFromArray(data_2d) - path = tmp_path / "flat.nrrd" - sitk.WriteImage(sitk_image, str(path)) - - image = ScalarImage(path) - with pytest.raises(ValueError, match="Expected 3D"): - image.shape - - -class TestSimpleITKReaderEdgeCases: - def test_multichannel_nrrd_loads_data(self, tmp_path: Path): - data = np.random.randn(10, 12, 14, 3).astype(np.float32) - sitk_image = sitk.GetImageFromArray( - rearrange(data, "i j k c -> k j i c"), - isVector=True, - ) - path = tmp_path / "multi.nrrd" - sitk.WriteImage(sitk_image, str(path)) - - image = ScalarImage(path) - _ = image.data # trigger actual data load - assert image.shape == (3, 10, 12, 14) - - def test_5d_vector_nifti_loads_data(self, tmp_path: Path): - # SimpleITK writes multichannel NIfTI as 5D vector: (I, J, K, 1, C) - tensor = torch.randn(3, 10, 10, 10) - image = ScalarImage(tensor) - path = tmp_path / "vector.nii.gz" - image.save(path) - - loaded = ScalarImage(path) - _ = loaded.data # trigger nibabel load of 5D vector NIfTI - assert loaded.shape == (3, 10, 10, 10) - - -class TestImageIO: - def test_save_and_load_nifti(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - image = ScalarImage(tensor, affine=affine) - path = tmp_path / "output.nii.gz" - image.save(path) - - loaded = ScalarImage(path) - assert loaded.shape == (1, 10, 10, 10) - np.testing.assert_allclose(loaded.spacing, (2.0, 2.0, 2.0)) - - def test_save_and_load_nrrd(self, tmp_path: Path): - tensor = torch.randn(1, 10, 12, 14) - affine = np.diag([0.5, 0.8, 1.2, 1.0]) - image = ScalarImage(tensor, affine=affine) - path = tmp_path / "output.nrrd" - image.save(path) - - loaded = ScalarImage(path) - assert loaded.shape == (1, 10, 12, 14) - np.testing.assert_allclose(loaded.spacing, (0.5, 0.8, 1.2)) - - @pytest.mark.parametrize("extension", [".nii.gz", ".nrrd"]) - def test_save_preserves_affine(self, tmp_path: Path, extension: str): - """Round-trip save/load must preserve the full affine matrix.""" - tensor = torch.randn(1, 8, 10, 12) - affine = np.array( - [ - [0.0, 0.0, 1.5, -10.0], - [0.5, 0.0, 0.0, -20.0], - [0.0, 0.8, 0.0, 30.0], - [0.0, 0.0, 0.0, 1.0], - ] - ) - image = ScalarImage(tensor, affine=affine) - path = tmp_path / f"output{extension}" - image.save(path) - - loaded = ScalarImage(path) - np.testing.assert_allclose( - loaded.affine.numpy(), - affine, - atol=1e-6, - ) - torch.testing.assert_close(loaded.data, tensor, atol=1e-4, rtol=1e-4) - - def test_save_preserves_lps_orientation(self, tmp_path: Path): - """NIfTI with LPS+ orientation must survive a save/load round-trip.""" - tensor = torch.randn(1, 8, 10, 12) - affine = np.array( - [ - [-0.5, 0.0, 0.0, 90.0], - [0.0, -0.5, 0.0, 126.0], - [0.0, 0.0, 0.5, -72.0], - [0.0, 0.0, 0.0, 1.0], - ] - ) - image = ScalarImage(tensor, affine=affine) - path = tmp_path / "lps.nii.gz" - image.save(path) - - loaded = ScalarImage(path) - np.testing.assert_allclose( - loaded.affine.numpy(), - affine, - atol=1e-6, - ) - assert loaded.affine.orientation == ("L", "P", "S") - - def test_save_multichannel(self, tmp_path: Path): - tensor = torch.randn(3, 10, 10, 10) - image = ScalarImage(tensor) - path = tmp_path / "multi.nii.gz" - image.save(path) - - loaded = ScalarImage(path) - assert loaded.shape == (3, 10, 10, 10) - - def test_custom_reader(self, tmp_path: Path): - path = tmp_path / "test.npy" - data = np.random.randn(1, 10, 10, 10).astype(np.float32) - np.save(path, data) - - def npy_reader(p): - arr = np.load(p) - return torch.from_numpy(arr), np.eye(4) - - image = ScalarImage(path, reader=npy_reader) - assert image.shape == (1, 10, 10, 10) - - def test_save_nii_zarr(self, tmp_path: Path): - tensor = torch.randn(1, 10, 12, 14) - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - image = ScalarImage(tensor, affine=affine) - path = tmp_path / "output.nii.zarr" - image.save(path) - - loaded = ScalarImage(path) - assert loaded.shape == (1, 10, 12, 14) - np.testing.assert_allclose(loaded.spacing, (2.0, 3.0, 4.0), atol=1e-5) - np.testing.assert_allclose( - loaded.data.numpy(), - tensor.numpy(), - atol=1e-5, - ) - - def test_save_nii_zarr_multichannel(self, tmp_path: Path): - tensor = torch.randn(3, 8, 8, 8) - image = ScalarImage(tensor) - path = tmp_path / "multi.nii.zarr" - image.save(path) - - loaded = ScalarImage(path) - assert loaded.shape == (3, 8, 8, 8) - - -class TestImageSlicing: - def test_slice_channel_int(self): - tensor = torch.randn(3, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[0] - assert sliced.shape == (1, 20, 20, 20) - torch.testing.assert_close(sliced.data, tensor[0:1]) - - def test_slice_channel_range(self): - tensor = torch.randn(5, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[1:3] - assert sliced.shape == (2, 20, 20, 20) - torch.testing.assert_close(sliced.data, tensor[1:3]) - - def test_slice_spatial_via_tuple(self): - tensor = torch.randn(1, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[:, 5:10] - assert sliced.shape == (1, 5, 20, 20) - torch.testing.assert_close(sliced.data, tensor[:, 5:10, :, :]) - - def test_slice_all_four_dims(self): - tensor = torch.randn(3, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[0:2, 2:8, 3:7, 4:10] - assert sliced.shape == (2, 6, 4, 6) - torch.testing.assert_close(sliced.data, tensor[0:2, 2:8, 3:7, 4:10]) - - def test_slice_preserves_class(self): - tensor = torch.randn(1, 20, 20, 20) - image = LabelMap(tensor) - sliced = image[:, 5:10] - assert isinstance(sliced, LabelMap) - - def test_slice_updates_affine_origin(self): - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - tensor = torch.randn(1, 20, 20, 20) - image = ScalarImage(tensor, affine=affine) - sliced = image[:, 5:10] - # New origin should be shifted by 5 voxels * 2mm spacing in I direction - expected_origin = (10.0, 0.0, 0.0) - np.testing.assert_allclose(sliced.origin, expected_origin) - - def test_slice_channel_does_not_affect_origin(self): - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - tensor = torch.randn(5, 20, 20, 20) - image = ScalarImage(tensor, affine=affine) - sliced = image[1:3] - np.testing.assert_allclose(sliced.origin, (0.0, 0.0, 0.0)) - - def test_slice_partial_dims(self): - tensor = torch.randn(1, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[:, 5:10, 3:7] - assert sliced.shape == (1, 5, 4, 20) - - def test_slice_negative_indices(self): - tensor = torch.randn(1, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[:, -5:] - assert sliced.shape == (1, 5, 20, 20) - torch.testing.assert_close(sliced.data, tensor[:, 15:, :, :]) - - def test_slice_with_step(self): - tensor = torch.randn(1, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[:, ::2] - assert sliced.shape == (1, 10, 20, 20) - torch.testing.assert_close(sliced.data, tensor[:, ::2, :, :]) - - def test_slice_ellipsis_trailing(self): - tensor = torch.randn(3, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[..., 5:10] - assert sliced.shape == (3, 20, 20, 5) - torch.testing.assert_close(sliced.data, tensor[..., 5:10]) - - def test_slice_ellipsis_leading(self): - tensor = torch.randn(3, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[0, ...] - assert sliced.shape == (1, 20, 20, 20) - torch.testing.assert_close(sliced.data, tensor[0:1]) - - def test_slice_ellipsis_middle(self): - tensor = torch.randn(3, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[0:2, ..., 5:10] - assert sliced.shape == (2, 20, 20, 5) - torch.testing.assert_close(sliced.data, tensor[0:2, :, :, 5:10]) - - def test_slice_bare_ellipsis(self): - tensor = torch.randn(1, 20, 20, 20) - image = ScalarImage(tensor) - sliced = image[...] - assert sliced.shape == (1, 20, 20, 20) - torch.testing.assert_close(sliced.data, tensor) - - def test_slice_double_ellipsis_raises(self): - image = ScalarImage(torch.randn(1, 20, 20, 20)) - with pytest.raises(IndexError, match="one ellipsis"): - image[..., ...] - - def test_slice_float_raises(self): - image = ScalarImage(torch.randn(1, 20, 20, 20)) - with pytest.raises(TypeError, match="not understood"): - image[1.5] - - def test_slice_too_many_dims_raises(self): - image = ScalarImage(torch.randn(1, 20, 20, 20)) - with pytest.raises(IndexError, match="Too many"): - image[:, :, :, :, :] - - def test_slice_lazy_does_not_load(self, tmp_path: Path): - data = np.random.randn(10, 12, 14).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - image = ScalarImage(path) - sliced = image[:, 2:5, 3:7, 4:8] - assert not image.is_loaded # parent not loaded - assert sliced.shape == (1, 3, 4, 4) - - def test_slice_preserves_metadata(self): - image = ScalarImage( - torch.randn(1, 20, 20, 20), - modality="T1", - ) - sliced = image[:, 5:10] - assert sliced.metadata["modality"] == "T1" - - -class TestReaderWriterKwargs: - def _make_nifti(self, tmp_path: Path) -> Path: - """Helper: write a small NIfTI file and return its path.""" - data = np.random.randn(10, 12, 14).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - return path - - def test_reader_kwargs_passed(self, tmp_path: Path): - """reader_kwargs are forwarded to the reader function.""" - from unittest.mock import patch - - path = self._make_nifti(tmp_path) - image = ScalarImage(path, reader_kwargs={"keep_file_open": True}) - with patch("torchio.data.image.nib.load", wraps=nib.load) as mock_load: - image.load() - mock_load.assert_called_once_with(path, keep_file_open=True) - - def test_writer_kwargs_passed(self, tmp_path: Path): - """writer kwargs are forwarded to SimpleITK.WriteImage.""" - from unittest.mock import patch - - image = ScalarImage(torch.randn(1, 4, 4, 4)) - out = tmp_path / "out.nii.gz" - with patch("torchio.data.image.sitk.WriteImage") as mock_write: - image.save(out, useCompression=True) - mock_write.assert_called_once() - _, call_kwargs = mock_write.call_args - assert call_kwargs["useCompression"] is True - - def test_reader_kwargs_deepcopy(self, tmp_path: Path): - """reader_kwargs survive deepcopy.""" - import copy - - path = self._make_nifti(tmp_path) - kw = {"keep_file_open": True} - image = ScalarImage(path, reader_kwargs=kw) - copied = copy.deepcopy(image) - assert copied._reader_kwargs == kw - assert copied._reader_kwargs is not image._reader_kwargs - - -class TestFromNiftiLazy: - """Tests for the lazy Image.from_nifti constructor.""" - - @staticmethod - def _make_nifti(shape: tuple[int, ...] = (10, 12, 14)) -> nib.Nifti1Image: - data = np.random.randn(*shape).astype(np.float32) - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - return nib.Nifti1Image(data, affine) - - def test_from_nifti_is_lazy(self): - nii = self._make_nifti() - image = ScalarImage(nii) - assert not image.is_loaded - - def test_data_triggers_load(self): - nii = self._make_nifti() - image = ScalarImage(nii) - _ = image.data - assert image.is_loaded - - def test_data_values_correct(self): - data = np.arange(24, dtype=np.float32).reshape(2, 3, 4) - nii = nib.Nifti1Image(data, np.eye(4)) - image = ScalarImage(nii) - expected = torch.tensor(data[np.newaxis], dtype=torch.float32) - torch.testing.assert_close(image.data, expected) - - def test_shape_without_load(self): - nii = self._make_nifti((10, 12, 14)) - image = ScalarImage(nii) - assert image.shape == (1, 10, 12, 14) - assert not image.is_loaded - - def test_affine_without_load(self): - nii = self._make_nifti() - image = ScalarImage(nii) - np.testing.assert_allclose( - image.affine.numpy(), - np.diag([2.0, 3.0, 4.0, 1.0]), - ) - assert not image.is_loaded - - def test_spacing_without_load(self): - nii = self._make_nifti() - image = ScalarImage(nii) - np.testing.assert_allclose(image.spacing, (2.0, 3.0, 4.0)) - assert not image.is_loaded - - def test_multichannel(self): - data = np.random.randn(10, 12, 14, 3).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - image = ScalarImage(nii) - assert image.shape == (3, 10, 12, 14) - assert not image.is_loaded - assert image.data.shape == (3, 10, 12, 14) - - def test_label_map_subclass(self): - nii = self._make_nifti() - image = LabelMap(nii) - assert isinstance(image, LabelMap) - assert not image.is_loaded - - def test_metadata_forwarded(self): - nii = self._make_nifti() - image = ScalarImage(nii, protocol="MPRAGE") - assert image.metadata["protocol"] == "MPRAGE" - assert not image.is_loaded - - def test_from_bytes_still_works(self, tmp_path: Path): - tensor = torch.randn(1, 8, 8, 8) - affine = np.diag([1.5, 1.5, 1.5, 1.0]) - nii = nib.Nifti1Image( - rearrange(tensor.numpy(), "c i j k -> i j k c"), - affine, - ) - path = tmp_path / "temp.nii.gz" - nib.save(nii, path) - raw = path.read_bytes() - image = ScalarImage(raw) - assert image.is_loaded - assert image.shape == (1, 8, 8, 8) - np.testing.assert_allclose(image.spacing, (1.5, 1.5, 1.5)) diff --git a/tests/test_image_annotations.py b/tests/test_image_annotations.py deleted file mode 100644 index 3ec87d6b7..000000000 --- a/tests/test_image_annotations.py +++ /dev/null @@ -1,307 +0,0 @@ -"""Tests for Image-level annotations (points and bounding boxes).""" - -from __future__ import annotations - -import copy - -import pytest -import torch - -from torchio import LabelMap -from torchio import ScalarImage -from torchio import Subject -from torchio.data.bboxes import BoundingBoxes -from torchio.data.bboxes import BoundingBoxFormat -from torchio.data.points import Points - - -class TestImageWithPoints: - def test_image_default_no_points(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - assert image.points == {} - - def test_image_with_points_kwarg(self): - pts = Points(torch.randn(5, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": pts}, - ) - assert "landmarks" in image.points - assert image.points["landmarks"] is pts - - def test_image_with_multiple_point_sets(self): - lm = Points(torch.randn(5, 3)) - fiducials = Points(torch.randn(3, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": lm, "fiducials": fiducials}, - ) - assert len(image.points) == 2 - - def test_image_points_validates_values(self): - """Points dict values must be Points instances.""" - with pytest.raises(TypeError, match="Points"): - ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": torch.randn(5, 3)}, - ) - - -class TestImageWithBoundingBoxes: - def test_image_default_no_bounding_boxes(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - assert image.bounding_boxes == {} - - def test_image_with_bboxes_kwarg(self): - boxes = BoundingBoxes( - torch.tensor([[10, 20, 30, 50, 60, 70]]), - format=BoundingBoxFormat.IJKIJK, - ) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - bounding_boxes={"tumors": boxes}, - ) - assert "tumors" in image.bounding_boxes - assert image.bounding_boxes["tumors"] is boxes - - def test_image_with_multiple_bbox_sets(self): - tumors = BoundingBoxes( - torch.randn(2, 6), - format=BoundingBoxFormat.IJKIJK, - ) - organs = BoundingBoxes( - torch.randn(5, 6), - format=BoundingBoxFormat.IJKWHD, - ) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - bounding_boxes={"tumors": tumors, "organs": organs}, - ) - assert len(image.bounding_boxes) == 2 - - def test_image_bboxes_validates_values(self): - """BoundingBoxes dict values must be BoundingBoxes instances.""" - with pytest.raises(TypeError, match="BoundingBoxes"): - ScalarImage( - torch.randn(1, 10, 10, 10), - bounding_boxes={"tumors": torch.randn(2, 6)}, - ) - - -class TestImageAnnotationsBothTypes: - def test_image_with_points_and_bboxes(self): - pts = Points(torch.randn(5, 3)) - boxes = BoundingBoxes( - torch.randn(2, 6), - format=BoundingBoxFormat.IJKIJK, - ) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": pts}, - bounding_boxes={"tumors": boxes}, - ) - assert len(image.points) == 1 - assert len(image.bounding_boxes) == 1 - - -class TestNewLikePreservesAnnotations: - def test_new_like_preserves_points(self): - pts = Points(torch.randn(5, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": pts}, - ) - new = image.new_like(data=torch.randn(1, 5, 5, 5)) - assert "landmarks" in new.points - assert new.points["landmarks"].num_points == 5 - # Should be a copy, not the same object - assert new.points["landmarks"] is not pts - - def test_new_like_preserves_bboxes(self): - boxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - ) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - bounding_boxes={"tumors": boxes}, - ) - new = image.new_like(data=torch.randn(1, 5, 5, 5)) - assert "tumors" in new.bounding_boxes - assert new.bounding_boxes["tumors"].num_boxes == 3 - assert new.bounding_boxes["tumors"] is not boxes - - def test_new_like_preserves_subclass_with_annotations(self): - pts = Points(torch.randn(5, 3)) - image = LabelMap( - torch.randint(0, 5, (1, 10, 10, 10)), - points={"landmarks": pts}, - ) - new = image.new_like(data=torch.randint(0, 5, (1, 5, 5, 5))) - assert isinstance(new, LabelMap) - assert "landmarks" in new.points - - -class TestDeepCopyPreservesAnnotations: - def test_deepcopy_copies_points(self): - pts = Points(torch.randn(5, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": pts}, - ) - copied = copy.deepcopy(image) - assert "landmarks" in copied.points - assert copied.points["landmarks"] is not pts - torch.testing.assert_close( - copied.points["landmarks"].data, - pts.data, - ) - - def test_deepcopy_copies_bboxes(self): - boxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - ) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - bounding_boxes={"tumors": boxes}, - ) - copied = copy.deepcopy(image) - assert "tumors" in copied.bounding_boxes - assert copied.bounding_boxes["tumors"] is not boxes - torch.testing.assert_close( - copied.bounding_boxes["tumors"].data, - boxes.data, - ) - - def test_deepcopy_independence(self): - """Modifying the copy's annotations doesn't affect the original.""" - pts = Points(torch.randn(5, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": pts}, - ) - copied = copy.deepcopy(image) - # Mutate the copy - copied.points["landmarks"]._data[0, 0] = 999.0 - assert image.points["landmarks"].data[0, 0] != 999.0 - - -class TestSlicingPreservesAnnotations: - def test_slice_preserves_points(self): - pts = Points(torch.randn(5, 3)) - image = ScalarImage( - torch.randn(1, 20, 20, 20), - points={"landmarks": pts}, - ) - sliced = image[:, 5:10] - assert "landmarks" in sliced.points - assert sliced.points["landmarks"].num_points == 5 - - def test_slice_preserves_bboxes(self): - boxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - ) - image = ScalarImage( - torch.randn(1, 20, 20, 20), - bounding_boxes={"tumors": boxes}, - ) - sliced = image[:, 5:10] - assert "tumors" in sliced.bounding_boxes - - -class TestSubjectWithImageLevelAnnotations: - def test_subject_image_level_points(self): - """Points on an Image are accessible through the Subject.""" - pts = Points(torch.randn(5, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": pts}, - ) - subject = Subject(t1=image) - assert "landmarks" in subject.t1.points - - def test_all_points_includes_both_levels(self): - """all_points() yields from both image-level and subject-level.""" - img_pts = Points(torch.randn(5, 3)) - subj_pts = Points(torch.randn(3, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"img_landmarks": img_pts}, - ) - subject = Subject( - t1=image, - subj_landmarks=subj_pts, - ) - all_pts = subject.all_points() - assert "subj_landmarks" in all_pts - assert ("t1", "img_landmarks") in all_pts - - def test_all_bounding_boxes_includes_both_levels(self): - """all_bounding_boxes() yields from both levels.""" - img_boxes = BoundingBoxes( - torch.randn(2, 6), - format=BoundingBoxFormat.IJKIJK, - ) - subj_boxes = BoundingBoxes( - torch.randn(1, 6), - format=BoundingBoxFormat.IJKIJK, - ) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - bounding_boxes={"img_tumors": img_boxes}, - ) - subject = Subject( - t1=image, - subj_tumors=subj_boxes, - ) - all_bb = subject.all_bounding_boxes() - assert "subj_tumors" in all_bb - assert ("t1", "img_tumors") in all_bb - - def test_all_points_no_overlap(self): - """Subject with only subject-level points.""" - subj_pts = Points(torch.randn(3, 3)) - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - landmarks=subj_pts, - ) - all_pts = subject.all_points() - assert "landmarks" in all_pts - assert len(all_pts) == 1 - - def test_all_points_only_image_level(self): - """Subject with only image-level points.""" - img_pts = Points(torch.randn(5, 3)) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": img_pts}, - ) - subject = Subject(t1=image) - all_pts = subject.all_points() - assert ("t1", "landmarks") in all_pts - assert len(all_pts) == 1 - - -class TestImageAnnotationsRepr: - def test_repr_includes_annotations(self): - pts = Points(torch.randn(5, 3)) - boxes = BoundingBoxes( - torch.randn(2, 6), - format=BoundingBoxFormat.IJKIJK, - ) - image = ScalarImage( - torch.randn(1, 10, 10, 10), - points={"landmarks": pts}, - bounding_boxes={"tumors": boxes}, - ) - r = repr(image) - assert "landmarks" in r - assert "tumors" in r - - def test_repr_no_annotations(self): - image = ScalarImage(torch.randn(1, 10, 10, 10)) - r = repr(image) - # Should not mention points/bboxes when empty - assert "points" not in r.lower() or "0" in r diff --git a/tests/test_inverse.py b/tests/test_inverse.py deleted file mode 100644 index ff2369b41..000000000 --- a/tests/test_inverse.py +++ /dev/null @@ -1,128 +0,0 @@ -"""Tests for the inverse transform module.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject() -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - seg=tio.LabelMap(torch.zeros(1, 10, 10, 10)), - ) - - -class TestApplyInverseTransform: - def test_flip_inverse(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - transformed = tio.Flip(axes=(0,))(subject) - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original) - - def test_does_not_mutate_subject(self) -> None: - subject = _make_subject() - transformed = tio.Flip(axes=(0,))(subject) - snapshot = transformed.t1.data.clone() - restored = transformed.apply_inverse_transform() - # Inverting must not modify the transformed input in place. - torch.testing.assert_close(transformed.t1.data, snapshot) - assert restored is not transformed - - def test_does_not_mutate_batch(self) -> None: - data = torch.rand(1, 16, 16, 16) - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(3)] - ) - transformed = tio.Affine(degrees=(0, 0, (10, 45)), default_pad_value=0.0)(batch) - snapshot = transformed.t1.data.clone() - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(transformed.t1.data, snapshot) - assert not torch.allclose(restored.t1.data, transformed.t1.data) - - def test_does_not_mutate_per_element_batch(self) -> None: - torch.manual_seed(0) - data = torch.rand(1, 16, 16, 16) - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(8)] - ) - transformed = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))])(batch) - snapshot = transformed.t1.data.clone() - transformed.apply_inverse_transform() - torch.testing.assert_close(transformed.t1.data, snapshot) - - def test_standalone_function_does_not_mutate(self) -> None: - data = torch.rand(1, 16, 16, 16) - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(3)] - ) - transformed = tio.Affine(degrees=(0, 0, (10, 45)), default_pad_value=0.0)(batch) - snapshot = transformed.t1.data.clone() - tio.apply_inverse_transform(transformed) - torch.testing.assert_close(transformed.t1.data, snapshot) - - def test_compose_inverse(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - pipeline = tio.Compose( - [ - tio.Flip(axes=(0,)), - tio.Flip(axes=(1,)), - ] - ) - transformed = pipeline(subject) - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original) - - def test_ignore_intensity(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - pipeline = tio.Compose( - [ - tio.Flip(axes=(0,)), - tio.Noise(std=0.1), - ] - ) - transformed = pipeline(subject) - restored = transformed.apply_inverse_transform( - ignore_intensity=True, - ) - # Shape restored, flip inverted, noise skipped. - assert restored.t1.data.shape == original.shape - - def test_get_inverse_transform(self) -> None: - subject = _make_subject() - transformed = tio.Flip(axes=(0,))(subject) - inverse = transformed.get_inverse_transform() - assert inverse is not None - - def test_standalone_function(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - transformed = tio.Flip(axes=(0,))(subject) - restored = tio.apply_inverse_transform(transformed) - torch.testing.assert_close(restored.t1.data, original) - - def test_no_history(self) -> None: - """Subject with no transforms should return itself.""" - subject = _make_subject() - original = subject.t1.data.clone() - restored = subject.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original) - - def test_missing_included_image_is_noop(self) -> None: - a = torch.arange(8.0).reshape(1, 2, 2, 2) - b = torch.arange(100.0, 108.0).reshape(1, 2, 2, 2) - subject = tio.Subject( - a=tio.ScalarImage(a.clone()), - b=tio.ScalarImage(b.clone()), - ) - - transformed = tio.Gamma(log_gamma=0.5, include=["a"])(subject) - current = tio.Subject(b=transformed.b) - current.applied_transforms = transformed.applied_transforms - restored = current.apply_inverse_transform() - - torch.testing.assert_close(restored.b.data, b) diff --git a/tests/test_keep_largest.py b/tests/test_keep_largest.py deleted file mode 100644 index ce945ea5d..000000000 --- a/tests/test_keep_largest.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Tests for KeepLargestComponent transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestKeepLargestComponent: - def test_keeps_largest_binary(self) -> None: - seg = torch.zeros(1, 20, 20, 20, dtype=torch.float32) - seg[0, 1:3, 1:3, 1:3] = 1 - seg[0, 10:18, 10:18, 10:18] = 1 - subject = tio.Subject(seg=tio.LabelMap(seg)) - result = tio.KeepLargestComponent()(subject) - assert result.seg.data[0, 2, 2, 2] == 0 - assert result.seg.data[0, 14, 14, 14] == 1 - - def test_multi_label(self) -> None: - seg = torch.zeros(1, 20, 20, 20, dtype=torch.float32) - seg[0, 0:2, 0:2, 0:2] = 1 - seg[0, 10:13, 10:13, 10:13] = 1 - seg[0, 5:8, 5:8, 5:8] = 2 - subject = tio.Subject(seg=tio.LabelMap(seg)) - result = tio.KeepLargestComponent()(subject) - assert result.seg.data[0, 1, 1, 1] == 0 - assert result.seg.data[0, 11, 11, 11] == 1 - assert result.seg.data[0, 6, 6, 6] == 2 - - def test_specific_labels(self) -> None: - seg = torch.zeros(1, 20, 20, 20, dtype=torch.float32) - seg[0, 0:2, 0:2, 0:2] = 1 - seg[0, 10:13, 10:13, 10:13] = 1 - seg[0, 5:7, 5:7, 5:7] = 2 - seg[0, 15:19, 15:19, 15:19] = 2 - subject = tio.Subject(seg=tio.LabelMap(seg)) - result = tio.KeepLargestComponent(labels=[1])(subject) - assert result.seg.data[0, 1, 1, 1] == 0 - assert result.seg.data[0, 6, 6, 6] == 2 - - def test_multichannel_raises(self) -> None: - seg = torch.zeros(2, 10, 10, 10, dtype=torch.float32) - subject = tio.Subject(seg=tio.LabelMap(seg)) - with pytest.raises(RuntimeError, match="single-channel"): - tio.KeepLargestComponent()(subject) - - def test_leaves_scalar_unchanged(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.KeepLargestComponent()(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_face_connectivity(self) -> None: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 0:3, 0:3, 0:3] = 1 - seg[0, 3:7, 3:7, 3:7] = 1 - subject = tio.Subject(seg=tio.LabelMap(seg)) - result_26 = tio.KeepLargestComponent(fully_connected=True)(subject) - assert result_26.seg.data[0, 1, 1, 1] == 1 - assert result_26.seg.data[0, 5, 5, 5] == 1 diff --git a/tests/test_labels_to_image.py b/tests/test_labels_to_image.py deleted file mode 100644 index 70a0942e4..000000000 --- a/tests/test_labels_to_image.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Tests for LabelsToImage transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestLabelsToImage: - def test_generates_image(self) -> None: - subject = _make_subject() - result = tio.LabelsToImage(label_key="seg")(subject) - assert "image_from_labels" in result - assert result.image_from_labels.data.shape[1:] == (10, 10, 10) - - def test_custom_key(self) -> None: - subject = _make_subject() - result = tio.LabelsToImage(label_key="seg", image_key="synth")(subject) - assert "synth" in result - - def test_auto_detect_label(self) -> None: - subject = _make_subject() - result = tio.LabelsToImage()(subject) - assert "image_from_labels" in result - - def test_ignore_background(self) -> None: - subject = _make_subject() - result = tio.LabelsToImage( - label_key="seg", - ignore_background=True, - )(subject) - bg_mask = subject.seg.data == 0 - bg_values = result.image_from_labels.data[0, bg_mask[0]] - assert bg_values.abs().max() < 1e-5 - - def test_no_label_raises(self) -> None: - subject = _make_subject(with_label=False) - with pytest.raises(KeyError, match="No LabelMap"): - tio.LabelsToImage()(subject) - - def test_missing_key_raises(self) -> None: - subject = _make_subject() - with pytest.raises(KeyError, match="nope"): - tio.LabelsToImage(label_key="nope")(subject) - - -class TestLabelsToImagePerInstance: - def _batch(self, batch_size: int = 5) -> tio.SubjectsBatch: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - subjects = [ - tio.Subject(seg=tio.LabelMap(seg.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_means_differ_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.LabelsToImage(label_key="seg", default_mean=(0.2, 0.9)) - result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["means"]) == batch.batch_size - means_for_label_1 = [m[1] for m in params["means"]] - assert len(set(means_for_label_1)) > 1 - assert result.image_from_labels.data.shape[0] == batch.batch_size - - def test_per_instance_false_shares_params(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.LabelsToImage( - label_key="seg", - default_mean=(0.2, 0.9), - per_instance=False, - ) - result = transform(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["means"], dict) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = _make_subject() - result = tio.LabelsToImage(label_key="seg", default_mean=(0.2, 0.9))(subject) - assert isinstance(result.applied_transforms[-1].params["means"], dict) - - -class TestLabelsToImagePerElementVectorized: - def test_each_element_uses_its_own_label_stats(self) -> None: - # The vectorized per-element generation must give each batch element - # its own per-label mean (no cross-element contamination), even though - # the random draw order differs from the old per-element loop. - size = 16 - label = torch.zeros(1, size, size, size) - label[0, : size // 2] = 1 - label[0, size // 2 :] = 2 - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(seg=tio.LabelMap(label.clone())) for _ in range(3)] - ) - transform = tio.LabelsToImage( - label_key="seg", - image_key="img", - default_mean=(0.0, 100.0), - default_std=(0.0, 0.05), - ) - torch.manual_seed(1) - result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - image = result.img.data - for index in range(batch.batch_size): - region_one = image[index, 0, : size // 2] - region_two = image[index, 0, size // 2 :] - assert region_one.mean().item() == pytest.approx( - params["means"][index][1], abs=0.5 - ) - assert region_two.mean().item() == pytest.approx( - params["means"][index][2], abs=0.5 - ) - # Independent per-element sampling: means vary across the batch. - label_one_means = {round(params["means"][i][1], 3) for i in range(3)} - assert len(label_one_means) > 1 diff --git a/tests/test_lambda_transform.py b/tests/test_lambda_transform.py deleted file mode 100644 index b11b30fe9..000000000 --- a/tests/test_lambda_transform.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Tests for Lambda transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestLambda: - def test_double(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Lambda(lambda x: 2 * x)(subject) - torch.testing.assert_close(result.t1.data, 2 * original) - - def test_scalar_only(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.Lambda(lambda x: x * 0, types_to_apply="scalar")(subject) - assert result.t1.data.sum() == 0 - torch.testing.assert_close(result.seg.data, original_seg) - - def test_label_only(self) -> None: - subject = _make_subject() - original_t1 = subject.t1.data.clone() - result = tio.Lambda(lambda x: x * 0, types_to_apply="label")(subject) - assert result.seg.data.sum() == 0 - torch.testing.assert_close(result.t1.data, original_t1) - - def test_not_callable_raises(self) -> None: - with pytest.raises(TypeError, match="callable"): - tio.Lambda(42) # type: ignore[arg-type] - - def test_unknown_types_to_apply_applies_all(self) -> None: - subject = _make_subject() - result = tio.Lambda(lambda x: x * 0, types_to_apply="unknown")(subject) - assert result.t1.data.sum() == 0 - assert result.seg.data.sum() == 0 diff --git a/tests/test_mask.py b/tests/test_mask.py deleted file mode 100644 index 36590546d..000000000 --- a/tests/test_mask.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Tests for Mask transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestMask: - def test_mask_with_label_key(self) -> None: - subject = _make_subject() - result = tio.Mask(masking_method="seg")(subject) - outside = subject.seg.data == 0 - assert (result.t1.data[outside.expand_as(result.t1.data)] == 0).all() - - def test_mask_with_callable(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Mask(masking_method=lambda x: x > 50)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_mask_with_labels(self) -> None: - subject = _make_subject() - result = tio.Mask(masking_method="seg", labels=[1])(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_mask_key_not_found(self) -> None: - subject = _make_subject(with_label=False) - with pytest.raises(KeyError, match="brain"): - tio.Mask(masking_method="brain")(subject) - - def test_outside_value(self) -> None: - subject = _make_subject() - result = tio.Mask(masking_method="seg", outside_value=-1)(subject) - outside = subject.seg.data == 0 - assert (result.t1.data[outside.expand_as(result.t1.data)] == -1).all() diff --git a/tests/test_monai_adapter.py b/tests/test_monai_adapter.py deleted file mode 100644 index f2be5ac79..000000000 --- a/tests/test_monai_adapter.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Tests for MonaiAdapter transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _monai_available() -> bool: - try: - import monai # noqa: F401 - - return True - except ImportError: - return False - - -HAS_MONAI = _monai_available() - - -@pytest.mark.skipif(not HAS_MONAI, reason="MONAI not installed") -class TestMonaiAdapterArray: - def test_array_transform(self) -> None: - from monai.transforms import NormalizeIntensity - - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 1), - ) - adapter = tio.MonaiAdapter(NormalizeIntensity()) - result = adapter(subject) - # NormalizeIntensity should zero-mean the data - assert abs(result.t1.data.mean().item()) < 0.5 - - def test_array_respects_include(self) -> None: - from monai.transforms import NormalizeIntensity - - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 5), - t2=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 5), - ) - original_t2 = subject.t2.data.clone() - adapter = tio.MonaiAdapter(NormalizeIntensity(), include=["t1"]) - result = adapter(subject) - torch.testing.assert_close(result.t2.data, original_t2) - - def test_array_skips_label_maps(self) -> None: - from monai.transforms import NormalizeIntensity - - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 8, 8, 8))), - ) - original_seg = subject.seg.data.clone() - adapter = tio.MonaiAdapter(NormalizeIntensity()) - result = adapter(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - -@pytest.mark.skipif(not HAS_MONAI, reason="MONAI not installed") -class TestMonaiAdapterDict: - def test_dict_transform(self) -> None: - from monai.transforms import NormalizeIntensityd - - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 1), - ) - adapter = tio.MonaiAdapter(NormalizeIntensityd(keys=["t1"])) - result = adapter(subject) - assert abs(result.t1.data.mean().item()) < 0.5 - - def test_dict_only_modifies_specified_keys(self) -> None: - from monai.transforms import NormalizeIntensityd - - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 5), - t2=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 5), - ) - original_t2 = subject.t2.data.clone() - adapter = tio.MonaiAdapter(NormalizeIntensityd(keys=["t1"])) - result = adapter(subject) - torch.testing.assert_close(result.t2.data, original_t2) - - -@pytest.mark.skipif(not HAS_MONAI, reason="MONAI not installed") -class TestMonaiAdapterGeneral: - def test_history_not_recorded(self) -> None: - from monai.transforms import NormalizeIntensity - - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - ) - adapter = tio.MonaiAdapter(NormalizeIntensity()) - result = adapter(subject) - # MonaiAdapter should not record itself in history - assert len(result.applied_transforms) == 0 - - def test_accepts_image(self) -> None: - from monai.transforms import NormalizeIntensity - - image = tio.ScalarImage(torch.rand(1, 8, 8, 8) + 1) - result = tio.MonaiAdapter(NormalizeIntensity())(image) - assert isinstance(result, tio.Image) - - def test_not_callable_raises(self) -> None: - with pytest.raises(TypeError, match="callable"): - tio.MonaiAdapter("not a transform") # type: ignore[arg-type] - - def test_in_compose(self) -> None: - from monai.transforms import NormalizeIntensity - - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) + 1), - ) - pipeline = tio.Compose([tio.MonaiAdapter(NormalizeIntensity())]) - result = pipeline(subject) - assert isinstance(result, tio.Subject) diff --git a/tests/test_motion.py b/tests/test_motion.py deleted file mode 100644 index b3c5d1304..000000000 --- a/tests/test_motion.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Tests for Motion transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestMotion: - def test_changes_data(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Motion(degrees=15, translation=10)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_num_transforms_validation(self) -> None: - with pytest.raises(ValueError, match="num_transforms"): - tio.Motion(num_transforms=0) - - def test_leaves_labels_unchanged(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.Motion()(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_preserves_shape(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Motion()(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_single_transform(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Motion(num_transforms=1)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - -class TestMotionPerInstance: - def _batch(self, batch_size: int = 5) -> tio.SubjectsBatch: - data = torch.rand(1, 12, 12, 12) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Motion(degrees=(5, 15), translation=(5, 15), num_transforms=2)( - batch - ) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["transforms"]) == batch.batch_size - assert not torch.allclose(result.t1.data[0], result.t1.data[1]) - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.Motion( - degrees=(5, 15), - translation=(5, 15), - num_transforms=2, - per_instance=False, - ) - result = transform(batch) - torch.testing.assert_close(result.t1.data[0], result.t1.data[1]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 12, 12, 12))) - result = tio.Motion(degrees=15, translation=10)(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params - - -class TestMotionDegenerateSegments: - def test_too_many_transforms_for_first_axis_raises(self) -> None: - # num_transforms + 1 segments cannot exceed the first spatial axis - # size; the transform must raise a clear error rather than silently - # replacing the whole spectrum. - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 2, 8, 8))) - with pytest.raises(ValueError, match="motion segments"): - tio.Motion(degrees=5, translation=5, num_transforms=4)(subject) diff --git a/tests/test_noise.py b/tests/test_noise.py deleted file mode 100644 index e5c912603..000000000 --- a/tests/test_noise.py +++ /dev/null @@ -1,308 +0,0 @@ -"""Tests for the Noise intensity transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - -HAS_MPS = torch.backends.mps.is_available() - - -class TestNoise: - def test_adds_noise(self) -> None: - tensor = torch.zeros(1, 8, 8, 8) - subject = tio.Subject(t1=tio.ScalarImage(tensor)) - result = tio.Noise(std=1.0)(subject) - # Should no longer be all zeros - assert result.t1.data.abs().sum() > 0 - - def test_mean_param(self) -> None: - tensor = torch.zeros(1, 8, 8, 8) - subject = tio.Subject(t1=tio.ScalarImage(tensor)) - result = tio.Noise(mean=10.0, std=0.0)(subject) - torch.testing.assert_close( - result.t1.data.mean(), - torch.tensor(10.0), - atol=0.01, - rtol=0, - ) - - def test_zero_std_no_change(self) -> None: - tensor = torch.rand(1, 8, 8, 8) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - result = tio.Noise(std=0.0)(subject) - torch.testing.assert_close(result.t1.data, tensor) - - def test_only_scalar_images(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - seg=tio.LabelMap(torch.zeros(1, 8, 8, 8, dtype=torch.long)), - ) - original_seg = subject.seg.data.clone() - result = tio.Noise(std=1.0)(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_history_recorded(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - result = tio.Noise(std=0.5)(subject) - assert len(result.applied_transforms) == 1 - trace = result.applied_transforms[0] - assert trace.name == "Noise" - assert "mean" in trace.params - assert "std" in trace.params - assert "seed" in trace.params - - def test_probability(self) -> None: - tensor = torch.rand(1, 4, 4, 4) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - result = tio.Noise(std=1.0, p=0.0)(subject) - torch.testing.assert_close(result.t1.data, tensor) - - def test_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - result = tio.Noise(std=0.1)(image) - assert isinstance(result, tio.Image) - - def test_accepts_tensor(self) -> None: - tensor = torch.rand(1, 4, 4, 4) - result = tio.Noise(std=0.1)(tensor) - assert isinstance(result, torch.Tensor) - - def test_differentiable(self) -> None: - tensor = torch.rand(1, 4, 4, 4, requires_grad=True) - result = tio.Noise(std=0.1, copy=False)(tensor) - loss = result.sum() - loss.backward() - assert tensor.grad is not None - - def test_in_compose(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - pipeline = tio.Compose([tio.Noise(std=0.1), tio.Noise(std=0.2)]) - result = pipeline(subject) - assert len(result.applied_transforms) == 2 - - def test_include_exclude(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)), - t2=tio.ScalarImage(torch.zeros(1, 4, 4, 4)), - ) - original_t2 = subject.t2.data.clone() - result = tio.Noise(std=1.0, include=["t1"])(subject) - torch.testing.assert_close(result.t2.data, original_t2) - - @pytest.mark.skipif(not HAS_MPS, reason="No MPS") - def test_noise_on_mps(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)), - ) - subject.to("mps") - result = tio.Noise(std=1.0)(subject) - assert result.t1.device.type == "mps" - assert result.t1.data.abs().sum() > 0 - - def test_seed_reproducibility(self) -> None: - """Replaying with saved params reproduces the same noise.""" - from torchio.data.batch import SubjectsBatch - - subject1 = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - ) - subject2 = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - ) - noise = tio.Noise(std=1.0) - params = { - "mean": 0.0, - "std": 1.0, - "seed": 42, - } - batch1 = SubjectsBatch.from_subjects([subject1]) - batch2 = SubjectsBatch.from_subjects([subject2]) - result1 = noise.apply_transform(batch1, params) - result2 = noise.apply_transform(batch2, params) - r1 = result1.unbatch()[0] - r2 = result2.unbatch()[0] - torch.testing.assert_close(r1.t1.data, r2.t1.data) - - def test_negative_std_raises(self) -> None: - with pytest.raises(ValueError, match="non-negative"): - tio.Noise(std=-1.0) - - def test_random_std_range(self) -> None: - """std=(lo, hi) samples uniformly each call.""" - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - ) - noise = tio.Noise(std=(0.5, 1.5)) - stds = set() - for _ in range(10): - result = noise(subject) - sampled_std = result.applied_transforms[-1].params["std"] - assert 0.5 <= sampled_std <= 1.5 - stds.add(round(sampled_std, 4)) - # Should have sampled different values - assert len(stds) > 1 - - def test_random_mean_range(self) -> None: - noise = tio.Noise(mean=(-1.0, 1.0), std=0.0) - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - ) - means = set() - for _ in range(10): - result = noise(subject) - sampled_mean = result.applied_transforms[-1].params["mean"] - assert -1.0 <= sampled_mean <= 1.0 - means.add(round(sampled_mean, 4)) - assert len(means) > 1 - - def test_deterministic_scalar(self) -> None: - """Scalar std is always the same.""" - noise = tio.Noise(std=0.5) - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 4, 4, 4)), - ) - result = noise(subject) - assert result.applied_transforms[0].params["std"] == 0.5 - - def test_rician_noise(self) -> None: - tensor = torch.ones(1, 8, 8, 8) - subject = tio.Subject(t1=tio.ScalarImage(tensor)) - result = tio.Noise(std=0.5, rician=True)(subject) - # Rician noise is always non-negative - assert (result.t1.data >= 0).all() - # Should differ from the original - assert not torch.equal(result.t1.data, tensor) - - def test_rician_recorded_in_params(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - result = tio.Noise(std=0.1, rician=True)(subject) - assert result.applied_transforms[0].params["rician"] is True - - def test_gaussian_vs_rician_differ(self) -> None: - torch.manual_seed(42) - tensor = torch.ones(1, 8, 8, 8) - subject_g = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - subject_r = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - gaussian = tio.Noise(std=0.5, rician=False)(subject_g) - rician = tio.Noise(std=0.5, rician=True)(subject_r) - # They should produce different results - assert not torch.equal(gaussian.t1.data, rician.t1.data) - - def test_distribution_for_std(self) -> None: - from torch.distributions import Uniform - - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - ) - noise = tio.Noise(std=Uniform(0.1, 0.5)) - result = noise(subject) - sampled_std = result.applied_transforms[0].params["std"] - assert 0.1 <= sampled_std <= 0.5 - - def test_distribution_for_mean(self) -> None: - from torch.distributions import Normal - - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - ) - noise = tio.Noise(mean=Normal(0.0, 0.1), std=0.0) - result = noise(subject) - # Mean was sampled from N(0, 0.1), should be near 0 - sampled_mean = result.applied_transforms[0].params["mean"] - assert isinstance(sampled_mean, float) - - def test_lognormal_distribution(self) -> None: - from torch.distributions import LogNormal - - subject = tio.Subject( - t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8)), - ) - noise = tio.Noise(std=LogNormal(loc=-2.0, scale=0.5)) - result = noise(subject) - sampled_std = result.applied_transforms[0].params["std"] - assert sampled_std > 0 # LogNormal always positive - - -class TestNoisePerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8))) - for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_default_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Noise(std=(0.5, 1.5))(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["std"], list) - assert len(params["std"]) == batch.batch_size - assert len(set(params["std"])) > 1 - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Noise(std=(0.5, 1.5), per_instance=False)(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["std"], float) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.zeros(1, 8, 8, 8))) - result = tio.Noise(std=(0.5, 1.5))(subject) - assert isinstance(result.applied_transforms[-1].params["std"], float) - - def test_per_instance_mean_applied_per_element(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=5) - result = tio.Noise(mean=(5.0, 20.0), std=0.0)(batch) - means = result.applied_transforms[-1].params["mean"] - for i, mean in enumerate(means): - torch.testing.assert_close( - result.t1.data[i].mean(), - torch.tensor(float(mean)), - atol=0.01, - rtol=0, - ) - - def test_per_instance_p_gates_some_elements(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=64) - result = tio.Noise(std=1.0, p=0.5)(batch) - changed = [result.t1.data[i].abs().sum() > 0 for i in range(batch.batch_size)] - assert any(changed) - assert not all(changed) - - def test_per_instance_p_masked_elements_have_no_history(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=32) - result = tio.Noise(std=1.0, p=0.5)(batch) - for subject in result.unbatch(): - changed = subject.t1.data.abs().sum() > 0 - has_history = len(subject.applied_transforms) == 1 - assert bool(changed) == has_history - - def test_per_instance_p_rician_masked_elements_unchanged(self) -> None: - """Rician is non-linear, so gated-out elements must be restored.""" - torch.manual_seed(0) - data = torch.randn(1, 8, 8, 8) # signed data - subjects = [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(32)] - batch = tio.SubjectsBatch.from_subjects(subjects) - original = batch.t1.data.clone() - result = tio.Noise(std=1.0, rician=True, p=0.5)(batch) - unchanged = [ - torch.allclose(result.t1.data[i], original[i]) - for i in range(batch.batch_size) - ] - # Some elements were gated out and must be exactly the input. - assert any(unchanged) - assert not all(unchanged) diff --git a/tests/test_normalize.py b/tests/test_normalize.py deleted file mode 100644 index 788e605c0..000000000 --- a/tests/test_normalize.py +++ /dev/null @@ -1,331 +0,0 @@ -"""Tests for Normalize.""" - -from __future__ import annotations - -import numpy as np -import pytest -import torch - -import torchio as tio -from torchio.transforms._statistics import compute_quantile - - -def _make_subject( - values: torch.Tensor | None = None, - with_label: bool = False, -) -> tio.Subject: - if values is None: - values = torch.arange(1000, dtype=torch.float32).reshape(1, 10, 10, 10) - kwargs: dict = {"t1": tio.ScalarImage(values)} - if with_label: - mask = torch.zeros(1, 10, 10, 10) - mask[0, 2:8, 2:8, 2:8] = 1 - kwargs["brain"] = tio.LabelMap(mask) - return tio.Subject(**kwargs) - - -class TestBasic: - def test_default_rescales_to_minus1_1(self) -> None: - subject = _make_subject() - result = tio.Normalize()(subject) - data = result.t1.data - assert abs(data.min().item() - (-1.0)) < 1e-5 - assert abs(data.max().item() - 1.0) < 1e-5 - - def test_rescale_to_0_1(self) -> None: - subject = _make_subject() - result = tio.Normalize(out_min=0.0, out_max=1.0)(subject) - data = result.t1.data - assert abs(data.min().item()) < 1e-5 - assert abs(data.max().item() - 1.0) < 1e-5 - - def test_rescale_to_0_255(self) -> None: - subject = _make_subject() - result = tio.Normalize(out_min=0.0, out_max=255.0)(subject) - data = result.t1.data - assert abs(data.min().item()) < 1e-3 - assert abs(data.max().item() - 255.0) < 1e-3 - - def test_ct_windowing(self) -> None: - data = torch.tensor([-1500, -1000, 0, 500, 1000, 2000], dtype=torch.float32) - data = data.reshape(1, 1, 1, 6) - subject = tio.Subject(ct=tio.ScalarImage(data)) - result = tio.Normalize( - out_min=0.0, - out_max=1.0, - in_min=-1000.0, - in_max=1000.0, - )(subject) - out = result.ct.data.flatten() - # -1500 gets clipped to -1000 -> maps to 0 - assert abs(out[0].item()) < 1e-5 - # 0 maps to 0.5 - assert abs(out[2].item() - 0.5) < 1e-5 - # 2000 gets clipped to 1000 -> maps to 1 - assert abs(out[5].item() - 1.0) < 1e-5 - - -class TestPercentiles: - def test_percentile_clipping(self) -> None: - data = torch.cat( - [ - torch.zeros(1, 5, 10, 10), - torch.ones(1, 5, 10, 10) * 100, - ], - dim=1, - ) - # First 50% is 0, second 50% is 100 - subject = tio.Subject(t1=tio.ScalarImage(data)) - result = tio.Normalize( - out_min=0.0, - out_max=1.0, - percentile_low=1.0, - percentile_high=99.0, - )(subject) - assert result.t1.data.min() >= -0.01 - assert result.t1.data.max() <= 1.01 - - def test_nnunet_percentiles(self) -> None: - torch.manual_seed(42) - data = torch.randn(1, 20, 20, 20) * 100 - subject = tio.Subject(t1=tio.ScalarImage(data)) - result = tio.Normalize( - out_min=0.0, - out_max=1.0, - percentile_low=0.5, - percentile_high=99.5, - )(subject) - # Most values should be in [0, 1], outliers clipped - in_range = (result.t1.data >= -0.01) & (result.t1.data <= 1.01) - assert in_range.float().mean() > 0.98 - - -class TestMasking: - def test_masking_with_label_key(self) -> None: - subject = _make_subject(with_label=True) - result = tio.Normalize( - out_min=0.0, - out_max=1.0, - masking_method="brain", - )(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_masking_with_callable(self) -> None: - subject = _make_subject() - result = tio.Normalize( - out_min=0.0, - out_max=1.0, - masking_method=lambda x: x > 500, - )(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_masking_key_not_found_raises(self) -> None: - subject = _make_subject() - with pytest.raises(KeyError, match="nonexistent"): - tio.Normalize( - masking_method="nonexistent", - )(subject) - - def test_masking_key_not_labelmap_raises(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - t2=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - with pytest.raises(TypeError, match="LabelMap"): - tio.Normalize(masking_method="t2")(subject) - - -class TestRandom: - def test_random_out_range(self) -> None: - subject = _make_subject() - transform = tio.Normalize( - out_min=(-2.0, -0.5), - out_max=(0.5, 2.0), - ) - results = [transform(subject).t1.data.min().item() for _ in range(5)] - # With random sampling, not all results should be identical - assert len({f"{v:.2f}" for v in results}) > 1 - - def test_random_percentiles(self) -> None: - torch.manual_seed(0) - subject = _make_subject() - transform = tio.Normalize( - out_min=0.0, - out_max=1.0, - percentile_low=(0.0, 5.0), - percentile_high=(95.0, 100.0), - ) - # Random percentiles change the clipping bounds, so - # values above the high percentile get clamped to 1.0 and - # the interior distribution shifts. - results = [transform(subject).t1.data.mean().item() for _ in range(10)] - assert len({f"{v:.4f}" for v in results}) > 1 - - -class TestEdgeCases: - def test_constant_value_warns(self) -> None: - data = torch.ones(1, 4, 4, 4) * 42.0 - subject = tio.Subject(t1=tio.ScalarImage(data)) - with pytest.warns(RuntimeWarning, match="zero"): - result = tio.Normalize()(subject) - # Data unchanged - torch.testing.assert_close(result.t1.data, data) - - def test_empty_mask_warns(self) -> None: - subject = _make_subject() - with pytest.warns(RuntimeWarning, match="empty"): - tio.Normalize( - out_min=0.0, - out_max=1.0, - masking_method=lambda x: torch.zeros_like(x, dtype=torch.bool), - )(subject) - - def test_leaves_label_maps_unchanged(self) -> None: - subject = _make_subject(with_label=True) - original_label = subject.brain.data.clone() - result = tio.Normalize()(subject) - torch.testing.assert_close(result.brain.data, original_label) - - -class TestInverse: - def test_inverse_restores_values(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - transformed = tio.Normalize( - out_min=0.0, - out_max=1.0, - )(subject) - restored = transformed.apply_inverse_transform() - np.testing.assert_allclose( - restored.t1.data.numpy(), - original.numpy(), - atol=1e-4, - ) - - def test_inverse_with_ct_windowing(self) -> None: - data = torch.linspace(-500, 500, 1000).reshape(1, 10, 10, 10) - subject = tio.Subject(ct=tio.ScalarImage(data)) - transformed = tio.Normalize( - out_min=0.0, - out_max=1.0, - in_min=-1000.0, - in_max=1000.0, - )(subject) - restored = transformed.apply_inverse_transform() - # Clipped values can't be restored, but the linear map is reversed - np.testing.assert_allclose( - restored.ct.data.numpy(), - data.numpy(), - atol=1e-4, - ) - - -class TestExports: - def test_available_at_top_level(self) -> None: - assert hasattr(tio, "RescaleIntensity") - - -class TestAlias: - def test_rescale_intensity_alias(self) -> None: - assert tio.RescaleIntensity is tio.Normalize - - -class TestQuantile: - """Tests for the `torch.kthvalue`-based quantile helper.""" - - @pytest.mark.parametrize("q", [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]) - def test_matches_torch_quantile(self, q: float) -> None: - values = torch.linspace(-3.0, 7.0, 101) - expected = torch.quantile(values, q) - result = compute_quantile(values, q) - assert torch.allclose(result, expected, atol=1e-5) - - def test_invalid_q_raises(self) -> None: - values = torch.arange(10, dtype=torch.float32) - with pytest.raises(ValueError, match="0 <= q <= 1"): - compute_quantile(values, 1.5) - - def test_large_tensor_interior_quantile(self) -> None: - # torch.quantile raises for more than 2**24 elements; kthvalue does not. - values = torch.arange(2**24 + 1, dtype=torch.float32) - result = compute_quantile(values, 0.5) - assert result.item() == pytest.approx(2**23) - - def test_rescale_intensity_large_image(self) -> None: - # Exceeds torch.quantile's 2**24-element limit; uses min/max endpoints. - data = torch.zeros(1, 2**24 + 1, 1, 1, dtype=torch.float32) - # A single non-zero voxel becomes the input maximum; everything else - # is the minimum, so the output spans the full [0, 1] range. - input_max = 4.0 - data[0, -1] = input_max - image = tio.ScalarImage(data) - transform = tio.RescaleIntensity(out_min=0.0, out_max=1.0, copy=False) - result = transform(image) - assert result.data[0, 0, 0, 0].item() == pytest.approx(0.0) - assert result.data[0, -1, 0, 0].item() == pytest.approx(1.0) - - -class TestNormalizePerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) * 100)) - for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_out_range_differs(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.RescaleIntensity(out_min=(-1.0, 0.0), out_max=(0.5, 1.0)) - result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["out_min"]) == batch.batch_size - assert len(set(params["out_min"])) > 1 - # Each element rescaled to its own output range. - for i in range(batch.batch_size): - data = result.t1.data[i] - assert data.min() >= params["out_min"][i] - 1e-4 - assert data.max() <= params["out_max"][i] + 1e-4 - - def test_per_instance_false_shares_params(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.RescaleIntensity( - out_min=(-1.0, 0.0), - out_max=(0.5, 1.0), - per_instance=False, - ) - result = transform(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["out_min"], float) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 8, 8, 8) * 100)) - result = tio.RescaleIntensity(out_min=(-1.0, 0.0), out_max=(0.5, 1.0))(subject) - assert isinstance(result.applied_transforms[-1].params["out_min"], float) - - def test_per_instance_inverse_zero_range_no_nan(self) -> None: - # A degenerate out_min == out_max (zero output range) must not - # produce NaNs on the per-element inverse. - torch.manual_seed(0) - data = torch.rand(1, 8, 8, 8) * 100 - subjects = [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(4)] - batch = tio.SubjectsBatch.from_subjects(subjects) - transform = tio.RescaleIntensity(out_min=0.0, out_max=0.0) - result = transform(batch) - assert "_batched_keys" in result.applied_transforms[-1].params - restored = result.apply_inverse_transform() - assert not torch.isnan(restored.t1.data).any() - # Identical inputs: the batch-shared input range covers every - # element, so only the per-element output range varies and the - # round-trip is exact. - torch.manual_seed(0) - data = torch.rand(1, 8, 8, 8) * 100 - subjects = [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(6)] - batch = tio.SubjectsBatch.from_subjects(subjects) - original = batch.t1.data.clone() - transform = tio.RescaleIntensity(out_min=(-1.0, 0.0), out_max=(0.5, 1.0)) - result = transform(batch) - restored = result.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original, atol=1e-3, rtol=0) diff --git a/tests/test_one_hot.py b/tests/test_one_hot.py deleted file mode 100644 index 8efb97497..000000000 --- a/tests/test_one_hot.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Tests for OneHot transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestOneHot: - def test_one_hot_encoding(self) -> None: - subject = _make_subject() - result = tio.OneHot()(subject) - assert result.seg.data.shape[0] == 3 - assert (result.seg.data.sum(dim=0) == 1).all() - - def test_num_classes(self) -> None: - subject = _make_subject() - result = tio.OneHot(num_classes=5)(subject) - assert result.seg.data.shape[0] == 5 - - def test_inverse(self) -> None: - subject = _make_subject() - original = subject.seg.data.clone() - transformed = tio.OneHot()(subject) - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(restored.seg.data, original) - - def test_leaves_scalar_images_unchanged(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.OneHot()(subject) - torch.testing.assert_close(result.t1.data, original) diff --git a/tests/test_one_of.py b/tests/test_one_of.py deleted file mode 100644 index 1d03b5fee..000000000 --- a/tests/test_one_of.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Tests for OneOf transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject() -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - - -class TestOneOf: - def test_applies_one(self) -> None: - subject = _make_subject() - transform = tio.OneOf( - [ - tio.Flip(axes=(0,)), - tio.Gamma(log_gamma=0.3), - ] - ) - result = transform(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_single_transform(self) -> None: - subject = _make_subject() - result = tio.OneOf([tio.Flip(axes=(0,))])(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_with_weights(self) -> None: - subject = _make_subject() - transform = tio.OneOf( - { - tio.Flip(axes=(0,)): 1.0, - tio.Gamma(log_gamma=0.0): 0.0, - } - ) - result = transform(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_history_recorded(self) -> None: - subject = _make_subject() - result = tio.OneOf([tio.Flip(axes=(0,))])(subject) - assert len(result.applied_transforms) > 0 - - -class TestOneOfPerInstance: - def _batch(self, batch_size: int = 32) -> tio.SubjectsBatch: - data = torch.rand(1, 8, 8, 8) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_element_chooses_different_transforms(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.OneOf([tio.Gamma(log_gamma=0.5), tio.Flip(axes=(0,))]) - result = transform(batch) - names = set() - for subject in result.unbatch(): - assert len(subject.applied_transforms) == 1 - names.add(subject.applied_transforms[-1].name) - assert names == {"Gamma", "Flip"} - - def test_per_instance_false_is_batch_wide(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=8) - transform = tio.OneOf( - [tio.Gamma(log_gamma=0.5), tio.Flip(axes=(0,))], - per_instance=False, - ) - result = transform(batch) - names = {subject.applied_transforms[-1].name for subject in result.unbatch()} - assert len(names) == 1 - - def test_single_subject_unaffected(self) -> None: - subject = _make_subject() - result = tio.OneOf([tio.Gamma(log_gamma=0.5), tio.Flip(axes=(0,))])(subject) - assert len(result.applied_transforms) == 1 - - def test_history_composes_after_oneof(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=16) - pipeline = tio.Compose( - [ - tio.OneOf([tio.Gamma(log_gamma=0.5), tio.Flip(axes=(0,))]), - tio.Gamma(log_gamma=0.2), - ] - ) - result = pipeline(batch) - for subject in result.unbatch(): - names = [trace.name for trace in subject.applied_transforms] - assert len(names) == 2 - assert names[-1] == "Gamma" - - def test_p_gates_some_elements(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=64) - transform = tio.OneOf([tio.Gamma(log_gamma=0.8)], p=0.5) - result = transform(batch) - applied = [len(subject.applied_transforms) == 1 for subject in result.unbatch()] - assert any(applied) - assert not all(applied) - - def test_per_element_inverse_restores(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=16) - original = batch.t1.data.clone() - transform = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))]) - result = transform(batch) - restored = result.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original) - - def test_functional_inverse_restores_per_element(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=16) - original = batch.t1.data.clone() - transform = tio.OneOf([tio.Flip(axes=(0,)), tio.Flip(axes=(1,))]) - result = transform(batch) - restored = tio.apply_inverse_transform(result) - torch.testing.assert_close(restored.t1.data, original) - - def test_get_inverse_transform_raises_for_per_element(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=4) - result = tio.OneOf([tio.Flip(axes=(0,))])(batch) - with pytest.raises(RuntimeError, match="per-element"): - result.get_inverse_transform() - - def test_clear_history_clears_per_element(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=4) - result = tio.OneOf([tio.Flip(axes=(0,))])(batch) - result.clear_history() - assert result._per_element_history is None - for subject in result.unbatch(): - assert subject.applied_transforms == [] - - def test_p_zero_is_noop_preserving_history(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=4) - original = batch.t1.data.clone() - # A prior shared, invertible transform. - flipped = tio.Flip(axes=(0,))(batch) - result = tio.OneOf([tio.Flip(axes=(1,))], p=0.0)(flipped) - torch.testing.assert_close(result.t1.data, flipped.t1.data) - assert result._per_element_history is None - # The shared Flip history is intact and still invertible as a batch. - restored = result.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original) - - -class TestOneOfCopy: - def test_does_not_mutate_input(self) -> None: - subject = _make_subject() - snapshot = subject.t1.data.clone() - tio.OneOf([tio.Gamma(log_gamma=0.5)])(subject) - torch.testing.assert_close(subject.t1.data, snapshot) - - def test_restores_child_copy_flag(self) -> None: - child = tio.Gamma(log_gamma=0.5) - assert child.copy is True - tio.OneOf([child])(_make_subject()) - assert child.copy is True - - def test_does_not_double_copy_children(self) -> None: - # The child must run with copy=False (the parent already copied), - # avoiding a redundant deep copy of the whole input. - seen: list[bool] = [] - - class _Spy(tio.IntensityTransform): - def apply_transform(self, batch, params): - seen.append(self.copy) - return batch - - tio.OneOf([_Spy()])(_make_subject()) - assert seen == [False] diff --git a/tests/test_pad.py b/tests/test_pad.py deleted file mode 100644 index 92e195f24..000000000 --- a/tests/test_pad.py +++ /dev/null @@ -1,264 +0,0 @@ -"""Tests for the Pad transform and Crop/Pad invertibility.""" - -from __future__ import annotations - -import warnings - -import pytest -import torch - -import torchio as tio -from torchio.transforms.spatial._padding import pad_tensor - - -class TestPad: - def test_pad_tensor_rejects_invalid_number_of_dimensions(self) -> None: - with pytest.raises(ValueError, match="4D or 5D"): - pad_tensor( - torch.ones(2, 2, 2), - (0, 0, 0, 0, 0, 0), - "mean", - 0, - ) - - def test_pad_uniform(self) -> None: - image = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - subject = tio.Subject(t1=image) - result = tio.Pad(padding=5)(subject) - assert result.t1.shape == (1, 20, 20, 20) - - def test_pad_per_axis(self) -> None: - image = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - subject = tio.Subject(t1=image) - result = tio.Pad(padding=(2, 4, 6))(subject) - assert result.t1.shape == (1, 14, 18, 22) - - def test_pad_six_values(self) -> None: - image = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - subject = tio.Subject(t1=image) - result = tio.Pad(padding=(1, 2, 3, 4, 5, 6))(subject) - assert result.t1.shape == (1, 13, 17, 21) - - def test_pad_constant_fill(self) -> None: - tensor = torch.ones(1, 4, 4, 4) - image = tio.ScalarImage(tensor) - subject = tio.Subject(t1=image) - result = tio.Pad(padding=1, fill=0)(subject) - # Corners should be 0 (padded) - assert result.t1.data[0, 0, 0, 0] == 0 - # Interior should be 1 (original) - assert result.t1.data[0, 1, 1, 1] == 1 - - def test_pad_reflect(self) -> None: - image = tio.ScalarImage(torch.rand(1, 4, 4, 4)) - subject = tio.Subject(t1=image) - result = tio.Pad(padding=1, padding_mode="reflect")(subject) - assert result.t1.shape == (1, 6, 6, 6) - - def test_invalid_padding_mode(self) -> None: - with pytest.raises(ValueError, match="padding_mode"): - tio.Pad(padding=1, padding_mode="maximum") - - @pytest.mark.parametrize( - ("padding_mode", "expected"), - [ - ("mean", 1.5), - ("median", 1.5), - ("minimum", 0), - ], - ) - def test_pad_statistic_mode( - self, - padding_mode: str, - expected: float, - ) -> None: - tensor = torch.arange(4, dtype=torch.float32).reshape(1, 1, 2, 2) - result = tio.Pad( - padding=(0, 0, 0, 1, 0, 0), - padding_mode=padding_mode, - )(tensor) - torch.testing.assert_close( - result[0, 0, 2], - torch.full((2,), expected, dtype=tensor.dtype), - ) - - @pytest.mark.parametrize("padding_mode", ["mean", "median", "minimum"]) - def test_pad_statistic_mode_per_batch_element( - self, - padding_mode: str, - ) -> None: - subjects = [ - tio.Subject(t1=tio.ScalarImage(torch.full((1, 2, 2, 2), value))) - for value in (1.0, 3.0) - ] - batch = tio.SubjectsBatch.from_subjects(subjects) - result = tio.Pad(padding=1, padding_mode=padding_mode)(batch) - torch.testing.assert_close( - result.t1.data[:, 0, 0, 0, 0], - torch.tensor([1.0, 3.0]), - ) - - @pytest.mark.parametrize( - ("padding_mode", "expected"), - [ - ("mean", 0), - ("median", 1), - ], - ) - def test_pad_statistic_mode_warns_for_integer_truncation( - self, - padding_mode: str, - expected: int, - ) -> None: - tensor = torch.tensor([0, 1, 1, 1]).reshape(1, 1, 2, 2) - with pytest.warns(RuntimeWarning, match="might be truncated"): - result = tio.Pad( - padding=(0, 0, 0, 1, 0, 0), - padding_mode=padding_mode, - )(tensor) - assert result.dtype == tensor.dtype - assert result[0, 0, 2, 0].item() == expected - - def test_pad_minimum_does_not_warn_for_integer_input(self) -> None: - tensor = torch.tensor([0, 1, 1, 1]).reshape(1, 1, 2, 2) - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - result = tio.Pad( - padding=(0, 0, 0, 1, 0, 0), - padding_mode="minimum", - )(tensor) - assert result[0, 0, 2, 0].item() == 0 - assert not any(issubclass(item.category, RuntimeWarning) for item in caught) - - @pytest.mark.parametrize("padding_mode", ["mean", "median", "minimum"]) - def test_pad_statistic_mode_is_differentiable( - self, - padding_mode: str, - ) -> None: - tensor = torch.rand(1, 2, 2, 2, requires_grad=True) - result = tio.Pad( - padding=1, - padding_mode=padding_mode, - copy=False, - )(tensor) - result.sum().backward() - assert tensor.grad is not None - - @pytest.mark.parametrize("padding_mode", ["mean", "median"]) - def test_pad_statistic_mode_preserves_float64_precision( - self, - padding_mode: str, - ) -> None: - values = torch.tensor([1.0, 1.0 + 2**-40], dtype=torch.float64) - tensor = values.reshape(1, 1, 1, 2) - result = tio.Pad( - padding=(0, 0, 0, 1, 0, 0), - padding_mode=padding_mode, - )(tensor) - expected = values.mean() - torch.testing.assert_close( - result[0, 0, 1, 0], - expected, - rtol=0, - atol=0, - ) - - def test_pad_all_images(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 10, 10, 10))), - ) - result = tio.Pad(padding=5)(subject) - assert result.t1.shape == (1, 20, 20, 20) - assert result.seg.shape == (1, 20, 20, 20) - - def test_pad_affine_updated(self) -> None: - image = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - subject = tio.Subject(t1=image) - original_origin = subject.t1.affine.origin - result = tio.Pad(padding=(5, 0, 0, 0, 0, 0))(subject) - # Origin should shift back along first axis - assert result.t1.affine.origin[0] != original_origin[0] - - def test_pad_history(self) -> None: - image = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - subject = tio.Subject(t1=image) - result = tio.Pad(padding=5)(subject) - assert len(result.applied_transforms) == 1 - assert result.applied_transforms[0].name == "Pad" - - def test_pad_accepts_tensor(self) -> None: - tensor = torch.rand(1, 10, 10, 10) - result = tio.Pad(padding=2)(tensor) - assert isinstance(result, torch.Tensor) - assert result.shape == (1, 14, 14, 14) - - def test_pad_batch(self) -> None: - from torchio.data.batch import SubjectsBatch - - subjects = [ - tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.Pad(padding=2)(batch) - assert result.t1.data.shape == (3, 1, 14, 14, 14) - - -class TestCropPadInvertibility: - def test_pad_invertible(self) -> None: - assert tio.Pad(padding=5).invertible - - def test_crop_invertible(self) -> None: - assert tio.Crop(cropping=5).invertible - - def test_pad_then_inverse_gives_original_shape(self) -> None: - tensor = torch.rand(1, 10, 10, 10) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - pad = tio.Pad(padding=5) - padded = pad(subject) - assert padded.t1.shape == (1, 20, 20, 20) - restored = padded.apply_inverse_transform() - assert restored.t1.shape == (1, 10, 10, 10) - torch.testing.assert_close(restored.t1.data, tensor) - - def test_crop_then_inverse_gives_original_shape(self) -> None: - tensor = torch.rand(1, 20, 20, 20) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - crop = tio.Crop(cropping=5) - cropped = crop(subject) - assert cropped.t1.shape == (1, 10, 10, 10) - restored = cropped.apply_inverse_transform() - # Shape restored, but cropped data is lost (filled with 0) - assert restored.t1.shape == (1, 20, 20, 20) - - def test_crop_pad_compose_inverse(self) -> None: - tensor = torch.rand(1, 20, 20, 20) - subject = tio.Subject(t1=tio.ScalarImage(tensor.clone())) - pipeline = tio.Compose( - [ - tio.Crop(cropping=2), - tio.Pad(padding=3), - ] - ) - transformed = pipeline(subject) - assert transformed.t1.shape == (1, 22, 22, 22) - restored = transformed.apply_inverse_transform() - assert restored.t1.shape == (1, 20, 20, 20) - - def test_crop_or_pad_inverse_respects_include_scope(self) -> None: - a = torch.ones(1, 4, 4, 4) - b = torch.ones(1, 4, 4, 4) * 2 - subject = tio.Subject( - a=tio.ScalarImage(a.clone()), - b=tio.ScalarImage(b.clone()), - ) - - transformed = tio.CropOrPad(6, include=["a"])(subject) - restored = transformed.apply_inverse_transform() - - assert restored.a.shape == (1, 4, 4, 4) - torch.testing.assert_close(restored.a.data, a) - torch.testing.assert_close(restored.b.data, b) diff --git a/tests/test_parameter_range.py b/tests/test_parameter_range.py deleted file mode 100644 index c1f2487b0..000000000 --- a/tests/test_parameter_range.py +++ /dev/null @@ -1,265 +0,0 @@ -"""Tests for _ParameterRange.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio -from torchio.transforms.parameter_range import _ParameterRange - - -class TestParameterRangeParsing: - def test_scalar_is_deterministic(self) -> None: - pr = _ParameterRange(0.5) - assert pr.is_deterministic - assert pr.sample() == (0.5, 0.5, 0.5) - - def test_two_tuple_is_range(self) -> None: - """(lo, hi) → sample from U(lo, hi) per axis.""" - pr = _ParameterRange((0.8, 1.2)) - assert not pr.is_deterministic - for _ in range(50): - values = pr.sample() - assert len(values) == 3 - for v in values: - assert 0.8 <= v <= 1.2 - - def test_three_tuple_is_fixed(self) -> None: - """(a, b, c) → deterministic per-axis values.""" - pr = _ParameterRange((1.0, 2.0, 3.0)) - assert pr.is_deterministic - assert pr.sample() == (1.0, 2.0, 3.0) - - def test_six_tuple_is_per_axis_ranges(self) -> None: - """(lo0, hi0, lo1, hi1, lo2, hi2) → per-axis ranges.""" - pr = _ParameterRange((0.0, 1.0, 10.0, 20.0, 100.0, 200.0)) - assert not pr.is_deterministic - for _ in range(50): - v0, v1, v2 = pr.sample() - assert 0.0 <= v0 <= 1.0 - assert 10.0 <= v1 <= 20.0 - assert 100.0 <= v2 <= 200.0 - - def test_zero_scalar_is_deterministic(self) -> None: - pr = _ParameterRange(0.0) - assert pr.is_deterministic - assert pr.sample() == (0.0, 0.0, 0.0) - - def test_invalid_tuple_length(self) -> None: - with pytest.raises(ValueError, match="1, 2, 3, or 6"): - _ParameterRange((1.0, 2.0, 3.0, 4.0)) - - -class TestParameterRangeSampling: - def test_reproducible_with_generator(self) -> None: - pr = _ParameterRange((0.0, 100.0)) - g1 = torch.Generator().manual_seed(42) - g2 = torch.Generator().manual_seed(42) - assert pr.sample(generator=g1) == pr.sample(generator=g2) - - def test_different_seeds_differ(self) -> None: - pr = _ParameterRange((0.0, 100.0)) - g1 = torch.Generator().manual_seed(1) - g2 = torch.Generator().manual_seed(2) - # Very unlikely to be equal - assert pr.sample(generator=g1) != pr.sample(generator=g2) - - def test_scalar_sample_ignores_generator(self) -> None: - pr = _ParameterRange(5.0) - g = torch.Generator().manual_seed(42) - assert pr.sample(generator=g) == (5.0, 5.0, 5.0) - - def test_sample_1d(self) -> None: - """sample_1d returns a single float.""" - pr = _ParameterRange((0.0, 10.0)) - for _ in range(50): - v = pr.sample_1d() - assert isinstance(v, float) - assert 0.0 <= v <= 10.0 - - def test_sample_1d_deterministic(self) -> None: - pr = _ParameterRange(3.14) - assert pr.sample_1d() == 3.14 - - -class TestParameterRangeBatchedSampling: - """Batched (per-instance) sampling via the ``n`` argument.""" - - def test_sample_1d_none_returns_float(self) -> None: - """``n=None`` keeps the legacy float return type.""" - pr = _ParameterRange((0.0, 10.0)) - value = pr.sample_1d() - assert isinstance(value, float) - - def test_sample_none_returns_tuple(self) -> None: - """``n=None`` keeps the legacy 3-tuple return type.""" - pr = _ParameterRange((0.0, 10.0)) - value = pr.sample() - assert isinstance(value, tuple) - assert len(value) == 3 - - def test_sample_1d_batched_shape(self) -> None: - pr = _ParameterRange((0.0, 10.0)) - values = pr.sample_1d(n=5) - assert isinstance(values, torch.Tensor) - assert values.shape == (5,) - assert ((values >= 0.0) & (values <= 10.0)).all() - - def test_sample_batched_shape(self) -> None: - pr = _ParameterRange((0.0, 10.0)) - values = pr.sample(n=4) - assert isinstance(values, torch.Tensor) - assert values.shape == (4, 3) - - def test_batched_uniform_values_differ(self) -> None: - """Independent draws across the batch are (almost surely) distinct.""" - pr = _ParameterRange((0.0, 100.0)) - values = pr.sample_1d(n=8) - assert values.unique().numel() > 1 - - def test_batched_deterministic_is_constant(self) -> None: - pr = _ParameterRange(2.5) - values = pr.sample_1d(n=6) - assert values.shape == (6,) - torch.testing.assert_close(values, torch.full((6,), 2.5)) - - def test_batched_deterministic_per_axis(self) -> None: - pr = _ParameterRange((1.0, 2.0, 3.0)) - values = pr.sample(n=4) - expected = torch.tensor([1.0, 2.0, 3.0]).expand(4, 3) - torch.testing.assert_close(values, expected) - - def test_batched_six_tuple_per_axis_ranges(self) -> None: - pr = _ParameterRange((0.0, 1.0, 10.0, 20.0, 100.0, 200.0)) - values = pr.sample(n=16) - assert values.shape == (16, 3) - assert ((values[:, 0] >= 0.0) & (values[:, 0] <= 1.0)).all() - assert ((values[:, 1] >= 10.0) & (values[:, 1] <= 20.0)).all() - assert ((values[:, 2] >= 100.0) & (values[:, 2] <= 200.0)).all() - - def test_batched_choice(self) -> None: - pr = _ParameterRange(tio.Choice([-10.0, 0.0, 10.0])) - values = pr.sample_1d(n=32) - assert values.shape == (32,) - allowed = torch.tensor([-10.0, 0.0, 10.0]) - assert torch.isin(values, allowed).all() - - def test_batched_distribution(self) -> None: - from torch.distributions import Uniform - - pr = _ParameterRange(Uniform(5.0, 10.0)) - values = pr.sample_1d(n=10) - assert values.shape == (10,) - assert ((values >= 5.0) & (values <= 10.0)).all() - - def test_batched_reproducible_with_generator(self) -> None: - pr = _ParameterRange((0.0, 100.0)) - g1 = torch.Generator().manual_seed(42) - g2 = torch.Generator().manual_seed(42) - torch.testing.assert_close( - pr.sample_1d(n=7, generator=g1), pr.sample_1d(n=7, generator=g2) - ) - - def test_batched_n_one_returns_length_one_tensor(self) -> None: - pr = _ParameterRange((0.0, 10.0)) - values = pr.sample_1d(n=1) - assert isinstance(values, torch.Tensor) - assert values.shape == (1,) - - -class TestParameterRangeRepr: - def test_scalar_repr(self) -> None: - pr = _ParameterRange(0.5) - assert repr(pr) == "0.5" - - def test_range_repr(self) -> None: - pr = _ParameterRange((1.0, 2.0)) - assert repr(pr) == "(1.0, 2.0)" - - def test_three_tuple_repr(self) -> None: - pr = _ParameterRange((1.0, 2.0, 3.0)) - assert repr(pr) == "(1.0, 2.0, 3.0)" - - -class TestParameterRangeDistribution: - def test_distribution_not_deterministic(self) -> None: - from torch.distributions import Normal - - pr = _ParameterRange(Normal(0.0, 1.0)) - assert not pr.is_deterministic - - def test_distribution_sample_1d(self) -> None: - from torch.distributions import Uniform - - pr = _ParameterRange(Uniform(5.0, 10.0)) - for _ in range(50): - v = pr.sample_1d() - assert 5.0 <= v <= 10.0 - - def test_distribution_sample_3d(self) -> None: - from torch.distributions import Normal - - pr = _ParameterRange(Normal(0.0, 1.0)) - v0, v1, v2 = pr.sample() - assert isinstance(v0, float) - assert isinstance(v1, float) - assert isinstance(v2, float) - - def test_distribution_repr(self) -> None: - from torch.distributions import Normal - - pr = _ParameterRange(Normal(0.0, 1.0)) - assert "Normal" in repr(pr) - - -# ── Coverage gap tests ─────────────────────────────────────────────── - - -class TestChoiceEdgeCases: - def test_empty_choice_raises(self) -> None: - with pytest.raises(ValueError, match="at least one"): - tio.Choice([]) - - def test_mismatched_probabilities_raises(self) -> None: - with pytest.raises(ValueError, match="probabilities"): - tio.Choice([1, 2, 3], probabilities=[0.5, 0.5]) - - def test_repr_uniform(self) -> None: - c = tio.Choice([1.0, 2.0, 3.0]) - r = repr(c) - assert "Choice(" in r - assert "p=" not in r - - def test_repr_custom_probs(self) -> None: - c = tio.Choice([1.0, 2.0], probabilities=[0.3, 0.7]) - r = repr(c) - assert "p=" in r - - -class TestParameterRangeEdgeCases: - def test_invalid_type_raises(self) -> None: - with pytest.raises(TypeError, match="Expected float"): - _ParameterRange("bad") # type: ignore[arg-type] - - def test_ranges_for_choice_axis(self) -> None: - pr = _ParameterRange(tio.Choice([1.0, 2.0])) - lo, hi = pr._ranges[0] - assert lo == 0.0 - assert hi == 0.0 - - def test_mixed_specs_wrong_count_raises(self) -> None: - with pytest.raises(ValueError, match="Mixed per-axis"): - _ParameterRange((tio.Choice([1.0]), tio.Choice([2.0]))) - - def test_single_element_tuple(self) -> None: - pr = _ParameterRange((5.0,)) - assert pr._ranges == ((5.0, 5.0), (5.0, 5.0), (5.0, 5.0)) - - def test_invalid_axis_spec_raises(self) -> None: - with pytest.raises(TypeError, match="Per-axis spec"): - _ParameterRange(("a", "b", "c")) # type: ignore[arg-type] - - def test_invalid_tuple_length_raises(self) -> None: - with pytest.raises(ValueError, match="1, 2, 3, or 6"): - _ParameterRange((1.0, 2.0, 3.0, 4.0)) diff --git a/tests/test_patches.py b/tests/test_patches.py deleted file mode 100644 index d5dee82df..000000000 --- a/tests/test_patches.py +++ /dev/null @@ -1,320 +0,0 @@ -"""Tests for patch samplers and aggregator.""" - -from __future__ import annotations - -import torch - -import torchio as tio -from torchio.data.patch import PatchLocation - - -def _make_subject( - shape: tuple[int, int, int] = (20, 20, 20), -) -> tio.Subject: - data = torch.arange( - shape[0] * shape[1] * shape[2], - dtype=torch.float32, - ).reshape(1, *shape) - return tio.Subject(t1=tio.ScalarImage(data)) - - -def _make_subject_with_label( - shape: tuple[int, int, int] = (20, 20, 20), -) -> tio.Subject: - data = torch.rand(1, *shape) - label = torch.zeros(1, *shape, dtype=torch.float32) - label[0, 8:12, 8:12, 8:12] = 1 - return tio.Subject( - t1=tio.ScalarImage(data), - seg=tio.LabelMap(label), - ) - - -# --------------------------------------------------------------------------- -# PatchLocation -# --------------------------------------------------------------------------- - - -class TestPatchLocation: - def test_index_fin(self) -> None: - loc = PatchLocation(index=(0, 0, 0), size=(10, 10, 10)) - assert loc.index_fin == (10, 10, 10) - - def test_to_slices(self) -> None: - loc = PatchLocation(index=(5, 10, 15), size=(3, 4, 5)) - si, sj, sk = loc.to_slices() - assert si == slice(5, 8) - assert sj == slice(10, 14) - assert sk == slice(15, 20) - - def test_scaled(self) -> None: - loc = PatchLocation(index=(10, 20, 30), size=(8, 8, 8)) - scaled = loc.scaled((0.5, 0.5, 0.5)) - assert scaled.index == (5, 10, 15) - assert scaled.size == (4, 4, 4) - - -# --------------------------------------------------------------------------- -# GridSampler -# --------------------------------------------------------------------------- - - -class TestGridSampler: - def test_no_overlap(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=10) - assert len(sampler) == 8 # 2x2x2 grid - patch = sampler[0] - assert patch.t1.spatial_shape == (10, 10, 10) - - def test_with_overlap(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=12, patch_overlap=4) - assert len(sampler) > 1 - for i in range(len(sampler)): - assert sampler[i].t1.spatial_shape == (12, 12, 12) - - def test_with_padding(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler( - subject, - patch_size=12, - patch_overlap=4, - padding_mode="constant", - ) - assert len(sampler) > 1 - - def test_patch_has_location(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=10) - patch = sampler[0] - loc = patch.patch_location - assert isinstance(loc, PatchLocation) - - def test_covers_volume(self) -> None: - """All voxels should be covered by at least one patch.""" - subject = _make_subject((15, 15, 15)) - sampler = tio.GridSampler(subject, patch_size=10) - covered = torch.zeros(15, 15, 15) - for i in range(len(sampler)): - loc = sampler[i].patch_location - si, sj, sk = loc.to_slices() - covered[si, sj, sk] = 1 - assert covered.all() - - def test_works_with_dataloader(self) -> None: - from torchio.loader import SubjectsLoader - - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=10) - loader = SubjectsLoader(sampler, batch_size=4) - total = 0 - for batch in loader: - total += batch.batch_size - assert total == 8 - - -# --------------------------------------------------------------------------- -# UniformSampler -# --------------------------------------------------------------------------- - - -class TestUniformSampler: - def test_yields_correct_count(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.UniformSampler(subject, patch_size=8, num_patches=5) - patches = list(sampler) - assert len(patches) == 5 - - def test_correct_shape(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.UniformSampler(subject, patch_size=8, num_patches=1) - patch = next(iter(sampler)) - assert patch.t1.spatial_shape == (8, 8, 8) - - def test_patches_vary(self) -> None: - torch.manual_seed(42) - subject = _make_subject((50, 50, 50)) - sampler = tio.UniformSampler(subject, patch_size=8, num_patches=3) - locations = [p.patch_location.index for p in sampler] - assert len(set(locations)) > 1 - - def test_works_with_dataloader(self) -> None: - from torchio.loader import SubjectsLoader - - subject = _make_subject((20, 20, 20)) - sampler = tio.UniformSampler(subject, patch_size=8, num_patches=10) - loader = SubjectsLoader(sampler, batch_size=4) - total = sum(batch.batch_size for batch in loader) - assert total == 10 - - -# --------------------------------------------------------------------------- -# WeightedSampler -# --------------------------------------------------------------------------- - - -class TestWeightedSampler: - def test_samples_from_high_probability(self) -> None: - subject = _make_subject_with_label() - sampler = tio.WeightedSampler( - subject, - patch_size=4, - probability_map="seg", - num_patches=10, - ) - patches = list(sampler) - assert len(patches) == 10 - for p in patches: - assert p.t1.spatial_shape == (4, 4, 4) - - -# --------------------------------------------------------------------------- -# LabelSampler -# --------------------------------------------------------------------------- - - -class TestLabelSampler: - def test_samples_near_labels(self) -> None: - subject = _make_subject_with_label() - sampler = tio.LabelSampler( - subject, - patch_size=4, - label_name="seg", - num_patches=10, - ) - patches = list(sampler) - assert len(patches) == 10 - - def test_custom_probabilities(self) -> None: - subject = _make_subject_with_label() - sampler = tio.LabelSampler( - subject, - patch_size=4, - label_name="seg", - label_probabilities={0: 0.0, 1: 1.0}, - num_patches=5, - ) - patches = list(sampler) - assert len(patches) == 5 - - -# --------------------------------------------------------------------------- -# PatchAggregator -# --------------------------------------------------------------------------- - - -class TestAggregatorCrop: - def test_reconstruct_identity(self) -> None: - """Crop mode with no overlap should perfectly reconstruct.""" - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=10) - aggregator = tio.PatchAggregator( - spatial_shape=(20, 20, 20), - overlap_mode="crop", - ) - for i in range(len(sampler)): - patch = sampler[i] - loc = patch.patch_location - aggregator.add_batch( - patch.t1.data.unsqueeze(0), - [loc], - ) - output = aggregator.get_output() - torch.testing.assert_close(output, subject.t1.data) - - def test_with_overlap(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=12, patch_overlap=4) - aggregator = tio.PatchAggregator( - spatial_shape=(20, 20, 20), - overlap_mode="crop", - patch_overlap=4, - ) - for i in range(len(sampler)): - patch = sampler[i] - loc = patch.patch_location - aggregator.add_batch( - patch.t1.data.unsqueeze(0), - [loc], - ) - output = aggregator.get_output() - assert output.shape == (1, 20, 20, 20) - - -class TestAggregatorAverage: - def test_average_mode(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=12, patch_overlap=4) - aggregator = tio.PatchAggregator( - spatial_shape=(20, 20, 20), - overlap_mode="average", - ) - for i in range(len(sampler)): - patch = sampler[i] - loc = patch.patch_location - aggregator.add_batch( - patch.t1.data.unsqueeze(0), - [loc], - ) - output = aggregator.get_output() - assert output.shape == (1, 20, 20, 20) - - -class TestAggregatorHann: - def test_hann_mode(self) -> None: - subject = _make_subject((20, 20, 20)) - sampler = tio.GridSampler(subject, patch_size=12, patch_overlap=4) - aggregator = tio.PatchAggregator( - spatial_shape=(20, 20, 20), - overlap_mode="hann", - ) - for i in range(len(sampler)): - patch = sampler[i] - loc = patch.patch_location - aggregator.add_batch( - patch.t1.data.unsqueeze(0), - [loc], - ) - output = aggregator.get_output() - assert output.shape == (1, 20, 20, 20) - - -class TestAggregatorOutputShape: - def test_downsampled_output(self) -> None: - """Aggregator with output smaller than input.""" - aggregator = tio.PatchAggregator( - spatial_shape=(20, 20, 20), - overlap_mode="average", - output_shape=(10, 10, 10), - ) - loc = PatchLocation(index=(0, 0, 0), size=(20, 20, 20)) - patch = torch.rand(1, 1, 10, 10, 10) - aggregator.add_batch(patch, [loc]) - output = aggregator.get_output() - assert output.shape == (1, 10, 10, 10) - - -class TestAggregatorMultiKey: - def test_dict_output(self) -> None: - aggregator = tio.PatchAggregator( - spatial_shape=(10, 10, 10), - overlap_mode="average", - ) - loc = PatchLocation(index=(0, 0, 0), size=(10, 10, 10)) - seg = torch.rand(1, 2, 10, 10, 10) - emb = torch.rand(1, 64, 10, 10, 10) - aggregator.add_batch({"seg": seg, "emb": emb}, [loc]) - assert aggregator.get_output("seg").shape == (2, 10, 10, 10) - assert aggregator.get_output("emb").shape == (64, 10, 10, 10) - - -class TestAggregatorValidation: - def test_invalid_mode(self) -> None: - import pytest - - with pytest.raises(ValueError, match="overlap_mode"): - tio.PatchAggregator( - spatial_shape=(10, 10, 10), - overlap_mode="invalid", - ) diff --git a/tests/test_pca.py b/tests/test_pca.py deleted file mode 100644 index 65022075d..000000000 --- a/tests/test_pca.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Tests for PCA transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestPCA: - def test_reduces_channels(self) -> None: - data = torch.rand(8, 10, 10, 10) - subject = tio.Subject(emb=tio.ScalarImage(data)) - result = tio.PCA(num_components=3)(subject) - assert result.emb.data.shape[0] == 3 - - def test_output_range(self) -> None: - data = torch.randn(16, 10, 10, 10) - subject = tio.Subject(emb=tio.ScalarImage(data)) - result = tio.PCA(num_components=3, clip=True)(subject) - assert result.emb.data.min() >= 0.0 - assert result.emb.data.max() <= 1.0 - - def test_too_few_channels_raises(self) -> None: - data = torch.rand(2, 10, 10, 10) - subject = tio.Subject(emb=tio.ScalarImage(data)) - with pytest.raises(ValueError, match="channels"): - tio.PCA(num_components=5)(subject) - - def test_invalid_num_components_raises(self) -> None: - with pytest.raises(ValueError, match="num_components"): - tio.PCA(num_components=0) - - def test_no_whitening(self) -> None: - data = torch.randn(8, 10, 10, 10) - subject = tio.Subject(emb=tio.ScalarImage(data)) - result = tio.PCA(num_components=3, whiten=False, normalize=False)(subject) - assert result.emb.data.shape[0] == 3 diff --git a/tests/test_per_instance.py b/tests/test_per_instance.py deleted file mode 100644 index d5e0bdef6..000000000 --- a/tests/test_per_instance.py +++ /dev/null @@ -1,190 +0,0 @@ -"""Cross-cutting tests for per-instance batch augmentation.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _identical_batch(batch_size: int = 4) -> tio.SubjectsBatch: - data = torch.rand(1, 8, 8, 8) + 0.1 - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - -class TestCapabilityFlags: - def test_base_defaults_false(self) -> None: - transform = tio.transforms.Transform() - assert transform.supports_per_instance_params is False - assert transform.supports_per_instance_p is False - - @pytest.mark.parametrize( - "cls_kwargs", - [ - (tio.Noise, {"std": 0.1}), - (tio.Gamma, {"log_gamma": 0.1}), - ], - ) - def test_converted_intensity_opt_in(self, cls_kwargs) -> None: - cls, kwargs = cls_kwargs - transform = cls(**kwargs) - assert transform.supports_per_instance_params - assert transform.supports_per_instance_p - - def test_spatial_opts_in(self) -> None: - transform = tio.Affine(degrees=10.0) - assert transform.supports_per_instance_params - assert transform.supports_per_instance_p - - def test_resample_disables_per_instance_p(self) -> None: - # A resampling target is shape-changing, so per-element gating is off. - transform = tio.Resample(2) - assert transform.supports_per_instance_params - assert not transform.supports_per_instance_p - - -class TestUnconvertedTransforms: - """A transform that does not opt in stays batch-shared on a batch.""" - - def test_unconverted_transform_resolves_no_batch(self) -> None: - class Plain(tio.transforms.IntensityTransform): - # Does not override the capability flags, so it never samples - # per-instance parameters even for a batch. - def make_params(self, batch): - return {"n": self._resolve_n(batch)} - - def apply_transform(self, batch, params): - return batch - - batch = _identical_batch() - result = Plain()(batch) - params = result.applied_transforms[-1].params - assert params["n"] is None - assert "_batched_keys" not in params - - def test_unconverted_flags_default_false(self) -> None: - class Plain(tio.transforms.IntensityTransform): - def apply_transform(self, batch, params): - return batch - - transform = Plain() - assert not transform.supports_per_instance_params - assert not transform.supports_per_instance_p - - -class TestComposePerInstance: - def test_compose_child_is_per_instance(self) -> None: - torch.manual_seed(0) - batch = _identical_batch() - pipeline = tio.Compose([tio.Gamma(log_gamma=(0.2, 0.8))]) - result = pipeline(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], list) - assert len(set(params["log_gamma"])) > 1 - - def test_compose_respects_per_instance_false(self) -> None: - torch.manual_seed(0) - batch = _identical_batch() - pipeline = tio.Compose([tio.Gamma(log_gamma=(0.2, 0.8), per_instance=False)]) - result = pipeline(batch) - params = result.applied_transforms[-1].params - assert isinstance(params["log_gamma"], float) - - -class TestPerInstanceHistory: - def test_unbatch_slices_history(self) -> None: - torch.manual_seed(0) - batch = _identical_batch(batch_size=4) - result = tio.Gamma(log_gamma=(0.2, 0.8))(batch) - batch_log_gammas = result.applied_transforms[-1].params["log_gamma"] - for i, subject in enumerate(result.unbatch()): - trace = subject.applied_transforms[-1] - assert trace.params["log_gamma"] == batch_log_gammas[i] - assert "_batched_keys" not in trace.params - - -class TestSpatialBatchSizeValidation: - def test_mismatched_batch_size_raises(self) -> None: - torch.manual_seed(0) - batch = _identical_batch(batch_size=4) - transform = tio.Affine(degrees=(20.0, 80.0), default_pad_value=0.0) - result = transform(batch) - params = result.applied_transforms[-1].params - smaller = _identical_batch(batch_size=2) - with pytest.raises(RuntimeError, match="Per-instance spatial parameters"): - transform.apply_transform(smaller, params) - - def test_history_slice_out_of_range_raises(self) -> None: - # Slicing per-instance history for an element beyond the recorded - # batch size must fail with a clear error rather than an opaque one. - from torchio.data.batch import _slice_history - - torch.manual_seed(0) - batch = _identical_batch(batch_size=4) - result = tio.Noise(std=(0.1, 0.5))(batch) - history = result.applied_transforms - with pytest.raises(IndexError, match="batch of size 4"): - _slice_history(history, 4) - - -class TestPerInstanceDtypePreservation: - """Per-instance gating must not produce mixed-dtype batch outputs. - - Float-domain intensity transforms compute in ``float32`` internally. - When per-element gating skips some elements, the skipped ones keep the - input dtype while applied ones must be cast back, so ``torch.cat`` over - the batch does not fail on a dtype mismatch for non-``float32`` inputs. - """ - - @staticmethod - def _mixed_dtype_batch(dtype: torch.dtype, batch_size: int = 8): - data = (torch.rand(1, 8, 8, 8) + 0.5).to(dtype) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - @pytest.mark.parametrize( - "transform", - [ - tio.Ghosting(num_ghosts=4, intensity=1.0, p=0.5), - tio.Spike(num_spikes=2, intensity=1.0, p=0.5), - tio.Motion(degrees=10.0, translation=10.0, num_transforms=2, p=0.5), - ], - ) - def test_fft_transforms_preserve_float64(self, transform) -> None: - torch.manual_seed(0) - batch = self._mixed_dtype_batch(torch.float64) - result = transform(batch) - assert result.t1.data.dtype == torch.float64 - - def test_bias_field_preserves_float16(self) -> None: - torch.manual_seed(0) - batch = self._mixed_dtype_batch(torch.float16) - result = tio.BiasField(std=0.5, p=0.5)(batch) - assert result.t1.data.dtype == torch.float16 - - -class TestFullyGatedNoHistory: - def test_fully_gated_records_no_history(self) -> None: - # p=0 gates out every element, an exact no-op, so no history trace - # is recorded (and replaying it cannot trigger a spurious resample). - torch.manual_seed(0) - batch = _identical_batch(batch_size=4) - result = tio.Affine(degrees=20.0, p=0.0)(batch) - assert result.applied_transforms == [] - - def test_fully_gated_inverse_preserves_float64(self) -> None: - torch.manual_seed(0) - data = torch.rand(1, 8, 8, 8, dtype=torch.float64) - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(4)] - ) - original = batch.t1.data.clone() - result = tio.Affine(degrees=20.0, p=0.0)(batch) - restored = result.apply_inverse_transform() - assert torch.equal(restored.t1.data, original) diff --git a/tests/test_points.py b/tests/test_points.py deleted file mode 100644 index 3fd75ab04..000000000 --- a/tests/test_points.py +++ /dev/null @@ -1,220 +0,0 @@ -"""Tests for Points.""" - -from __future__ import annotations - -import copy - -import numpy as np -import pytest -import torch - -from torchio.data.points import Points - - -class TestPointsCreation: - def test_from_tensor(self): - coords = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - pts = Points(coords) - assert pts.data.shape == (2, 3) - - def test_from_numpy(self): - coords = np.array([[1.0, 2.0, 3.0]]) - pts = Points(coords) - assert isinstance(pts.data, torch.Tensor) - assert pts.data.shape == (1, 3) - - def test_with_affine(self): - coords = torch.tensor([[1.0, 2.0, 3.0]]) - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - pts = Points(coords, affine=affine) - np.testing.assert_array_equal(pts.affine.numpy(), affine) - - def test_default_affine_is_identity(self): - pts = Points(torch.tensor([[0.0, 0.0, 0.0]])) - np.testing.assert_array_equal(pts.affine.numpy(), np.eye(4)) - - def test_with_metadata(self): - pts = Points( - torch.tensor([[1.0, 2.0, 3.0]]), - metadata={"structure": "hippocampus"}, - ) - assert pts.metadata["structure"] == "hippocampus" - - def test_empty_points(self): - pts = Points(torch.zeros(0, 3)) - assert len(pts) == 0 - assert pts.data.shape == (0, 3) - - def test_wrong_shape_raises(self): - with pytest.raises(ValueError, match="N, 3"): - Points(torch.tensor([1.0, 2.0, 3.0])) - - def test_wrong_columns_raises(self): - with pytest.raises(ValueError, match="N, 3"): - Points(torch.tensor([[1.0, 2.0]])) - - def test_default_axes_ijk(self): - pts = Points(torch.tensor([[1.0, 2.0, 3.0]])) - assert pts.axes == "IJK" - - def test_custom_axes(self): - pts = Points(torch.tensor([[1.0, 2.0, 3.0]]), axes="RAS") - assert pts.axes == "RAS" - - def test_invalid_axes_raises(self): - with pytest.raises(ValueError, match="Invalid"): - Points(torch.tensor([[1.0, 2.0, 3.0]]), axes="XYZ") - - -class TestPointsProperties: - def test_len(self): - pts = Points(torch.randn(5, 3)) - assert len(pts) == 5 - - def test_num_points(self): - pts = Points(torch.randn(7, 3)) - assert pts.num_points == 7 - - -class TestPointsNewLike: - def test_new_like_preserves_affine(self): - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - pts = Points(torch.randn(3, 3), affine=affine) - new = pts.new_like(data=torch.randn(5, 3)) - np.testing.assert_array_equal(new.affine.numpy(), affine) - - def test_new_like_preserves_metadata(self): - pts = Points(torch.randn(3, 3), metadata={"label": "tumor"}) - new = pts.new_like(data=torch.randn(2, 3)) - assert new.metadata["label"] == "tumor" - - def test_new_like_with_new_affine(self): - pts = Points(torch.randn(3, 3)) - new_affine = np.diag([3.0, 3.0, 3.0, 1.0]) - new = pts.new_like(data=torch.randn(2, 3), affine=new_affine) - np.testing.assert_array_equal(new.affine.numpy(), new_affine) - - def test_new_like_preserves_axes(self): - pts = Points(torch.randn(3, 3), axes="RAS") - new = pts.new_like(data=torch.randn(2, 3)) - assert new.axes == "RAS" - - -class TestPointsToAxes: - """Test conversion between axis conventions.""" - - def test_ijk_to_kji(self): - pts = Points(torch.tensor([[1.0, 2.0, 3.0]])) - converted = pts.to_axes("KJI") - expected = torch.tensor([[3.0, 2.0, 1.0]]) - torch.testing.assert_close(converted.data, expected) - assert converted.axes == "KJI" - - def test_ijk_to_jki(self): - pts = Points(torch.tensor([[10.0, 20.0, 30.0]])) - converted = pts.to_axes("JKI") - expected = torch.tensor([[20.0, 30.0, 10.0]]) - torch.testing.assert_close(converted.data, expected) - - def test_ras_to_lpi(self): - pts = Points(torch.tensor([[10.0, 20.0, 30.0]]), axes="RAS") - converted = pts.to_axes("LPI") - expected = torch.tensor([[-10.0, -20.0, -30.0]]) - torch.testing.assert_close(converted.data, expected) - - def test_ras_to_asr(self): - pts = Points(torch.tensor([[10.0, 20.0, 30.0]]), axes="RAS") - converted = pts.to_axes("ASR") - expected = torch.tensor([[20.0, 30.0, 10.0]]) - torch.testing.assert_close(converted.data, expected) - - def test_roundtrip_ijk_kji(self): - data = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - pts = Points(data) - roundtrip = pts.to_axes("KJI").to_axes("IJK") - torch.testing.assert_close(roundtrip.data, data) - - def test_same_axes_noop(self): - data = torch.tensor([[1.0, 2.0, 3.0]]) - pts = Points(data) - converted = pts.to_axes("IJK") - torch.testing.assert_close(converted.data, data) - - def test_voxel_to_anatomical(self): - """IJK → RAS uses the affine.""" - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - pts = Points(torch.tensor([[10.0, 10.0, 10.0]]), affine=affine) - converted = pts.to_axes("RAS") - expected = torch.tensor([[20.0, 30.0, 40.0]]) - torch.testing.assert_close(converted.data, expected) - - def test_anatomical_to_voxel(self): - """RAS → IJK uses inverse affine.""" - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - pts = Points( - torch.tensor([[20.0, 30.0, 40.0]]), - axes="RAS", - affine=affine, - ) - converted = pts.to_axes("IJK") - expected = torch.tensor([[10.0, 10.0, 10.0]]) - torch.testing.assert_close(converted.data, expected, atol=1e-5, rtol=1e-5) - - def test_voxel_to_anatomical_roundtrip(self): - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - affine[:3, 3] = [10, 20, 30] - data = torch.tensor([[5.0, 10.0, 15.0]]) - pts = Points(data, affine=affine) - roundtrip = pts.to_axes("RAS").to_axes("IJK") - torch.testing.assert_close(roundtrip.data, data, atol=1e-5, rtol=1e-5) - - def test_cross_type_raises_without_matching(self): - """IJK → RAS should work using the affine.""" - pts = Points(torch.tensor([[1.0, 2.0, 3.0]])) - converted = pts.to_axes("RAS") - assert converted.axes == "RAS" - - -class TestPointsTransform: - def test_to_world(self): - # Points at voxel (1, 0, 0), affine scales by 2mm - affine = np.diag([2.0, 3.0, 4.0, 1.0]) - pts = Points(torch.tensor([[1.0, 0.0, 0.0]]), affine=affine) - world = pts.to_world() - expected = torch.tensor([[2.0, 0.0, 0.0]]) - torch.testing.assert_close(world, expected) - - def test_to_world_with_origin(self): - affine = np.eye(4) - affine[:3, 3] = [10.0, 20.0, 30.0] - pts = Points(torch.tensor([[0.0, 0.0, 0.0]]), affine=affine) - world = pts.to_world() - expected = torch.tensor([[10.0, 20.0, 30.0]]) - torch.testing.assert_close(world, expected) - - -class TestPointsRepr: - def test_repr(self): - pts = Points(torch.randn(5, 3)) - r = repr(pts) - assert "Points" in r - assert "5" in r - - def test_repr_with_axes(self): - pts = Points(torch.randn(5, 3), axes="RAS") - r = repr(pts) - assert "RAS" in r - - -class TestPointsCopy: - def test_copy(self): - pts = Points(torch.randn(3, 3), metadata={"a": 1}) - copied = copy.deepcopy(pts) - assert torch.equal(copied.data, pts.data) - copied.metadata["a"] = 2 - assert pts.metadata["a"] == 1 # original unchanged - - def test_copy_preserves_axes(self): - pts = Points(torch.randn(3, 3), axes="RAS") - copied = copy.deepcopy(pts) - assert copied.axes == "RAS" diff --git a/tests/test_queue.py b/tests/test_queue.py deleted file mode 100644 index a299185fd..000000000 --- a/tests/test_queue.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Tests for Queue.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subjects( - n: int = 4, - shape: tuple[int, int, int] = (20, 20, 20), -) -> list[tio.Subject]: - return [ - tio.Subject( - t1=tio.ScalarImage( - torch.rand(1, *shape) + i, - ), - ) - for i in range(n) - ] - - -class TestQueueBasic: - def test_yields_correct_total(self) -> None: - subjects = _make_subjects(4) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=5, - max_length=50, - ) - patches = list(queue) - assert len(patches) == 4 * 5 - - def test_shuffle_patches(self) -> None: - subjects = _make_subjects(4) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=5, - shuffle_patches=True, - ) - patches = list(queue) - assert len(patches) == 20 - - def test_no_shuffle(self) -> None: - subjects = _make_subjects(2) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=3, - shuffle_subjects=False, - shuffle_patches=False, - ) - patches = list(queue) - assert len(patches) == 6 - - -class TestQueueTransform: - def test_transform_applied(self) -> None: - subjects = _make_subjects(2) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - transform = tio.Flip(axes=0, p=1) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=2, - transform=transform, - ) - patches = list(queue) - assert len(patches) == 4 - - -class TestQueueThreaded: - def test_num_workers(self) -> None: - subjects = _make_subjects(4) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=3, - num_workers=2, - ) - patches = list(queue) - assert len(patches) == 12 - - -class TestQueueDistributed: - def test_disjoint_subsets(self) -> None: - """Two simulated ranks receive disjoint subject subsets.""" - from torch.utils.data import SubsetRandomSampler - - subjects = _make_subjects(6) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - - rank0_indices = [0, 2, 4] - rank1_indices = [1, 3, 5] - sampler0 = SubsetRandomSampler(rank0_indices) - sampler1 = SubsetRandomSampler(rank1_indices) - - queue0 = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=2, - shuffle_subjects=False, - subject_sampler=sampler0, - ) - queue1 = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=2, - shuffle_subjects=False, - subject_sampler=sampler1, - ) - - patches0 = list(queue0) - patches1 = list(queue1) - assert len(patches0) == 3 * 2 - assert len(patches1) == 3 * 2 - - def test_shuffle_with_sampler_raises(self) -> None: - import pytest - from torch.utils.data import SubsetRandomSampler - - subjects = _make_subjects(2) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - with pytest.raises(ValueError, match="shuffle_subjects"): - tio.Queue( - subjects, - patch_sampler=sampler, - shuffle_subjects=True, - subject_sampler=SubsetRandomSampler([0, 1]), - ) - - -class TestQueueMemory: - def test_max_memory(self) -> None: - subjects = _make_subjects(2, shape=(10, 10, 10)) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - max_length=100, - ) - # 1 channel * 8^3 voxels * 4 bytes * 100 patches - assert queue.max_memory == 1 * 512 * 4 * 100 - - def test_max_memory_pretty(self) -> None: - subjects = _make_subjects(2, shape=(10, 10, 10)) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - max_length=100, - ) - assert "KiB" in queue.max_memory_pretty - - -class TestQueueProperties: - def test_num_subjects(self) -> None: - subjects = _make_subjects(5) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue(subjects, patch_sampler=sampler) - assert queue.num_subjects == 5 - - def test_patches_per_epoch(self) -> None: - subjects = _make_subjects(4) - sampler = tio.UniformSampler(subjects[0], patch_size=8) - queue = tio.Queue( - subjects, - patch_sampler=sampler, - patches_per_volume=10, - ) - assert queue.patches_per_epoch == 40 diff --git a/tests/test_remap_labels.py b/tests/test_remap_labels.py deleted file mode 100644 index 5bd96adef..000000000 --- a/tests/test_remap_labels.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Tests for RemapLabels transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestRemapLabels: - def test_basic_remap(self) -> None: - subject = _make_subject() - result = tio.RemapLabels({1: 10, 2: 20})(subject) - assert 10 in result.seg.data.unique().tolist() - assert 20 in result.seg.data.unique().tolist() - assert 1 not in result.seg.data.unique().tolist() - assert 2 not in result.seg.data.unique().tolist() - - def test_merge_labels(self) -> None: - subject = _make_subject() - result = tio.RemapLabels({2: 1})(subject) - assert 2 not in result.seg.data.unique().tolist() - assert 1 in result.seg.data.unique().tolist() - - def test_swap_labels(self) -> None: - subject = _make_subject() - original_1_count = (subject.seg.data == 1).sum().item() - original_2_count = (subject.seg.data == 2).sum().item() - result = tio.RemapLabels({1: 2, 2: 1})(subject) - assert (result.seg.data == 1).sum().item() == original_2_count - assert (result.seg.data == 2).sum().item() == original_1_count - - def test_leaves_unlisted_labels(self) -> None: - subject = _make_subject() - original_0_count = (subject.seg.data == 0).sum().item() - result = tio.RemapLabels({1: 10})(subject) - assert (result.seg.data == 0).sum().item() == original_0_count - - def test_inverse(self) -> None: - subject = _make_subject() - original = subject.seg.data.clone() - transformed = tio.RemapLabels({1: 10, 2: 20})(subject) - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(restored.seg.data, original) - - def test_leaves_scalar_unchanged(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.RemapLabels({1: 10})(subject) - torch.testing.assert_close(result.t1.data, original) diff --git a/tests/test_remote_loading.py b/tests/test_remote_loading.py deleted file mode 100644 index b62c12a41..000000000 --- a/tests/test_remote_loading.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Tests for remote/URL/file-like Image loading via fsspec.""" - -from __future__ import annotations - -import io -from pathlib import Path - -import nibabel as nib -import numpy as np -import pytest - -import torchio as tio - - -@pytest.fixture -def nifti_path(tmp_path: Path) -> Path: - path = tmp_path / "test.nii.gz" - data = np.random.rand(8, 8, 8).astype(np.float32) - nib.save(nib.Nifti1Image(data, np.eye(4)), path) - return path - - -class TestLocalPath: - def test_str_path(self, nifti_path: Path) -> None: - image = tio.ScalarImage(str(nifti_path)) - assert image.shape == (1, 8, 8, 8) - - def test_path_object(self, nifti_path: Path) -> None: - image = tio.ScalarImage(nifti_path) - assert image.shape == (1, 8, 8, 8) - - -class TestFileUrl: - def test_file_uri(self, nifti_path: Path) -> None: - """file:// URIs should work via fsspec.""" - uri = f"file://{nifti_path}" - image = tio.ScalarImage(uri) - assert image.shape == (1, 8, 8, 8) - - -class TestFileLike: - def test_bytes_io(self, nifti_path: Path) -> None: - """BytesIO should work by materializing to a temp file.""" - buf = io.BytesIO(nifti_path.read_bytes()) - image = tio.ScalarImage(buf, suffix=".nii.gz") - assert image.shape == (1, 8, 8, 8) - - def test_open_file(self, nifti_path: Path) -> None: - """An open binary file should work.""" - with open(nifti_path, "rb") as f: - image = tio.ScalarImage(f, suffix=".nii.gz") - assert image.shape == (1, 8, 8, 8) - - -class TestFsspec: - def test_local_fsspec(self, nifti_path: Path) -> None: - """fsspec local filesystem should work.""" - import fsspec - - of = fsspec.open(str(nifti_path), mode="rb") - image = tio.ScalarImage(of) - assert image.shape == (1, 8, 8, 8) diff --git a/tests/test_remote_zarr.py b/tests/test_remote_zarr.py deleted file mode 100644 index 4b5f4cffa..000000000 --- a/tests/test_remote_zarr.py +++ /dev/null @@ -1,250 +0,0 @@ -"""Tests for remote NIfTI-Zarr streaming (lazy reads without full download).""" - -from __future__ import annotations - -from pathlib import Path -from unittest.mock import MagicMock -from unittest.mock import patch - -import nibabel as nib -import numpy as np -import pytest -import torch - -import torchio as tio -from torchio.data.io import is_remote_nifti_zarr - -# ── Detection helper ──────────────────────────────────────────────── - - -class TestIsRemoteNiftiZarr: - """Test the is_remote_nifti_zarr() detection function.""" - - @pytest.mark.parametrize( - "uri", - [ - "az://container/image.nii.zarr", - "s3://bucket/image.nii.zarr", - "gs://bucket/image.nii.zarr", - "https://example.com/image.nii.zarr", - "abfs://container/path/to/image.nii.zarr", - ], - ) - def test_remote_zarr_detected(self, uri: str) -> None: - assert is_remote_nifti_zarr(uri) is True - - @pytest.mark.parametrize( - "uri", - [ - "az://container/image.nii.gz", - "s3://bucket/image.nii", - "/local/path/image.nii.zarr", - "relative/path/image.nii.zarr", - "az://container/image.nrrd", - "https://example.com/image.nii.gz", - ], - ) - def test_non_remote_zarr_not_detected(self, uri: str) -> None: - assert is_remote_nifti_zarr(uri) is False - - def test_trailing_slash_handled(self) -> None: - assert is_remote_nifti_zarr("az://container/image.nii.zarr/") is True - - -# ── Image construction ────────────────────────────────────────────── - - -class TestRemoteZarrImageConstruction: - """Verify that remote .nii.zarr URIs are stored without downloading.""" - - def test_remote_zarr_stores_uri(self) -> None: - """The URI should be preserved, not fetched to a temp file.""" - uri = "az://my-container/brain.nii.zarr" - with patch("torchio.data.io._fetch_remote") as mock_fetch: - image = tio.ScalarImage(uri) - mock_fetch.assert_not_called() - assert image._remote_zarr_uri == uri - assert image._path is None - - def test_non_zarr_remote_still_fetched(self) -> None: - """A remote .nii.gz should still be downloaded (existing behaviour).""" - uri = "az://my-container/brain.nii.gz" - with patch("torchio.data.io._fetch_remote") as mock_fetch: - mock_fetch.return_value = Path("/tmp/fake.nii.gz") - _image = tio.ScalarImage(uri) - mock_fetch.assert_called_once() - - def test_local_zarr_not_treated_as_remote(self, tmp_path: Path) -> None: - """A local .nii.zarr path should go through the normal code path.""" - try: - import niizarr - except ImportError: - pytest.skip("nifti-zarr not installed") - data = np.random.rand(8, 8, 8).astype(np.float32) - nii = nib.Nifti1Image(data, np.eye(4)) - nii_path = tmp_path / "test.nii" - nib.save(nii, nii_path) - zarr_path = tmp_path / "test.nii.zarr" - niizarr.nii2zarr(str(nii_path), str(zarr_path)) - - image = tio.ScalarImage(zarr_path) - assert image._remote_zarr_uri is None - assert image._path == zarr_path - - -# ── Lazy backend from remote URI ──────────────────────────────────── - - -class TestRemoteZarrBackend: - """Verify that _ensure_backend passes the URI to ZarrBackend.""" - - def test_ensure_backend_uses_uri(self) -> None: - """ZarrBackend should receive the remote URI, not a local path.""" - uri = "az://container/brain.nii.zarr" - - with patch("torchio.data.io._fetch_remote"): - image = tio.ScalarImage(uri) - - with patch( - "torchio.data.backends.ZarrBackend.__init__", return_value=None - ) as mock_init: - # Provide a mock backend so _ensure_backend finishes - mock_backend = MagicMock() - mock_backend.shape = (1, 8, 8, 8) - mock_backend.affine = torch.eye(4, dtype=torch.float64) - - def side_effect(path, **kwargs): - image._backend = mock_backend - - mock_init.side_effect = side_effect - image._ensure_backend() - mock_init.assert_called_once_with(uri, affine=None) - - def test_reader_kwargs_forwarded_to_backend(self) -> None: - """store_opt and other kwargs should reach ZarrBackend.""" - uri = "az://container/brain.nii.zarr" - kwargs = {"account_name": "myaccount", "account_key": "secret"} - - with patch("torchio.data.io._fetch_remote"): - image = tio.ScalarImage(uri, reader_kwargs=kwargs) - - with patch( - "torchio.data.backends.ZarrBackend.__init__", return_value=None - ) as mock_init: - mock_backend = MagicMock() - mock_backend.shape = (1, 8, 8, 8) - mock_backend.affine = torch.eye(4, dtype=torch.float64) - - def side_effect(path, **kw): - image._backend = mock_backend - - mock_init.side_effect = side_effect - image._ensure_backend() - mock_init.assert_called_once_with(uri, affine=None, **kwargs) - - def test_shape_via_remote_backend(self) -> None: - """image.shape should work through the remote ZarrBackend.""" - uri = "az://container/brain.nii.zarr" - - with patch("torchio.data.io._fetch_remote"): - image = tio.ScalarImage(uri) - - mock_backend = MagicMock() - mock_backend.shape = (1, 32, 32, 32) - mock_backend.affine = torch.eye(4, dtype=torch.float64) - image._backend = mock_backend - - assert image.shape == (1, 32, 32, 32) - - def test_load_via_remote_backend(self) -> None: - """image.load() should materialize from the remote ZarrBackend.""" - uri = "az://container/brain.nii.zarr" - - with patch("torchio.data.io._fetch_remote"): - image = tio.ScalarImage(uri) - - tensor = torch.randn(1, 8, 8, 8) - mock_backend = MagicMock() - mock_backend.shape = (1, 8, 8, 8) - mock_backend.affine = torch.eye(4, dtype=torch.float64) - mock_backend.to_tensor.return_value = tensor - image._backend = mock_backend - - image.load() - assert image._data is not None - assert image._data.shape == (1, 8, 8, 8) - - -# ── repr and deepcopy ─────────────────────────────────────────────── - - -class TestRemoteZarrReprAndCopy: - def test_repr_shows_uri(self) -> None: - uri = "az://container/brain.nii.zarr" - with patch("torchio.data.io._fetch_remote"): - image = tio.ScalarImage(uri) - - mock_backend = MagicMock() - mock_backend.shape = (1, 8, 8, 8) - mock_backend.affine = torch.eye(4, dtype=torch.float64) - mock_backend.dtype = np.dtype(np.float32) - mock_backend.to_tensor.return_value = torch.randn(1, 8, 8, 8) - image._backend = mock_backend - - r = repr(image) - assert "az://container/brain.nii.zarr" in r - - def test_deepcopy_preserves_uri(self) -> None: - import copy - - uri = "az://container/brain.nii.zarr" - with patch("torchio.data.io._fetch_remote"): - image = tio.ScalarImage(uri) - - tensor = torch.randn(1, 8, 8, 8) - mock_backend = MagicMock() - mock_backend.shape = (1, 8, 8, 8) - mock_backend.affine = torch.eye(4, dtype=torch.float64) - mock_backend.to_tensor.return_value = tensor - image._backend = mock_backend - image._data = tensor - - with patch("torchio.data.io._fetch_remote"): - copied = copy.deepcopy(image) - assert copied._remote_zarr_uri == uri - - -# ── End-to-end with local Zarr store ──────────────────────────────── - - -class TestRemoteZarrEndToEnd: - """Integration test using a local .nii.zarr but going through remote path.""" - - @pytest.fixture - def zarr_path(self, tmp_path: Path) -> Path: - try: - import niizarr - except ImportError: - pytest.skip("nifti-zarr not installed") - data = np.arange(16**3, dtype=np.float32).reshape(16, 16, 16) - nii = nib.Nifti1Image(data, np.eye(4)) - nii_path = tmp_path / "test.nii" - nib.save(nii, nii_path) - zarr_path = tmp_path / "test.nii.zarr" - niizarr.nii2zarr(str(nii_path), str(zarr_path)) - return zarr_path - - def test_zarr_backend_accepts_kwargs(self, zarr_path: Path) -> None: - """ZarrBackend should forward kwargs to niizarr.zarr2nii().""" - from torchio.data.backends import ZarrBackend - - backend = ZarrBackend(str(zarr_path)) - assert backend.shape == (1, 16, 16, 16) - - def test_slice_without_full_load(self, zarr_path: Path) -> None: - """Slicing a remote-style zarr should not load the full tensor.""" - from torchio.data.backends import ZarrBackend - - backend = ZarrBackend(str(zarr_path)) - roi = backend[:, 4:8, 4:8, 4:8] - assert roi.shape == (1, 4, 4, 4) diff --git a/tests/test_remove_labels.py b/tests/test_remove_labels.py deleted file mode 100644 index d27f65289..000000000 --- a/tests/test_remove_labels.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Tests for RemoveLabels transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestRemoveLabels: - def test_removes_specified_labels(self) -> None: - subject = _make_subject() - result = tio.RemoveLabels([2])(subject) - assert 2 not in result.seg.data.unique().tolist() - assert 1 in result.seg.data.unique().tolist() - - def test_removes_multiple_labels(self) -> None: - subject = _make_subject() - result = tio.RemoveLabels([1, 2])(subject) - unique = result.seg.data.unique().tolist() - assert unique == [0.0] - - def test_custom_background(self) -> None: - subject = _make_subject() - result = tio.RemoveLabels([1], background_label=99)(subject) - assert 1 not in result.seg.data.unique().tolist() - assert 99 in result.seg.data.unique().tolist() - - def test_leaves_scalar_unchanged(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.RemoveLabels([1])(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_noop_when_label_absent(self) -> None: - subject = _make_subject() - original = subject.seg.data.clone() - result = tio.RemoveLabels([99])(subject) - torch.testing.assert_close(result.seg.data, original) diff --git a/tests/test_reorient.py b/tests/test_reorient.py deleted file mode 100644 index 3c7d33966..000000000 --- a/tests/test_reorient.py +++ /dev/null @@ -1,258 +0,0 @@ -"""Tests for the Reorient transform.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio -from torchio.data.affine import AffineMatrix - - -def _make_subject( - shape: tuple[int, int, int] = (10, 20, 30), - orientation: str = "RAS", -) -> tio.Subject: - """Create a subject with a known orientation.""" - import nibabel as nib - - # Build an affine for the desired orientation - target_ornt = nib.orientations.axcodes2ornt(tuple(orientation)) - # Start from identity and apply the orientation - base_affine = nib.orientations.inv_ornt_aff(target_ornt, shape) - affine = AffineMatrix(base_affine) - data = torch.arange(shape[0] * shape[1] * shape[2], dtype=torch.float32) - data = data.reshape(1, *shape) - image = tio.ScalarImage(data, affine=affine) - return tio.Subject(t1=image) - - -# --------------------------------------------------------------------------- -# Basic reorientation -# --------------------------------------------------------------------------- - - -class TestReorientBasic: - def test_no_op_when_already_target(self) -> None: - subject = _make_subject(orientation="RAS") - result = tio.Reorient(orientation="RAS")(subject) - assert result.t1.affine.orientation == ("R", "A", "S") - torch.testing.assert_close(result.t1.data, subject.t1.data) - - def test_ras_to_las(self) -> None: - subject = _make_subject(orientation="RAS") - result = tio.Reorient(orientation="LAS")(subject) - assert result.t1.affine.orientation == ("L", "A", "S") - - def test_ras_to_pls(self) -> None: - subject = _make_subject(orientation="RAS") - result = tio.Reorient(orientation="PLS")(subject) - assert result.t1.affine.orientation == ("P", "L", "S") - - def test_default_is_ras(self) -> None: - subject = _make_subject(orientation="LAS") - result = tio.Reorient()(subject) - assert result.t1.affine.orientation == ("R", "A", "S") - - def test_shape_changes_with_permutation(self) -> None: - subject = _make_subject(shape=(10, 20, 30), orientation="RAS") - result = tio.Reorient(orientation="ASR")(subject) - assert result.t1.affine.orientation == ("A", "S", "R") - # Spatial shape should be a permutation of (10, 20, 30) - assert sorted(result.t1.spatial_shape) == [10, 20, 30] - - -# --------------------------------------------------------------------------- -# Round-trip: reorient and back preserves data -# --------------------------------------------------------------------------- - - -class TestRoundTrip: - def test_ras_to_las_and_back(self) -> None: - subject = _make_subject(orientation="RAS") - original_data = subject.t1.data.clone() - to_las = tio.Reorient(orientation="LAS") - to_ras = tio.Reorient(orientation="RAS") - result = to_ras(to_las(subject)) - torch.testing.assert_close(result.t1.data, original_data) - - def test_ras_to_spl_and_back(self) -> None: - subject = _make_subject(orientation="RAS") - original_data = subject.t1.data.clone() - result = tio.Reorient(orientation="RAS")( - tio.Reorient(orientation="SPL")(subject) - ) - torch.testing.assert_close(result.t1.data, original_data) - - -# --------------------------------------------------------------------------- -# Validation -# --------------------------------------------------------------------------- - - -class TestValidation: - def test_invalid_length(self) -> None: - with pytest.raises(ValueError, match="3-letter"): - tio.Reorient(orientation="RA") - - def test_invalid_characters(self) -> None: - with pytest.raises(ValueError, match="distinct"): - tio.Reorient(orientation="XYZ") - - def test_missing_axis(self) -> None: - with pytest.raises(ValueError, match="each axis"): - tio.Reorient(orientation="RAA") - - def test_case_insensitive(self) -> None: - transform = tio.Reorient(orientation="ras") - assert transform.orientation == "RAS" - - -# --------------------------------------------------------------------------- -# All images in subject -# --------------------------------------------------------------------------- - - -class TestAllImages: - def test_reorients_all_images(self) -> None: - import nibabel as nib - - ornt = nib.orientations.axcodes2ornt(tuple("RAS")) - affine = AffineMatrix(nib.orientations.inv_ornt_aff(ornt, (10, 20, 30))) - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 20, 30), affine=affine), - seg=tio.LabelMap( - torch.randint(0, 3, (1, 10, 20, 30)), - affine=affine, - ), - ) - result = tio.Reorient(orientation="LAS")(subject) - assert result.t1.affine.orientation == ("L", "A", "S") - assert result.seg.affine.orientation == ("L", "A", "S") - - -# --------------------------------------------------------------------------- -# Invertibility -# --------------------------------------------------------------------------- - - -class TestInvertibility: - def test_invertible(self) -> None: - assert tio.Reorient().invertible - - def test_inverse_restores_shape(self) -> None: - subject = _make_subject(shape=(10, 20, 30), orientation="RAS") - original_data = subject.t1.data.clone() - transformed = tio.Reorient(orientation="SPL")(subject) - restored = transformed.apply_inverse_transform() - assert restored.t1.spatial_shape == (10, 20, 30) - torch.testing.assert_close(restored.t1.data, original_data) - - -# --------------------------------------------------------------------------- -# Input types -# --------------------------------------------------------------------------- - - -class TestInputTypes: - def test_accepts_image(self) -> None: - import nibabel as nib - - ornt = nib.orientations.axcodes2ornt(tuple("RAS")) - affine = AffineMatrix(nib.orientations.inv_ornt_aff(ornt, (10, 20, 30))) - image = tio.ScalarImage( - torch.rand(1, 10, 20, 30), - affine=affine, - ) - result = tio.Reorient(orientation="LAS")(image) - assert isinstance(result, tio.Image) - assert result.affine.orientation == ("L", "A", "S") - - def test_accepts_subject(self) -> None: - subject = _make_subject(orientation="RAS") - result = tio.Reorient(orientation="LAS")(subject) - assert isinstance(result, tio.Subject) - - -# --------------------------------------------------------------------------- -# Batch mode -# --------------------------------------------------------------------------- - - -class TestBatch: - def test_batch_reorient(self) -> None: - import nibabel as nib - - from torchio.data.batch import SubjectsBatch - - ornt = nib.orientations.axcodes2ornt(tuple("RAS")) - affine = AffineMatrix(nib.orientations.inv_ornt_aff(ornt, (10, 20, 30))) - subjects = [ - tio.Subject( - t1=tio.ScalarImage( - torch.rand(1, 10, 20, 30), - affine=affine, - ), - ) - for _ in range(3) - ] - batch = SubjectsBatch.from_subjects(subjects) - result = tio.Reorient(orientation="LAS")(batch) - assert result.t1.data.shape == (3, 1, 10, 20, 30) - - -# --------------------------------------------------------------------------- -# Probability -# --------------------------------------------------------------------------- - - -class TestProbability: - def test_p_zero_is_no_op(self) -> None: - subject = _make_subject(orientation="RAS") - result = tio.Reorient(orientation="LAS", p=0)(subject) - assert result.t1.affine.orientation == ("R", "A", "S") - - -# --------------------------------------------------------------------------- -# World coordinate preservation -# --------------------------------------------------------------------------- - - -class TestWorldCoordinates: - """Reorientation must not change the physical position of voxels.""" - - @staticmethod - def _world_bbox(image: tio.Image) -> torch.Tensor: - s = image.spatial_shape - corners = torch.tensor( - [[0, 0, 0, 1], [s[0] - 1, s[1] - 1, s[2] - 1, 1]], - dtype=torch.float64, - ) - aff = torch.as_tensor(image.affine.data, dtype=torch.float64) - return ((aff @ corners.T).T)[:, :3] - - def _assert_bbox_preserved( - self, - original: tio.Image, - reoriented: tio.Image, - ) -> None: - orig = self._world_bbox(original).sort(dim=0).values - new = self._world_bbox(reoriented).sort(dim=0).values - torch.testing.assert_close(orig, new, atol=1e-5, rtol=0) - - @pytest.mark.parametrize("target", ["PSR", "LPS", "SLA", "AIR", "RAS"]) - def test_bbox_preserved_identity_spacing(self, target: str) -> None: - subject = _make_subject(shape=(10, 12, 14), orientation="RAS") - result = tio.Reorient(target)(subject) - self._assert_bbox_preserved(subject.t1, result.t1) - - def test_bbox_preserved_anisotropic_spacing(self) -> None: - import numpy as np - - affine = np.diag([2.0, 0.5, 1.5, 1.0]) - affine[:3, 3] = [10, 20, 30] - data = torch.rand(1, 10, 12, 14) - img = tio.ScalarImage(data, affine=AffineMatrix(affine)) - subject = tio.Subject(t1=img) - result = tio.Reorient("PSR")(subject) - self._assert_bbox_preserved(subject.t1, result.t1) diff --git a/tests/test_repr_html.py b/tests/test_repr_html.py deleted file mode 100644 index 639af7424..000000000 --- a/tests/test_repr_html.py +++ /dev/null @@ -1,189 +0,0 @@ -"""Tests for _repr_html_ on Image and Subject.""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio -from torchio.data.bboxes import BoundingBoxes -from torchio.data.bboxes import BoundingBoxFormat -from torchio.data.points import Points - - -class TestImageReprHtml: - def test_returns_html_string(self) -> None: - image = tio.ScalarImage(torch.rand(1, 16, 16, 16)) - html = image._repr_html_() - assert isinstance(html, str) - assert " None: - image = tio.ScalarImage(torch.rand(1, 8, 10, 12)) - html = image._repr_html_() - assert "(8, 10, 12)" in html - - def test_contains_spacing(self) -> None: - image = tio.ScalarImage(torch.rand(1, 16, 16, 16)) - html = image._repr_html_() - assert "1.00" in html - - def test_contains_orientation(self) -> None: - image = tio.ScalarImage(torch.rand(1, 16, 16, 16)) - html = image._repr_html_() - assert "RAS" in html - - def test_contains_class_name(self) -> None: - image = tio.ScalarImage(torch.rand(1, 16, 16, 16)) - html = image._repr_html_() - assert "ScalarImage" in html - - def test_label_map_class_name(self) -> None: - image = tio.LabelMap(torch.randint(0, 3, (1, 8, 8, 8))) - html = image._repr_html_() - assert "LabelMap" in html - - def test_contains_dtype(self) -> None: - image = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - html = image._repr_html_() - assert "float32" in html - - def test_shows_points(self) -> None: - pts = Points(torch.tensor([[1.0, 2.0, 3.0]])) - image = tio.ScalarImage( - torch.rand(1, 16, 16, 16), - points={"landmarks": pts}, - ) - html = image._repr_html_() - assert "landmarks" in html - assert "1 point" in html - - def test_shows_bounding_boxes(self) -> None: - boxes = BoundingBoxes( - torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]), - format=BoundingBoxFormat.IJKIJK, - ) - image = tio.ScalarImage( - torch.rand(1, 16, 16, 16), - bounding_boxes={"tumors": boxes}, - ) - html = image._repr_html_() - assert "tumors" in html - assert "1 box" in html - - def test_shows_memory(self) -> None: - image = tio.ScalarImage(torch.rand(1, 16, 16, 16)) - html = image._repr_html_() - assert "Memory" in html - - def test_unloaded_shows_dtype_and_memory(self, tmp_path) -> None: - """Unloaded image shows dtype and memory from header.""" - import nibabel as nib - import numpy as np - - path = tmp_path / "test.nii.gz" - nib.save( - nib.Nifti1Image(np.zeros((8, 8, 8)), np.eye(4)), - path, - ) - image = tio.ScalarImage(path) - html = image._repr_html_() - assert "dtype" in html - assert "Memory" in html - - -class TestSubjectReprHtml: - def test_returns_html_string(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - ) - html = subject._repr_html_() - assert isinstance(html, str) - assert " None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 16, 16, 16))), - ) - html = subject._repr_html_() - assert "t1" in html - assert "seg" in html - - def test_contains_image_types(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 16, 16, 16))), - ) - html = subject._repr_html_() - assert "ScalarImage" in html - assert "LabelMap" in html - - def test_shows_metadata(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - age=42, - diagnosis="healthy", - ) - html = subject._repr_html_() - assert "age" in html - assert "42" in html - assert "diagnosis" in html - assert "healthy" in html - - def test_shows_points(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - landmarks=Points(torch.rand(5, 3)), - ) - html = subject._repr_html_() - assert "landmarks" in html - assert "5 points" in html - - def test_shows_bboxes(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - tumors=BoundingBoxes( - torch.rand(3, 6), - format=BoundingBoxFormat.IJKIJK, - ), - ) - html = subject._repr_html_() - assert "tumors" in html - assert "3 boxes" in html - - def test_shows_shapes(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 10, 12)), - ) - html = subject._repr_html_() - assert "(1, 8, 10, 12)" in html - - def test_metadata_only_subject(self) -> None: - subject = tio.Subject(age=42, name="test") - html = subject._repr_html_() - assert "age" in html - assert "42" in html - - -class TestPlotInteractive: - def test_plot_interactive_returns_widget(self) -> None: - pytest.importorskip("ipyniivue") - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - widget = img.plot_interactive() - assert widget is not None - assert hasattr(widget, "_repr_mimebundle_") - - def test_plot_interactive_radiological(self) -> None: - pytest.importorskip("ipyniivue") - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - widget = img.plot_interactive() - assert widget.opts.is_radiological_convention is True - - def test_plot_interactive_no_ipyniivue_raises(self, monkeypatch) -> None: - from torchio.external import imports - - monkeypatch.setattr(imports, "find_spec", lambda m: None) - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - with pytest.raises(ImportError, match="ipyniivue"): - img.plot_interactive() diff --git a/tests/test_resize.py b/tests/test_resize.py deleted file mode 100644 index 4706326ea..000000000 --- a/tests/test_resize.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Tests for Resize transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestResize: - def test_resize_to_target(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Resize(5)(subject) - assert result.t1.data.shape[1:] == (5, 5, 5) - - def test_resize_anisotropic(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Resize((8, 6, 4))(subject) - assert result.t1.data.shape[1:] == (8, 6, 4) - - def test_resize_preserves_dtype(self) -> None: - subject = _make_subject() - result = tio.Resize(5)(subject) - assert result.t1.data.dtype == subject.t1.data.dtype - - def test_resize_labels_nearest(self) -> None: - subject = _make_subject() - result = tio.Resize(5)(subject) - unique = result.seg.data.unique().tolist() - for v in unique: - assert v == int(v) - - def test_resize_with_labels(self) -> None: - subject = _make_subject() - result = tio.Resize(5)(subject) - assert result.seg.data.shape[1:] == (5, 5, 5) diff --git a/tests/test_sequential_labels.py b/tests/test_sequential_labels.py deleted file mode 100644 index 91ddd7a4e..000000000 --- a/tests/test_sequential_labels.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Tests for SequentialLabels transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestSequentialLabels: - def test_basic_sequential(self) -> None: - seg = torch.zeros(1, 5, 5, 5, dtype=torch.float32) - seg[0, 0:2, :, :] = 5 - seg[0, 3:5, :, :] = 10 - subject = tio.Subject(seg=tio.LabelMap(seg)) - result = tio.SequentialLabels()(subject) - unique = sorted(result.seg.data.unique().tolist()) - assert unique == [0.0, 1.0, 2.0] - - def test_already_sequential(self) -> None: - subject = _make_subject() - result = tio.SequentialLabels()(subject) - torch.testing.assert_close(result.seg.data, subject.seg.data) - - def test_inverse(self) -> None: - seg = torch.zeros(1, 5, 5, 5, dtype=torch.float32) - seg[0, 0:2, :, :] = 5 - seg[0, 3:5, :, :] = 10 - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 5, 5, 5)), - seg=tio.LabelMap(seg), - ) - original = subject.seg.data.clone() - transformed = tio.SequentialLabels()(subject) - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(restored.seg.data, original) - - def test_leaves_scalar_unchanged(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - result = tio.SequentialLabels()(subject) - torch.testing.assert_close(result.t1.data, original) diff --git a/tests/test_some_of.py b/tests/test_some_of.py deleted file mode 100644 index d66eb7e86..000000000 --- a/tests/test_some_of.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Tests for SomeOf transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject() -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - - -class TestSomeOf: - def test_applies_subset(self) -> None: - subject = _make_subject() - transform = tio.SomeOf( - [tio.Flip(axes=(0,)), tio.Gamma(log_gamma=0.0)], - num_transforms=1, - ) - result = transform(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_all_transforms(self) -> None: - subject = _make_subject() - transform = tio.SomeOf( - [tio.Flip(axes=(0,)), tio.Gamma(log_gamma=0.0)], - num_transforms=2, - ) - result = transform(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_range_num_transforms(self) -> None: - subject = _make_subject() - transform = tio.SomeOf( - [ - tio.Flip(axes=(0,)), - tio.Gamma(log_gamma=0.0), - tio.Blur(std=0.0), - ], - num_transforms=(1, 3), - ) - result = transform(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_history_recorded(self) -> None: - subject = _make_subject() - transform = tio.SomeOf( - [tio.Flip(axes=(0,))], - num_transforms=1, - ) - result = transform(subject) - assert len(result.applied_transforms) > 0 - - -class TestSomeOfPerInstance: - def _batch(self, batch_size: int = 16) -> tio.SubjectsBatch: - data = torch.rand(1, 8, 8, 8) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_element_subsets_differ(self) -> None: - torch.manual_seed(0) - batch = self._batch() - transform = tio.SomeOf( - [tio.Gamma(log_gamma=0.5), tio.Flip(axes=(0,)), tio.Noise(std=0.3)], - num_transforms=2, - ) - result = transform(batch) - counts = {len(subject.applied_transforms) for subject in result.unbatch()} - # Each element applies exactly two transforms. - assert counts == {2} - name_sets = { - tuple(sorted(t.name for t in subject.applied_transforms)) - for subject in result.unbatch() - } - assert len(name_sets) > 1 - - def test_per_instance_false_is_batch_wide(self) -> None: - torch.manual_seed(0) - batch = self._batch(batch_size=8) - transform = tio.SomeOf( - [tio.Gamma(log_gamma=0.5), tio.Flip(axes=(0,)), tio.Noise(std=0.3)], - num_transforms=2, - per_instance=False, - ) - result = transform(batch) - name_sets = { - tuple(sorted(t.name for t in subject.applied_transforms)) - for subject in result.unbatch() - } - assert len(name_sets) == 1 - - def test_single_subject_unaffected(self) -> None: - subject = _make_subject() - transform = tio.SomeOf( - [tio.Gamma(log_gamma=0.5), tio.Flip(axes=(0,))], - num_transforms=1, - ) - result = transform(subject) - assert len(result.applied_transforms) == 1 - - -class TestSomeOfCopy: - def test_does_not_mutate_input(self) -> None: - subject = _make_subject() - snapshot = subject.t1.data.clone() - tio.SomeOf([tio.Gamma(log_gamma=0.5)], num_transforms=1)(subject) - torch.testing.assert_close(subject.t1.data, snapshot) - - def test_restores_child_copy_flag(self) -> None: - child = tio.Gamma(log_gamma=0.5) - assert child.copy is True - tio.SomeOf([child], num_transforms=1)(_make_subject()) - assert child.copy is True - - def test_children_applied_without_copy(self) -> None: - seen: list[bool] = [] - - class _Spy(tio.IntensityTransform): - def apply_transform(self, batch, params): - seen.append(self.copy) - return batch - - tio.SomeOf([_Spy()], num_transforms=1)(_make_subject()) - assert seen == [False] diff --git a/tests/test_spatial.py b/tests/test_spatial.py deleted file mode 100644 index 4c98800cc..000000000 --- a/tests/test_spatial.py +++ /dev/null @@ -1,1217 +0,0 @@ -"""Tests for spatial transforms.""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest -import torch - -import torchio as tio -from torchio import AffineMatrix -from torchio.transforms import Affine as AffineTransform -from torchio.transforms.spatial.spatial import _border_mean -from torchio.transforms.spatial.spatial import _build_sampling_grid -from torchio.transforms.spatial.spatial import _check_folding -from torchio.transforms.spatial.spatial import _check_shared_space -from torchio.transforms.spatial.spatial import _compute_channel_pad_value -from torchio.transforms.spatial.spatial import _is_spacing_list -from torchio.transforms.spatial.spatial import _is_spacing_tuple -from torchio.transforms.spatial.spatial import _is_target_space_tuple -from torchio.transforms.spatial.spatial import _normalize_parameter_value -from torchio.transforms.spatial.spatial import _otsu_threshold -from torchio.transforms.spatial.spatial import _parse_center -from torchio.transforms.spatial.spatial import _parse_control_points -from torchio.transforms.spatial.spatial import _parse_default_pad_value -from torchio.transforms.spatial.spatial import _parse_interpolation -from torchio.transforms.spatial.spatial import _parse_locked_borders -from torchio.transforms.spatial.spatial import _parse_num_control_points -from torchio.transforms.spatial.spatial import _parse_spacing -from torchio.transforms.spatial.spatial import _parse_target_space_tuple -from torchio.transforms.spatial.spatial import _prepare_fill_value -from torchio.transforms.spatial.spatial import _to_nonnegative_parameter_range -from torchio.transforms.spatial.spatial import _to_positive_range -from torchio.transforms.spatial.spatial import _validate_isotropic - - -def _make_subject( - shape: tuple[int, int, int] = (11, 11, 11), - spacing: tuple[float, float, float] = (1.0, 1.0, 1.0), -) -> tio.Subject: - data = torch.arange( - shape[0] * shape[1] * shape[2], - dtype=torch.float32, - ).reshape(1, *shape) - label = torch.zeros(1, *shape, dtype=torch.float32) - label[ - 0, - shape[0] // 4 : shape[0] // 4 + 2, - shape[1] // 4 : shape[1] // 4 + 2, - shape[2] // 4 : shape[2] // 4 + 2, - ] = 1 - affine = np.diag([spacing[0], spacing[1], spacing[2], 1.0]) - return tio.Subject( - t1=tio.ScalarImage(data, affine=affine), - seg=tio.LabelMap(label, affine=affine), - ) - - -class TestSpatial: - def test_identity(self) -> None: - subject = _make_subject() - transformed = tio.Spatial()(subject) - - torch.testing.assert_close(transformed.t1.data, subject.t1.data) - torch.testing.assert_close(transformed.seg.data, subject.seg.data) - np.testing.assert_allclose( - transformed.t1.affine.numpy(), - subject.t1.affine.numpy(), - ) - - def test_affine_first_changes_result(self) -> None: - subject = _make_subject() - control_points = torch.zeros(5, 5, 5, 3) - control_points[2, 2, 2, 0] = 2.0 - - kwargs = { - "scales": 1.0, - "degrees": (0.0, 0.0, 45.0), - "translation": 0.0, - "control_points": control_points, - "default_pad_value": 0.0, - "default_pad_label": 0.0, - } - - first = tio.Spatial(affine_first=True, **kwargs)(subject) - second = tio.Spatial(affine_first=False, **kwargs)(subject) - - assert not torch.allclose(first.t1.data, second.t1.data) - - def test_2d_suppresses_out_of_plane(self) -> None: - data = torch.rand(1, 8, 8, 1) - subject = tio.Subject(t1=tio.ScalarImage(data)) - transformed = AffineTransform( - scales=1.0, - degrees=(0.0, 0.0, 10.0), - translation=0.0, - center="image", - default_pad_value=0.0, - )(subject) - assert transformed.t1.spatial_shape[-1] == 1 - - -class TestResample: - def test_spacing_target_changes_shape_and_affine(self) -> None: - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - transformed = tio.Resample(2)(subject) - - assert transformed.t1.spatial_shape == (6, 6, 6) - assert transformed.seg.spatial_shape == (6, 6, 6) - np.testing.assert_allclose(transformed.t1.spacing, (2.0, 2.0, 2.0)) - np.testing.assert_allclose( - transformed.t1.affine.numpy(), - transformed.seg.affine.numpy(), - ) - - def test_named_image_target_uses_reference_space(self) -> None: - reference = tio.ScalarImage( - torch.ones(1, 6, 6, 6), - affine=np.diag([2.0, 2.0, 2.0, 1.0]), - ) - moving = tio.ScalarImage( - torch.ones(1, 12, 12, 12), - affine=np.diag([1.0, 1.0, 1.0, 1.0]), - ) - subject = tio.Subject(t1=reference, t2=moving) - - transformed = tio.Resample("t1")(subject) - - assert transformed.t2.spatial_shape == transformed.t1.spatial_shape - np.testing.assert_allclose( - transformed.t2.affine.numpy(), - transformed.t1.affine.numpy(), - ) - - def test_inverse_restores_geometry(self) -> None: - subject = _make_subject(shape=(12, 12, 12)) - transformed = tio.Resample(2)(subject) - - restored = transformed.apply_inverse_transform() - - assert restored.t1.spatial_shape == subject.t1.spatial_shape - np.testing.assert_allclose( - restored.t1.affine.numpy(), - subject.t1.affine.numpy(), - ) - - def test_target_image_object(self) -> None: - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - reference = tio.ScalarImage( - torch.ones(1, 6, 6, 6), - affine=np.diag([2.0, 2.0, 2.0, 1.0]), - ) - transformed = tio.Resample(target=reference)(subject) - assert transformed.t1.spatial_shape == (6, 6, 6) - - def test_target_tuple_spacing(self) -> None: - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - transformed = tio.Resample(target=(2.0, 2.0, 2.0))(subject) - assert transformed.t1.spatial_shape == (6, 6, 6) - - def test_target_shape_affine_pair(self) -> None: - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - target_affine = np.diag([2.0, 2.0, 2.0, 1.0]) - transformed = tio.Resample(target=((6, 6, 6), target_affine))(subject) - assert transformed.t1.spatial_shape == (6, 6, 6) - - def test_target_file_path(self, tmp_path: Path) -> None: - import nibabel as nib - - ref_data = np.zeros((6, 6, 6), dtype=np.float32) - ref_affine = np.diag([2.0, 2.0, 2.0, 1.0]) - nib.save(nib.Nifti1Image(ref_data, ref_affine), tmp_path / "ref.nii") - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - transformed = tio.Resample(target=str(tmp_path / "ref.nii"))(subject) - assert transformed.t1.spatial_shape == (6, 6, 6) - - def test_target_ndarray_spacing(self) -> None: - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - transformed = tio.Resample(target=np.array([2.0, 2.0, 2.0]))(subject) - assert transformed.t1.spatial_shape == (6, 6, 6) - - -class TestResampleTargetRange: - """The spacing form of `target` accepts random ranges and distributions. - - See discussion #1472: `Resample(target=range)` should sample a spacing at - apply time, consistent with how `degrees`/`scales` accept ranges. - """ - - def test_deterministic_forms_unchanged(self) -> None: - # Scalar and 3-tuple spacings stay deterministic. - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - np.testing.assert_allclose(tio.Resample(2)(subject).t1.spacing, (2.0, 2.0, 2.0)) - np.testing.assert_allclose( - tio.Resample(target=(2.0, 3.0, 4.0))(subject).t1.spacing, - (2.0, 3.0, 4.0), - ) - - def test_two_tuple_uniform_range_within_bounds(self) -> None: - subject = _make_subject(shape=(40, 40, 40), spacing=(1.0, 1.0, 1.0)) - torch.manual_seed(0) - for _ in range(10): - spacing = tio.Resample(target=(2.0, 4.0))(subject).t1.spacing - assert all(2.0 <= s <= 4.0 for s in spacing) - - def test_six_tuple_per_axis_ranges_within_bounds(self) -> None: - # The exact form from discussion #1472. - subject = _make_subject(shape=(40, 40, 40), spacing=(1.0, 1.0, 1.0)) - bounds = [(2.0, 4.0), (2.0, 4.0), (3.0, 6.0)] - torch.manual_seed(0) - for _ in range(10): - spacing = tio.Resample(target=(2, 4, 2, 4, 3, 6))(subject).t1.spacing - for value, (low, high) in zip(spacing, bounds, strict=True): - assert low <= value <= high - - def test_choice_target(self) -> None: - from torchio.transforms.parameter_range import Choice - - subject = _make_subject(shape=(40, 40, 40), spacing=(1.0, 1.0, 1.0)) - torch.manual_seed(0) - for _ in range(10): - spacing = tio.Resample(target=Choice([2.0, 4.0]))(subject).t1.spacing - assert all(s in (2.0, 4.0) for s in spacing) - - def test_distribution_target(self) -> None: - from torch.distributions import Uniform - - subject = _make_subject(shape=(40, 40, 40), spacing=(1.0, 1.0, 1.0)) - torch.manual_seed(0) - spacing = tio.Resample(target=Uniform(2.0, 4.0))(subject).t1.spacing - assert all(2.0 <= s <= 4.0 for s in spacing) - - def test_two_tuple_not_treated_as_shape_affine(self) -> None: - # Regression: a 2-tuple of numbers must be a spacing range, not an - # attempted (shape, affine) pair (which used to raise). - subject = _make_subject(shape=(20, 20, 20), spacing=(1.0, 1.0, 1.0)) - transformed = tio.Resample(target=(2.0, 4.0))(subject) - assert all(2.0 <= s <= 4.0 for s in transformed.t1.spacing) - - def test_nonpositive_range_raises(self) -> None: - subject = _make_subject(shape=(12, 12, 12), spacing=(1.0, 1.0, 1.0)) - with pytest.raises(ValueError, match="positive"): - tio.Resample(target=(-2.0, -1.0))(subject) - - def test_seed_reproducible(self) -> None: - subject = _make_subject(shape=(40, 40, 40), spacing=(1.0, 1.0, 1.0)) - torch.manual_seed(123) - first = tio.Resample(target=(2, 4, 2, 4, 3, 6))(subject).t1.spacing - torch.manual_seed(123) - second = tio.Resample(target=(2, 4, 2, 4, 3, 6))(subject).t1.spacing - np.testing.assert_allclose(first, second) - - def test_antialias_smooths_before_downsample(self) -> None: - subject = _make_subject(shape=(20, 20, 20), spacing=(0.5, 0.5, 0.5)) - no_aa = tio.Resample(2)(subject) - with_aa = tio.Resample(2, antialias=True)(subject) - assert with_aa.t1.spatial_shape == no_aa.t1.spatial_shape - # Antialiased result should be smoother (lower high-freq energy) - assert not torch.allclose(with_aa.t1.data, no_aa.t1.data) - - def test_antialias_skips_label_maps(self) -> None: - subject = _make_subject(shape=(20, 20, 20), spacing=(0.5, 0.5, 0.5)) - transformed = tio.Resample(2, antialias=True)(subject) - unique = set(transformed.seg.data.unique().tolist()) - assert unique <= {0.0, 1.0} - - def test_antialias_noop_on_upsample(self) -> None: - subject = _make_subject(shape=(6, 6, 6), spacing=(2.0, 2.0, 2.0)) - no_aa = tio.Resample(1)(subject) - with_aa = tio.Resample(1, antialias=True)(subject) - torch.testing.assert_close(with_aa.t1.data, no_aa.t1.data) - - -class TestAffine: - def test_transform_changes_data(self) -> None: - subject = _make_subject() - transform = AffineTransform( - scales=1.0, - degrees=(0.0, 0.0, 90.0), - translation=0.0, - center="image", - default_pad_value=0.0, - default_pad_label=0.0, - ) - - transformed = transform(subject) - - assert not torch.allclose(transformed.t1.data, subject.t1.data) - np.testing.assert_allclose( - transformed.t1.affine.numpy(), - subject.t1.affine.numpy(), - ) - - def test_inverse_restores_geometry(self) -> None: - subject = _make_subject() - transform = AffineTransform( - scales=(1.1, 0.9, 1.0), - degrees=(0.0, 0.0, 20.0), - translation=(1.0, -2.0, 0.5), - center="image", - default_pad_value=0.0, - default_pad_label=0.0, - ) - - restored = transform(subject).apply_inverse_transform() - - assert restored.t1.spatial_shape == subject.t1.spatial_shape - np.testing.assert_allclose( - restored.t1.affine.numpy(), - subject.t1.affine.numpy(), - ) - - def test_inverse_leaves_excluded_images_untouched(self) -> None: - # The forward pass records the images it transformed; the inverse - # must only resample those, so an excluded image (here the label - # map) is bit-for-bit identical after a forward + inverse round - # trip. - subject = _make_subject() - transform = AffineTransform( - degrees=(0.0, 0.0, 20.0), - translation=(1.0, -2.0, 0.5), - default_pad_value=0.0, - include=["t1"], - ) - original_seg = subject.seg.data.clone() - - result = transform(subject) - assert result.applied_transforms[-1].params["selected_images"] == ["t1"] - assert torch.equal(result.seg.data, original_seg) - - restored = result.apply_inverse_transform() - assert torch.equal(restored.seg.data, original_seg) - subject = _make_subject() - transform = AffineTransform( - scales=(0.9, 1.1), - isotropic=True, - degrees=0.0, - translation=0.0, - default_pad_value=0.0, - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_center_origin(self) -> None: - subject = _make_subject() - transform = AffineTransform( - scales=1.0, - degrees=(0.0, 0.0, 10.0), - translation=0.0, - center="origin", - default_pad_value=0.0, - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_choice_degrees(self) -> None: - subject = _make_subject() - transform = AffineTransform( - scales=1.0, - degrees=tio.Choice([-90, 0, 90, 180]), - translation=0.0, - default_pad_value=0.0, - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_per_axis_mixed_specs(self) -> None: - subject = _make_subject() - transform = AffineTransform( - scales=1.0, - degrees=(0, 0, tio.Choice([-90, 0, 90])), - translation=0.0, - default_pad_value=0.0, - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_distribution_parameter(self) -> None: - from torch.distributions import Normal - - subject = _make_subject() - transform = AffineTransform( - scales=1.0, - degrees=Normal(0.0, 5.0), - translation=0.0, - default_pad_value=0.0, - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - -class TestSpatialPerInstance: - def _identical_batch( - self, - batch_size: int = 4, - shape: tuple[int, int, int] = (12, 12, 12), - ) -> tio.SubjectsBatch: - data = torch.arange( - shape[0] * shape[1] * shape[2], - dtype=torch.float32, - ).reshape(1, *shape) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_rotations_differ(self) -> None: - torch.manual_seed(0) - batch = self._identical_batch() - transform = AffineTransform(degrees=(20.0, 80.0), default_pad_value=0.0) - result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["affine_matrix"]) == batch.batch_size - data = result.t1.data - assert not torch.allclose(data[0], data[1]) - assert not torch.allclose(data[1], data[2]) - - def test_per_instance_cubic_interpolation(self) -> None: - # Exercises the per-sample high-order interpolation path - # (interpol.grid_pull) with batched grids. - pytest.importorskip("interpol") - torch.manual_seed(0) - batch = self._identical_batch() - transform = AffineTransform( - degrees=(20.0, 80.0), - default_pad_value=0.0, - image_interpolation="cubic", - ) - result = transform(batch) - assert result.t1.data.shape == batch.t1.data.shape - assert not torch.allclose(result.t1.data[0], result.t1.data[1]) - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._identical_batch() - transform = AffineTransform( - degrees=(20.0, 80.0), - default_pad_value=0.0, - per_instance=False, - ) - result = transform(batch) - data = result.t1.data - torch.testing.assert_close(data[0], data[1]) - torch.testing.assert_close(data[1], data[2]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = _make_subject() - result = AffineTransform(degrees=(20.0, 80.0), default_pad_value=0.0)(subject) - params = result.applied_transforms[-1].params - assert "_batched_keys" not in params - - def test_per_instance_inverse_restores_geometry(self) -> None: - torch.manual_seed(0) - batch = self._identical_batch() - transform = AffineTransform( - scales=(0.9, 1.1), - degrees=(20.0, 80.0), - translation=(-2.0, 2.0), - default_pad_value=0.0, - ) - result = transform(batch) - restored = result.apply_inverse_transform() - assert restored.t1.data.shape == batch.t1.data.shape - for affine in restored.t1.affines: - np.testing.assert_allclose( - affine.numpy(), - batch.t1.affines[0].numpy(), - atol=1e-5, - ) - - def test_per_instance_p_gates_some_elements(self) -> None: - torch.manual_seed(0) - batch = self._identical_batch(batch_size=64) - original = batch.t1.data.clone() - transform = AffineTransform(degrees=(40.0, 80.0), default_pad_value=0.0, p=0.5) - result = transform(batch) - changed = [ - not torch.allclose(result.t1.data[i], original[i]) - for i in range(batch.batch_size) - ] - assert any(changed) - assert not all(changed) - - def test_per_instance_p_masked_elements_unchanged(self) -> None: - torch.manual_seed(0) - batch = self._identical_batch(batch_size=32) - original = batch.t1.data.clone() - transform = AffineTransform(degrees=(40.0, 80.0), default_pad_value=0.0, p=0.5) - result = transform(batch) - for i, subject in enumerate(result.unbatch()): - changed = not torch.allclose(subject.t1.data, original[i]) - has_history = len(subject.applied_transforms) == 1 - assert changed == has_history - - def test_per_instance_elastic_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._identical_batch() - transform = tio.ElasticDeformation( - num_control_points=5, - max_displacement=(1.0, 3.0), - ) - result = transform(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["control_points"]) == batch.batch_size - data = result.t1.data - assert not torch.allclose(data[0], data[1]) - assert not torch.allclose(data[1], data[2]) - - def test_fully_gated_noop_preserves_per_sample_affines(self) -> None: - # When every element is gated out (p=0) and there is no target - # resampling, the transform must be a true no-op: the data and the - # distinct per-sample affines must be left untouched rather than - # rebuilt from an identity grid using the first element's affine. - torch.manual_seed(0) - subjects = [] - for index in range(4): - affine = np.eye(4) - affine[0, 3] = float(index * 10) - image = tio.ScalarImage(torch.rand(1, 8, 8, 8), affine=affine) - subjects.append(tio.Subject(t1=image)) - batch = tio.SubjectsBatch.from_subjects(subjects) - original_data = batch.t1.data.clone() - original_affines = [affine.numpy().copy() for affine in batch.t1.affines] - - result = AffineTransform(degrees=20.0, p=0.0)(batch) - - assert torch.equal(result.t1.data, original_data) - for original, new in zip( - original_affines, - result.t1.affines, - strict=True, - ): - assert np.allclose(original, new.numpy()) - - def test_partially_gated_elements_are_exact_noops(self) -> None: - # In a mixed batch (0 < p < 1), gated-out elements must be exact, - # bit-for-bit no-ops even for float64 inputs: the per-sample - # resample runs an identity grid in float32, so those rows are - # restored from the input afterwards. - torch.manual_seed(0) - data = torch.rand(1, 8, 8, 8, dtype=torch.float64) - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(8)] - ) - original = batch.t1.data.clone() - torch.manual_seed(1) - result = AffineTransform(degrees=30.0, p=0.5)(batch) - assert result.t1.data.dtype == torch.float64 - exact = [ - torch.equal(result.t1.data[i], original[i]) for i in range(batch.batch_size) - ] - changed = [ - not torch.allclose(result.t1.data[i], original[i], atol=1e-6) - for i in range(batch.batch_size) - ] - # Every element is either an exact no-op or genuinely augmented. - assert all( - is_exact ^ is_changed - for is_exact, is_changed in zip(exact, changed, strict=True) - ) - assert any(exact) - assert any(changed) - - -class TestElasticDeformation: - def test_accepts_tensor_control_points(self) -> None: - subject = _make_subject() - control_points = torch.zeros(5, 5, 5, 3) - control_points[2, 2, 2, 0] = 2.0 - - transformed = tio.ElasticDeformation(control_points=control_points)(subject) - - assert not torch.allclose(transformed.t1.data, subject.t1.data) - - def test_label_interpolation_preserves_label_values(self) -> None: - subject = _make_subject() - transformed = AffineTransform( - scales=(1.1, 1.0, 1.0), - degrees=(0.0, 0.0, 15.0), - translation=0.0, - center="image", - default_pad_value=0.0, - default_pad_label=0.0, - )(subject) - - unique = set(transformed.seg.data.unique().tolist()) - assert unique <= {0.0, 1.0} - - def test_sampled_max_displacement(self) -> None: - subject = _make_subject() - transformed = tio.ElasticDeformation( - max_displacement=2.0, - num_control_points=5, - locked_borders=1, - )(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_inverse_with_elastic(self) -> None: - subject = _make_subject() - control_points = torch.zeros(5, 5, 5, 3) - control_points[2, 2, 2, 0] = 1.0 - transformed = tio.ElasticDeformation( - control_points=control_points, - )(subject) - restored = transformed.apply_inverse_transform() - assert restored.t1.spatial_shape == subject.t1.spatial_shape - - def test_folding_warning(self) -> None: - with pytest.warns(RuntimeWarning, match="folding"): - _check_folding( - control_points=np.ones((5, 5, 5, 3)), - max_displacement=(100.0, 100.0, 100.0), - shape=(10, 10, 10), - spacing=np.array([1.0, 1.0, 1.0]), - ) - - -class TestPadValue: - def test_pad_value_mean(self) -> None: - subject = _make_subject() - transform = tio.Spatial( - degrees=(0.0, 0.0, 30.0), - scales=1.0, - translation=0.0, - default_pad_value="mean", - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_pad_value_otsu(self) -> None: - subject = _make_subject() - transform = tio.Spatial( - degrees=(0.0, 0.0, 30.0), - scales=1.0, - translation=0.0, - default_pad_value="otsu", - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_pad_value_numeric_nonzero(self) -> None: - subject = _make_subject() - transform = tio.Spatial( - degrees=(0.0, 0.0, 30.0), - scales=1.0, - translation=0.0, - default_pad_value=42.0, - ) - transformed = transform(subject) - assert transformed.t1.spatial_shape == subject.t1.spatial_shape - - def test_compute_channel_pad_minimum(self) -> None: - t = torch.arange(27, dtype=torch.float32).reshape(3, 3, 3) - assert _compute_channel_pad_value(t, "minimum") == 0.0 - - def test_compute_channel_pad_mean(self) -> None: - t = torch.ones(3, 3, 3) - val = _compute_channel_pad_value(t, "mean") - assert abs(val - 1.0) < 1e-5 - - def test_compute_channel_pad_otsu(self) -> None: - t = torch.ones(3, 3, 3) - val = _compute_channel_pad_value(t, "otsu") - assert isinstance(val, float) - - -class TestBorderMeanAndOtsu: - def test_border_mean_no_otsu(self) -> None: - t = torch.ones(5, 5, 5) * 3.0 - assert abs(_border_mean(t, filter_otsu=False) - 3.0) < 1e-5 - - def test_border_mean_with_otsu(self) -> None: - t = torch.ones(5, 5, 5) - result = _border_mean(t, filter_otsu=True) - assert isinstance(result, float) - - def test_otsu_threshold_empty(self) -> None: - assert _otsu_threshold(torch.tensor([])) == 0.0 - - def test_otsu_threshold_basic(self) -> None: - values = torch.tensor([0.0, 0.0, 0.0, 10.0, 10.0, 10.0]) - threshold = _otsu_threshold(values) - assert 0.0 <= threshold <= 10.0 - - -class TestValidation: - def test_locked_borders_invalid(self) -> None: - with pytest.raises(ValueError, match="locked_borders"): - _parse_locked_borders(5) - - def test_locked_borders_2_with_4_control_points(self) -> None: - with pytest.raises(ValueError, match="identity elastic field"): - tio.Spatial( - num_control_points=4, - locked_borders=2, - ) - - def test_invalid_default_pad_label(self) -> None: - with pytest.raises(TypeError, match="default_pad_label"): - tio.Spatial(default_pad_label="bad") # type: ignore[arg-type] - - def test_negative_scales(self) -> None: - with pytest.raises(ValueError, match="strictly positive"): - _to_positive_range(-1.0) - - def test_negative_max_displacement(self) -> None: - with pytest.raises(ValueError, match="non-negative"): - _to_nonnegative_parameter_range(-1.0) - - def test_isotropic_with_per_axis(self) -> None: - with pytest.raises(ValueError, match="isotropic"): - _validate_isotropic((1.0, 1.0, 1.0), isotropic=True) - - def test_parse_num_control_points_too_small(self) -> None: - with pytest.raises(ValueError, match="greater than 3"): - _parse_num_control_points(3) - - def test_parse_control_points_bad_shape(self) -> None: - with pytest.raises(ValueError, match="n_i, n_j, n_k, 3"): - _parse_control_points(torch.zeros(5, 5, 2)) - - def test_parse_control_points_axis_too_small(self) -> None: - with pytest.raises(ValueError, match="at least 4"): - _parse_control_points(torch.zeros(3, 5, 5, 3)) - - def test_parse_interpolation_invalid(self) -> None: - with pytest.raises(ValueError, match="not supported"): - _parse_interpolation("bicubic") # type: ignore[arg-type] - - def test_parse_interpolation_int(self) -> None: - assert _parse_interpolation(3) == "cubic" - assert _parse_interpolation(0) == "nearest" - - def test_parse_interpolation_int_invalid(self) -> None: - with pytest.raises(ValueError, match="not supported"): - _parse_interpolation(99) - - def test_parse_interpolation_not_string(self) -> None: - with pytest.raises(TypeError, match="string or int"): - _parse_interpolation(42.5) # type: ignore[arg-type] - - def test_parse_default_pad_value_invalid_string(self) -> None: - with pytest.raises(ValueError, match="minimum"): - _parse_default_pad_value("bad") # type: ignore[arg-type] - - def test_parse_center_invalid(self) -> None: - with pytest.raises(ValueError, match="center"): - _parse_center("middle") # type: ignore[arg-type] - - def test_parse_spacing_negative(self) -> None: - with pytest.raises(ValueError, match="positive"): - _parse_spacing(-1.0) - - def test_parse_spacing_wrong_length(self) -> None: - with pytest.raises(ValueError, match="3 values"): - _parse_spacing((1.0, 2.0)) - - def test_parse_spacing_ndarray_wrong_size(self) -> None: - with pytest.raises(ValueError, match="3 values"): - _parse_spacing(np.array([1.0, 2.0])) - - def test_parse_spacing_tuple_3(self) -> None: - result = _parse_spacing((1.0, 2.0, 3.0)) - assert result == (1.0, 2.0, 3.0) - - def test_parse_spacing_ndarray(self) -> None: - result = _parse_spacing(np.array([1.0, 2.0, 3.0])) - assert result == (1.0, 2.0, 3.0) - - def test_target_not_understood(self) -> None: - subject = _make_subject() - with pytest.raises(ValueError, match="not understood"): - tio.Resample(target=object())(subject) # type: ignore[arg-type] - - def test_target_unknown_string(self) -> None: - subject = _make_subject() - with pytest.raises(ValueError, match="Unknown target"): - tio.Resample(target="nonexistent_image")(subject) - - def test_shared_space_shape_mismatch(self) -> None: - from torchio.data.batch import ImagesBatch - - batch_a = ImagesBatch( - torch.rand(1, 1, 8, 8, 8), - [AffineMatrix()], - ) - batch_b = ImagesBatch( - torch.rand(1, 1, 10, 10, 10), - [AffineMatrix()], - ) - with pytest.raises(RuntimeError, match="shape"): - _check_shared_space( - {"a": batch_a, "b": batch_b}, - (8, 8, 8), - AffineMatrix(), - ) - - def test_shared_space_affine_mismatch(self) -> None: - from torchio.data.batch import ImagesBatch - - batch_a = ImagesBatch( - torch.rand(1, 1, 8, 8, 8), - [AffineMatrix()], - ) - batch_b = ImagesBatch( - torch.rand(1, 1, 8, 8, 8), - [AffineMatrix(np.diag([2.0, 2.0, 2.0, 1.0]))], - ) - with pytest.raises(RuntimeError, match="same affine"): - _check_shared_space( - {"a": batch_a, "b": batch_b}, - (8, 8, 8), - AffineMatrix(), - ) - - -class TestTypeGuards: - def test_is_spacing_tuple(self) -> None: - assert _is_spacing_tuple((1.0, 2.0, 3.0)) - assert not _is_spacing_tuple((1.0, 2.0)) - assert not _is_spacing_tuple([1.0, 2.0, 3.0]) - - def test_is_spacing_list(self) -> None: - assert _is_spacing_list([1.0, 2.0, 3.0]) - assert not _is_spacing_list((1.0, 2.0, 3.0)) - - def test_is_target_space_tuple(self) -> None: - assert _is_target_space_tuple(((6, 6, 6), np.eye(4))) - assert not _is_target_space_tuple((1.0, 2.0, 3.0)) - # A 2-tuple of plain numbers is a spacing range, not a (shape, affine). - assert not _is_target_space_tuple((2.0, 4.0)) - - -class TestEdgeCases: - """Tests for defensive branches and edge cases.""" - - def test_exclude_all_images(self) -> None: - """make_params returns empty when all images are excluded (line 210).""" - subject = _make_subject() - transform = tio.Spatial(exclude=["t1", "seg"]) - from torchio.data.batch import SubjectsBatch - - batch = SubjectsBatch.from_subjects([subject]) - params = transform.make_params(batch) - assert params == {"selected_images": []} - - def test_apply_transform_empty_selection(self) -> None: - """apply_transform returns batch unchanged (line 278).""" - subject = _make_subject() - from torchio.data.batch import SubjectsBatch - - batch = SubjectsBatch.from_subjects([subject]) - transform = tio.Spatial() - original_data = batch.images["t1"].data.clone() - result = transform.apply_transform(batch, {"selected_images": []}) - torch.testing.assert_close(result.images["t1"].data, original_data) - - def test_inverse_missing_original_space(self) -> None: - """inverse raises when original is None (lines 330-331).""" - transform = tio.Spatial() - params = { - "affine_matrix": None, - "control_points": None, - "original": None, - "affine_first": True, - "image_interpolation": "linear", - "label_interpolation": "nearest", - "default_pad_value": "minimum", - "default_pad_label": 0, - } - with pytest.raises(RuntimeError, match="original output space"): - transform.inverse(params) - - def test_apply_spatial_empty_names(self) -> None: - """_apply_spatial_to_batch returns early for empty names (line 556).""" - from torchio.transforms.spatial.spatial import _apply_spatial_to_batch - - subject = _make_subject() - from torchio.data.batch import SubjectsBatch - - batch = SubjectsBatch.from_subjects([subject]) - original = batch.images["t1"].data.clone() - _apply_spatial_to_batch( - batch=batch, - image_names=[], - target_space=None, - affine_matrix=None, - control_points=None, - max_displacement=None, - affine_first=True, - image_interpolation="linear", - label_interpolation="nearest", - antialias=False, - default_pad_value="minimum", - default_pad_label=0.0, - ) - torch.testing.assert_close(batch.images["t1"].data, original) - - def test_build_grid_elastic_no_max_displacement(self) -> None: - """_build_sampling_grid computes max_displacement from field (line 742).""" - control_points = torch.zeros(5, 5, 5, 3) - control_points[2, 2, 2, 0] = 0.5 - grid = _build_sampling_grid( - input_shape=(11, 11, 11), - input_affine=AffineMatrix(), - output_shape=(11, 11, 11), - output_affine=AffineMatrix(), - affine_matrix=None, - control_points=control_points, - max_displacement=None, - affine_first=True, - device=torch.device("cpu"), - ) - assert grid.shape == (11, 11, 11, 3) - - def test_batch_fill_value_bad_type(self) -> None: - """_batch_fill_value raises TypeError for non-str non-number (lines 907-911).""" - from torchio.data.batch import ImagesBatch - from torchio.transforms.spatial.spatial import _batch_fill_value - - batch = ImagesBatch( - torch.rand(1, 1, 4, 4, 4), - [AffineMatrix()], - ) - with pytest.raises(TypeError, match="string or number"): - _batch_fill_value( - batch, - default_pad_value=object(), # type: ignore[arg-type] - default_pad_label=0.0, - ) - - def test_prepare_fill_value_multidim(self) -> None: - """_prepare_fill_value returns a >1D tensor unchanged (line 945).""" - fill = torch.ones(1, 2, 1, 1, 1) - ref = torch.zeros(1, 2, 3, 3, 3) - result = _prepare_fill_value(fill, ref) - assert result is not None - assert result.shape == (1, 2, 1, 1, 1) - - def test_compute_channel_pad_unknown(self) -> None: - """_compute_channel_pad_value raises for unknown strategy (lines 959-960).""" - with pytest.raises(ValueError, match="Unknown"): - _compute_channel_pad_value( - torch.ones(3, 3, 3), - "bad_strategy", # type: ignore[arg-type] - ) - - def test_parse_target_space_tuple_wrong_length(self) -> None: - """_parse_target_space_tuple raises for non-3 shape (lines 1386-1387).""" - with pytest.raises(ValueError, match="length 3"): - _parse_target_space_tuple([6, 6], np.eye(4)) - - def test_normalize_parameter_value_distribution(self) -> None: - """_normalize_parameter_value returns Distribution unchanged (line 1548).""" - from torch.distributions import Normal - - dist = Normal(0.0, 1.0) - assert _normalize_parameter_value(dist) is dist - - -class TestExports: - def test_root_exports_expose_transform_and_matrix(self) -> None: - assert hasattr(tio, "Spatial") - assert hasattr(tio, "Resample") - assert hasattr(tio, "ElasticDeformation") - assert tio.Affine is AffineTransform - assert tio.AffineMatrix is AffineMatrix - - -# --------------------------------------------------------------------------- -# High-order interpolation (torch-interpol) -# --------------------------------------------------------------------------- - - -class TestHighOrderInterpolation: - def test_cubic_produces_different_result_from_linear(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 16, 16, 16))) - linear = tio.Affine( - degrees=10, - image_interpolation="linear", - )(subject) - cubic = tio.Affine( - degrees=10, - image_interpolation="cubic", - )(subject) - # Same params won't be sampled, so compare shapes at least - assert linear.t1.data.shape == cubic.t1.data.shape - - def test_cubic_resample(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 16, 16, 16))) - result = tio.Resample( - target=2.0, - image_interpolation="cubic", - )(subject) - assert result.t1.data.shape[1:] == (8, 8, 8) - - def test_quadratic_interpolation(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 10, 10, 10))) - result = tio.Affine( - degrees=5, - image_interpolation="quadratic", - )(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_int_order_3(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 10, 10, 10))) - result = tio.Affine( - degrees=5, - image_interpolation=3, - )(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_order_0_uses_fast_path(self) -> None: - """Nearest interpolation should still work via F.grid_sample.""" - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 10, 10, 10))) - result = tio.Affine( - degrees=5, - image_interpolation="nearest", - )(subject) - assert result.t1.data.shape == subject.t1.data.shape - - -def _sphere_label( - n: int = 64, - radius: float = 20.0, - value: float = 1.0, -) -> torch.Tensor: - center = (n - 1) / 2 - zz, yy, xx = torch.meshgrid( - torch.arange(n), - torch.arange(n), - torch.arange(n), - indexing="ij", - ) - distance = ((xx - center) ** 2 + (yy - center) ** 2 + (zz - center) ** 2).sqrt() - sphere = (distance <= radius).float() * value - return sphere[None] - - -def _dice(a: torch.Tensor, b: torch.Tensor) -> float: - mask_a = a > 0 - mask_b = b > 0 - intersection = (mask_a & mask_b).sum() - return (2 * intersection / (mask_a.sum() + mask_b.sum())).item() - - -class TestLabelInterpolation: - def test_parse_interpolation_accepts_label(self) -> None: - assert _parse_interpolation("label") == "label" - assert _parse_interpolation("LABEL") == "label" - - def test_image_interpolation_label_raises(self) -> None: - with pytest.raises(ValueError, match="image_interpolation"): - tio.Resample(2, image_interpolation="label") - - def test_no_invalid_labels_when_downsampling(self) -> None: - data = torch.zeros(1, 32, 32, 32) - data[0, 8:24, 8:24, 8:24] = 2 - data[0, 12:20, 12:20, 12:20] = 5 # non-contiguous label values - subject = tio.Subject(seg=tio.LabelMap(data, affine=np.eye(4))) - result = tio.Resample(4, label_interpolation="label")(subject) - unique = set(result.seg.data.unique().tolist()) - assert unique <= {0.0, 2.0, 5.0} - - def test_no_invalid_labels_when_upsampling(self) -> None: - data = torch.zeros(1, 16, 16, 16) - data[0, 4:12, 4:12, 4:12] = 3 - subject = tio.Subject(seg=tio.LabelMap(data, affine=np.eye(4))) - result = tio.Resample(0.5, label_interpolation="label")(subject) - unique = set(result.seg.data.unique().tolist()) - assert unique <= {0.0, 3.0} - - def test_roundtrip_dice_beats_nearest(self) -> None: - original = _sphere_label() - subject = tio.Subject(seg=tio.LabelMap(original, affine=np.eye(4))) - - def roundtrip(mode: str) -> torch.Tensor: - down = tio.Resample(4, label_interpolation=mode)(subject) - back = tio.Resample(subject.seg, label_interpolation=mode)(down) - return back.seg.data - - dice_label = _dice(roundtrip("label"), original) - dice_nearest = _dice(roundtrip("nearest"), original) - # "label" is reliably better for a compact sphere, but assert only - # "not worse" to stay robust across grid alignment and library - # versions (a tie should not fail the test). - assert dice_label >= dice_nearest - - def test_default_pad_label_fills_out_of_bounds(self) -> None: - data = torch.ones(1, 16, 16, 16) - subject = tio.Subject(seg=tio.LabelMap(data, affine=np.eye(4))) - # Shift the whole volume out of the field of view along one axis. - transformed = AffineTransform( - translation=(100.0, 0.0, 0.0), - label_interpolation="label", - default_pad_label=7.0, - )(subject) - assert (transformed.seg.data == 7.0).any() - - def test_antialias_label_runs_and_keeps_valid_labels(self) -> None: - original = _sphere_label(value=4.0) - subject = tio.Subject(seg=tio.LabelMap(original, affine=np.eye(4))) - result = tio.Resample( - 4, - label_interpolation="label", - antialias=True, - )(subject) - unique = set(result.seg.data.unique().tolist()) - assert unique <= {0.0, 4.0} - assert result.seg.data.shape[1:] == (16, 16, 16) - - def test_multichannel_label_resamples_without_argmax(self) -> None: - data = torch.zeros(2, 16, 16, 16) - data[0] = 1.0 - data[0, 4:12, 4:12, 4:12] = 0.0 - data[1, 4:12, 4:12, 4:12] = 1.0 - subject = tio.Subject(seg=tio.LabelMap(data, affine=np.eye(4))) - result = tio.Resample(2, label_interpolation="label")(subject) - assert result.seg.data.shape[0] == 2 - - def test_multichannel_integer_input_preserves_partial_volumes(self) -> None: - # An integer one-hot encoding must not be truncated back to 0/1: - # linear resampling should yield fractional partial volumes. - data = torch.zeros(2, 16, 16, 16, dtype=torch.uint8) - data[0] = 1 - data[0, :8] = 0 - data[1, :8] = 1 - subject = tio.Subject(seg=tio.LabelMap(data, affine=np.eye(4))) - result = tio.Resample((1.5, 1.0, 1.0), label_interpolation="label")(subject) - assert result.seg.data.dtype.is_floating_point - fractional = (result.seg.data > 0) & (result.seg.data < 1) - assert fractional.any() - - def _three_label_junction(self) -> tio.Subject: - # Three labels meeting at a wavy junction, where the per-channel - # interpolation order changes the argmax outcome. - n = 40 - yy, xx, zz = torch.meshgrid( - torch.arange(n), - torch.arange(n), - torch.arange(n), - indexing="ij", - ) - seg = torch.zeros(n, n, n) - boundary = n / 2 + 3 * torch.sin(xx.float() / 3) - seg[yy > boundary] = 1 - seg[(yy <= boundary) & (zz > n / 2)] = 2 - return tio.Subject(seg=tio.LabelMap(seg[None], affine=np.eye(4))) - - def test_one_hot_label_interpolation_label_raises(self) -> None: - with pytest.raises(ValueError, match="one_hot_label_interpolation"): - tio.Resample( - 2, - label_interpolation="label", - one_hot_label_interpolation="label", - ) - - def test_one_hot_label_interpolation_default_is_linear(self) -> None: - subject = self._three_label_junction() - default = tio.Resample(0.5, label_interpolation="label")(subject) - explicit = tio.Resample( - 0.5, - label_interpolation="label", - one_hot_label_interpolation="linear", - )(subject) - torch.testing.assert_close(default.seg.data, explicit.seg.data) - - def test_one_hot_label_interpolation_higher_order_differs(self) -> None: - subject = self._three_label_junction() - linear = tio.Resample( - 0.5, - label_interpolation="label", - one_hot_label_interpolation="linear", - )(subject) - cubic = tio.Resample( - 0.5, - label_interpolation="label", - one_hot_label_interpolation="cubic", - )(subject) - # The order changes the result, but never invents labels. - assert not torch.equal(linear.seg.data, cubic.seg.data) - assert set(cubic.seg.data.unique().tolist()) <= {0.0, 1.0, 2.0} - - def test_one_hot_label_interpolation_accepts_integer_order(self) -> None: - subject = self._three_label_junction() - result = tio.Resample( - 0.5, - label_interpolation="label", - one_hot_label_interpolation=3, - )(subject) - assert set(result.seg.data.unique().tolist()) <= {0.0, 1.0, 2.0} - - def test_label_mode_per_instance_batch(self) -> None: - # Per-instance Affine on a batch must run the label partial-volume - # mode through the per-sample sampling grid, one geometry per element. - seg = torch.zeros(1, 24, 24, 24) - seg[0, 6:18, 6:18, 6:18] = 1 - seg[0, 10:14, 10:14, 10:14] = 2 - subjects = [ - tio.Subject(seg=tio.LabelMap(seg.clone(), affine=np.eye(4))) - for _ in range(4) - ] - batch = tio.SubjectsBatch.from_subjects(subjects) - torch.manual_seed(0) - transform = AffineTransform( - degrees=(-25, 25), - scales=(0.8, 1.2), - label_interpolation="label", - ) - data = transform(batch).images["seg"].data - assert data.shape[0] == 4 - assert set(data.unique().tolist()) <= {0.0, 1.0, 2.0} - # Per-instance sampling: elements receive different geometry. - assert not torch.equal(data[0], data[1]) diff --git a/tests/test_spike.py b/tests/test_spike.py deleted file mode 100644 index 3d1fc1729..000000000 --- a/tests/test_spike.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Tests for Spike transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestSpike: - def test_changes_data(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Spike(num_spikes=3, intensity=2.0)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_zero_intensity_is_identity(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Spike(intensity=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - def test_leaves_labels_unchanged(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = tio.Spike(num_spikes=3, intensity=2.0)(subject) - torch.testing.assert_close(result.seg.data, original_seg) - - def test_single_spike(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Spike(num_spikes=1, intensity=1.0)(subject) - assert not torch.allclose(result.t1.data, original) - - -class TestSpikePerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - data = torch.rand(1, 12, 12, 12) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Spike(intensity=(1.0, 3.0))(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["intensity"]) == batch.batch_size - assert not torch.allclose(result.t1.data[0], result.t1.data[1]) - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Spike(intensity=(1.0, 3.0), per_instance=False)(batch) - torch.testing.assert_close(result.t1.data[0], result.t1.data[1]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 12, 12, 12))) - result = tio.Spike(intensity=(1.0, 3.0))(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params diff --git a/tests/test_standardize.py b/tests/test_standardize.py deleted file mode 100644 index 711d7c85e..000000000 --- a/tests/test_standardize.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Tests for Standardize (z-score normalization).""" - -from __future__ import annotations - -import numpy as np -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = False) -> tio.Subject: - torch.manual_seed(42) - data = torch.randn(1, 10, 10, 10) * 50 + 100 # mean ~100, std ~50 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - mask = torch.zeros(1, 10, 10, 10) - mask[0, 2:8, 2:8, 2:8] = 1 - kwargs["brain"] = tio.LabelMap(mask) - return tio.Subject(**kwargs) - - -class TestBasic: - def test_output_has_zero_mean_unit_std(self) -> None: - subject = _make_subject() - result = tio.Standardize()(subject) - data = result.t1.data - assert abs(data.mean().item()) < 0.01 - assert abs(data.std().item() - 1.0) < 0.01 - - def test_leaves_label_maps_unchanged(self) -> None: - subject = _make_subject(with_label=True) - original = subject.brain.data.clone() - result = tio.Standardize()(subject) - torch.testing.assert_close(result.brain.data, original) - - -class TestMasking: - def test_masking_with_label_key(self) -> None: - subject = _make_subject(with_label=True) - result = tio.Standardize(masking_method="brain")(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_masking_with_callable(self) -> None: - subject = _make_subject() - result = tio.Standardize(masking_method=lambda x: x > 100)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_masking_key_not_found(self) -> None: - subject = _make_subject() - with pytest.raises(KeyError, match="nonexistent"): - tio.Standardize(masking_method="nonexistent")(subject) - - def test_masking_not_labelmap(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - t2=tio.ScalarImage(torch.rand(1, 4, 4, 4)), - ) - with pytest.raises(TypeError, match="LabelMap"): - tio.Standardize(masking_method="t2")(subject) - - -class TestEdgeCases: - def test_zero_std_raises(self) -> None: - data = torch.ones(1, 4, 4, 4) * 42.0 - subject = tio.Subject(t1=tio.ScalarImage(data)) - with pytest.raises(RuntimeError, match="zero"): - tio.Standardize()(subject) - - def test_empty_mask_warns(self) -> None: - subject = _make_subject() - with pytest.warns(RuntimeWarning, match="empty"): - tio.Standardize( - masking_method=lambda x: torch.zeros_like(x, dtype=torch.bool), - )(subject) - - -class TestInverse: - def test_inverse_restores_values(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - transformed = tio.Standardize()(subject) - restored = transformed.apply_inverse_transform() - np.testing.assert_allclose( - restored.t1.data.numpy(), - original.numpy(), - atol=1e-4, - ) - - -class TestExports: - def test_available_at_top_level(self) -> None: - assert hasattr(tio, "Standardize") - - def test_znormalization_alias(self) -> None: - assert tio.ZNormalization is tio.Standardize diff --git a/tests/test_subject.py b/tests/test_subject.py deleted file mode 100644 index 2be27b49c..000000000 --- a/tests/test_subject.py +++ /dev/null @@ -1,456 +0,0 @@ -"""Tests for Subject.""" - -from __future__ import annotations - -import copy -from pathlib import Path - -import nibabel as nib -import numpy as np -import pytest -import torch -from einops import rearrange - -from torchio import LabelMap -from torchio import ScalarImage -from torchio import Subject -from torchio.data.bboxes import BoundingBoxes -from torchio.data.bboxes import BoundingBoxFormat -from torchio.data.points import Points - - -class TestSubjectCreation: - def test_create_with_kwargs(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - seg=LabelMap(torch.randint(0, 5, (1, 10, 10, 10))), - ) - assert len(subject.images) == 2 - - def test_create_from_unpacked_dict(self): - data = { - "t1": ScalarImage(torch.randn(1, 10, 10, 10)), - "seg": LabelMap(torch.randint(0, 5, (1, 10, 10, 10))), - } - subject = Subject(**data) - assert len(subject.images) == 2 - - def test_metadata_from_kwargs(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - age=45, - name="John", - ) - assert subject.metadata["age"] == 45 - assert subject.metadata["name"] == "John" - - def test_empty_subject_raises(self): - with pytest.raises(ValueError, match="at least one"): - Subject() - - def test_metadata_only_subject(self): - subject = Subject(age=45, name="patient_1") - assert subject.age == 45 - assert subject.name == "patient_1" - assert len(subject.images) == 0 - - def test_points_only_subject(self): - pts = Points(torch.randn(5, 3)) - subject = Subject(landmarks=pts) - assert len(subject.points) == 1 - - -class TestSubjectAccess: - @pytest.fixture - def subject(self) -> Subject: - return Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - seg=LabelMap(torch.randint(0, 5, (1, 10, 10, 10))), - age=45, - ) - - def test_getattr_image(self, subject: Subject): - assert isinstance(subject.t1, ScalarImage) - assert isinstance(subject.seg, LabelMap) - - def test_getattr_metadata(self, subject: Subject): - assert subject.age == 45 - - def test_getitem(self, subject: Subject): - assert isinstance(subject["t1"], ScalarImage) - assert isinstance(subject["seg"], LabelMap) - - def test_getattr_missing_raises(self, subject: Subject): - with pytest.raises(AttributeError): - subject.nonexistent - - def test_getitem_missing_raises(self, subject: Subject): - with pytest.raises(KeyError): - subject["nonexistent"] - - def test_images_returns_only_images(self, subject: Subject): - images = subject.images - assert len(images) == 2 - assert "t1" in images - assert "seg" in images - - def test_metadata_access(self, subject: Subject): - assert subject.metadata["age"] == 45 - - def test_contains(self, subject: Subject): - assert "t1" in subject - assert "nonexistent" not in subject - - -class TestSubjectProperties: - @pytest.fixture - def subject(self) -> Subject: - return Subject( - t1=ScalarImage(torch.randn(1, 10, 20, 30)), - seg=LabelMap(torch.randint(0, 5, (1, 10, 20, 30))), - ) - - def test_spatial_shape(self, subject: Subject): - assert subject.spatial_shape == (10, 20, 30) - - def test_shape(self, subject: Subject): - assert subject.shape == (1, 10, 20, 30) - - def test_spacing(self, subject: Subject): - assert subject.spacing == (1.0, 1.0, 1.0) - - def test_inconsistent_shapes_raises(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - t2=ScalarImage(torch.randn(1, 20, 20, 20)), - ) - with pytest.raises(RuntimeError, match="Inconsistent"): - subject.spatial_shape - - def test_inconsistent_spacing_raises(self): - subject = Subject( - t1=ScalarImage( - torch.randn(1, 10, 10, 10), - affine=np.diag([1.0, 1.0, 1.0, 1.0]), - ), - t2=ScalarImage( - torch.randn(1, 10, 10, 10), - affine=np.diag([2.0, 2.0, 2.0, 1.0]), - ), - ) - with pytest.raises(RuntimeError, match="Inconsistent"): - subject.spacing - - def test_single_image_properties(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - ) - assert subject.spatial_shape == (10, 10, 10) - assert subject.spacing == (1.0, 1.0, 1.0) - - def test_applied_transforms_starts_empty(self, subject: Subject): - assert subject.applied_transforms == [] - - -class TestSubjectHistory: - def test_add_transform(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - ) - subject.applied_transforms.append( - {"name": "Affine", "parameters": {"scales": (1.1, 1.1, 1.1)}} - ) - assert len(subject.applied_transforms) == 1 - - def test_clear_history(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - ) - subject.applied_transforms.append( - {"name": "Affine", "parameters": {"scales": (1.1, 1.1, 1.1)}} - ) - subject.clear_history() - assert len(subject.applied_transforms) == 0 - - -class TestSubjectLoad: - def test_load_all(self, tmp_path: Path): - tensor = torch.randn(1, 10, 10, 10) - array = rearrange(tensor.numpy(), "c i j k -> i j k c") - nii = nib.Nifti1Image(array, np.eye(4)) - path = tmp_path / "test.nii.gz" - nib.save(nii, path) - - subject = Subject(t1=ScalarImage(path)) - assert not subject.t1.is_loaded - subject.load() - assert subject.t1.is_loaded - - -class TestSubjectCopy: - def test_copy(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - age=45, - ) - copied = copy.deepcopy(subject) - assert isinstance(copied, Subject) - assert isinstance(copied.t1, ScalarImage) - assert copied.metadata["age"] == 45 - # Verify it's a deep copy - copied.t1.set_data(torch.zeros(1, 10, 10, 10)) - assert not torch.equal(subject.t1.data, copied.t1.data) - - -class TestSubjectRepr: - def test_repr(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - seg=LabelMap(torch.randint(0, 5, (1, 10, 10, 10))), - ) - r = repr(subject) - assert "Subject" in r - assert "t1" in r - assert "seg" in r - - -class TestSubjectIteration: - def test_iter_yields_image_keys(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - seg=LabelMap(torch.randint(0, 5, (1, 10, 10, 10))), - age=45, - ) - keys = list(subject) - assert "t1" in keys - assert "seg" in keys - assert "age" not in keys - - def test_iter_yields_all_spatial_keys(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - landmarks=Points(torch.randn(5, 3)), - tumors=BoundingBoxes( - torch.randn(2, 6), - format=BoundingBoxFormat.IJKIJK, - ), - age=45, - ) - keys = list(subject) - assert "t1" in keys - assert "landmarks" in keys - assert "tumors" in keys - assert "age" not in keys - - -class TestSubjectWithPoints: - def test_points_access(self): - pts = Points(torch.randn(5, 3)) - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - landmarks=pts, - ) - assert subject.landmarks is pts - assert subject["landmarks"] is pts - - def test_points_dict(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - lm1=Points(torch.randn(3, 3)), - lm2=Points(torch.randn(7, 3)), - ) - pts = subject.points - assert len(pts) == 2 - assert "lm1" in pts - assert "lm2" in pts - - def test_contains_points(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - landmarks=Points(torch.randn(5, 3)), - ) - assert "landmarks" in subject - - def test_len_includes_points(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - landmarks=Points(torch.randn(5, 3)), - ) - assert len(subject) == 2 - - -class TestSubjectWithBoundingBoxes: - def test_bboxes_access(self): - boxes = BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - ) - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - tumors=boxes, - ) - assert subject.tumors is boxes - assert subject["tumors"] is boxes - - def test_bboxes_dict(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - tumors=BoundingBoxes( - torch.randn(3, 6), - format=BoundingBoxFormat.IJKIJK, - ), - organs=BoundingBoxes( - torch.randn(5, 6), - format=BoundingBoxFormat.IJKWHD, - ), - ) - bb = subject.bounding_boxes - assert len(bb) == 2 - assert "tumors" in bb - assert "organs" in bb - - def test_contains_bboxes(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - tumors=BoundingBoxes( - torch.randn(1, 6), - format=BoundingBoxFormat.IJKIJK, - ), - ) - assert "tumors" in subject - - def test_len_includes_bboxes(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - tumors=BoundingBoxes( - torch.randn(1, 6), - format=BoundingBoxFormat.IJKIJK, - ), - ) - assert len(subject) == 2 - - -class TestSubjectMixed: - def test_all_types(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - seg=LabelMap(torch.randint(0, 5, (1, 10, 10, 10))), - landmarks=Points(torch.randn(5, 3)), - tumors=BoundingBoxes( - torch.randn(2, 6), - format=BoundingBoxFormat.IJKIJK, - ), - age=45, - ) - assert len(subject.images) == 2 - assert len(subject.points) == 1 - assert len(subject.bounding_boxes) == 1 - assert subject.metadata["age"] == 45 - assert len(subject) == 4 # 2 images + 1 points + 1 bboxes - - def test_repr_with_all_types(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 10, 10, 10)), - landmarks=Points(torch.randn(5, 3)), - tumors=BoundingBoxes( - torch.randn(2, 6), - format=BoundingBoxFormat.IJKIJK, - ), - ) - r = repr(subject) - assert "images" in r - assert "points" in r - assert "bboxes" in r - - -class TestSubjectSlicing: - @pytest.fixture - def subject(self) -> Subject: - return Subject( - t1=ScalarImage(torch.randn(1, 20, 30, 40)), - seg=LabelMap(torch.randint(0, 3, (1, 20, 30, 40))), - age=42, - ) - - def test_slice_single_dim(self, subject: Subject): - cropped = subject[5:15] - assert cropped.t1.spatial_shape == (10, 30, 40) - assert cropped.seg.spatial_shape == (10, 30, 40) - - def test_slice_two_dims(self, subject: Subject): - cropped = subject[5:15, 10:20] - assert cropped.t1.spatial_shape == (10, 10, 40) - assert cropped.seg.spatial_shape == (10, 10, 40) - - def test_slice_three_dims(self, subject: Subject): - cropped = subject[2:12, 5:25, 10:30] - assert cropped.t1.spatial_shape == (10, 20, 20) - - def test_slice_with_ellipsis(self, subject: Subject): - cropped = subject[..., 10:30] - assert cropped.t1.spatial_shape == (20, 30, 20) - - def test_slice_with_int(self, subject: Subject): - cropped = subject[5] - assert cropped.t1.spatial_shape == (1, 30, 40) - - def test_preserves_metadata(self, subject: Subject): - cropped = subject[5:15] - assert cropped.age == 42 - - def test_preserves_channels(self): - subject = Subject( - rgb=ScalarImage(torch.randn(3, 20, 30, 40)), - ) - cropped = subject[5:15] - assert cropped.rgb.shape == (3, 10, 30, 40) - - def test_preserves_points(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 20, 30, 40)), - landmarks=Points(torch.randn(5, 3)), - ) - cropped = subject[5:15] - assert len(cropped.points) == 1 - assert "landmarks" in cropped.points - - def test_preserves_bboxes(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 20, 30, 40)), - tumors=BoundingBoxes( - torch.tensor([[1, 2, 3, 4, 5, 6]]), - format=BoundingBoxFormat.IJKIJK, - ), - ) - cropped = subject[5:15] - assert len(cropped.bounding_boxes) == 1 - - def test_preserves_transform_history(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 20, 30, 40)), - ) - subject.applied_transforms.append({"name": "RandomFlip"}) - cropped = subject[5:15] - assert len(cropped.applied_transforms) == 1 - - def test_inconsistent_shapes_raises(self): - subject = Subject( - t1=ScalarImage(torch.randn(1, 20, 30, 40)), - t2=ScalarImage(torch.randn(1, 10, 30, 40)), - ) - with pytest.raises(RuntimeError, match="Inconsistent"): - subject[5:10] - - def test_string_key_still_works(self, subject: Subject): - img = subject["t1"] - assert isinstance(img, ScalarImage) - - def test_no_images_slice_raises(self): - subject = Subject(landmarks=Points(torch.randn(5, 3))) - with pytest.raises(RuntimeError, match="no images"): - subject[5:15] - - def test_is_new_subject(self, subject: Subject): - cropped = subject[5:15] - assert isinstance(cropped, Subject) - assert cropped is not subject diff --git a/tests/test_swap.py b/tests/test_swap.py deleted file mode 100644 index 45aba36d5..000000000 --- a/tests/test_swap.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Tests for Swap transform.""" - -from __future__ import annotations - -import warnings - -import pytest -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestSwap: - def test_changes_data(self) -> None: - subject = _make_subject(with_label=False) - original = subject.t1.data.clone() - result = tio.Swap(patch_size=3, num_iterations=10)(subject) - assert not torch.allclose(result.t1.data, original) - - def test_preserves_shape(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Swap(patch_size=3, num_iterations=5)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - def test_warns_with_labels(self) -> None: - subject = _make_subject() - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - tio.Swap(patch_size=3, num_iterations=1)(subject) - assert any("LabelMap" in str(warning.message) for warning in w) - - def test_patch_too_large_raises(self) -> None: - subject = _make_subject(with_label=False) - with pytest.raises(ValueError, match="cannot be larger"): - tio.Swap(patch_size=100, num_iterations=1)(subject) - - def test_single_iteration(self) -> None: - subject = _make_subject(with_label=False) - result = tio.Swap(patch_size=3, num_iterations=1)(subject) - assert result.t1.data.shape == subject.t1.data.shape - - -class TestSwapPerInstance: - def _batch(self, batch_size: int = 6) -> tio.SubjectsBatch: - data = torch.rand(1, 16, 16, 16) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - def test_per_instance_differs_across_batch(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Swap(patch_size=4, num_iterations=20)(batch) - params = result.applied_transforms[-1].params - assert "_batched_keys" in params - assert len(params["locations"]) == batch.batch_size - assert not torch.allclose(result.t1.data[0], result.t1.data[1]) - - def test_per_instance_false_is_shared(self) -> None: - torch.manual_seed(0) - batch = self._batch() - result = tio.Swap(patch_size=4, num_iterations=20, per_instance=False)(batch) - torch.testing.assert_close(result.t1.data[0], result.t1.data[1]) - - def test_single_subject_keeps_scalar_params(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 16, 16, 16))) - result = tio.Swap(patch_size=4, num_iterations=20)(subject) - assert "_batched_keys" not in result.applied_transforms[-1].params - - -class TestSwapGatedOut: - def test_gated_out_elements_are_exact_no_ops(self) -> None: - torch.manual_seed(0) - data = torch.rand(1, 16, 16, 16) - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(32)] - ) - original = batch.t1.data.clone() - result = tio.Swap(patch_size=4, num_iterations=20, p=0.5)(batch) - unchanged = [ - torch.equal(result.t1.data[i], original[i]) for i in range(batch.batch_size) - ] - assert any(unchanged) - assert not all(unchanged) diff --git a/tests/test_tensordict.py b/tests/test_tensordict.py deleted file mode 100644 index 80136c7bd..000000000 --- a/tests/test_tensordict.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Tests for data loaders and batch collation.""" - -from __future__ import annotations - -import torch -from torch.utils.data import Dataset - -import torchio as tio - - -def _make_subject(idx: int = 0) -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 16, 16, 16))), - age=42 + idx, - name=f"subject_{idx}", - ) - - -# ── collate_subjects / collate_images ────────────────────────────────── - - -class TestCollate: - def test_collate_subjects(self) -> None: - subjects = [_make_subject(i) for i in range(4)] - batch = tio.collate_subjects(subjects) - assert batch.batch_size == 4 - assert batch.t1.data.shape == (4, 1, 16, 16, 16) - - def test_collate_images(self) -> None: - images = [tio.ScalarImage(torch.rand(1, 8, 8, 8)) for _ in range(4)] - batch = tio.collate_images(images) - assert batch.batch_size == 4 - assert batch.data.shape == (4, 1, 8, 8, 8) - - -# ── SubjectsLoader ──────────────────────────────────────────────────── - - -class _SimpleSubjectsDataset(Dataset): - def __init__(self, n: int = 8) -> None: - self.n = n - - def __len__(self) -> int: - return self.n - - def __getitem__(self, idx: int) -> tio.Subject: - return _make_subject(idx) - - -class TestSubjectsLoader: - def test_basic_iteration(self) -> None: - dataset = _SimpleSubjectsDataset(8) - loader = tio.SubjectsLoader(dataset, batch_size=4) - batch = next(iter(loader)) - - assert batch.batch_size == 4 - assert batch.t1.data.shape == (4, 1, 16, 16, 16) - - def test_all_batches(self) -> None: - dataset = _SimpleSubjectsDataset(8) - loader = tio.SubjectsLoader(dataset, batch_size=4) - batches = list(loader) - assert len(batches) == 2 - - def test_metadata_in_batch(self) -> None: - dataset = _SimpleSubjectsDataset(4) - loader = tio.SubjectsLoader(dataset, batch_size=4) - batch = next(iter(loader)) - assert batch.metadata["age"] == [42, 43, 44, 45] - - def test_passes_dataloader_kwargs(self) -> None: - dataset = _SimpleSubjectsDataset(8) - loader = tio.SubjectsLoader( - dataset, - batch_size=2, - shuffle=False, - num_workers=0, - ) - batches = list(loader) - assert len(batches) == 4 - - -# ── ImagesLoader ────────────────────────────────────────────────────── - - -class _SimpleImagesDataset(Dataset): - def __init__(self, n: int = 8) -> None: - self.n = n - - def __len__(self) -> int: - return self.n - - def __getitem__(self, idx: int) -> tio.ScalarImage: - return tio.ScalarImage(torch.rand(1, 8, 8, 8)) - - -class TestImagesLoader: - def test_basic_iteration(self) -> None: - dataset = _SimpleImagesDataset(8) - loader = tio.ImagesLoader(dataset, batch_size=4) - batch = next(iter(loader)) - - assert batch.batch_size == 4 - assert batch.data.shape == (4, 1, 8, 8, 8) - - def test_all_batches(self) -> None: - dataset = _SimpleImagesDataset(8) - loader = tio.ImagesLoader(dataset, batch_size=4) - batches = list(loader) - assert len(batches) == 2 - - def test_affines_in_batch(self) -> None: - dataset = _SimpleImagesDataset(4) - loader = tio.ImagesLoader(dataset, batch_size=4) - batch = next(iter(loader)) - assert len(batch.affines) == 4 - assert isinstance(batch.affines[0], tio.AffineMatrix) diff --git a/tests/test_to.py b/tests/test_to.py deleted file mode 100644 index d05c71380..000000000 --- a/tests/test_to.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Tests for To transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject() -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 5, 5, 5)), - ) - - -class TestTo: - def test_cast_dtype(self) -> None: - subject = _make_subject() - result = tio.To(torch.float64)(subject) - assert result.t1.data.dtype == torch.float64 - - def test_cast_to_half(self) -> None: - subject = _make_subject() - result = tio.To(torch.float16)(subject) - assert result.t1.data.dtype == torch.float16 - - def test_device_cpu(self) -> None: - subject = _make_subject() - result = tio.To("cpu")(subject) - assert result.t1.data.device.type == "cpu" - - def test_in_compose(self) -> None: - subject = _make_subject() - pipeline = tio.Compose( - [ - tio.To(torch.float64), - tio.Gamma(log_gamma=0.0), - ] - ) - result = pipeline(subject) - assert result.t1.data.shape == subject.t1.data.shape diff --git a/tests/test_to_reference_space.py b/tests/test_to_reference_space.py deleted file mode 100644 index 73dfdd420..000000000 --- a/tests/test_to_reference_space.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Tests for ToReferenceSpace transform.""" - -from __future__ import annotations - -import numpy as np -import pytest -import torch - -import torchio as tio - - -def _reference(shape=(64, 64, 64), spacing=2.0, origin=(10, 20, 30)) -> tio.ScalarImage: - affine = np.diag([spacing, spacing, spacing, 1.0]) - affine[:3, 3] = origin - return tio.ScalarImage( - torch.rand(1, *shape), - affine=tio.AffineMatrix(affine), - ) - - -def _fov_center(image: tio.Image) -> np.ndarray: - shape = np.array(image.spatial_shape) - matrix = image.affine.data.cpu().numpy() - corner0 = matrix @ np.array([0, 0, 0, 1.0]) - corner1 = matrix @ np.array([*(shape - 1), 1.0]) - return (corner0[:3] + corner1[:3]) / 2 - - -class TestFromTensor: - def test_shape_preserved(self) -> None: - ref = _reference() - embedding = torch.rand(8, 16, 16, 16) - image = tio.ToReferenceSpace.from_tensor(embedding, ref) - assert image.spatial_shape == (16, 16, 16) - assert image.data.shape[0] == 8 - - def test_spacing_scaled(self) -> None: - ref = _reference(shape=(64, 64, 64), spacing=2.0) - embedding = torch.rand(1, 16, 16, 16) # downsample x4 - image = tio.ToReferenceSpace.from_tensor(embedding, ref) - np.testing.assert_allclose(image.spacing, (8.0, 8.0, 8.0), atol=1e-5) - - def test_center_preserved(self) -> None: - ref = _reference() - embedding = torch.rand(1, 16, 16, 16) - image = tio.ToReferenceSpace.from_tensor(embedding, ref) - np.testing.assert_allclose(_fov_center(image), _fov_center(ref), atol=1e-4) - - def test_class_preserved(self) -> None: - ref = tio.LabelMap(torch.zeros(1, 32, 32, 32)) - embedding = torch.rand(1, 8, 8, 8) - image = tio.ToReferenceSpace.from_tensor(embedding, ref) - assert isinstance(image, tio.LabelMap) - - def test_same_shape_keeps_affine(self) -> None: - ref = _reference(shape=(32, 32, 32), spacing=1.5) - embedding = torch.rand(1, 32, 32, 32) - image = tio.ToReferenceSpace.from_tensor(embedding, ref) - np.testing.assert_allclose( - image.affine.data.cpu().numpy(), - ref.affine.data.cpu().numpy(), - atol=1e-5, - ) - - def test_anisotropic_shape(self) -> None: - ref = _reference(shape=(64, 32, 16), spacing=1.0) - embedding = torch.rand(1, 16, 16, 16) - image = tio.ToReferenceSpace.from_tensor(embedding, ref) - np.testing.assert_allclose(image.spacing, (4.0, 2.0, 1.0), atol=1e-5) - - -class TestTransform: - def test_data_unchanged(self) -> None: - ref = _reference() - subject = tio.Subject(emb=tio.ScalarImage(torch.rand(8, 16, 16, 16))) - original = subject.emb.data.clone() - result = tio.ToReferenceSpace(ref)(subject) - torch.testing.assert_close(result.emb.data, original) - - def test_affine_updated(self) -> None: - ref = _reference() - subject = tio.Subject(emb=tio.ScalarImage(torch.rand(1, 16, 16, 16))) - result = tio.ToReferenceSpace(ref)(subject) - np.testing.assert_allclose(result.emb.spacing, (8.0, 8.0, 8.0), atol=1e-5) - - def test_applies_to_all_images(self) -> None: - ref = _reference() - subject = tio.Subject( - a=tio.ScalarImage(torch.rand(1, 16, 16, 16)), - b=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - ) - result = tio.ToReferenceSpace(ref)(subject) - np.testing.assert_allclose(result.a.spacing, (8.0, 8.0, 8.0), atol=1e-5) - np.testing.assert_allclose(result.b.spacing, (16.0, 16.0, 16.0), atol=1e-5) - - def test_invalid_reference_raises(self) -> None: - with pytest.raises(TypeError, match="reference must be a TorchIO Image"): - tio.ToReferenceSpace("not an image") # type: ignore[arg-type] - - -class TestExport: - def test_top_level(self) -> None: - assert hasattr(tio, "ToReferenceSpace") diff --git a/tests/test_transforms_base.py b/tests/test_transforms_base.py deleted file mode 100644 index 8907b2bce..000000000 --- a/tests/test_transforms_base.py +++ /dev/null @@ -1,451 +0,0 @@ -"""Tests for the transform base classes and composition.""" - -from __future__ import annotations - -from dataclasses import asdict -from typing import Any - -import nibabel as nib -import numpy as np -import pytest -import SimpleITK as sitk -import torch - -import torchio as tio -from torchio.data.bboxes import BoundingBoxes -from torchio.data.bboxes import BoundingBoxFormat -from torchio.data.points import Points - -# ── Helpers ────────────────────────────────────────────────────────── - - -def _make_subject() -> tio.Subject: - return tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 8, 8, 8)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 8, 8, 8))), - landmarks=Points(torch.rand(3, 3)), - tumors=BoundingBoxes( - torch.rand(2, 6), - format=BoundingBoxFormat.IJKIJK, - ), - age=42, - ) - - -class _IdentityTransform(tio.Transform): - """Transform that does nothing (for testing the base flow).""" - - def apply_transform(self, batch: Any, params: dict) -> Any: - return batch - - -class _DoubleIntensity(tio.IntensityTransform): - """Doubles intensity of ScalarImages (for testing IntensityTransform).""" - - def apply_transform(self, batch: Any, params: dict) -> Any: - for _name, img_batch in self._get_images(batch).items(): - img_batch.data = img_batch.data * 2 - return batch - - -class _FlipSpatial(tio.SpatialTransform): - """Flips along axis 0 (for testing SpatialTransform).""" - - def apply_transform(self, batch: Any, params: dict) -> Any: - for _name, img_batch in self._get_images(batch).items(): - img_batch.data = torch.flip(img_batch.data, [-3]) - return batch - - -# ── Transform base ─────────────────────────────────────────────────── - - -class TestTransformBase: - def test_forward_returns_subject(self) -> None: - subject = _make_subject() - result = _IdentityTransform()(subject) - assert isinstance(result, tio.Subject) - - def test_forward_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - result = _IdentityTransform()(image) - assert isinstance(result, tio.Image) - - def test_forward_accepts_tensor(self) -> None: - tensor = torch.rand(1, 8, 8, 8) - result = _IdentityTransform()(tensor) - assert isinstance(result, torch.Tensor) - assert result.shape == (1, 8, 8, 8) - - def test_forward_accepts_ndarray(self) -> None: - array = np.random.rand(1, 8, 8, 8).astype(np.float32) - result = _IdentityTransform()(array) - assert isinstance(result, np.ndarray) - assert result.shape == (1, 8, 8, 8) - - def test_forward_accepts_ndarray_3d(self) -> None: - array = np.random.rand(8, 8, 8).astype(np.float32) - result = _IdentityTransform()(array) - assert isinstance(result, np.ndarray) - - def test_forward_accepts_sitk(self) -> None: - sitk_image = sitk.Image(8, 8, 8, sitk.sitkFloat32) - result = _IdentityTransform()(sitk_image) - assert isinstance(result, sitk.Image) - - def test_forward_accepts_nifti(self) -> None: - nifti = nib.Nifti1Image(np.zeros((8, 8, 8)), np.eye(4)) - result = _IdentityTransform()(nifti) - assert isinstance(result, nib.Nifti1Image) - - def test_sitk_preserves_spacing(self) -> None: - sitk_image = sitk.Image(8, 8, 8, sitk.sitkFloat32) - sitk_image.SetSpacing((2.0, 2.0, 2.0)) - result = _IdentityTransform()(sitk_image) - assert result.GetSpacing() == pytest.approx((2.0, 2.0, 2.0)) - - def test_nifti_preserves_affine(self) -> None: - affine = np.diag([2.0, 2.0, 2.0, 1.0]) - nifti = nib.Nifti1Image(np.zeros((8, 8, 8)), affine) - result = _IdentityTransform()(nifti) - np.testing.assert_array_almost_equal(result.affine, affine) - - def test_forward_accepts_dict(self) -> None: - data = { - "t1": torch.rand(1, 8, 8, 8), - "seg": torch.randint(0, 3, (1, 8, 8, 8)), - } - result = _IdentityTransform()(data) - assert isinstance(result, dict) - assert set(result.keys()) == {"t1", "seg"} - assert isinstance(result["t1"], torch.Tensor) - - def test_dict_transform_modifies_data(self) -> None: - data = { - "t1": torch.rand(1, 8, 8, 8), - } - original = data["t1"].clone() - result = _DoubleIntensity()(data) - torch.testing.assert_close(result["t1"], original * 2) - - def test_dict_metadata_passthrough(self) -> None: - data = { - "t1": torch.rand(1, 8, 8, 8), - "age": 42, - } - result = _IdentityTransform()(data) - assert result["age"] == 42 - - def test_probability_zero_skips(self) -> None: - subject = _make_subject() - original_data = subject.t1.data.clone() - result = _DoubleIntensity(p=0.0)(subject) - torch.testing.assert_close(result.t1.data, original_data) - - def test_probability_zero_skips_when_rand_is_zero( - self, - monkeypatch: pytest.MonkeyPatch, - ) -> None: - """p=0 must be a no-op even if torch.rand returns exactly 0.""" - subject = _make_subject() - original_data = subject.t1.data.clone() - monkeypatch.setattr(torch, "rand", lambda *a, **k: torch.zeros(1)) - result = _DoubleIntensity(p=0.0)(subject) - torch.testing.assert_close(result.t1.data, original_data) - - def test_probability_one_applies(self) -> None: - subject = _make_subject() - original_data = subject.t1.data.clone() - result = _DoubleIntensity(p=1.0)(subject) - torch.testing.assert_close(result.t1.data, original_data * 2) - - def test_history_recorded(self) -> None: - subject = _make_subject() - result = _IdentityTransform()(subject) - assert len(result.applied_transforms) == 1 - assert result.applied_transforms[0].name == "_IdentityTransform" - - def test_history_has_params(self) -> None: - subject = _make_subject() - result = _IdentityTransform()(subject) - trace = result.applied_transforms[0] - assert isinstance(trace.params, dict) - - def test_history_serializable(self) -> None: - subject = _make_subject() - result = _IdentityTransform()(subject) - d = asdict(result.applied_transforms[0]) - assert "name" in d - assert "params" in d - - def test_is_nn_module(self) -> None: - t = _IdentityTransform() - assert isinstance(t, torch.nn.Module) - - def test_invalid_input_type(self) -> None: - with pytest.raises(TypeError): - _IdentityTransform()("not a valid input") - - -# ── include/exclude ────────────────────────────────────────────────── - - -class TestIncludeExclude: - def test_include_filters(self) -> None: - subject = _make_subject() - t = _DoubleIntensity(include=["t1"]) - result = t(subject) - # t1 should be doubled - assert result.t1.data.mean() > 0 - # seg should be unchanged (it's a LabelMap so IntensityTransform - # skips it anyway, but include also restricts) - - def test_exclude_filters(self) -> None: - subject = _make_subject() - original_t1 = subject.t1.data.clone() - t = _DoubleIntensity(exclude=["t1"]) - result = t(subject) - # t1 should be unchanged (excluded) - torch.testing.assert_close(result.t1.data, original_t1) - - -# ── IntensityTransform ─────────────────────────────────────────────── - - -class TestIntensityTransform: - def test_only_scalar_images(self) -> None: - subject = _make_subject() - original_seg = subject.seg.data.clone() - result = _DoubleIntensity()(subject) - # seg (LabelMap) should not be modified - torch.testing.assert_close(result.seg.data, original_seg) - - def test_scalar_image_modified(self) -> None: - subject = _make_subject() - original_t1 = subject.t1.data.clone() - result = _DoubleIntensity()(subject) - torch.testing.assert_close(result.t1.data, original_t1 * 2) - - -# ── SpatialTransform ───────────────────────────────────────────────── - - -class TestSpatialTransform: - def test_all_images_modified(self) -> None: - subject = _make_subject() - result = _FlipSpatial()(subject) - # Both t1 and seg should be flipped - assert result.t1.data.shape == (1, 8, 8, 8) - assert result.seg.data.shape == (1, 8, 8, 8) - - -# ── Compose ────────────────────────────────────────────────────────── - - -class TestCompose: - def test_sequential_application(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - composed = tio.Compose([_DoubleIntensity(), _DoubleIntensity()]) - result = composed(subject) - torch.testing.assert_close(result.t1.data, original * 4) - - def test_copy_true_preserves_original(self) -> None: - subject = _make_subject() - original = subject.t1.data.clone() - composed = tio.Compose([_DoubleIntensity()], copy=True) - composed(subject) - # Original should be unchanged - torch.testing.assert_close(subject.t1.data, original) - - def test_copy_false_no_deepcopy(self) -> None: - """copy=False skips deepcopy (used inside Compose).""" - subject = _make_subject() - composed = tio.Compose([_DoubleIntensity()], copy=False) - result = composed(subject) - # Result should be transformed - assert result.t1.data.mean() > 0 - - def test_empty_compose(self) -> None: - subject = _make_subject() - composed = tio.Compose([]) - result = composed(subject) - assert isinstance(result, tio.Subject) - - def test_history_from_children(self) -> None: - subject = _make_subject() - composed = tio.Compose([_IdentityTransform(), _DoubleIntensity()]) - result = composed(subject) - assert len(result.applied_transforms) == 2 - - def test_accepts_image(self) -> None: - image = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - composed = tio.Compose([_IdentityTransform()]) - result = composed(image) - assert isinstance(result, tio.Image) - - def test_accepts_tensor(self) -> None: - tensor = torch.rand(1, 8, 8, 8) - composed = tio.Compose([_IdentityTransform()]) - result = composed(tensor) - assert isinstance(result, torch.Tensor) - - -# ── OneOf ───────────────────────────────────────────────────────────── - - -class TestOneOf: - def test_applies_exactly_one(self) -> None: - subject = _make_subject() - one_of = tio.OneOf([_DoubleIntensity(), _IdentityTransform()]) - result = one_of(subject) - # Exactly one transform should be in history - assert len(result.applied_transforms) == 1 - - def test_with_weights(self) -> None: - subject = _make_subject() - # Weight 1.0 on identity, 0.0 on double: should always pick identity - one_of = tio.OneOf( - {_IdentityTransform(): 1.0, _DoubleIntensity(): 0.0}, - ) - result = one_of(subject) - assert result.applied_transforms[0].name == "_IdentityTransform" - - -# ── SomeOf ──────────────────────────────────────────────────────────── - - -class TestSomeOf: - def test_applies_n_transforms(self) -> None: - subject = _make_subject() - some_of = tio.SomeOf( - [_IdentityTransform(), _DoubleIntensity(), _IdentityTransform()], - num_transforms=2, - ) - result = some_of(subject) - assert len(result.applied_transforms) == 2 - - def test_num_transforms_range(self) -> None: - subject = _make_subject() - some_of = tio.SomeOf( - [_IdentityTransform(), _DoubleIntensity(), _IdentityTransform()], - num_transforms=(1, 3), - ) - result = some_of(subject) - assert 1 <= len(result.applied_transforms) <= 3 - - -# ── Operator sugar ─────────────────────────────────────────────────── - - -class TestAddOperator: - def test_add_creates_compose(self) -> None: - t1 = tio.Flip(axes=(0,)) - t2 = tio.Noise(std=0.1) - result = t1 + t2 - assert isinstance(result, tio.Compose) - assert len(result.transforms) == 2 - - def test_add_flattens_compose(self) -> None: - t1 = tio.Flip(axes=(0,)) - t2 = tio.Noise(std=0.1) - t3 = tio.BiasField() - result = t1 + t2 + t3 - assert isinstance(result, tio.Compose) - assert len(result.transforms) == 3 - - def test_add_compose_plus_transform(self) -> None: - c = tio.Compose([tio.Flip(axes=(0,)), tio.Noise(std=0.1)]) - t = tio.BiasField() - result = c + t - assert isinstance(result, tio.Compose) - assert len(result.transforms) == 3 - - def test_add_not_implemented_for_non_transform(self) -> None: - with pytest.raises(TypeError): - tio.Flip(axes=(0,)) + 42 # type: ignore[operator] - - def test_add_produces_working_pipeline(self) -> None: - subject = _make_subject() - pipeline = tio.Flip(axes=(0,)) + tio.Noise(std=0.01) - result = pipeline(subject) - assert result.t1.shape == subject.t1.shape - - -class TestOrOperator: - def test_or_creates_oneof(self) -> None: - t1 = tio.Flip(axes=(0,)) - t2 = tio.Noise(std=0.1) - result = t1 | t2 - assert isinstance(result, tio.OneOf) - assert len(result.transforms) == 2 - - def test_or_flattens_oneof(self) -> None: - t1 = tio.Flip(axes=(0,)) - t2 = tio.Noise(std=0.1) - t3 = tio.BiasField() - result = t1 | t2 | t3 - assert isinstance(result, tio.OneOf) - assert len(result.transforms) == 3 - - def test_or_not_implemented_for_non_transform(self) -> None: - with pytest.raises(TypeError): - tio.Flip(axes=(0,)) | "bad" # type: ignore[operator] - - def test_or_produces_working_pipeline(self) -> None: - subject = _make_subject() - pipeline = tio.Flip(axes=(0,)) | tio.Noise(std=0.01) - result = pipeline(subject) - assert result.t1.shape == subject.t1.shape - - -# ── Coverage gap tests ─────────────────────────────────────────────── - - -class TestTransformEdgeCases: - def test_invalid_probability_raises(self) -> None: - with pytest.raises(ValueError, match="Probability"): - tio.Flip(axes=(0,), p=1.5) - - def test_repr_shows_params(self) -> None: - t = tio.Flip(axes=(0,)) - r = repr(t) - assert "Flip" in r - - def test_base_apply_transform_raises(self) -> None: - t = tio.Transform() - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - with pytest.raises(NotImplementedError): - t(subject) - - def test_non_invertible_raises(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - result = tio.Blur(std=1.0)(subject) - # Blur is not invertible; inverse should warn/skip - restored = result.apply_inverse_transform() - assert restored.t1.data.shape == subject.t1.data.shape - - def test_exclude_images(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 5, 5, 5)), - t2=tio.ScalarImage(torch.rand(1, 5, 5, 5)), - ) - original_t2 = subject.t2.data.clone() - result = tio.Gamma(log_gamma=0.5, exclude=["t2"])(subject) - torch.testing.assert_close(result.t2.data, original_t2) - - -class TestOneOfSkip: - def test_one_of_with_p_zero_is_identity(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - original = subject.t1.data.clone() - result = tio.OneOf([tio.Flip(axes=(0,))], p=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) - - -class TestSomeOfSkip: - def test_some_of_with_p_zero_is_identity(self) -> None: - subject = tio.Subject(t1=tio.ScalarImage(torch.rand(1, 5, 5, 5))) - original = subject.t1.data.clone() - result = tio.SomeOf([tio.Flip(axes=(0,))], num_transforms=1, p=0.0)(subject) - torch.testing.assert_close(result.t1.data, original) diff --git a/tests/test_transpose.py b/tests/test_transpose.py deleted file mode 100644 index 164f4f0c4..000000000 --- a/tests/test_transpose.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Tests for Transpose transform.""" - -from __future__ import annotations - -import torch - -import torchio as tio - - -def _make_subject(with_label: bool = True) -> tio.Subject: - data = torch.rand(1, 10, 10, 10) * 100 - kwargs: dict = {"t1": tio.ScalarImage(data)} - if with_label: - seg = torch.zeros(1, 10, 10, 10, dtype=torch.float32) - seg[0, 2:5, 2:5, 2:5] = 1 - seg[0, 6:9, 6:9, 6:9] = 2 - kwargs["seg"] = tio.LabelMap(seg) - return tio.Subject(**kwargs) - - -class TestTranspose: - def test_swaps_axes(self) -> None: - data = torch.rand(1, 8, 10, 12) - subject = tio.Subject(t1=tio.ScalarImage(data)) - result = tio.Transpose()(subject) - assert result.t1.data.shape == (1, 12, 10, 8) - - def test_double_transpose_restores_shape(self) -> None: - data = torch.rand(1, 8, 10, 12) - subject = tio.Subject(t1=tio.ScalarImage(data)) - result = tio.Transpose()(tio.Transpose()(subject)) - assert result.t1.data.shape == (1, 8, 10, 12) - - def test_inverse(self) -> None: - data = torch.rand(1, 8, 10, 12) - subject = tio.Subject(t1=tio.ScalarImage(data)) - original = subject.t1.data.clone() - transformed = tio.Transpose()(subject) - restored = transformed.apply_inverse_transform() - torch.testing.assert_close(restored.t1.data, original) - - def test_is_invertible(self) -> None: - t = tio.Transpose() - assert t.invertible is True - - def test_symmetric_shape_unchanged(self) -> None: - data = torch.rand(1, 10, 10, 10) - subject = tio.Subject(t1=tio.ScalarImage(data)) - result = tio.Transpose()(subject) - assert result.t1.data.shape == (1, 10, 10, 10) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 000000000..0215ef322 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,87 @@ +import copy + +import pytest +import torch + +import torchio as tio + +from .utils import TorchioTestCase + + +class TestUtils(TorchioTestCase): + """Tests for `utils` module.""" + + def test_to_tuple(self): + assert tio.utils.to_tuple(1) == (1,) + assert tio.utils.to_tuple((1,)) == (1,) + assert tio.utils.to_tuple(1, length=3) == (1, 1, 1) + assert tio.utils.to_tuple((1, 2)) == (1, 2) + assert tio.utils.to_tuple((1, 2), length=3) == (1, 2) + assert tio.utils.to_tuple([1, 2], length=3) == (1, 2) + + def test_get_stem(self): + assert tio.utils.get_stem('/home/image.nii.gz') == 'image' + assert tio.utils.get_stem('/home/image.nii') == 'image' + assert tio.utils.get_stem('/home/image.nrrd') == 'image' + + def test_guess_type(self): + assert tio.utils.guess_type('None') is None + assert isinstance(tio.utils.guess_type('1'), int) + assert isinstance(tio.utils.guess_type('1.5'), float) + assert isinstance(tio.utils.guess_type('(1, 3, 5)'), tuple) + assert isinstance(tio.utils.guess_type('(1,3,5)'), tuple) + assert isinstance(tio.utils.guess_type('[1,3,5]'), list) + assert isinstance(tio.utils.guess_type('test'), str) + + def test_apply_transform_to_file(self): + transform = tio.RandomFlip() + tio.utils.apply_transform_to_file( + self.get_image_path('input'), + transform, + self.get_image_path('output'), + verbose=True, + ) + + def test_subjects_from_batch(self): + dataset = tio.SubjectsDataset(4 * [self.sample_subject]) + loader = tio.SubjectsLoader(dataset, batch_size=4) + batch = tio.utils.get_first_item(loader) + subjects = tio.utils.get_subjects_from_batch(batch) + assert isinstance(subjects[0], tio.Subject) + + def test_subjects_from_batch_with_string_metadata(self): + subject_c_with_string_metadata = tio.Subject( + name='John Doe', + label=tio.LabelMap(self.get_image_path('label_c', binary=True)), + ) + + dataset = tio.SubjectsDataset(4 * [subject_c_with_string_metadata]) + loader = tio.SubjectsLoader(dataset, batch_size=4) + batch = tio.utils.get_first_item(loader) + subjects = tio.utils.get_subjects_from_batch(batch) + assert isinstance(subjects[0], tio.Subject) + assert 'label' in subjects[0] + assert 'name' in subjects[0] + + def test_subjects_from_batch_with_int_metadata(self): + subject_c_with_int_metadata = tio.Subject( + age=45, + label=tio.LabelMap(self.get_image_path('label_c', binary=True)), + ) + dataset = tio.SubjectsDataset(4 * [subject_c_with_int_metadata]) + loader = tio.SubjectsLoader(dataset, batch_size=4) + batch = tio.utils.get_first_item(loader) + subjects = tio.utils.get_subjects_from_batch(batch) + assert isinstance(subjects[0], tio.Subject) + assert 'label' in subjects[0] + assert 'age' in subjects[0] + + def test_add_images_from_batch(self): + subject = copy.deepcopy(self.sample_subject) + subjects = 4 * [subject] + preds = torch.rand(4, *subject.shape) + tio.utils.add_images_from_batch(subjects, preds) + + def test_empty_batch(self): + with pytest.raises(RuntimeError): + tio.utils.get_batch_images_and_size({}) diff --git a/tests/test_vectorization.py b/tests/test_vectorization.py deleted file mode 100644 index ea7364d82..000000000 --- a/tests/test_vectorization.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Equivalence gate: per-instance batch transforms must be vectorized. - -Each transform here has an `apply_transform` that is deterministic given the -recorded parameters (the randomness lives in `make_params`). The -`assert_vectorized` fixture checks that applying the transform to a batch -produces, for every element, the same result as applying the element's own -parameters to that element alone. A correct vectorized implementation passes; -any cross-element contamination or broadcasting mistake fails. - -Transforms that sample inside `apply_transform` (Noise, LabelsToImage) are -excluded here and covered by their own tests. -""" - -from __future__ import annotations - -import pytest -import torch - -import torchio as tio - - -def _batch( - batch_size: int = 4, dtype: torch.dtype = torch.float32 -) -> tio.SubjectsBatch: - data = torch.rand(1, 12, 12, 12, dtype=dtype) - subjects = [ - tio.Subject(t1=tio.ScalarImage(data.clone() + index)) - for index in range(batch_size) - ] - return tio.SubjectsBatch.from_subjects(subjects) - - -@pytest.mark.parametrize( - "transform", - [ - tio.Ghosting(num_ghosts=(2, 5), intensity=(0.5, 1.0)), - tio.Spike(num_spikes=(1, 3), intensity=(0.3, 0.8)), - tio.Blur(std=(0.5, 2.0)), - tio.BiasField(std=(0.3, 0.8)), - tio.Flip(axes=(0, 1, 2), flip_probability=0.5), - tio.Motion(degrees=10.0, translation=10.0, num_transforms=2), - tio.Swap(patch_size=3, num_iterations=5), - tio.Anisotropy(downsampling=(1.5, 4.0)), - ], -) -def test_vectorized_matches_per_element(transform, assert_vectorized) -> None: - torch.manual_seed(0) - assert_vectorized(transform, _batch()) - - -@pytest.mark.parametrize( - "transform", - [ - tio.Ghosting(num_ghosts=4, intensity=1.0, p=0.5), - tio.Spike(num_spikes=2, intensity=1.0, p=0.5), - tio.Blur(std=1.5, p=0.5), - tio.BiasField(std=0.5, p=0.5), - tio.Flip(axes=(0, 1, 2), flip_probability=1.0, p=0.5), - tio.Motion(degrees=10.0, translation=10.0, num_transforms=2, p=0.5), - tio.Swap(patch_size=3, num_iterations=5, p=0.5), - ], -) -def test_vectorized_matches_per_element_with_gating( - transform, - assert_vectorized, -) -> None: - torch.manual_seed(0) - assert_vectorized(transform, _batch(batch_size=6)) - - -def test_anisotropy_tie_rounding_matches_scalar(assert_vectorized) -> None: - # An odd spatial length with factor 2.0 makes length/factor land on a .5 - # tie (e.g. 9/2 = 4.5). The per-instance path uses torch.round and the - # scalar path uses Python round; both must agree (round-half-to-even), so - # the vectorized per-element result must match the scalar application. - torch.manual_seed(0) - data = torch.rand(1, 9, 9, 9) - batch = tio.SubjectsBatch.from_subjects( - [tio.Subject(t1=tio.ScalarImage(data.clone() + index)) for index in range(4)] - ) - assert_vectorized(tio.Anisotropy(downsampling=2.0), batch) diff --git a/tests/test_visualization.py b/tests/test_visualization.py deleted file mode 100644 index 7748b1dcb..000000000 --- a/tests/test_visualization.py +++ /dev/null @@ -1,389 +0,0 @@ -"""Tests for visualization and Image repr.""" - -from __future__ import annotations - -import shutil - -import matplotlib -import matplotlib.pyplot as plt -import nibabel as nib -import numpy as np -import pytest -import torch -from matplotlib.figure import Figure - -import torchio as tio -from torchio.data.affine import AffineMatrix - -matplotlib.use("Agg") - -requires_ffmpeg = pytest.mark.skipif( - shutil.which("ffmpeg") is None, - reason="ffmpeg binary not available", -) - - -class TestEulerAngles: - def test_identity_gives_zeros(self) -> None: - a = AffineMatrix() - angles = a.euler_angles - assert all(abs(v) < 1e-6 for v in angles) - - def test_rotation_around_z(self) -> None: - theta = np.radians(15) - m = np.eye(4) - m[0, 0] = np.cos(theta) - m[0, 1] = -np.sin(theta) - m[1, 0] = np.sin(theta) - m[1, 1] = np.cos(theta) - a = AffineMatrix(m) - x, y, z = a.euler_angles - assert abs(z - 15.0) < 0.1 - assert abs(x) < 0.1 - assert abs(y) < 0.1 - - def test_rotation_around_x(self) -> None: - theta = np.radians(30) - m = np.eye(4) - m[1, 1] = np.cos(theta) - m[1, 2] = -np.sin(theta) - m[2, 1] = np.sin(theta) - m[2, 2] = np.cos(theta) - a = AffineMatrix(m) - x, _y, _z = a.euler_angles - assert abs(x - 30.0) < 0.1 - - -class TestImageRepr: - def test_multiline_format(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 20, 30)) - r = repr(img) - assert "ScalarImage(\n" in r - assert "channels:" in r - assert "spatial:" in r - assert "spacing:" in r - assert "orientation:" in r - assert "angles:" in r - assert "dtype:" in r - assert "memory:" in r - - def test_lazy_shows_backend(self, tmp_path) -> None: - path = tmp_path / "test.nii" - nib.save(nib.Nifti1Image(np.zeros((10, 10, 10)), np.eye(4)), path) - img = tio.ScalarImage(path) - r = repr(img) - assert "lazy" in r - assert "NIfTI" in r - - def test_loaded_shows_loaded(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - r = repr(img) - assert "in memory" in r - - def test_origin_shown(self) -> None: - affine = AffineMatrix.from_spacing((1, 1, 1), origin=(10.0, 20.0, 30.0)) - img = tio.ScalarImage( - torch.rand(1, 5, 5, 5), - affine=affine, - ) - r = repr(img) - assert "10.00" in r - assert "origin:" in r - - -class TestPlotImage: - def test_returns_figure(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 20, 30)) - fig = img.plot(show=False) - assert isinstance(fig, Figure) - - def test_custom_indices(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 20, 30)) - fig = img.plot(indices=(5, 10, 15), show=False) - axes = fig.axes - assert len(axes) == 3 - # Titles show the slice index for the sliced axis - titles = [ax.get_title() for ax in axes] - assert any("5" in t for t in titles) - assert any("10" in t for t in titles) - assert any("15" in t for t in titles) - - def test_views_are_sagittal_coronal_axial(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 20, 30)) - fig = img.plot(show=False) - titles = [ax.get_title() for ax in fig.axes] - assert "Sagittal" in titles[0] - assert "Coronal" in titles[1] - assert "Axial" in titles[2] - - def test_orientation_labels_show_tensor_axis(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - fig = img.plot(show=False) - ax = fig.axes[0] - xlabel = ax.get_xlabel() - ylabel = ax.get_ylabel() - # Default (voxels=False): "Anterior [mm] (j)" format - assert any(c in xlabel for c in ("i", "j", "k")) - assert any(c in ylabel for c in ("i", "j", "k")) - - def test_voxel_labels_show_arrow(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - fig = img.plot(show=False, voxels=True) - ax = fig.axes[0] - xlabel = ax.get_xlabel() - ylabel = ax.get_ylabel() - # voxels=True: "j (A ↔ P)" format - assert "↔" in xlabel - assert "↔" in ylabel - - def test_save_to_file(self, tmp_path) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - out = tmp_path / "test.png" - img.plot(output_path=out, show=False) - assert out.exists() - assert out.stat().st_size > 0 - - def test_external_axes(self) -> None: - fig, axes = plt.subplots(1, 3) - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - result = img.plot(axes=axes, show=False) - assert result is fig - - def test_label_map_uses_nearest(self) -> None: - label = tio.LabelMap(torch.randint(0, 3, (1, 10, 10, 10))) - fig = label.plot(show=False) - ax = fig.axes[0] - im = ax.images[0] - assert im.get_interpolation() == "none" - - def test_voxels_mode(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - fig = img.plot(show=False, voxels=True) - assert isinstance(fig, Figure) - - def test_consistent_views_across_orientations(self) -> None: - """Sagittal/Coronal/Axial order is the same for RAS and LPS.""" - img_ras = tio.ScalarImage(torch.rand(1, 10, 20, 30)) - fig_ras = img_ras.plot(show=False) - - ornt = nib.orientations.axcodes2ornt(("L", "P", "S")) - affine = tio.AffineMatrix(nib.orientations.inv_ornt_aff(ornt, (10, 20, 30))) - img_lps = tio.ScalarImage( - torch.rand(1, 10, 20, 30), - affine=affine, - ) - fig_lps = img_lps.plot(show=False) - - titles_ras = [ax.get_title().split("[")[0].strip() for ax in fig_ras.axes] - titles_lps = [ax.get_title().split("[")[0].strip() for ax in fig_lps.axes] - assert titles_ras == titles_lps == ["Sagittal", "Coronal", "Axial"] - - def test_coordinates_kwarg(self) -> None: - """Passing world coordinates resolves to the correct voxel.""" - affine = tio.AffineMatrix.from_spacing((2.0, 2.0, 2.0)) - img = tio.ScalarImage( - torch.rand(1, 50, 50, 50), - affine=affine, - ) - # origin = (0,0,0), spacing = 2mm → coord 20mm = voxel 10 - fig = img.plot(coordinates=(20.0, 20.0, 20.0), show=False) - assert isinstance(fig, Figure) - # Title should show index 10 - titles = [ax.get_title() for ax in fig.axes] - assert any("10" in t for t in titles) - - def test_coordinates_and_indices_exclusive(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - with pytest.raises(ValueError, match="mutually exclusive"): - img.plot(indices=(5, 5, 5), coordinates=(0.0, 0.0, 0.0), show=False) - - def test_coordinates_with_none(self) -> None: - """None entries in coordinates default to mid-slice.""" - img = tio.ScalarImage(torch.rand(1, 20, 20, 20)) - fig = img.plot(coordinates=(None, None, None), show=False) - assert isinstance(fig, Figure) - - -class TestReprHtml: - def test_contains_table(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - html = img._repr_html_() - assert "tio-table" in html - assert "Channels" in html - assert "Spatial shape" in html - assert "Euler angles" in html - - def test_contains_plot(self) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - html = img._repr_html_() - assert "data:image/png;base64" in html - - -class TestPlotSubject: - def test_returns_figure(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - seg=tio.LabelMap(torch.randint(0, 3, (1, 10, 10, 10))), - ) - fig = subject.plot(show=False) - assert isinstance(fig, Figure) - - def test_many_images_transposes(self) -> None: - """With >3 images, layout should transpose to rows=views, cols=images.""" - subject = tio.Subject( - **{f"img{i}": tio.ScalarImage(torch.rand(1, 10, 10, 10)) for i in range(4)} - ) - fig = subject.plot(show=False) - assert isinstance(fig, Figure) - # 3 rows (views) x 4 cols (images) = 12 axes - assert len(fig.axes) == 12 - - def test_few_images_rows(self) -> None: - """With ≤3 images, rows=images, cols=views.""" - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - t2=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - fig = subject.plot(show=False) - # 2 rows (images) x 3 cols (views) = 6 axes - assert len(fig.axes) == 6 - - def test_cmap_dict(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - fig = subject.plot(show=False, cmap_dict={"t1": "hot"}) - assert isinstance(fig, Figure) - - def test_save_to_file(self, tmp_path) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - out = tmp_path / "subject.png" - subject.plot(output_path=out, show=False) - assert out.exists() - - def test_subject_repr_html_contains_plot(self) -> None: - subject = tio.Subject( - t1=tio.ScalarImage(torch.rand(1, 10, 10, 10)), - ) - html = subject._repr_html_() - assert "data:image/png;base64" in html - - -class TestMakeGif: - def test_to_gif_creates_file(self, tmp_path) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - out = tmp_path / "test.gif" - img.to_gif(out, direction="I") - assert out.exists() - assert out.stat().st_size > 0 - - def test_to_gif_reverse(self, tmp_path) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 10, 10)) - out = tmp_path / "rev.gif" - img.to_gif(out, direction="S", reverse=True) - assert out.exists() - - def test_to_gif_no_rescale(self, tmp_path) -> None: - data = torch.randint(0, 256, (1, 8, 8, 8), dtype=torch.uint8) - img = tio.ScalarImage(data.float()) - out = tmp_path / "noscale.gif" - img.to_gif(out, direction="A", rescale=False) - assert out.exists() - - def test_to_gif_multichannel(self, tmp_path) -> None: - img = tio.ScalarImage(torch.rand(3, 8, 8, 8)) - out = tmp_path / "rgb.gif" - img.to_gif(out, direction="R") - assert out.exists() - - def test_to_gif_warns_on_quantization(self, tmp_path) -> None: - # 200 slices at 0.01s → <1ms/frame → clamped to 20ms → big mismatch - img = tio.ScalarImage(torch.rand(1, 200, 4, 4)) - out = tmp_path / "fast.gif" - with pytest.warns(RuntimeWarning, match="quantized"): - img.to_gif(out, direction="I", seconds=0.01) - - def test_to_gif_all_directions(self, tmp_path) -> None: - img = tio.ScalarImage(torch.rand(1, 10, 12, 14)) - for direction in ("I", "S", "A", "P", "R", "L"): - out = tmp_path / f"{direction}.gif" - img.to_gif(out, direction=direction, seconds=1.0) - assert out.exists() - - -class TestJupyterReturn: - def test_to_gif_returns_none_outside_jupyter(self, tmp_path) -> None: - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - out = tmp_path / "test.gif" - result = img.to_gif(out, direction="I") - assert result is None - - def test_to_gif_returns_ipy_image_in_jupyter(self, tmp_path, monkeypatch) -> None: - pytest.importorskip("IPython") - from torchio.data import image as image_module - - monkeypatch.setattr(image_module, "_in_jupyter", lambda: True) - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - out = tmp_path / "test.gif" - result = img.to_gif(out, direction="I") - from IPython.display import Image as IPyImage - - assert isinstance(result, IPyImage) - - def test_to_gif_no_path_outside_jupyter_raises(self) -> None: - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - with pytest.raises(ValueError, match="output_path is required"): - img.to_gif() - - def test_to_gif_no_path_in_jupyter(self, monkeypatch) -> None: - pytest.importorskip("IPython") - from torchio.data import image as image_module - - monkeypatch.setattr(image_module, "_in_jupyter", lambda: True) - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - result = img.to_gif() - from IPython.display import Image as IPyImage - - assert isinstance(result, IPyImage) - - @requires_ffmpeg - def test_to_video_returns_none_outside_jupyter(self, tmp_path) -> None: - pytest.importorskip("ffmpeg") - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - out = tmp_path / "test.mp4" - result = img.to_video(out, direction="I") - assert result is None - - @requires_ffmpeg - def test_to_video_returns_ipy_video_in_jupyter(self, tmp_path, monkeypatch) -> None: - pytest.importorskip("ffmpeg") - pytest.importorskip("IPython") - from torchio.data import image as image_module - - monkeypatch.setattr(image_module, "_in_jupyter", lambda: True) - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - out = tmp_path / "test.mp4" - result = img.to_video(out, direction="I") - from IPython.display import Video - - assert isinstance(result, Video) - - def test_to_video_no_path_outside_jupyter_raises(self) -> None: - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - with pytest.raises(ValueError, match="output_path is required"): - img.to_video() - - @requires_ffmpeg - def test_to_video_no_path_in_jupyter(self, monkeypatch) -> None: - pytest.importorskip("ffmpeg") - pytest.importorskip("IPython") - from torchio.data import image as image_module - - monkeypatch.setattr(image_module, "_in_jupyter", lambda: True) - img = tio.ScalarImage(torch.rand(1, 8, 8, 8)) - result = img.to_video() - from IPython.display import Video - - assert isinstance(result, Video) diff --git a/tests/transforms/__init__.py b/tests/transforms/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/transforms/augmentation/__init__.py b/tests/transforms/augmentation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/transforms/augmentation/test_oneof.py b/tests/transforms/augmentation/test_oneof.py new file mode 100644 index 000000000..1f7ab6d45 --- /dev/null +++ b/tests/transforms/augmentation/test_oneof.py @@ -0,0 +1,44 @@ +from typing import Any +from typing import cast + +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestOneOf(TorchioTestCase): + """Tests for `OneOf`.""" + + def test_wrong_input_type(self): + with pytest.raises(ValueError): + tio.OneOf(cast(Any, 1)) + + def test_negative_probabilities(self): + transforms: dict[tio.Transform, float] = { + tio.RandomAffine(): -1, + tio.RandomElasticDeformation(): 1, + } + with pytest.raises(ValueError): + tio.OneOf(transforms) + + def test_zero_probabilities(self): + with pytest.raises(ValueError): + transforms: dict[tio.Transform, float] = { + tio.RandomAffine(): 0, + tio.RandomElasticDeformation(): 0, + } + tio.OneOf(transforms) + + def test_not_transform(self): + with pytest.raises(ValueError): + tio.OneOf(cast(Any, {tio.RandomAffine: 1, tio.RandomElasticDeformation: 2})) + + def test_one_of(self): + transforms: dict[tio.Transform, float] = { + tio.RandomAffine(): 0.2, + tio.RandomElasticDeformation(max_displacement=0.5): 0.8, + } + transform = tio.OneOf(transforms) + transform(self.sample_subject) diff --git a/tests/transforms/augmentation/test_random_affine.py b/tests/transforms/augmentation/test_random_affine.py new file mode 100644 index 000000000..ecf03f710 --- /dev/null +++ b/tests/transforms/augmentation/test_random_affine.py @@ -0,0 +1,263 @@ +from typing import Any +from typing import cast + +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRandomAffine(TorchioTestCase): + """Tests for `RandomAffine`.""" + + def setUp(self): + # Set image origin far from center + super().setUp() + affine = self.sample_subject.t1.affine + affine[:3, 3] = 1e5 + + def test_rotation_image(self): + # Rotation around image center + transform = tio.RandomAffine( + degrees=(90, 90), + default_pad_value=0, + center='image', + ) + transformed = transform(self.sample_subject) + total = transformed.t1.data.sum() + self.assertNotEqual(total, 0) + + def test_rotation_origin(self): + # Rotation around far away point, image should be empty + transform = tio.RandomAffine( + degrees=(90, 90), + default_pad_value=0, + center='origin', + ) + transformed = transform(self.sample_subject) + total = transformed.t1.data.sum() + assert total == 0 + + def test_no_rotation(self): + transform = tio.RandomAffine( + scales=(1, 1), + degrees=(0, 0), + default_pad_value=0, + center='image', + ) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + transform = tio.RandomAffine( + scales=(1, 1), + degrees=(180, 180), + default_pad_value=0, + center='image', + ) + transformed = transform(self.sample_subject) + transformed = transform(transformed) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_isotropic(self): + tio.RandomAffine(isotropic=True)(self.sample_subject) + + def test_mean(self): + tio.RandomAffine(default_pad_value='mean')(self.sample_subject) + + def test_otsu(self): + tio.RandomAffine(default_pad_value='otsu')(self.sample_subject) + + def test_bad_center(self): + with pytest.raises(ValueError): + tio.RandomAffine(center='bad') + + def test_negative_scales(self): + with pytest.raises(ValueError): + tio.RandomAffine(scales=(-1, 1)) + + def test_scale_too_large(self): + with pytest.raises(ValueError): + tio.RandomAffine(scales=1.5) + + def test_scales_range_with_negative_min(self): + with pytest.raises(ValueError): + tio.RandomAffine(scales=(-1, 4)) + + def test_wrong_scales_type(self): + with pytest.raises(ValueError): + tio.RandomAffine(scales=cast(Any, 'wrong')) + + def test_wrong_degrees_type(self): + with pytest.raises(ValueError): + tio.RandomAffine(degrees=cast(Any, 'wrong')) + + def test_too_many_translation_values(self): + with pytest.raises(ValueError): + tio.RandomAffine(translation=(-10, 4, 42)) + + def test_wrong_translation_type(self): + with pytest.raises(ValueError): + tio.RandomAffine(translation=cast(Any, 'wrong')) + + def test_wrong_center(self): + with pytest.raises(ValueError): + tio.RandomAffine(center=cast(Any, 0)) + + def test_wrong_default_pad_value(self): + with pytest.raises(ValueError): + tio.RandomAffine(default_pad_value='wrong') + + def test_wrong_image_interpolation_type(self): + with pytest.raises(TypeError): + tio.RandomAffine(image_interpolation=cast(Any, 0)) + + def test_wrong_image_interpolation_value(self): + with pytest.raises(ValueError): + tio.RandomAffine(image_interpolation='wrong') + + def test_incompatible_args_isotropic(self): + with pytest.raises(ValueError): + tio.RandomAffine(scales=(0.8, 0.5, 0.1), isotropic=True) + + def test_parse_scales(self): + def do_assert(transform: tio.RandomAffine) -> None: + assert transform.scales == 3 * (0.9, 1.1) + + triplet_scales: tuple[float, float, float] = (0.1, 0.1, 0.1) + sextet_scales: tuple[float, float, float, float, float, float] = ( + 0.9, + 1.1, + 0.9, + 1.1, + 0.9, + 1.1, + ) + do_assert(tio.RandomAffine(scales=0.1)) + do_assert(tio.RandomAffine(scales=(0.9, 1.1))) + do_assert(tio.RandomAffine(scales=triplet_scales)) + do_assert(tio.RandomAffine(scales=sextet_scales)) + + def test_parse_degrees(self): + def do_assert(transform: tio.RandomAffine) -> None: + assert transform.degrees == 3 * (-10, 10) + + triplet_degrees: tuple[int, int, int] = (10, 10, 10) + sextet_degrees: tuple[int, int, int, int, int, int] = ( + -10, + 10, + -10, + 10, + -10, + 10, + ) + do_assert(tio.RandomAffine(degrees=10)) + do_assert(tio.RandomAffine(degrees=(-10, 10))) + do_assert(tio.RandomAffine(degrees=triplet_degrees)) + do_assert(tio.RandomAffine(degrees=sextet_degrees)) + + def test_parse_translation(self): + def do_assert(transform: tio.RandomAffine) -> None: + assert transform.translation == 3 * (-10, 10) + + triplet_translation: tuple[int, int, int] = (10, 10, 10) + sextet_translation: tuple[int, int, int, int, int, int] = ( + -10, + 10, + -10, + 10, + -10, + 10, + ) + do_assert(tio.RandomAffine(translation=10)) + do_assert(tio.RandomAffine(translation=(-10, 10))) + do_assert(tio.RandomAffine(translation=triplet_translation)) + do_assert(tio.RandomAffine(translation=sextet_translation)) + + def test_default_value_label_map(self): + # From https://github.com/TorchIO-project/torchio/issues/626 + a = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).reshape(1, 3, 3, 1) + image = tio.LabelMap(tensor=a) + aff = tio.RandomAffine(translation=(0, 1, 1), default_pad_value='otsu') + transformed = aff(image) + assert all(n in (0, 1) for n in transformed.data.flatten()) + + def test_default_pad_label_parameter(self): + # Test for issue #1304: Using default_pad_value if image is of type LABEL + # Create a simple label map + label_data = torch.ones((1, 2, 2, 2)) + subject = tio.Subject(label=tio.LabelMap(tensor=label_data)) + + # Test 1: default_pad_label should be respected + transform = tio.RandomAffine( + translation=(10, 10), + default_pad_label=250, + ) + transformed_subject = transform(subject) + + # Should contain the specified pad value for labels + message = 'default_pad_label=250 should be respected for LABEL images' + transformed_label = transformed_subject.get_label_map('label') + has_expected_value = (transformed_label.tensor == 250).any() + assert has_expected_value, message + + # Test 2: backward compatibility - default_pad_value should still be ignored for labels + message = 'default_pad_value should still be ignored for LABEL images (backward compatibility)' + aff_old = tio.RandomAffine( + translation=(-10, 10, -10, 10, -10, 10), + default_pad_value=250, # This should be ignored for labels + ) + s_aug_old = aff_old.apply_transform(subject) + + # Should still use 0 (default for labels), not the default_pad_value + augmented_label = s_aug_old.get_label_map('label') + non_one_values = augmented_label.data[augmented_label.data != 1] + all_zeros = (non_one_values == 0).all() if len(non_one_values) > 0 else True + assert all_zeros, message + + # Test 3: Test direct Affine class with default_pad_label + affine_transform = tio.Affine( + scales=(1, 1, 1), + degrees=(0, 0, 0), + translation=(5, 0, 0), + default_pad_label=123, + ) + s_affine = affine_transform.apply_transform(subject) + affine_label = s_affine.get_label_map('label') + has_affine_value = (affine_label.tensor == 123).any() + assert has_affine_value, 'Direct Affine class should respect default_pad_label' + + def test_wrong_default_pad_label(self): + with pytest.raises(ValueError): + tio.RandomAffine(default_pad_label=cast(Any, 'minimum')) + + def test_no_inverse(self): + tensor = torch.zeros((1, 2, 2, 2)) + tensor[0, 1, 1, 1] = 1 # most RAS voxel + expected = torch.zeros((1, 2, 2, 2)) + expected[0, 0, 1, 1] = 1 + scales = 1, 1, 1 + degrees = 0, 0, 90 # anterior should go left + translation = 0, 0, 0 + apply_affine = tio.Affine( + scales, + degrees, + translation, + ) + transformed = apply_affine(tensor) + self.assert_tensor_almost_equal(transformed, expected) + + def test_different_spaces(self): + t1 = self.sample_subject.t1 + label = tio.Resample(2)(self.sample_subject.label) + new_subject = tio.Subject(t1=t1, label=label) + with pytest.raises(RuntimeError): + tio.RandomAffine()(new_subject) + tio.RandomAffine(check_shape=False)(new_subject) diff --git a/tests/transforms/augmentation/test_random_affine_elastic_deformation.py b/tests/transforms/augmentation/test_random_affine_elastic_deformation.py new file mode 100644 index 000000000..500e4244a --- /dev/null +++ b/tests/transforms/augmentation/test_random_affine_elastic_deformation.py @@ -0,0 +1,353 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRandomAffineElasticDeformation(TorchioTestCase): + """Tests for `RandomAffineElasticDeformation`.""" + + def setUp(self): + # Set image origin far from center + super().setUp() + affine = self.sample_subject.t1.affine + affine[:3, 3] = 1e5 + + def test_inputs_pta_gt_one(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(p=1.5) + + def test_inputs_pta_lt_zero(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(p=-1) + + def test_inputs_interpolation_int(self): + with pytest.raises(TypeError): + tio.RandomAffineElasticDeformation(image_interpolation=1) + + def test_inputs_interpolation(self): + with pytest.raises(TypeError): + tio.RandomAffineElasticDeformation(image_interpolation=0) + + def test_num_control_points_noint(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation( + elastic_kwargs={'num_control_points': 2.5} + ) + + def test_num_control_points_small(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(elastic_kwargs={'num_control_points': 3}) + + def test_max_displacement_no_num(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation( + elastic_kwargs={'max_displacement': None} + ) + + def test_max_displacement_negative(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(elastic_kwargs={'max_displacement': -1}) + + def test_wrong_locked_borders(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(elastic_kwargs={'locked_borders': -1}) + + def test_coarse_grid_removed(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation( + elastic_kwargs={'num_control_points': (4, 5, 6), 'locked_borders': 2} + ) + + def test_folding(self): + # Assume shape is (10, 20, 30) and spacing is (1, 1, 1) + # Then grid spacing is (10/(12-2), 20/(5-2), 30/(5-2)) + # or (1, 6.7, 10), and half is (0.5, 3.3, 5) + transform = tio.RandomAffineElasticDeformation( + elastic_kwargs={'num_control_points': (12, 5, 5), 'max_displacement': 6} + ) + with pytest.warns(RuntimeWarning): + transform(self.sample_subject) + + def test_num_control_points(self): + tio.RandomAffineElasticDeformation(elastic_kwargs={'num_control_points': 5}) + tio.RandomAffineElasticDeformation( + elastic_kwargs={'num_control_points': (5, 6, 7)} + ) + + def test_max_displacement(self): + tio.RandomAffineElasticDeformation(elastic_kwargs={'max_displacement': 5}) + tio.RandomAffineElasticDeformation( + elastic_kwargs={'max_displacement': (5, 6, 7)} + ) + + def test_no_displacement(self): + transform = tio.RandomAffineElasticDeformation( + affine_kwargs={'scales': 0, 'degrees': 0, 'translation': 0}, + elastic_kwargs={'max_displacement': 0}, + ) + transformed = transform(self.sample_subject) + self.assert_tensor_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + self.assert_tensor_equal( + self.sample_subject.label.data, + transformed.label.data, + ) + + def test_rotation_image(self): + # Rotation around image center + transform = tio.RandomAffineElasticDeformation( + affine_kwargs={ + 'degrees': (90, 90), + 'default_pad_value': 0, + 'center': 'image', + } + ) + transformed = transform(self.sample_subject) + total = transformed.t1.data.sum() + self.assertNotEqual(total, 0) + + def test_rotation_origin(self): + # Rotation around far away point, image should be empty + transform = tio.RandomAffineElasticDeformation( + affine_kwargs={ + 'degrees': (90, 90), + 'default_pad_value': 0, + 'center': 'origin', + } + ) + transformed = transform(self.sample_subject) + total = transformed.t1.data.sum() + assert total == 0 + + def test_no_rotation(self): + transform = tio.RandomAffineElasticDeformation( + affine_kwargs={ + 'scales': (1, 1), + 'degrees': (0, 0), + 'default_pad_value': 0, + 'center': 'image', + }, + elastic_kwargs={'max_displacement': 0}, + ) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + transform = tio.RandomAffineElasticDeformation( + affine_kwargs={ + 'scales': (1, 1), + 'degrees': (180, 180), + 'default_pad_value': 0, + 'center': 'image', + }, + elastic_kwargs={'max_displacement': 0}, + ) + transformed = transform(self.sample_subject) + transformed = transform(transformed) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_isotropic(self): + tio.RandomAffineElasticDeformation(affine_kwargs={'isotropic': True})( + self.sample_subject + ) + + def test_mean(self): + tio.RandomAffineElasticDeformation(affine_kwargs={'default_pad_value': 'mean'})( + self.sample_subject + ) + + def test_otsu(self): + tio.RandomAffineElasticDeformation(affine_kwargs={'default_pad_value': 'otsu'})( + self.sample_subject + ) + + def test_bad_center(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'center': 'bad'}) + + def test_negative_scales(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'scales': (-1, 1)}) + + def test_scale_too_large(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 1.5}) + + def test_scales_range_with_negative_min(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'scales': (-1, 4)}) + + def test_wrong_scales_type(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 'wrong'}) + + def test_wrong_degrees_type(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 'wrong'}) + + def test_too_many_translation_values(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation( + affine_kwargs={'translation': (-10, 4, 42)} + ) + + def test_wrong_translation_type(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'translation': 'wrong'}) + + def test_wrong_center(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation(affine_kwargs={'center': 0}) + + def test_wrong_default_pad_value(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation( + affine_kwargs={'default_pad_value': 'wrong'} + ) + + def test_wrong_image_interpolation_type(self): + with pytest.raises(TypeError): + tio.RandomAffineElasticDeformation(affine_kwargs={'image_interpolation': 0}) + + def test_wrong_image_interpolation_value(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation( + affine_kwargs={'image_interpolation': 'wrong'} + ) + + def test_incompatible_args_isotropic(self): + with pytest.raises(ValueError): + tio.RandomAffineElasticDeformation( + affine_kwargs={'scales': (0.8, 0.5, 0.1), 'isotropic': True} + ) + + def test_parse_scales(self): + def do_assert(transform): + assert transform.random_affine.scales == 3 * (0.9, 1.1) + + do_assert(tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 0.1})) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'scales': (0.9, 1.1)}) + ) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 3 * (0.1,)}) + ) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'scales': 3 * [0.9, 1.1]}) + ) + + def test_parse_degrees(self): + def do_assert(transform): + assert transform.random_affine.degrees == 3 * (-10, 10) + + do_assert(tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 10})) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': (-10, 10)}) + ) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 3 * (10,)}) + ) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'degrees': 3 * [-10, 10]}) + ) + + def test_parse_translation(self): + def do_assert(transform): + assert transform.random_affine.translation == 3 * (-10, 10) + + do_assert(tio.RandomAffineElasticDeformation(affine_kwargs={'translation': 10})) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'translation': (-10, 10)}) + ) + do_assert( + tio.RandomAffineElasticDeformation(affine_kwargs={'translation': 3 * (10,)}) + ) + do_assert( + tio.RandomAffineElasticDeformation( + affine_kwargs={'translation': 3 * [-10, 10]} + ) + ) + + def test_default_value_label_map(self): + # From https://github.com/fepegar/torchio/issues/626 + a = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).reshape(1, 3, 3, 1) + image = tio.LabelMap(tensor=a) + aff = tio.RandomAffineElasticDeformation( + affine_kwargs={'translation': (0, 1, 1), 'default_pad_value': 'otsu'} + ) + transformed = aff(image) + assert all(n in (0, 1) for n in transformed.data.flatten()) + + def test_no_inverse(self): + tensor = torch.zeros((1, 2, 2, 2)) + tensor[0, 1, 1, 1] = 1 # most RAS voxel + expected = torch.zeros((1, 2, 2, 2)) + expected[0, 0, 1, 1] = 1 + scales = 1, 1, 1 + degrees = 0, 0, 90 # anterior should go left + translation = 0, 0, 0 + apply_affine = tio.Affine( + scales, + degrees, + translation, + ) + transformed = apply_affine(tensor) + self.assert_tensor_almost_equal(transformed, expected) + + def test_different_spaces(self): + t1 = self.sample_subject.t1 + label = tio.Resample(2)(self.sample_subject.label) + new_subject = tio.Subject(t1=t1, label=label) + with pytest.raises(RuntimeError): + tio.RandomAffineElasticDeformation()(new_subject) + tio.RandomAffineElasticDeformation(affine_kwargs={'check_shape': False})( + new_subject + ) + + def test_transform_order(self): + src_transform = tio.RandomAffineElasticDeformation( + affine_kwargs={'scales': 0, 'degrees': 0, 'translation': 1}, + elastic_kwargs={'num_control_points': 5, 'max_displacement': 1}, + ) + + (scales, degrees, translation), control_points = src_transform.get_params() + + max_displacement = src_transform.random_elastic.max_displacement + + transform1 = tio.AffineElasticDeformation( + affine_first=True, + affine_params={ + 'scales': scales, + 'degrees': degrees, + 'translation': translation, + }, + elastic_params={ + 'control_points': control_points, + 'max_displacement': max_displacement, + }, + ) + transform2 = tio.AffineElasticDeformation( + affine_first=False, + affine_params={ + 'scales': scales, + 'degrees': degrees, + 'translation': translation, + }, + elastic_params={ + 'control_points': control_points, + 'max_displacement': max_displacement, + }, + ) + + transformed1 = transform1(self.sample_subject) + transformed2 = transform2(self.sample_subject) + self.assert_tensor_not_equal(transformed1.t1.data, transformed2.t1.data) diff --git a/tests/transforms/augmentation/test_random_anisotropy.py b/tests/transforms/augmentation/test_random_anisotropy.py new file mode 100644 index 000000000..ec81db6e6 --- /dev/null +++ b/tests/transforms/augmentation/test_random_anisotropy.py @@ -0,0 +1,46 @@ +from typing import Any +from typing import cast + +import pytest +import torch + +from torchio import RandomAnisotropy +from torchio import ScalarImage + +from ...utils import TorchioTestCase + + +class TestRandomAnisotropy(TorchioTestCase): + """Tests for `RandomAnisotropy`.""" + + def test_downsample(self): + transform = RandomAnisotropy( + axes=1, + downsampling=(2, 2), + ) + transformed = transform(self.sample_subject) + assert self.sample_subject.spacing[1] == transformed.spacing[1] + + def test_out_of_range_axis(self): + with pytest.raises(ValueError): + RandomAnisotropy(axes=3) + + def test_out_of_range_axis_in_tuple(self): + with pytest.raises(ValueError): + RandomAnisotropy(axes=(0, -1, 2)) + + def test_wrong_axes_type(self): + with pytest.raises(ValueError): + RandomAnisotropy(axes=cast(Any, 'wrong')) + + def test_wrong_downsampling_type(self): + with pytest.raises(ValueError): + RandomAnisotropy(downsampling=cast(Any, 'wrong')) + + def test_below_one_downsampling(self): + with pytest.raises(ValueError): + RandomAnisotropy(downsampling=0.2) + + def test_2d_rgb(self): + image = ScalarImage(tensor=torch.rand(3, 4, 5, 6)) + RandomAnisotropy()(image) diff --git a/tests/transforms/augmentation/test_random_bias_field.py b/tests/transforms/augmentation/test_random_bias_field.py new file mode 100644 index 000000000..54c351020 --- /dev/null +++ b/tests/transforms/augmentation/test_random_bias_field.py @@ -0,0 +1,43 @@ +from typing import Any +from typing import cast + +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRandomBiasField(TorchioTestCase): + def test_no_bias(self): + transform = tio.RandomBiasField(coefficients=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_bias(self): + transform = tio.RandomBiasField(coefficients=0.1) + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_wrong_coefficient_type(self): + with pytest.raises(ValueError): + tio.RandomBiasField(coefficients=cast(Any, 'wrong')) + + def test_negative_order(self): + with pytest.raises(ValueError): + tio.RandomBiasField(order=-1) + + def test_wrong_order_type(self): + with pytest.raises(TypeError): + tio.RandomBiasField(order=cast(Any, 'wrong')) + + def test_small_image(self): + # https://github.com/TorchIO-project/torchio/issues/300 + tio.RandomBiasField()(torch.rand(1, 2, 3, 4)) diff --git a/tests/transforms/augmentation/test_random_blur.py b/tests/transforms/augmentation/test_random_blur.py new file mode 100644 index 000000000..b87dc706d --- /dev/null +++ b/tests/transforms/augmentation/test_random_blur.py @@ -0,0 +1,51 @@ +from typing import Any +from typing import cast + +import pytest + +from torchio import RandomBlur + +from ...utils import TorchioTestCase + + +class TestRandomBlur(TorchioTestCase): + """Tests for `RandomBlur`.""" + + def test_no_blurring(self): + transform = RandomBlur(std=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_blurring(self): + transform = RandomBlur(std=(1, 3)) + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_negative_std(self): + with pytest.raises(ValueError): + RandomBlur(std=-2) + + def test_std_range_with_negative_min(self): + with pytest.raises(ValueError): + RandomBlur(std=(-0.5, 4)) + + def test_wrong_std_type(self): + with pytest.raises(ValueError): + RandomBlur(std=cast(Any, 'wrong')) + + def test_parse_stds(self): + def do_assert(transform: RandomBlur) -> None: + assert transform.std_ranges == 3 * (0, 1) + + triplet_std: tuple[int, int, int] = (1, 1, 1) + sextet_std: tuple[int, int, int, int, int, int] = (0, 1, 0, 1, 0, 1) + do_assert(RandomBlur(std=1)) + do_assert(RandomBlur(std=(0, 1))) + do_assert(RandomBlur(std=cast(Any, triplet_std))) + do_assert(RandomBlur(std=cast(Any, sextet_std))) diff --git a/tests/transforms/augmentation/test_random_elastic_deformation.py b/tests/transforms/augmentation/test_random_elastic_deformation.py new file mode 100644 index 000000000..801b6504b --- /dev/null +++ b/tests/transforms/augmentation/test_random_elastic_deformation.py @@ -0,0 +1,86 @@ +from typing import Any +from typing import cast + +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRandomElasticDeformation(TorchioTestCase): + """Tests for `RandomElasticDeformation`.""" + + def test_inputs_pta_gt_one(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation(p=1.5) + + def test_inputs_pta_lt_zero(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation(p=-1) + + def test_inputs_interpolation_int(self): + with pytest.raises(TypeError): + tio.RandomElasticDeformation(image_interpolation=cast(Any, 1)) + + def test_inputs_interpolation(self): + with pytest.raises(TypeError): + tio.RandomElasticDeformation(image_interpolation=cast(Any, 0)) + + def test_num_control_points_noint(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation(num_control_points=cast(Any, 2.5)) + + def test_num_control_points_small(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation(num_control_points=3) + + def test_max_displacement_no_num(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation(max_displacement=cast(Any, None)) + + def test_max_displacement_negative(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation(max_displacement=-1) + + def test_wrong_locked_borders(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation(locked_borders=-1) + + def test_coarse_grid_removed(self): + with pytest.raises(ValueError): + tio.RandomElasticDeformation( + num_control_points=(4, 5, 6), + locked_borders=2, + ) + + def test_folding(self): + # Assume shape is (10, 20, 30) and spacing is (1, 1, 1) + # Then grid spacing is (10/(12-2), 20/(5-2), 30/(5-2)) + # or (1, 6.7, 10), and half is (0.5, 3.3, 5) + transform = tio.RandomElasticDeformation( + num_control_points=(12, 5, 5), + max_displacement=6, + ) + with pytest.warns(RuntimeWarning): + transform(self.sample_subject) + + def test_num_control_points(self): + tio.RandomElasticDeformation(num_control_points=5) + tio.RandomElasticDeformation(num_control_points=(5, 6, 7)) + + def test_max_displacement(self): + tio.RandomElasticDeformation(max_displacement=5) + tio.RandomElasticDeformation(max_displacement=(5, 6, 7)) + + def test_no_displacement(self): + transform = tio.RandomElasticDeformation(max_displacement=0) + transformed = transform(self.sample_subject) + self.assert_tensor_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + self.assert_tensor_equal( + self.sample_subject.label.data, + transformed.label.data, + ) diff --git a/tests/transforms/augmentation/test_random_flip.py b/tests/transforms/augmentation/test_random_flip.py new file mode 100644 index 000000000..5c7b35f69 --- /dev/null +++ b/tests/transforms/augmentation/test_random_flip.py @@ -0,0 +1,45 @@ +from typing import Any +from typing import cast + +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRandomFlip(TorchioTestCase): + def test_2d(self): + subject = self.make_2d(self.sample_subject) + transform = tio.RandomFlip(axes=(1, 2), flip_probability=1) + transformed = transform(subject) + self.assert_tensor_equal( + torch.from_numpy(subject.t1.data.numpy()[..., ::-1, ::-1].copy()), + transformed.t1.data, + ) + + def test_out_of_range_axis(self): + with pytest.raises(ValueError): + tio.RandomFlip(axes=3) + + def test_out_of_range_axis_in_tuple(self): + with pytest.raises(ValueError): + tio.RandomFlip(axes=(0, -1, 2)) + + def test_wrong_axes_type(self): + with pytest.raises(ValueError): + tio.RandomFlip(axes=cast(Any, None)) + + def test_wrong_flip_probability_type(self): + with pytest.raises(ValueError): + tio.RandomFlip(flip_probability=cast(Any, 'wrong')) + + def test_anatomical_axis(self): + transform = tio.RandomFlip(axes=cast(Any, ['i']), flip_probability=1) + tensor = torch.rand(1, 2, 3, 4) + transformed = transform(tensor) + self.assert_tensor_equal( + torch.from_numpy(tensor.numpy()[..., ::-1].copy()), + transformed, + ) diff --git a/tests/transforms/augmentation/test_random_gamma.py b/tests/transforms/augmentation/test_random_gamma.py new file mode 100644 index 000000000..221a01fa1 --- /dev/null +++ b/tests/transforms/augmentation/test_random_gamma.py @@ -0,0 +1,50 @@ +from typing import Any +from typing import cast + +import pytest +import torch + +from torchio import RandomGamma + +from ...utils import TorchioTestCase + + +class TestRandomGamma(TorchioTestCase): + """Tests for `RandomGamma`.""" + + def get_random_tensor_zero_one(self): + return torch.rand(4, 5, 6, 7) + + def test_with_zero_gamma(self): + transform = RandomGamma(log_gamma=0) + tensor = self.get_random_tensor_zero_one() + transformed = transform(tensor) + self.assert_tensor_almost_equal(tensor, transformed) + + def test_with_non_zero_gamma(self): + transform = RandomGamma(log_gamma=(0.1, 0.3)) + tensor = self.get_random_tensor_zero_one() + transformed = transform(tensor) + self.assert_tensor_not_equal(tensor, transformed) + + def test_with_high_gamma(self): + transform = RandomGamma(log_gamma=(100, 100)) + tensor = self.get_random_tensor_zero_one() + transformed = transform(tensor) + self.assert_tensor_almost_equal( + tensor == 1, + transformed, + ) + + def test_with_low_gamma(self): + transform = RandomGamma(log_gamma=(-100, -100)) + tensor = self.get_random_tensor_zero_one() + transformed = transform(tensor) + self.assert_tensor_almost_equal( + tensor > 0, + transformed, + ) + + def test_wrong_gamma_type(self): + with pytest.raises(ValueError): + RandomGamma(log_gamma=cast(Any, 'wrong')) diff --git a/tests/transforms/augmentation/test_random_ghosting.py b/tests/transforms/augmentation/test_random_ghosting.py new file mode 100644 index 000000000..7ba42eca4 --- /dev/null +++ b/tests/transforms/augmentation/test_random_ghosting.py @@ -0,0 +1,80 @@ +from typing import Any +from typing import cast + +import pytest + +from torchio import RandomGhosting + +from ...utils import TorchioTestCase + + +class TestRandomGhosting(TorchioTestCase): + """Tests for `RandomGhosting`.""" + + def test_with_zero_intensity(self): + transform = RandomGhosting(intensity=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_zero_ghost(self): + transform = RandomGhosting(num_ghosts=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_ghosting(self): + transform = RandomGhosting() + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_intensity_range_with_negative_min(self): + with pytest.raises(ValueError): + RandomGhosting(intensity=(-0.5, 4)) + + def test_wrong_intensity_type(self): + with pytest.raises(ValueError): + RandomGhosting(intensity=cast(Any, 'wrong')) + + def test_negative_num_ghosts(self): + with pytest.raises(ValueError): + RandomGhosting(num_ghosts=-1) + + def test_num_ghosts_range_with_negative_min(self): + with pytest.raises(ValueError): + RandomGhosting(num_ghosts=(-1, 4)) + + def test_not_integer_num_ghosts(self): + with pytest.raises(ValueError): + RandomGhosting(num_ghosts=cast(Any, (0.7, 4))) + + def test_wrong_num_ghosts_type(self): + with pytest.raises(ValueError): + RandomGhosting(num_ghosts=cast(Any, 'wrong')) + + def test_out_of_range_axis(self): + with pytest.raises(ValueError): + RandomGhosting(axes=3) + + def test_out_of_range_axis_in_tuple(self): + with pytest.raises(ValueError): + RandomGhosting(axes=(0, -1, 2)) + + def test_wrong_axes_type(self): + with pytest.raises(ValueError): + RandomGhosting(axes=cast(Any, None)) + + def test_out_of_range_restore(self): + with pytest.raises(ValueError): + RandomGhosting(restore=-1) + + def test_wrong_restore_type(self): + with pytest.raises(ValueError): + RandomGhosting(restore=cast(Any, 'wrong')) diff --git a/tests/transforms/augmentation/test_random_labels_to_image.py b/tests/transforms/augmentation/test_random_labels_to_image.py new file mode 100644 index 000000000..b5416e016 --- /dev/null +++ b/tests/transforms/augmentation/test_random_labels_to_image.py @@ -0,0 +1,302 @@ +from typing import Any +from typing import cast + +import pytest + +from torchio.transforms import RandomLabelsToImage + +from ...utils import TorchioTestCase + + +class TestRandomLabelsToImage(TorchioTestCase): + """Tests for `RandomLabelsToImage`.""" + + def test_random_simulation(self): + """The transform runs without error and an 'image_from_labels' key is + present in the transformed subject.""" + transform = RandomLabelsToImage(label_key='label') + transformed = transform(self.sample_subject) + self.assertIn('image_from_labels', transformed) + + def test_deterministic_simulation(self): + """The transform creates an image where values are equal to given mean + if standard deviation is zero. + + Using a label map. + """ + transform = RandomLabelsToImage( + label_key='label', + mean=[0.5, 2], + std=[0, 0], + ) + transformed = transform(self.sample_subject) + sample_label = self.sample_subject.get_label_map('label') + self.assert_tensor_equal( + transformed['image_from_labels'].data == 0.5, + sample_label.data == 0, + ) + self.assert_tensor_equal( + transformed['image_from_labels'].data == 2, + sample_label.data == 1, + ) + + def test_deterministic_simulation_with_discretized_label_map(self): + """The transform creates an image where values are equal to given mean + if standard deviation is zero. + + Using a discretized label map. + """ + transform = RandomLabelsToImage( + label_key='label', + mean=[0.5, 2], + std=[0, 0], + discretize=True, + ) + transformed = transform(self.sample_subject) + sample_label = self.sample_subject.get_label_map('label') + self.assert_tensor_equal( + transformed['image_from_labels'].data == 0.5, + sample_label.data == 0, + ) + self.assert_tensor_equal( + transformed['image_from_labels'].data == 2, + sample_label.data == 1, + ) + + def test_deterministic_simulation_with_pv_map(self): + """The transform creates an image where values are equal to given mean + weighted by partial-volume if standard deviation is zero.""" + subject = self.get_subject_with_partial_volume_label_map(components=2) + transform = RandomLabelsToImage( + label_key='label', + mean=[0.5, 1], + std=[0, 0], + ) + transformed = transform(subject) + self.assert_tensor_almost_equal( + transformed['image_from_labels'].data[0], + subject['label'].data[0] * 0.5 + subject['label'].data[1] * 1, + ) + assert transformed['image_from_labels'].data.shape == (1, 10, 20, 30) + + def test_deterministic_simulation_with_discretized_pv_map(self): + """The transform creates an image where values are equal to given mean + if standard deviation is zero. + + Using a discretized partial-volume label map. + """ + subject = self.get_subject_with_partial_volume_label_map() + transform = RandomLabelsToImage( + label_key='label', + mean=[0.5], + std=[0], + discretize=True, + ) + transformed = transform(subject) + self.assert_tensor_almost_equal( + transformed['image_from_labels'].data, + (subject['label'].data > 0) * 0.5, + ) + + def test_filling(self): + """The transform can fill in the generated image with an already + existing image. + + Using a label map. + """ + transform = RandomLabelsToImage( + label_key='label', + image_key='t1', + used_labels=[1], + ) + sample_label = self.sample_subject.get_label_map('label') + sample_t1 = self.sample_subject.get_scalar_image('t1') + t1_indices = sample_label.data == 0 + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + transformed['t1'].data[t1_indices], + sample_t1.data[t1_indices], + ) + + def test_filling_with_discretized_label_map(self): + """The transform can fill in the generated image with an already + existing image. + + Using a discretized label map. + """ + transform = RandomLabelsToImage( + label_key='label', + image_key='t1', + discretize=True, + used_labels=[1], + ) + sample_label = self.sample_subject.get_label_map('label') + sample_t1 = self.sample_subject.get_scalar_image('t1') + t1_indices = sample_label.data < 0.5 + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + transformed['t1'].data[t1_indices], + sample_t1.data[t1_indices], + ) + + def test_filling_with_discretized_pv_label_map(self): + """The transform can fill in the generated image with an already + existing image. + + Using a discretized partial-volume label map. + """ + subject = self.get_subject_with_partial_volume_label_map(components=2) + transform = RandomLabelsToImage( + label_key='label', + image_key='t1', + discretize=True, + used_labels=[1], + ) + t1_indices = subject['label'].data.argmax(dim=0) == 0 + transformed = transform(subject) + self.assert_tensor_almost_equal( + transformed['t1'].data[0][t1_indices], + subject['t1'].data[0][t1_indices], + ) + + def test_filling_without_any_hole(self): + """The transform does not fill anything if there is no hole.""" + transform = RandomLabelsToImage( + label_key='label', + image_key='t1', + default_std=0, + default_mean=-1, + ) + original_t1 = self.sample_subject.t1.data.clone() + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal(original_t1, transformed.t1.data) + + def test_with_bad_default_mean_range(self): + """The transform raises an error if default_mean is not a single value + nor a tuple of two values.""" + with pytest.raises(ValueError): + RandomLabelsToImage( + label_key='label', + default_mean=cast(Any, (0, 1, 2)), + ) + + def test_with_bad_default_mean_type(self): + """The transform raises an error if default_mean has the wrong type.""" + with pytest.raises(ValueError): + RandomLabelsToImage( + label_key='label', + default_mean=cast(Any, 'wrong'), + ) + + def test_with_bad_default_std_range(self): + """The transform raises an error if default_std is not a single value + nor a tuple of two values.""" + with pytest.raises(ValueError): + RandomLabelsToImage( + label_key='label', + default_std=cast(Any, (0, 1, 2)), + ) + + def test_with_bad_default_std_type(self): + """The transform raises an error if default_std has the wrong type.""" + with pytest.raises(ValueError): + RandomLabelsToImage( + label_key='label', + default_std=cast(Any, 'wrong'), + ) + + def test_with_wrong_label_key_type(self): + """The transform raises an error if a wrong type is given for + label_key.""" + with pytest.raises(TypeError): + RandomLabelsToImage(label_key=cast(Any, 42)) + + def test_with_wrong_used_labels_type(self): + """The transform raises an error if a wrong type is given for + used_labels.""" + with pytest.raises(TypeError): + RandomLabelsToImage( + label_key='label', + used_labels=cast(Any, 42), + ) + + def test_with_wrong_used_labels_elements_type(self): + """The transform raises an error if wrong type are given for + used_labels elements.""" + with pytest.raises(ValueError): + RandomLabelsToImage( + label_key='label', + used_labels=cast(Any, ['wrong']), + ) + + def test_with_wrong_mean_type(self): + """The transform raises an error if wrong type is given for mean.""" + with pytest.raises(TypeError): + RandomLabelsToImage(label_key='label', mean=cast(Any, 42)) + + def test_with_wrong_mean_elements_type(self): + """The transform raises an error if wrong type are given for mean + elements.""" + with pytest.raises(ValueError): + RandomLabelsToImage( + label_key='label', + mean=cast(Any, ['wrong']), + ) + + def test_with_wrong_std_type(self): + """The transform raises an error if wrong type is given for std.""" + with pytest.raises(TypeError): + RandomLabelsToImage(label_key='label', std=cast(Any, 42)) + + def test_with_wrong_std_elements_type(self): + """The transform raises an error if wrong type are given for std + elements.""" + with pytest.raises(ValueError): + RandomLabelsToImage( + label_key='label', + std=cast(Any, ['wrong']), + ) + + def test_mean_and_std_len_not_matching(self): + """The transform raises an error if mean and std length don't match.""" + with pytest.raises(AssertionError): + RandomLabelsToImage(label_key='label', mean=[0], std=[0, 1]) + + def test_mean_and_used_labels_len_not_matching(self): + """The transform raises an error if mean and used_labels length don't + match.""" + with pytest.raises(AssertionError): + RandomLabelsToImage( + label_key='label', + mean=[0], + used_labels=[0, 1], + ) + + def test_std_and_used_labels_len_not_matching(self): + """The transform raises an error if std and used_labels length don't + match.""" + with pytest.raises(AssertionError): + RandomLabelsToImage(label_key='label', std=[0], used_labels=[0, 1]) + + def test_mean_not_matching_number_of_labels(self): + """The transform raises an error at runtime if mean length does not + match label numbers.""" + transform = RandomLabelsToImage(label_key='label', mean=[0]) + with pytest.raises(RuntimeError): + transform(self.sample_subject) + + def test_std_not_matching_number_of_labels(self): + """The transform raises an error at runtime if std length does not + match label numbers.""" + transform = RandomLabelsToImage(label_key='label', std=[1, 2, 3]) + with pytest.raises(RuntimeError): + transform(self.sample_subject) + + def test_bad_range(self): + with pytest.raises(ValueError): + RandomLabelsToImage(default_mean=(2, 1)) + + def test_no_labels(self): + transform = RandomLabelsToImage() + with pytest.raises(RuntimeError): + transform(self.sample_subject.t1) diff --git a/tests/transforms/augmentation/test_random_motion.py b/tests/transforms/augmentation/test_random_motion.py new file mode 100644 index 000000000..02d526a95 --- /dev/null +++ b/tests/transforms/augmentation/test_random_motion.py @@ -0,0 +1,64 @@ +from typing import Any +from typing import cast + +import pytest + +from torchio import RandomMotion + +from ...utils import TorchioTestCase + + +class TestRandomMotion(TorchioTestCase): + """Tests for `RandomMotion`.""" + + def test_bad_num_transforms_value(self): + with pytest.raises(ValueError): + RandomMotion(num_transforms=0) + + def test_no_movement(self): + transform = RandomMotion( + degrees=0, + translation=0, + num_transforms=1, + ) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + atol=1e-4, + rtol=0, + ) + + def test_with_movement(self): + transform = RandomMotion( + num_transforms=1, + ) + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_negative_degrees(self): + with pytest.raises(ValueError): + RandomMotion(degrees=-10) + + def test_wrong_degrees_type(self): + with pytest.raises(ValueError): + RandomMotion(degrees=cast(Any, 'wrong')) + + def test_negative_translation(self): + with pytest.raises(ValueError): + RandomMotion(translation=-10) + + def test_wrong_translation_type(self): + with pytest.raises(ValueError): + RandomMotion(translation=cast(Any, 'wrong')) + + def test_wrong_image_interpolation_type(self): + with pytest.raises(TypeError): + RandomMotion(image_interpolation=cast(Any, 0)) + + def test_wrong_image_interpolation_value(self): + with pytest.raises(ValueError): + RandomMotion(image_interpolation='wrong') diff --git a/tests/transforms/augmentation/test_random_noise.py b/tests/transforms/augmentation/test_random_noise.py new file mode 100644 index 000000000..ed764e5b4 --- /dev/null +++ b/tests/transforms/augmentation/test_random_noise.py @@ -0,0 +1,52 @@ +from typing import Any +from typing import cast + +import pytest + +from torchio import RandomNoise + +from ...utils import TorchioTestCase + + +class TestRandomNoise(TorchioTestCase): + """Tests for `RandomNoise`.""" + + def test_no_noise(self): + transform = RandomNoise(mean=0, std=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_noise(self): + transform = RandomNoise() + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_constant_noise(self): + transform = RandomNoise(mean=(5, 5), std=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data + 5, + transformed.t1.data, + ) + + def test_negative_std(self): + with pytest.raises(ValueError): + RandomNoise(std=-2) + + def test_std_range_with_negative_min(self): + with pytest.raises(ValueError): + RandomNoise(std=(-0.5, 4)) + + def test_wrong_std_type(self): + with pytest.raises(ValueError): + RandomNoise(std=cast(Any, 'wrong')) + + def test_wrong_mean_type(self): + with pytest.raises(ValueError): + RandomNoise(mean=cast(Any, 'wrong')) diff --git a/tests/transforms/augmentation/test_random_spike.py b/tests/transforms/augmentation/test_random_spike.py new file mode 100644 index 000000000..d30b93498 --- /dev/null +++ b/tests/transforms/augmentation/test_random_spike.py @@ -0,0 +1,56 @@ +from typing import Any +from typing import cast + +import pytest + +from torchio import RandomSpike + +from ...utils import TorchioTestCase + + +class TestRandomSpike(TorchioTestCase): + """Tests for `RandomSpike`.""" + + def test_with_zero_intensity(self): + transform = RandomSpike(intensity=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_zero_spike(self): + transform = RandomSpike(num_spikes=0) + transformed = transform(self.sample_subject) + self.assert_tensor_almost_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_spikes(self): + transform = RandomSpike() + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_negative_num_spikes(self): + with pytest.raises(ValueError): + RandomSpike(num_spikes=-1) + + def test_num_spikes_range_with_negative_min(self): + with pytest.raises(ValueError): + RandomSpike(num_spikes=(-1, 4)) + + def test_not_integer_num_spikes(self): + with pytest.raises(ValueError): + RandomSpike(num_spikes=cast(Any, (0.7, 4))) + + def test_wrong_num_spikes_type(self): + with pytest.raises(ValueError): + RandomSpike(num_spikes=cast(Any, 'wrong')) + + def test_wrong_intensity_type(self): + with pytest.raises(ValueError): + RandomSpike(intensity=cast(Any, 'wrong')) diff --git a/tests/transforms/augmentation/test_random_swap.py b/tests/transforms/augmentation/test_random_swap.py new file mode 100644 index 000000000..b5c68d4eb --- /dev/null +++ b/tests/transforms/augmentation/test_random_swap.py @@ -0,0 +1,34 @@ +from typing import Any +from typing import cast + +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRandomSwap(TorchioTestCase): + def test_no_swap(self): + transform = tio.RandomSwap(patch_size=5, num_iterations=0) + transformed = transform(self.sample_subject) + self.assert_tensor_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_with_swap(self): + transform = tio.RandomSwap(patch_size=5) + transformed = transform(self.sample_subject) + self.assert_tensor_not_equal( + self.sample_subject.t1.data, + transformed.t1.data, + ) + + def test_wrong_num_iterations_type(self): + with pytest.raises(TypeError): + tio.RandomSwap(num_iterations=cast(Any, 'wrong')) + + def test_negative_num_iterations(self): + with pytest.raises(ValueError): + tio.RandomSwap(num_iterations=-1) diff --git a/tests/transforms/label/__init__.py b/tests/transforms/label/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/transforms/label/test_remap_labels.py b/tests/transforms/label/test_remap_labels.py new file mode 100644 index 000000000..458e8e2a3 --- /dev/null +++ b/tests/transforms/label/test_remap_labels.py @@ -0,0 +1,56 @@ +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +@pytest.mark.parametrize( + 'original_label_set', + ( + {0}, + {0, 1}, + {0, 1, 2}, + {0, 2}, + {1, 2, 5, 6}, # values from original @efirdc test + ), +) +@pytest.mark.parametrize( + 'remapping', + ( + {}, + {0: 10}, + {0: 10, 1: 11, 2: 12}, + {0: 1}, + {0: 1, 1: 0}, + {0: 1, 1: 2, 2: 0}, + {2: 1, 5: 1}, + {3: 4}, + {3: 1}, + {1: 2, 2: 1, 5: 10, 6: 11}, # values from original @efirdc test + ), +) +def test_remap(original_label_set, remapping): + source_label_set = set(remapping.keys()) + target_label_set = set(remapping.values()) + remap_labels = tio.RemapLabels(remapping=remapping) + tensor = TorchioTestCase.get_tensor_with_labels(original_label_set) + subject = tio.Subject(label=tio.LabelMap(tensor=tensor)) + transformed = remap_labels(subject) + + new_label_set = TorchioTestCase.get_unique_labels(transformed.label.data) + + if source_label_set.intersection(original_label_set): + assert new_label_set.intersection(target_label_set) + else: + assert new_label_set == original_label_set + + if len(target_label_set) < len(remapping.keys()): + with pytest.raises(RuntimeError): + _ = transformed.apply_inverse_transform() + else: + inverse_data = transformed.apply_inverse_transform().label.data + inverted_label_set = TorchioTestCase.get_unique_labels(inverse_data) + # Users are warned about this in the docs for the transform + if target_label_set.isdisjoint(original_label_set): + assert inverted_label_set == original_label_set diff --git a/tests/transforms/label/test_remove_labels.py b/tests/transforms/label/test_remove_labels.py new file mode 100644 index 000000000..1c5536f0b --- /dev/null +++ b/tests/transforms/label/test_remove_labels.py @@ -0,0 +1,29 @@ +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRemoveLabels(TorchioTestCase): + """Tests for `RemoveLabels`.""" + + def test_remove(self): + original_labels = (1, 2, 3, 4, 5, 6, 7) + labels_to_remove = (1, 2, 5, 6) + remaining_labels = (3, 4, 7) + + remove_labels = tio.RemoveLabels(labels_to_remove) + + tensor = TorchioTestCase.get_tensor_with_labels(original_labels) + subject = tio.Subject(label=tio.LabelMap(tensor=tensor)) + transformed = remove_labels(subject) + + for removed_label in labels_to_remove: + original_mask = subject.label.data == removed_label + new_values = transformed.label.data[original_mask] + self.assert_tensor_all_zeros(new_values) + + for remaining_label in remaining_labels: + original_mask = subject.label.data == remaining_label + original_values = subject.label.data[original_mask] + output_values = transformed.label.data[original_mask] + self.assert_tensor_equal(original_values, output_values) diff --git a/tests/transforms/label/test_sequential_labels.py b/tests/transforms/label/test_sequential_labels.py new file mode 100644 index 000000000..d06771e98 --- /dev/null +++ b/tests/transforms/label/test_sequential_labels.py @@ -0,0 +1,34 @@ +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +@pytest.mark.parametrize( + 'original_labels', + ( + (0,), + (0, 1), + (0, 1, 2), + (0, 2), + (0, 4, 8), + (1,), + (1, 2), + (3, 5, 9, 15, 16, 23), # values from original @efirdc docstring + (0, 3, 5, 9, 15, 16, 23), + (2, 8, 9, 10, 15, 20, 100), # values from original @efirdc test + (0, 2, 8, 9, 10, 15, 20, 100), + ), +) +def test_sequential(original_labels): + remap_labels = tio.SequentialLabels() + tensor = TorchioTestCase.get_tensor_with_labels(original_labels) + subject = tio.Subject(label=tio.LabelMap(tensor=tensor)) + transformed = remap_labels(subject) + for i, label in enumerate(original_labels): + original_mask = tensor == label + new_mask = transformed.label.data == i + TorchioTestCase.assert_tensor_equal(original_mask, new_mask) + inverted = transformed.apply_inverse_transform() + TorchioTestCase.assert_tensor_equal(tensor, inverted.label.data) diff --git a/tests/transforms/preprocessing/__init__.py b/tests/transforms/preprocessing/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/transforms/preprocessing/test_clamp.py b/tests/transforms/preprocessing/test_clamp.py new file mode 100644 index 000000000..3ae8f5bf6 --- /dev/null +++ b/tests/transforms/preprocessing/test_clamp.py @@ -0,0 +1,51 @@ +from typing import cast + +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestClamp(TorchioTestCase): + """Tests for :class:`tio.Clamp` class.""" + + def test_out_min_max(self): + transform = tio.Clamp(out_min=0, out_max=1) + transformed = transform(self.sample_subject) + assert transformed.t1.data.min() == 0 + assert transformed.t1.data.max() == 1 + + def test_ct(self): + ct_max = 1500 + ct_min = -2000 + ct_range = ct_max - ct_min + tensor = torch.rand(1, 30, 30, 30) * ct_range + ct_min + ct = tio.ScalarImage(tensor=tensor) + ct_air = -1000 + ct_bone = 1000 + clamp = tio.Clamp(ct_air, ct_bone) + clamped = clamp(ct) + assert clamped.data.min() == ct_air + assert clamped.data.max() == ct_bone + + def test_too_many_values_for_out_min(self): + with pytest.raises(TypeError): + clamp = tio.Clamp(out_min=cast(float, (1, 2))) + clamp(self.sample_subject) + + def test_too_many_values_for_out_max(self): + with pytest.raises(TypeError): + clamp = tio.Clamp(out_max=cast(float, (1, 2))) + clamp(self.sample_subject) + + def test_wrong_out_min_type(self): + with pytest.raises(TypeError): + clamp = tio.Clamp(out_min=cast(float, 'foo')) + clamp(self.sample_subject) + + def test_wrong_out_max_type(self): + with pytest.raises(TypeError): + clamp = tio.Clamp(out_max=cast(float, 'foo')) + clamp(self.sample_subject) diff --git a/tests/transforms/preprocessing/test_contour.py b/tests/transforms/preprocessing/test_contour.py new file mode 100644 index 000000000..ae13459fc --- /dev/null +++ b/tests/transforms/preprocessing/test_contour.py @@ -0,0 +1,19 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestContour(TorchioTestCase): + """Tests for `Contour`.""" + + def test_one_hot(self): + image = self.sample_subject.label + tio.Contour()(image) + + def test_multichannel(self): + label_map = tio.LabelMap(tensor=torch.rand(2, 3, 3, 3) > 1) + with pytest.raises(RuntimeError): + tio.Contour()(label_map) diff --git a/tests/transforms/preprocessing/test_copy_affine.py b/tests/transforms/preprocessing/test_copy_affine.py new file mode 100644 index 000000000..b55e3ee98 --- /dev/null +++ b/tests/transforms/preprocessing/test_copy_affine.py @@ -0,0 +1,50 @@ +from typing import cast + +import numpy as np +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestCopyAffine(TorchioTestCase): + """Tests for `CopyAffine`.""" + + def test_missing_reference(self): + transform = tio.CopyAffine(target='missing') + with pytest.raises(RuntimeError): + transform(self.sample_subject) + + def test_wrong_target_type(self): + with pytest.raises(ValueError): + tio.CopyAffine(target=cast(str, [1])) + + def test_same_affine(self): + image = tio.ScalarImage(tensor=torch.rand(2, 2, 2, 2)) + mask = tio.LabelMap(tensor=torch.rand(2, 2, 2, 2)) + mask.affine *= 1.1 + subject = tio.Subject(t1=image, mask=mask) + transform = tio.CopyAffine('t1') + transformed = transform(subject) + transformed_t1 = transformed.get_image('t1') + transformed_mask = transformed.get_image('mask') + self.assert_tensor_equal( + transformed_t1.affine, + transformed_mask.affine, + ) + + def test_before_loading(self): + affine = 2 * np.eye(4) + image = tio.ScalarImage(tensor=torch.rand(2, 2, 2, 2), affine=affine) + mask = tio.LabelMap(self.get_image_path('mask')) + subject = tio.Subject(t1=image, mask=mask) + transform = tio.CopyAffine('t1') + transformed = transform(subject) + transformed_t1 = transformed.get_image('t1') + transformed_mask = transformed.get_image('mask') + self.assert_tensor_equal( + transformed_t1.affine, + transformed_mask.affine, + ) diff --git a/tests/transforms/preprocessing/test_crop.py b/tests/transforms/preprocessing/test_crop.py new file mode 100644 index 000000000..c1edaee93 --- /dev/null +++ b/tests/transforms/preprocessing/test_crop.py @@ -0,0 +1,35 @@ +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestCrop(TorchioTestCase): + def test_tensor_single_channel(self): + crop = tio.Crop(1) + assert crop(torch.rand(1, 10, 10, 10)).shape == (1, 8, 8, 8) + + def test_tensor_multi_channel(self): + crop = tio.Crop(1) + assert crop(torch.rand(3, 10, 10, 10)).shape == (3, 8, 8, 8) + + def test_subject_copy(self): + crop = tio.Crop(1, copy=True) + subject = tio.Subject(t1=tio.ScalarImage(tensor=torch.rand(1, 10, 10, 10))) + cropped_subject = crop(subject) + assert cropped_subject.t1.shape == (1, 8, 8, 8) + assert subject.t1.shape == (1, 10, 10, 10) + + cropped2_subject = crop(cropped_subject) + assert cropped2_subject.t1.shape == (1, 6, 6, 6) + assert cropped_subject.t1.shape == (1, 8, 8, 8) + assert len(cropped2_subject.applied_transforms) == 2 + assert len(cropped_subject.applied_transforms) == 1 + + def test_subject_no_copy(self): + crop = tio.Crop(1, copy=False) + subject = tio.Subject(t1=tio.ScalarImage(tensor=torch.rand(1, 10, 10, 10))) + cropped_subject = crop(subject) + assert cropped_subject.t1.shape == (1, 8, 8, 8) + assert subject.t1.shape == (1, 8, 8, 8) diff --git a/tests/transforms/preprocessing/test_crop_pad.py b/tests/transforms/preprocessing/test_crop_pad.py new file mode 100644 index 000000000..fb1914757 --- /dev/null +++ b/tests/transforms/preprocessing/test_crop_pad.py @@ -0,0 +1,255 @@ +from typing import cast + +import numpy as np +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestCropOrPad(TorchioTestCase): + """Tests for `CropOrPad`.""" + + def test_no_changes(self): + sample_t1 = self.sample_subject.get_scalar_image('t1') + shape = sample_t1.spatial_shape + transform = tio.CropOrPad(shape) + transformed = transform(self.sample_subject) + transformed_t1 = transformed.get_scalar_image('t1') + self.assert_tensor_equal(sample_t1.data, transformed_t1.data) + self.assert_tensor_equal(sample_t1.affine, transformed_t1.affine) + + def test_no_changes_mask(self): + sample_t1 = self.sample_subject.get_scalar_image('t1') + sample_mask = self.sample_subject.get_label_map('label').data + sample_mask *= 0 + shape = sample_t1.spatial_shape + transform = tio.CropOrPad(shape, mask_name='label') + with pytest.warns(RuntimeWarning): + transformed = transform(self.sample_subject) + for key, transformed_image in transformed.get_images_dict( + intensity_only=False + ).items(): + image = self.sample_subject.get_image(key) + self.assert_tensor_equal(image.data, transformed_image.data) + self.assert_tensor_equal(image.affine, transformed_image.affine) + + def test_different_shape(self): + shape = self.sample_subject.get_scalar_image('t1').spatial_shape + target_shape = 9, 21, 30 + transform = tio.CropOrPad(target_shape) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + self.assertNotEqual(shape, result_shape) + + def test_shape_right(self): + target_shape = 9, 21, 30 + transform = tio.CropOrPad(target_shape) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + assert target_shape == result_shape + + def test_only_pad(self): + target_shape = 11, 22, 30 + transform = tio.CropOrPad(target_shape) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + assert target_shape == result_shape + + def test_only_crop(self): + target_shape = 9, 18, 30 + transform = tio.CropOrPad(target_shape) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + assert target_shape == result_shape + + def test_shape_negative(self): + with pytest.raises(ValueError): + tio.CropOrPad(-1) + + def test_shape_float(self): + with pytest.raises(ValueError): + tio.CropOrPad(cast(int, 2.5)) + + def test_shape_string(self): + with pytest.raises(TypeError): + tio.CropOrPad(cast(int, '')) + + def test_shape_one(self): + transform = tio.CropOrPad(1) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + assert result_shape == (1, 1, 1) + + def test_wrong_mask_name(self): + cop = tio.CropOrPad(1, mask_name='wrong') + with pytest.warns(RuntimeWarning): + cop(self.sample_subject) + + def test_empty_mask(self): + target_shape = 8, 22, 30 + transform = tio.CropOrPad(target_shape, mask_name='label') + mask = self.sample_subject.get_label_map('label').data + mask *= 0 + with pytest.warns(RuntimeWarning): + transform(self.sample_subject) + + def mask_only(self, target_shape): + transform = tio.CropOrPad(target_shape, mask_name='label') + mask = self.sample_subject.get_label_map('label').data + mask *= 0 + mask[0, 4:6, 5:8, 3:7] = 1 + transformed = transform(self.sample_subject) + shapes = [] + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + shapes.append(result_shape) + set_shapes = set(shapes) + message = f'Images have different shapes: {set_shapes}' + assert len(set_shapes) == 1, message + for key, image in transformed.get_images_dict(intensity_only=False).items(): + result_shape = image.spatial_shape + assert target_shape == result_shape, f'Wrong shape for image: {key}' + + def test_mask_only_pad(self): + self.mask_only((11, 22, 30)) + + def test_mask_only_crop(self): + self.mask_only((9, 18, 30)) + + def test_center_mask(self): + """The mask bounding box and the input image have the same center.""" + target_shape = 8, 22, 30 + transform_center = tio.CropOrPad(target_shape) + transform_mask = tio.CropOrPad(target_shape, mask_name='label') + mask = self.sample_subject.get_label_map('label').data + mask *= 0 + mask[0, 4:6, 9:11, 14:16] = 1 + transformed_center = transform_center(self.sample_subject) + transformed_mask = transform_mask(self.sample_subject) + zipped = zip( + transformed_center.get_images(intensity_only=False), + transformed_mask.get_images(intensity_only=False), + strict=True, + ) + for image_center, image_mask in zipped: + self.assert_tensor_equal( + image_center.data, + image_mask.data, + msg='Data is different after cropping', + ) + self.assert_tensor_equal( + image_center.affine, + image_mask.affine, + msg='Physical position is different after cropping', + ) + + def test_mask_corners(self): + """The mask bounding box and the input image have the same center.""" + target_shape = 8, 22, 30 + transform_center = tio.CropOrPad(target_shape) + transform_mask = tio.CropOrPad( + target_shape, + mask_name='label', + ) + mask = self.sample_subject.get_label_map('label').data + mask *= 0 + mask[0, 0, 0, 0] = 1 + mask[0, -1, -1, -1] = 1 + transformed_center = transform_center(self.sample_subject) + transformed_mask = transform_mask(self.sample_subject) + zipped = zip( + transformed_center.get_images(intensity_only=False), + transformed_mask.get_images(intensity_only=False), + strict=True, + ) + for image_center, image_mask in zipped: + self.assert_tensor_equal( + image_center.data, + image_mask.data, + msg='Data is different after cropping', + ) + self.assert_tensor_equal( + image_center.affine, + image_mask.affine, + msg='Physical position is different after cropping', + ) + + def test_2d(self): + # https://github.com/TorchIO-project/torchio/issues/434 + image = np.random.rand(1, 16, 16, 1) + mask = np.zeros_like(image, dtype=bool) + mask[0, 7, 0] = True + subject = tio.Subject( + image=tio.ScalarImage(tensor=image), + mask=tio.LabelMap(tensor=mask), + ) + transform = tio.CropOrPad((12, 12, 1), mask_name='mask') + transformed = transform(subject) + assert transformed.shape == (1, 12, 12, 1) + + def test_no_target_no_mask(self): + with pytest.raises(ValueError): + tio.CropOrPad() + + def test_labels_but_no_mask(self): + with pytest.raises(ValueError): + tio.CropOrPad(target_shape=(3, 4, 5), labels=[2, 3]) + + def test_no_target(self): + crop_with_mask = tio.CropOrPad(mask_name='label') + crop_with_mask(self.sample_subject) + + def test_persistent_bounds_params(self): + # https://github.com/TorchIO-project/torchio/issues/757 + shape = (1, 5, 5, 5) + mask_a = np.zeros(shape) + mask_a[0, 2, 2, 2] = 1 + mask_b = mask_a.copy() + mask_b[0, 1:4, 1:4, 1:4] = 1 + tensor = np.ones(shape) + image_a = tio.ScalarImage(tensor=tensor) + mask_a = tio.LabelMap(tensor=mask_a) + subject_a = tio.Subject(image=image_a, mask=mask_a) + image_b = tio.ScalarImage(tensor=tensor) + mask_b = tio.LabelMap(tensor=mask_b) + subject_b = tio.Subject(image=image_b, mask=mask_b) + crop = tio.CropOrPad(mask_name='mask') + for _ in range(2): + shape_a = crop(subject_a).image.shape + shape_b = crop(subject_b).image.shape + assert shape_a != shape_b + + def test_only_crop_pad_true(self): + with pytest.raises(ValueError): + tio.CropOrPad((1, 2, 3), only_crop=True, only_pad=True) + + def test_only_pad_true(self): + target_shape = 9, 21, 30 + orig_shape = self.sample_subject.get_scalar_image('t1').spatial_shape + expected_shape = tuple( + t if t > o else o for o, t in zip(orig_shape, target_shape, strict=True) + ) + transform = tio.CropOrPad(target_shape, only_pad=True) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + assert result_shape == expected_shape + + def test_only_crop_true(self): + target_shape = 9, 21, 30 + orig_shape = self.sample_subject.get_scalar_image('t1').spatial_shape + expected_shape = tuple( + t if t < o else o for o, t in zip(orig_shape, target_shape, strict=True) + ) + transform = tio.CropOrPad(target_shape, only_crop=True) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + result_shape = image.spatial_shape + assert result_shape == expected_shape diff --git a/tests/transforms/preprocessing/test_ensure_shape_multiple.py b/tests/transforms/preprocessing/test_ensure_shape_multiple.py new file mode 100644 index 000000000..2515ec04c --- /dev/null +++ b/tests/transforms/preprocessing/test_ensure_shape_multiple.py @@ -0,0 +1,33 @@ +import pytest + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestEnsureShapeMultiple(TorchioTestCase): + def test_bad_method(self): + with pytest.raises(ValueError): + tio.EnsureShapeMultiple(1, method='bad') + + def test_pad(self): + sample_t1 = self.sample_subject.t1 + assert sample_t1.shape == (1, 10, 20, 30) + transform = tio.EnsureShapeMultiple(4, method='pad') + transformed = transform(sample_t1) + assert transformed.shape == (1, 12, 20, 32) + + def test_crop(self): + sample_t1 = self.sample_subject.t1 + assert sample_t1.shape == (1, 10, 20, 30) + transform = tio.EnsureShapeMultiple(4, method='crop') + transformed = transform(sample_t1) + assert transformed.shape == (1, 8, 20, 28) + + def test_2d(self): + sample_t1 = self.sample_subject.t1 + sample_2d = sample_t1.data[..., :1] + assert sample_2d.shape == (1, 10, 20, 1) + transform = tio.EnsureShapeMultiple(4, method='crop') + transformed = transform(sample_2d) + assert transformed.shape == (1, 8, 20, 1) diff --git a/tests/transforms/preprocessing/test_histogram_standardization.py b/tests/transforms/preprocessing/test_histogram_standardization.py new file mode 100644 index 000000000..a9c4896b4 --- /dev/null +++ b/tests/transforms/preprocessing/test_histogram_standardization.py @@ -0,0 +1,101 @@ +from pathlib import Path +from typing import cast + +import numpy as np +import pytest +import torch +from packaging.version import Version + +from torchio import LabelMap +from torchio import ScalarImage +from torchio import Subject +from torchio import SubjectsDataset +from torchio.transforms import HistogramStandardization + +from ...utils import TorchioTestCase + + +class TestHistogramStandardization(TorchioTestCase): + """Tests for :class:`HistogramStandardization` class.""" + + def setUp(self): + super().setUp() + subjects = [] + for i in range(5): + image = ScalarImage(self.get_image_path(f'hs_image_{i}')) + label_path = self.get_image_path( + f'hs_label_{i}', + binary=True, + force_binary_foreground=True, + ) + label = LabelMap(label_path) + subject = Subject(image=image, label=label) + subjects.append(subject) + self.subjects = subjects + self.dataset = SubjectsDataset(self.subjects) + + def test_train_histogram(self): + paths = [subject.image.path for subject in self.dataset] + # Use a function to mask + HistogramStandardization.train( + paths, + masking_function=HistogramStandardization.mean, + output_path=(self.dir / 'landmarks.txt'), + progress=False, + ) + # Use a file to mask + HistogramStandardization.train( + paths, + mask_path=self.dataset[0].label.path, + output_path=(self.dir / 'landmarks.npy'), + progress=False, + ) + # Use files to mask + masks = [subject.label.path for subject in self.dataset] + HistogramStandardization.train( + paths, + mask_path=masks, + output_path=(self.dir / 'landmarks_masks.npy'), + progress=False, + ) + + def test_bad_paths_lengths(self): + with pytest.raises(ValueError): + image_paths = cast(list[str], [1, 2]) + mask_paths = cast(list[str], [1, 2, 3]) + HistogramStandardization.train( + image_paths, + mask_path=mask_paths, + ) + + def test_normalize(self): + landmarks = np.linspace(0, 100, 13) + landmarks_dict: dict[str, str | Path | np.ndarray] = {'image': landmarks} + transform = HistogramStandardization(landmarks_dict) + transform(self.dataset[0]) + + def test_wrong_image_key(self): + landmarks = np.linspace(0, 100, 13) + landmarks_dict: dict[str, str | Path | np.ndarray] = {'wrong_key': landmarks} + transform = HistogramStandardization(landmarks_dict) + with pytest.raises(KeyError): + transform(self.dataset[0]) + + def test_with_saved_dict(self): + landmarks = np.linspace(0, 100, 13) + landmarks_dict = {'image': landmarks} + landmarks_path = self.dir / 'landmarks_dict.pth' + torch.save(landmarks_dict, landmarks_path) + kwargs = {} + if Version(torch.__version__) >= Version('1.13'): + kwargs['weights_only'] = False + landmarks_dict = torch.load(landmarks_path, **kwargs) + transform = HistogramStandardization(landmarks_dict) + transform(self.dataset[0]) + + def test_with_saved_array(self): + landmarks = np.linspace(0, 100, 13) + np.save(self.dir / 'landmarks.npy', landmarks) + landmarks_dict = {'image': self.dir / 'landmarks.npy'} + transform = HistogramStandardization(landmarks_dict) + transform(self.dataset[0]) diff --git a/tests/transforms/preprocessing/test_keep_largest.py b/tests/transforms/preprocessing/test_keep_largest.py new file mode 100644 index 000000000..4a86af150 --- /dev/null +++ b/tests/transforms/preprocessing/test_keep_largest.py @@ -0,0 +1,21 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestKeepLargestComponent(TorchioTestCase): + """Tests for `KeepLargestComponent`.""" + + def test_one_hot(self): + tensor = torch.as_tensor([1, 0, 1, 1, 0, 1]).reshape(1, 1, 1, 6) + label_map = tio.LabelMap(tensor=tensor) + largest = tio.KeepLargestComponent()(label_map) + assert largest.data.sum() == 2 + + def test_multichannel(self): + label_map = tio.LabelMap(tensor=torch.rand(2, 3, 3, 3) > 1) + with pytest.raises(RuntimeError): + tio.KeepLargestComponent()(label_map) diff --git a/tests/transforms/preprocessing/test_mask.py b/tests/transforms/preprocessing/test_mask.py new file mode 100644 index 000000000..d7faa22a5 --- /dev/null +++ b/tests/transforms/preprocessing/test_mask.py @@ -0,0 +1,71 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestMask(TorchioTestCase): + def test_single_mask(self): + negated_mask = self.sample_subject.label.data.logical_not() + masked_voxel_indices = negated_mask.nonzero(as_tuple=True) + transform = tio.Mask(masking_method='label') + transformed = transform(self.sample_subject) + assert (transformed.t1.data[masked_voxel_indices] == 0).all() + + def test_single_mask_nonzero_background(self): + background_value = 314159 + negated_mask = self.sample_subject.label.data.logical_not() + masked_voxel_indices = negated_mask.nonzero(as_tuple=True) + + transform = tio.Mask( + masking_method='label', + outside_value=background_value, + ) + transformed = transform(self.sample_subject) + + assert (transformed.t1.data[masked_voxel_indices] == background_value).all() + + def test_mask_specified_label(self): + mask_label = [1] + negated_mask = self.sample_subject.label.data.logical_not() + masked_voxel_indices = negated_mask.nonzero(as_tuple=True) + + transform = tio.Mask(masking_method='label', labels=mask_label) + transformed = transform(self.sample_subject) + + assert (transformed.t1.data[masked_voxel_indices] == 0).all() + + def test_mask_specified_label_small(self): + def to_image(*numbers): + return torch.as_tensor(numbers).reshape(1, 1, 1, len(numbers)) + + image_tensor = to_image(1, 6, 7, 3, 0) + label_tensor = to_image(0, 1, 2, 3, 4) + mask_labels = [1, 2] + subject = tio.Subject( + image=tio.ScalarImage(tensor=image_tensor), + label=tio.LabelMap(tensor=label_tensor), + ) + transform = tio.Mask(masking_method='label', labels=mask_labels) + transformed = transform(subject) + masked_list = transformed.image.data.flatten().tolist() + assert masked_list == [0, 6, 7, 0, 0] + + def test_mask_example(self): + subject = self.sample_subject + negated_mask = subject.label.data.logical_not() + masked_voxel_indices = negated_mask.nonzero(as_tuple=True) + transform = tio.Mask(masking_method='label') + transformed = transform(subject) + assert (transformed.t1.data[masked_voxel_indices] == 0).all() + + def test_4d(self): + image = tio.ScalarImage(tensor=torch.rand(3, 4, 5, 6)) + mask = tio.LabelMap(tensor=torch.ones(1, 4, 5, 6)) + subject = tio.Subject(image=image, mask_lm=mask) + transform = tio.Mask(masking_method='mask_lm') + with pytest.warns(RuntimeWarning, match='^Expanding.*'): + masked = transform(subject) + assert masked.image.shape == image.shape diff --git a/tests/transforms/preprocessing/test_one_hot.py b/tests/transforms/preprocessing/test_one_hot.py new file mode 100644 index 000000000..70fa08e11 --- /dev/null +++ b/tests/transforms/preprocessing/test_one_hot.py @@ -0,0 +1,29 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestOneHot(TorchioTestCase): + """Tests for `OneHot`.""" + + def test_one_hot(self): + image = self.sample_subject.label + one_hot = tio.OneHot(num_classes=3)(image) + assert one_hot.num_channels == 3 + + def test_multichannel(self): + label_map = tio.LabelMap(tensor=torch.rand(2, 3, 3, 3) > 1) + with pytest.raises(RuntimeError): + tio.OneHot()(label_map) + + def test_inverse(self): + one_hot = tio.OneHot() + subject_one_hot = one_hot(self.sample_subject) + subject_back = subject_one_hot.apply_inverse_transform() + self.assert_tensor_equal( + self.sample_subject.label.data, + subject_back.label.data, + ) diff --git a/tests/transforms/preprocessing/test_pad.py b/tests/transforms/preprocessing/test_pad.py new file mode 100644 index 000000000..e6264ea9f --- /dev/null +++ b/tests/transforms/preprocessing/test_pad.py @@ -0,0 +1,82 @@ +from typing import Any +from typing import cast + +import pytest +import SimpleITK as sitk +import torch + +import torchio as tio +from torchio.data.io import sitk_to_nib + +from ...utils import TorchioTestCase + + +class TestPad(TorchioTestCase): + """Tests for `Pad`.""" + + def test_pad(self): + image = self.sample_subject.t1 + padding = 1, 2, 3, 4, 5, 6 + sitk_image = image.as_sitk() + low, high = padding[::2], padding[1::2] + sitk_padded = sitk.ConstantPad(sitk_image, low, high, 0) + tio_padded = tio.Pad(padding, padding_mode=0)(image) + sitk_tensor, sitk_affine = sitk_to_nib(sitk_padded) + tio_tensor, tio_affine = sitk_to_nib(tio_padded.as_sitk()) + self.assert_tensor_equal(sitk_tensor, tio_tensor) + self.assert_tensor_equal(sitk_affine, tio_affine) + + def test_nans_history(self): + padded = tio.Pad(1, padding_mode=2)(self.sample_subject) + again = padded.history[0](self.sample_subject) + assert not torch.isnan(again.t1.data).any() + + def test_padding_modes(self): + def padding_func() -> None: + return None + + for padding_mode in [0, *tio.Pad.PADDING_MODES]: + tio.Pad(0, padding_mode=padding_mode) + tio.Pad(0, padding_mode=cast(Any, padding_func)) + + with self.assertRaises(KeyError): + tio.Pad(0, padding_mode='abc') + + def test_padding_mean_label_map(self): + with self.assertWarns(RuntimeWarning): + tio.Pad(1, padding_mode='mean')(self.sample_subject.label) + + def test_padding_modes_global(self): + x = torch.ones(1, 1, 2, 2, dtype=torch.int) + x[..., 0, 0] = 0 + # The image should look like this: + # 0 1 + # 1 1 + + add_bottom_row = 0, 0, 0, 1, 0, 0 + with_zeros = tio.Pad(add_bottom_row)(x) + assert with_zeros[0, 0, 2].tolist() == [0, 0] + + with_minimum = tio.Pad(add_bottom_row, padding_mode='minimum')(x) + assert with_minimum[0, 0, 2].tolist() == [0, 0] + + with_maximum = tio.Pad(add_bottom_row, padding_mode='maximum')(x) + assert with_maximum[0, 0, 2].tolist() == [1, 1] + + with_median = tio.Pad(add_bottom_row, padding_mode='median')(x) + assert with_median[0, 0, 2].tolist() == [1, 1] + + # This is a special case: as we instantiated the tensor with integers, + # the mean (3/4) will be trucated to 0. + with_mean = tio.Pad(add_bottom_row, padding_mode='mean')(x) + assert with_mean[0, 0, 2].tolist() == [0, 0] + # So let's test with floats too + x = x.float() + with_mean = tio.Pad(add_bottom_row, padding_mode='mean')(x) + assert with_mean[0, 0, 2].tolist() == [0.75, 0.75] + + def test_truncation_warning(self): + x = torch.ones(1, 1, 2, 2, dtype=torch.int) + pad = tio.Pad(1, padding_mode='mean') + with pytest.warns(RuntimeWarning): + pad(x) diff --git a/tests/transforms/preprocessing/test_resample.py b/tests/transforms/preprocessing/test_resample.py new file mode 100644 index 000000000..abbf884ef --- /dev/null +++ b/tests/transforms/preprocessing/test_resample.py @@ -0,0 +1,122 @@ +from typing import Any +from typing import cast + +import numpy as np +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +def _apply_resample(target: object, subject: tio.Subject) -> tio.Subject: + transformed = tio.Resample(cast(Any, target))(subject) + assert isinstance(transformed, tio.Subject) + return transformed + + +class TestResample(TorchioTestCase): + """Tests for `Resample`.""" + + def test_spacing(self): + # Should this raise an error if sizes are different? + spacing = 2 + transform = tio.Resample(spacing) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + assert image.spacing == 3 * (spacing,) + + def test_reference_name(self): + subject = self.get_inconsistent_shape_subject() + reference_name = 't1' + transform = tio.Resample(reference_name) + transformed = transform(subject) + reference_image = subject[reference_name] + for image in transformed.get_images(intensity_only=False): + assert reference_image.shape == image.shape + self.assert_tensor_almost_equal( + reference_image.affine, + image.affine, + ) + + def test_affine(self): + spacing = 1 + affine_name = 'pre_affine' + transform = tio.Resample(spacing, pre_affine_name=affine_name) + transformed = transform(self.sample_subject) + for image in transformed.values(): + if affine_name in image: + target_affine = np.eye(4) + target_affine[:3, 3] = 10, 0, -0.1 + self.assert_tensor_almost_equal(image.affine, target_affine) + else: + self.assert_tensor_equal(image.affine, np.eye(4)) + + def test_missing_affine(self): + transform = tio.Resample(1, pre_affine_name='missing') + with pytest.raises(ValueError): + transform(self.sample_subject) + + def test_reference_path(self): + reference_image, reference_path = self.get_reference_image_and_path() + transform = tio.Resample(reference_path) + transformed = transform(self.sample_subject) + for image in transformed.values(): + assert reference_image.shape == image.shape + self.assert_tensor_almost_equal( + reference_image.affine, + image.affine, + ) + + def test_wrong_spacing_length(self): + with pytest.raises(RuntimeError): + _apply_resample((1, 2), self.sample_subject) + + def test_wrong_spacing_value(self): + with pytest.raises(ValueError): + tio.Resample(0)(self.sample_subject) + + def test_wrong_target_type(self): + with pytest.raises(RuntimeError): + tio.Resample(None)(self.sample_subject) + + def test_missing_reference(self): + transform = tio.Resample('missing') + with pytest.raises(ValueError): + transform(self.sample_subject) + + def test_2d(self): + """Check that image is still 2D after resampling.""" + image = tio.ScalarImage(tensor=torch.rand(1, 2, 3, 1)) + transform = tio.Resample(0.5) + shape = transform(image).shape + assert shape == (1, 4, 6, 1) + + def test_input_list(self): + _apply_resample([1, 2, 3], self.sample_subject) + + def test_input_array(self): + _apply_resample(np.asarray([1, 2, 3]), self.sample_subject) + + def test_image_target(self): + tio.Resample(self.sample_subject.t1)(self.sample_subject) + + def test_bad_affine(self): + shape = 1, 2, 3 + affine = np.eye(3) + target = shape, affine + transform = tio.Resample(target) + with pytest.raises(RuntimeError): + transform(self.sample_subject) + + def test_resample_flip_consistent(self): + image = torch.rand(1, 10, 10, 10) + resample = tio.Resample(1.35) + flip = tio.Flip(0) + flipped_and_resampled = resample(flip(image)) + resampled_and_flipped = flip(resample(image)) + self.assert_tensor_almost_equal( + flipped_and_resampled.data, + resampled_and_flipped.data, + ) diff --git a/tests/transforms/preprocessing/test_rescale.py b/tests/transforms/preprocessing/test_rescale.py new file mode 100644 index 000000000..4b80ec9fb --- /dev/null +++ b/tests/transforms/preprocessing/test_rescale.py @@ -0,0 +1,131 @@ +import copy +from typing import cast + +import numpy as np +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestRescaleIntensity(TorchioTestCase): + def test_rescale_to_same_intentisy(self): + min_t1 = float(self.sample_subject.t1.data.min()) + max_t1 = float(self.sample_subject.t1.data.max()) + transform = tio.RescaleIntensity(out_min_max=(min_t1, max_t1)) + transformed = transform(self.sample_subject) + assert np.allclose( + transformed.t1.data, + self.sample_subject.t1.data, + rtol=0, + atol=1e-05, + ) + + def test_min_max(self): + transform = tio.RescaleIntensity(out_min_max=(0, 1)) + transformed = transform(self.sample_subject) + assert transformed.t1.data.min() == 0 + assert transformed.t1.data.max() == 1 + + def test_percentiles(self): + low_quantile = np.percentile(self.sample_subject.t1.data, 5) + high_quantile = np.percentile(self.sample_subject.t1.data, 95) + low_indices = (self.sample_subject.t1.data < low_quantile).nonzero( + as_tuple=True, + ) + high_indices = (self.sample_subject.t1.data > high_quantile).nonzero( + as_tuple=True, + ) + rescale = tio.RescaleIntensity(out_min_max=(0, 1), percentiles=(5, 95)) + transformed = rescale(self.sample_subject) + assert (transformed.t1.data[low_indices] == 0).all() + assert (transformed.t1.data[high_indices] == 1).all() + + def test_masking_using_label(self): + transform = tio.RescaleIntensity( + out_min_max=(0, 1), + percentiles=(5, 95), + masking_method='label', + ) + transformed = transform(self.sample_subject) + mask = self.sample_subject.label.data > 0 + low_quantile = np.percentile(self.sample_subject.t1.data[mask], 5) + high_quantile = np.percentile(self.sample_subject.t1.data[mask], 95) + low_indices = (self.sample_subject.t1.data < low_quantile).nonzero( + as_tuple=True, + ) + high_indices = (self.sample_subject.t1.data > high_quantile).nonzero( + as_tuple=True, + ) + assert transformed.t1.data.min() == 0 + assert transformed.t1.data.max() == 1 + assert (transformed.t1.data[low_indices] == 0).all() + assert (transformed.t1.data[high_indices] == 1).all() + + def test_ct(self): + ct_max = 1500 + ct_min = -2000 + ct_range = ct_max - ct_min + tensor = torch.rand(1, 30, 30, 30) * ct_range + ct_min + ct = tio.ScalarImage(tensor=tensor) + ct_air = -1000 + ct_bone = 1000 + rescale = tio.RescaleIntensity( + out_min_max=(-1, 1), + in_min_max=(ct_air, ct_bone), + ) + rescaled = rescale(ct) + assert rescaled.data.min() < -1 + assert rescaled.data.max() > 1 + + def test_out_min_higher_than_out_max(self): + with pytest.raises(ValueError): + tio.RescaleIntensity(out_min_max=(1, 0)) + + def test_too_many_values_for_out_min_max(self): + with pytest.raises(ValueError): + invalid_out_min_max = cast(tuple[float, float], (1, 2, 3)) + tio.RescaleIntensity(out_min_max=invalid_out_min_max) + + def test_wrong_out_min_max_type(self): + with pytest.raises(ValueError): + invalid_out_min_max = cast(tuple[float, float], 'wrong') + tio.RescaleIntensity(out_min_max=invalid_out_min_max) + + def test_min_percentile_higher_than_max_percentile(self): + with pytest.raises(ValueError): + tio.RescaleIntensity(out_min_max=(0, 1), percentiles=(1, 0)) + + def test_too_many_values_for_percentiles(self): + with pytest.raises(ValueError): + invalid_percentiles = cast(tuple[float, float], (1, 2, 3)) + tio.RescaleIntensity(out_min_max=(0, 1), percentiles=invalid_percentiles) + + def test_wrong_percentiles_type(self): + with pytest.raises(ValueError): + invalid_percentiles = cast(tuple[float, float], 'wrong') + tio.RescaleIntensity(out_min_max=(0, 1), percentiles=invalid_percentiles) + + def test_empty_mask(self): + subject = copy.deepcopy(self.sample_subject) + subject.label.set_data(subject.label.data * 0) + rescale = tio.RescaleIntensity(masking_method='label') + with pytest.warns(RuntimeWarning): + rescale(subject) + + def test_persistent_in_min_max(self): + # see https://github.com/TorchIO-project/torchio/issues/1115 + img1 = torch.tensor([[[[0, 1]]]]) + img2 = torch.tensor([[[[0, 10]]]]) + + rescale = tio.RescaleIntensity(out_min_max=(0, 1)) + + assert rescale(img1).data.flatten().tolist() == [0, 1] + assert rescale(img2).data.flatten().tolist() == [0, 1] + + rescale = tio.RescaleIntensity(out_min_max=(0, 1)) + + assert rescale(img2).data.flatten().tolist() == [0, 1] + assert rescale(img1).data.flatten().tolist() == [0, 1] diff --git a/tests/transforms/preprocessing/test_resize.py b/tests/transforms/preprocessing/test_resize.py new file mode 100644 index 000000000..7d4835a79 --- /dev/null +++ b/tests/transforms/preprocessing/test_resize.py @@ -0,0 +1,21 @@ +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestResize(TorchioTestCase): + """Tests for `Resize`.""" + + def test_one_dim(self): + target_shape = 5 + transform = tio.Resize(target_shape) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + assert image.spatial_shape == 3 * (target_shape,) + + def test_all_dims(self): + target_shape = 11, 6, 7 + transform = tio.Resize(target_shape) + transformed = transform(self.sample_subject) + for image in transformed.get_images(intensity_only=False): + assert image.spatial_shape == target_shape diff --git a/tests/transforms/preprocessing/test_to.py b/tests/transforms/preprocessing/test_to.py new file mode 100644 index 000000000..c1c987d25 --- /dev/null +++ b/tests/transforms/preprocessing/test_to.py @@ -0,0 +1,17 @@ +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestTo(TorchioTestCase): + """Tests for :class:`tio.To` class.""" + + def test_to(self): + transform = tio.To(torch.int) + tensor = 10 * torch.rand(2, 3, 4, 5) + image = tio.ScalarImage(tensor=tensor) + transformed = transform(image) + assert image.data.dtype == torch.float32 + assert transformed.data.dtype == torch.int diff --git a/tests/transforms/preprocessing/test_to_canonical.py b/tests/transforms/preprocessing/test_to_canonical.py new file mode 100644 index 000000000..4bdd0aa21 --- /dev/null +++ b/tests/transforms/preprocessing/test_to_canonical.py @@ -0,0 +1,36 @@ +import numpy as np +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestToCanonical(TorchioTestCase): + def test_no_changes(self): + transform = tio.ToCanonical() + transformed = transform(self.sample_subject) + self.assert_tensor_equal( + transformed.t1.data, + self.sample_subject.t1.data, + ) + self.assert_tensor_equal( + transformed.t1.affine, + self.sample_subject.t1.affine, + ) + + def test_las_to_ras(self): + self.sample_subject.t1.affine[0, 0] = -1 # Change orientation to 'LAS' + transform = tio.ToCanonical() + transformed = transform(self.sample_subject) + assert transformed.t1.orientation == ('R', 'A', 'S') + array_flip = self.sample_subject.t1.data.numpy()[:, ::-1, :, :].copy() + self.assert_tensor_almost_equal( + transformed.t1.data, + torch.from_numpy(array_flip), + check_stride=False, + ) + + fixture = np.eye(4) + fixture[0, -1] = -self.sample_subject.t1.spatial_shape[0] + 1 + self.assert_tensor_equal(transformed.t1.affine, fixture) diff --git a/tests/transforms/preprocessing/test_to_orientation.py b/tests/transforms/preprocessing/test_to_orientation.py new file mode 100644 index 000000000..09eaf6601 --- /dev/null +++ b/tests/transforms/preprocessing/test_to_orientation.py @@ -0,0 +1,84 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestToOrientation(TorchioTestCase): + def test_invalid_orientation_length(self): + with pytest.raises(ValueError, match='3-letter'): + tio.ToOrientation('RA') # Too short + + def test_invalid_orientation_characters(self): + with pytest.raises(ValueError, match='three distinct characters'): + tio.ToOrientation('XYZ') + + def test_missing_axis_direction(self): + match = 'must include one character for each axis' + with pytest.raises(ValueError, match=match): + tio.ToOrientation('RAA') # no S/I direction + + def test_no_change_if_already_correct(self): + transform = tio.ToOrientation('RAS') + subject = transform(self.sample_subject) + self.assert_tensor_equal(subject.t1.data, self.sample_subject.t1.data) + self.assert_tensor_equal(subject.t1.affine, self.sample_subject.t1.affine) + + def test_ras_to_las(self): + # Step 1: Set initial orientation to RAS (default) + ras_subject = self.sample_subject + + # Step 2: RAS -> LAS + to_las = tio.ToOrientation('LAS') + las_subject = to_las(ras_subject) + + self.assertEqual(las_subject.t1.orientation, ('L', 'A', 'S')) + + # Manually compute expected LAS affine + expected_affine = ras_subject.t1.affine.copy() + expected_affine[0, 0] = -ras_subject.t1.affine[0, 0] + expected_affine[0, 3] = ( + ras_subject.t1.affine[0, 0] * (ras_subject.t1.spatial_shape[0] - 1) + + ras_subject.t1.affine[0, 3] + ) + + # Check transformation validity + flipped_data = torch.flip(ras_subject.t1.data, dims=[1]) + self.assert_tensor_almost_equal( + las_subject.t1.data, + flipped_data, + check_stride=False, + ) + self.assert_tensor_almost_equal( + las_subject.t1.affine, + expected_affine, + ) + + def test_ras_to_las_to_ras(self): + # Step 1: Start with RAS orientation + original_subject = self.sample_subject + original_data = original_subject.t1.data.clone() + original_affine = original_subject.t1.affine.copy() + + # Step 2: RAS -> LAS + to_las = tio.ToOrientation('LAS') + las_subject = to_las(original_subject) + self.assertEqual(las_subject.t1.orientation, ('L', 'A', 'S')) + + # Step 3: LAS -> RAS + to_ras = tio.ToOrientation('RAS') + recovered_subject = to_ras(las_subject) + self.assertEqual(recovered_subject.t1.orientation, ('R', 'A', 'S')) + + # Step 4: Check if data and affine are restored + self.assert_tensor_almost_equal( + recovered_subject.t1.data, + original_data, + check_stride=False, + ) + self.assert_tensor_almost_equal( + recovered_subject.t1.affine, + original_affine, + ) diff --git a/tests/transforms/preprocessing/test_transpose.py b/tests/transforms/preprocessing/test_transpose.py new file mode 100644 index 000000000..e4f3162bf --- /dev/null +++ b/tests/transforms/preprocessing/test_transpose.py @@ -0,0 +1,21 @@ +import SimpleITK as sitk + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestTranspose(TorchioTestCase): + def test_transpose(self): + transform = tio.Transpose() + image = tio.ScalarImage(self.get_image_path('image')) + transformed = transform(image) + sitk_image = sitk.GetImageFromArray(image.numpy()[0]) + from_sitk = tio.ScalarImage.from_sitk(sitk_image) + self.assert_tensor_equal(transformed.data, from_sitk.data) + + def test_orientation_reversed(self): + transform = tio.Transpose() + image = tio.ScalarImage(self.get_image_path('image')) + transformed = transform(image) + self.assertEqual(transformed.orientation_str, image.orientation_str[::-1]) diff --git a/tests/transforms/preprocessing/test_z_normalization.py b/tests/transforms/preprocessing/test_z_normalization.py new file mode 100644 index 000000000..50fc264b2 --- /dev/null +++ b/tests/transforms/preprocessing/test_z_normalization.py @@ -0,0 +1,29 @@ +import pytest +import torch + +import torchio as tio + +from ...utils import TorchioTestCase + + +class TestZNormalization(TorchioTestCase): + """Tests for :class:`ZNormalization` class.""" + + def test_z_normalization(self): + transform = tio.ZNormalization() + transformed = transform(self.sample_subject) + assert float(transformed.t1.data.mean()) == pytest.approx(0, abs=1e-6) + assert float(transformed.t1.data.std()) == pytest.approx(1) + + def test_no_std(self): + image = tio.ScalarImage(tensor=torch.ones(1, 2, 2, 2)) + with pytest.raises(RuntimeError): + tio.ZNormalization()(image) + + def test_dtype(self): + # https://github.com/TorchIO-project/torchio/issues/407 + tensor_int = (100 * torch.rand(1, 2, 3, 4)).byte() + transform = tio.ZNormalization(masking_method=tio.ZNormalization.mean) + transform(tensor_int) + transform = tio.ZNormalization() + transform(tensor_int) diff --git a/tests/transforms/test_collate.py b/tests/transforms/test_collate.py new file mode 100644 index 000000000..25e95ebe9 --- /dev/null +++ b/tests/transforms/test_collate.py @@ -0,0 +1,42 @@ +import torchio as tio + +from ..utils import TorchioTestCase + + +class TestCollate(TorchioTestCase): + def get_heterogeneous_dataset(self): + # Keys missing in one of the samples will not be present in the batch + # This is relevant for the case in which a transform is applied to some + # samples only, according to its probability (p argument) + transform_no = tio.RandomElasticDeformation(p=0, max_displacement=1) + transform_yes = tio.RandomElasticDeformation(p=1, max_displacement=1) + sample_no = transform_no(self.sample_subject) + sample_yes = transform_yes(self.sample_subject) + data = sample_no, sample_yes + + class Dataset: + def __init__(self, data): + self.data = data + + def __len__(self): + return len(self.data) + + def __getitem__(self, index): + return self.data[index] + + return Dataset(data) + + def test_collate(self): + loader = tio.SubjectsLoader(self.get_heterogeneous_dataset(), batch_size=2) + tio.utils.get_first_item(loader) + + def test_history_collate(self): + loader = tio.SubjectsLoader( + self.get_heterogeneous_dataset(), + batch_size=4, + collate_fn=tio.utils.history_collate, + ) + batch = tio.utils.get_first_item(loader) + empty_history, one_history = batch['history'] + assert not empty_history + assert len(one_history) == 1 diff --git a/tests/transforms/test_invertibility.py b/tests/transforms/test_invertibility.py new file mode 100644 index 000000000..bda8b3a4c --- /dev/null +++ b/tests/transforms/test_invertibility.py @@ -0,0 +1,68 @@ +import copy +import warnings + +import torch + +import torchio as tio + +from ..utils import TorchioTestCase + + +class TestInvertibility(TorchioTestCase): + def test_all_random_transforms(self): + transform = self.get_large_composed_transform() + # Remove RandomLabelsToImage as it will add a new image to the subject + for t in transform.transforms: + if t.name == 'RandomLabelsToImage': + transform.transforms.remove(t) # noqa: B038 + break + # Ignore elastic deformation and gamma warnings during execution + # Ignore some transforms not invertible + with warnings.catch_warnings(): + warnings.simplefilter('ignore', RuntimeWarning) + transformed = transform(self.sample_subject) + inverting_transform = transformed.get_inverse_transform() + transformed_back = inverting_transform(transformed) + assert transformed.t1.shape == transformed_back.t1.shape + self.assert_tensor_equal( + transformed.label.affine, + transformed_back.label.affine, + ) + + def test_different_interpolation(self): + def model_probs(subject): + subject = copy.deepcopy(subject) + subject.im.set_data(torch.rand_like(subject.im.data)) + return subject + + def model_label(subject): + subject = model_probs(subject) + subject.im.set_data(torch.bernoulli(subject.im.data)) + return subject + + transform = tio.RandomAffine(image_interpolation='bspline') + subject = copy.deepcopy(self.sample_subject) + tensor = (torch.rand(1, 20, 20, 20) > 0.5).float() # 0s and 1s + subject = tio.Subject(im=tio.ScalarImage(tensor=tensor)) + transformed = transform(subject) + assert transformed.im.data.min() < 0 + assert transformed.im.data.max() > 1 + + subject_probs = model_probs(transformed) + transformed_back = subject_probs.apply_inverse_transform() + assert transformed_back.im.data.min() < 0 + assert transformed_back.im.data.max() > 1 + transformed_back_linear = subject_probs.apply_inverse_transform( + image_interpolation='linear', + ) + assert transformed_back_linear.im.data.min() >= 0 + assert transformed_back_linear.im.data.max() <= 1 + + subject_label = model_label(transformed) + transformed_back = subject_label.apply_inverse_transform() + assert transformed_back.im.data.min() < 0 + assert transformed_back.im.data.max() > 1 + transformed_back_linear = subject_label.apply_inverse_transform( + image_interpolation='nearest', + ) + assert transformed_back_linear.im.data.unique().tolist() == [0, 1] diff --git a/tests/transforms/test_lambda_transform.py b/tests/transforms/test_lambda_transform.py new file mode 100644 index 000000000..2c2d401ce --- /dev/null +++ b/tests/transforms/test_lambda_transform.py @@ -0,0 +1,70 @@ +import pytest +import torch + +from torchio import LABEL +from torchio.transforms import Lambda + +from ..utils import TorchioTestCase + + +class TestLambda(TorchioTestCase): + """Tests for :class:`Lambda` class.""" + + def test_wrong_return_type(self): + transform = Lambda(lambda x: 'Not a tensor') + with pytest.raises(ValueError): + transform(self.sample_subject) + + def test_wrong_return_data_type(self): + transform = Lambda(lambda x: torch.rand(1) > 0) + with pytest.raises(ValueError): + transform(self.sample_subject) + + def test_wrong_return_shape(self): + transform = Lambda(lambda x: torch.rand(1)) + with pytest.raises(ValueError): + transform(self.sample_subject) + + def test_lambda(self): + transform = Lambda(lambda x: x + 1) + transformed = transform(self.sample_subject) + assert torch.all( + torch.eq( + transformed.t1.data, + self.sample_subject.t1.data + 1, + ), + ) + assert torch.all( + torch.eq( + transformed.t2.data, + self.sample_subject.t2.data + 1, + ), + ) + assert torch.all( + torch.eq( + transformed.label.data, + self.sample_subject.label.data + 1, + ), + ) + + def test_image_types(self): + transform = Lambda(lambda x: x + 1, types_to_apply=[LABEL]) + transformed = transform(self.sample_subject) + assert torch.all( + torch.eq( + transformed.t1.data, + self.sample_subject.t1.data, + ), + ) + assert torch.all( + torch.eq( + transformed.t2.data, + self.sample_subject.t2.data, + ), + ) + assert torch.all( + torch.eq( + transformed.label.data, + self.sample_subject.label.data + 1, + ), + ) diff --git a/tests/transforms/test_monai_adapter.py b/tests/transforms/test_monai_adapter.py new file mode 100644 index 000000000..0b45f51c2 --- /dev/null +++ b/tests/transforms/test_monai_adapter.py @@ -0,0 +1,370 @@ +import warnings + +import numpy as np +import pytest +import torch +from monai.transforms import MapTransform +from monai.transforms import NormalizeIntensity +from monai.transforms import NormalizeIntensityd +from monai.transforms import ScaleIntensity +from monai.transforms import ScaleIntensityd +from monai.transforms import SpatialCrop +from monai.transforms import SpatialCropd + +import torchio as tio + +from ..utils import TorchioTestCase + + +class TestMonaiAdapterDict(TorchioTestCase): + """Tests for :class:`MonaiAdapter` with dictionary transforms.""" + + def test_not_callable(self): + with pytest.raises(TypeError, match='callable'): + tio.MonaiAdapter('not a callable') + + def test_intensity_transform(self): + """Verify that a MONAI intensity transform modifies the data.""" + original_data = self.sample_subject.t1.data.clone() + transform = tio.MonaiAdapter(ScaleIntensityd(keys=['t1'], factor=2.0)) + transformed = transform(self.sample_subject) + assert not torch.equal(transformed.t1.data, original_data) + + def test_multiple_keys(self): + """Verify that a transform is applied to multiple keys.""" + original_t1 = self.sample_subject.t1.data.clone() + original_t2 = self.sample_subject.t2.data.clone() + transform = tio.MonaiAdapter( + ScaleIntensityd(keys=['t1', 't2'], factor=2.0), + ) + transformed = transform(self.sample_subject) + assert not torch.equal(transformed.t1.data, original_t1) + assert not torch.equal(transformed.t2.data, original_t2) + + def test_non_image_entries_preserved(self): + """Non-image entries in the Subject should be preserved.""" + subject = tio.Subject( + t1=tio.ScalarImage(tensor=torch.randn(1, 10, 10, 10)), + age=42, + name='test', + ) + transform = tio.MonaiAdapter(ScaleIntensityd(keys=['t1'], factor=1.0)) + transformed = transform(subject) + assert transformed['age'] == 42 + assert transformed['name'] == 'test' + + def test_affine_preserved_intensity(self): + """Affine should be unchanged for intensity-only transforms.""" + original_affine = self.sample_subject.t1.affine.copy() + transform = tio.MonaiAdapter( + NormalizeIntensityd(keys=['t1']), + ) + transformed = transform(self.sample_subject) + np.testing.assert_array_equal( + transformed.t1.affine, + original_affine, + ) + + def test_affine_updated_spatial(self): + """Cropping should shift the affine origin but keep the rotation/scale.""" + subject = tio.Subject( + t1=tio.ScalarImage(tensor=torch.randn(1, 32, 32, 32)), + ) + original_affine = subject.t1.affine.copy() + transform = tio.MonaiAdapter( + SpatialCropd(keys=['t1'], roi_start=[4, 4, 4], roi_end=[20, 20, 20]), + ) + transformed = transform(subject) + assert transformed.t1.spatial_shape == (16, 16, 16) + new_affine = transformed.t1.affine + # Rotation/scale (3×3) should be unchanged + np.testing.assert_array_equal(new_affine[:3, :3], original_affine[:3, :3]) + # Translation (origin) should differ due to cropping offset + assert not np.array_equal(new_affine[:3, 3], original_affine[:3, 3]) + + def test_spatial_shape_changes(self): + """Tensor shape should be updated by spatial transforms.""" + subject = tio.Subject( + t1=tio.ScalarImage(tensor=torch.randn(1, 32, 32, 32)), + seg=tio.LabelMap(tensor=torch.ones(1, 32, 32, 32)), + ) + transform = tio.MonaiAdapter( + SpatialCropd( + keys=['t1', 'seg'], + roi_start=[0, 0, 0], + roi_end=[16, 16, 16], + ), + ) + transformed = transform(subject) + assert transformed.t1.spatial_shape == (16, 16, 16) + assert transformed.seg.spatial_shape == (16, 16, 16) + + def test_compose_integration(self): + """MonaiAdapter should work inside tio.Compose.""" + pipeline = tio.Compose( + [ + tio.MonaiAdapter(ScaleIntensityd(keys=['t1'], factor=2.0)), + tio.RandomFlip(p=0), + ] + ) + original_data = self.sample_subject.t1.data.clone() + transformed = pipeline(self.sample_subject) + assert not torch.equal(transformed.t1.data, original_data) + + def test_compose_monai_between_torchio(self): + """MONAI transforms should chain correctly with TorchIO transforms.""" + pipeline = tio.Compose( + [ + tio.RandomFlip(p=0), + tio.MonaiAdapter(ScaleIntensityd(keys=['t1'], factor=1.0)), + tio.RandomFlip(p=0), + ] + ) + pipeline(self.sample_subject) + + def test_tensor_input(self): + """Dict MonaiAdapter should work with 4D tensor input via DataParser.""" + tensor = torch.randn(1, 10, 10, 10) + transform = tio.MonaiAdapter( + ScaleIntensityd(keys=['default_image_name'], factor=2.0), + include=['default_image_name'], + ) + result = transform(tensor) + assert isinstance(result, torch.Tensor) + assert result.shape == tensor.shape + + def test_image_input(self): + """Dict MonaiAdapter should work with Image input via DataParser.""" + image = tio.ScalarImage(tensor=torch.randn(1, 10, 10, 10)) + transform = tio.MonaiAdapter( + ScaleIntensityd(keys=['default_image_name'], factor=2.0), + include=['default_image_name'], + ) + result = transform(image) + assert isinstance(result, tio.ScalarImage) + + def test_probability(self): + """The p parameter should control application probability.""" + transform = tio.MonaiAdapter( + ScaleIntensityd(keys=['t1'], factor=100.0), + p=0, + ) + original_data = self.sample_subject.t1.data.clone() + transformed = transform(self.sample_subject) + assert torch.equal(transformed.t1.data, original_data) + + def test_label_not_modified_when_not_in_keys(self): + """Images not in MONAI keys should not be modified.""" + original_label = self.sample_subject.label.data.clone() + transform = tio.MonaiAdapter( + ScaleIntensityd(keys=['t1'], factor=2.0), + ) + transformed = transform(self.sample_subject) + assert torch.equal(transformed.label.data, original_label) + + +class TestMonaiAdapterArray(TorchioTestCase): + """Tests for :class:`MonaiAdapter` with array transforms.""" + + def test_array_intensity_transform(self): + """Array transform should modify all images.""" + original_t1 = self.sample_subject.t1.data.clone() + original_t2 = self.sample_subject.t2.data.clone() + transform = tio.MonaiAdapter(ScaleIntensity(factor=2.0)) + transformed = transform(self.sample_subject) + assert not torch.equal(transformed.t1.data, original_t1) + assert not torch.equal(transformed.t2.data, original_t2) + + def test_array_include(self): + """Array transform should respect the include parameter.""" + original_t2 = self.sample_subject.t2.data.clone() + original_label = self.sample_subject.label.data.clone() + transform = tio.MonaiAdapter( + ScaleIntensity(factor=2.0), + include=['t1'], + ) + transformed = transform(self.sample_subject) + assert not torch.equal( + transformed.t1.data, + self.sample_subject.t1.data, + ) + assert torch.equal(transformed.t2.data, original_t2) + assert torch.equal(transformed.label.data, original_label) + + def test_array_exclude(self): + """Array transform should respect the exclude parameter.""" + original_label = self.sample_subject.label.data.clone() + transform = tio.MonaiAdapter( + ScaleIntensity(factor=2.0), + exclude=['label'], + ) + transformed = transform(self.sample_subject) + assert not torch.equal( + transformed.t1.data, + self.sample_subject.t1.data, + ) + assert torch.equal(transformed.label.data, original_label) + + def test_array_affine_preserved_intensity(self): + """Affine should be unchanged for array intensity transforms.""" + original_affine = self.sample_subject.t1.affine.copy() + transform = tio.MonaiAdapter(NormalizeIntensity()) + transformed = transform(self.sample_subject) + np.testing.assert_array_equal( + transformed.t1.affine, + original_affine, + ) + + def test_array_affine_updated_spatial(self): + """Cropping should shift the affine origin but keep the rotation/scale.""" + subject = tio.Subject( + t1=tio.ScalarImage(tensor=torch.randn(1, 32, 32, 32)), + ) + original_affine = subject.t1.affine.copy() + transform = tio.MonaiAdapter( + SpatialCrop(roi_start=[4, 4, 4], roi_end=[20, 20, 20]), + ) + transformed = transform(subject) + assert transformed.t1.spatial_shape == (16, 16, 16) + new_affine = transformed.t1.affine + np.testing.assert_array_equal(new_affine[:3, :3], original_affine[:3, :3]) + assert not np.array_equal(new_affine[:3, 3], original_affine[:3, 3]) + + def test_array_tensor_input(self): + """Array MonaiAdapter should work with 4D tensor input.""" + tensor = torch.randn(1, 10, 10, 10) + transform = tio.MonaiAdapter(ScaleIntensity(factor=2.0)) + result = transform(tensor) + assert isinstance(result, torch.Tensor) + assert result.shape == tensor.shape + + def test_array_image_input(self): + """Array MonaiAdapter should work with Image input.""" + image = tio.ScalarImage(tensor=torch.randn(1, 10, 10, 10)) + transform = tio.MonaiAdapter(ScaleIntensity(factor=2.0)) + result = transform(image) + assert isinstance(result, tio.ScalarImage) + + def test_array_compose_integration(self): + """Array MonaiAdapter should work inside tio.Compose.""" + pipeline = tio.Compose( + [ + tio.MonaiAdapter(NormalizeIntensity()), + tio.RandomFlip(p=0), + ] + ) + pipeline(self.sample_subject) + + def test_array_probability(self): + """The p parameter should work with array transforms.""" + transform = tio.MonaiAdapter( + ScaleIntensity(factor=100.0), + p=0, + ) + original_data = self.sample_subject.t1.data.clone() + transformed = transform(self.sample_subject) + assert torch.equal(transformed.t1.data, original_data) + + +class TestMonaiAdapterEdgeCases(TorchioTestCase): + """Tests for edge cases and robustness of :class:`MonaiAdapter`.""" + + def test_history_not_recorded(self): + """MonaiAdapter should not be added to subject history. + + MONAI transform objects are not serializable, so recording them + in the history would cause failures during replay via + ``Subject.get_applied_transforms()``. + """ + transform = tio.MonaiAdapter(ScaleIntensityd(keys=['t1'], factor=2.0)) + transformed = transform(self.sample_subject) + assert len(transformed.applied_transforms) == 0 + + def test_random_array_warns_multi_image(self): + """A randomizable array transform on a multi-image subject should warn. + + When a MONAI ``Randomizable`` array transform is applied per-image, + each image receives different random parameters, breaking spatial + alignment. + """ + from monai.transforms import RandScaleIntensity + + transform = tio.MonaiAdapter(RandScaleIntensity(factors=0.5, prob=1.0)) + with pytest.warns(UserWarning, match='Randomizable'): + transform(self.sample_subject) + + def test_random_array_no_warn_single_image(self): + """No warning for a randomizable array transform on a single image.""" + from monai.transforms import RandScaleIntensity + + subject = tio.Subject( + t1=tio.ScalarImage(tensor=torch.randn(1, 10, 10, 10)), + ) + transform = tio.MonaiAdapter(RandScaleIntensity(factors=0.5, prob=1.0)) + with warnings.catch_warnings(): + warnings.simplefilter('error', UserWarning) + transform(subject) + + def test_new_dict_key_wrapped_as_image(self): + """New MetaTensor keys from MONAI dict transforms should become Images. + + When a MONAI dictionary transform creates a new key whose value + is a ``MetaTensor`` or ``torch.Tensor``, the adapter should wrap + it as a ``ScalarImage`` so it is visible to downstream TorchIO + processing. + """ + from monai.transforms import CopyItemsd + + subject = tio.Subject( + t1=tio.ScalarImage(tensor=torch.randn(1, 10, 10, 10)), + ) + transform = tio.MonaiAdapter( + CopyItemsd(keys=['t1'], names=['t1_copy']), + ) + transformed = transform(subject) + assert 't1_copy' in transformed + assert isinstance(transformed['t1_copy'], tio.ScalarImage) + # New keys should be accessible via attribute syntax + assert isinstance(transformed.t1_copy, tio.ScalarImage) + + def test_multi_sample_transform_raises(self): + """MONAI transforms returning list[dict] should raise a clear error.""" + + class FakeMultiSampleTransform(MapTransform): + """Simulates a MONAI dict transform returning multiple samples.""" + + def __init__(self): + super().__init__(keys=['t1']) + + def __call__(self, data): + return [data, data] + + transform = tio.MonaiAdapter(FakeMultiSampleTransform()) + with pytest.raises(TypeError, match='single mapping'): + transform(self.sample_subject) + + def test_new_non_image_tensor_kept_raw(self): + """Non-image tensors (0D/1D) from MONAI should not become Images.""" + + class FakeStatsTransform(MapTransform): + """Adds a scalar stat key to the output dict.""" + + def __init__(self): + super().__init__(keys=['t1']) + + def __call__(self, data): + data = dict(data) + data['t1'] = data['t1'] + data['mean'] = torch.tensor(0.5) + data['indices'] = torch.tensor([0, 1, 2]) + return data + + subject = tio.Subject( + t1=tio.ScalarImage(tensor=torch.randn(1, 10, 10, 10)), + ) + transform = tio.MonaiAdapter(FakeStatsTransform()) + transformed = transform(subject) + assert isinstance(transformed['mean'], torch.Tensor) + assert not isinstance(transformed['mean'], tio.Image) + assert isinstance(transformed['indices'], torch.Tensor) + assert not isinstance(transformed['indices'], tio.Image) diff --git a/tests/transforms/test_reproducibility.py b/tests/transforms/test_reproducibility.py new file mode 100644 index 000000000..93022ced1 --- /dev/null +++ b/tests/transforms/test_reproducibility.py @@ -0,0 +1,26 @@ +import warnings + +from ..utils import TorchioTestCase + + +def _ignore_reproducibility_warnings() -> None: + for category in (RuntimeWarning, UserWarning): + warnings.simplefilter('ignore', category) + + +class TestReproducibility(TorchioTestCase): + def test_all_random_transforms(self): + transform = self.get_large_composed_transform() + # Ignore elastic deformation and gamma warnings during execution + with warnings.catch_warnings(): # ignore elastic deformation warning + _ignore_reproducibility_warnings() + transformed = transform(self.sample_subject) + reproducing_transform = transformed.get_composed_history() + with warnings.catch_warnings(): # ignore elastic deformation warning + _ignore_reproducibility_warnings() + new_transformed = reproducing_transform(self.sample_subject) + self.assert_tensor_equal(transformed.t1.data, new_transformed.t1.data) + self.assert_tensor_equal( + transformed.label.data, + new_transformed.label.data, + ) diff --git a/tests/transforms/test_transforms.py b/tests/transforms/test_transforms.py new file mode 100644 index 000000000..697acbb36 --- /dev/null +++ b/tests/transforms/test_transforms.py @@ -0,0 +1,444 @@ +import copy +from collections.abc import Sequence +from pathlib import Path +from typing import Any +from typing import cast + +import numpy as np +import pytest +import SimpleITK as sitk +import torch +from nibabel.nifti1 import Nifti1Image + +import torchio as tio +from torchio.data.io import nib_to_sitk + +from ..utils import TorchioTestCase + + +def _parse_range_with_invalid_value( + transform: tio.Transform, + value: object, +) -> None: + transform._parse_range(cast(Any, value), 'name') + + +class TestTransforms(TorchioTestCase): + """Tests for all transforms.""" + + def get_transform( + self, + channels: Sequence[str], + is_3d: bool = True, + labels: bool = True, + ) -> tio.Compose: + landmarks_dict: dict[str, str | Path | np.ndarray] = { + channel: np.linspace(0, 100, 13) for channel in channels + } + disp = 1 if is_3d else (1, 1, 0.01) + elastic = tio.RandomElasticDeformation(max_displacement=disp) + affine_elastic = tio.RandomAffineElasticDeformation( + elastic_kwargs={'max_displacement': disp} + ) + cp_args = (9, 21, 30) if is_3d else (21, 30, 1) + resize_args = (10, 20, 30) if is_3d else (10, 20, 1) + flip_axes = axes_downsample = (0, 1, 2) if is_3d else (0, 1) + swap_patch = (2, 3, 4) if is_3d else (3, 4, 1) + pad_args = (1, 2, 3, 0, 5, 6) if is_3d else (0, 0, 3, 0, 5, 6) + crop_args = (3, 2, 8, 0, 1, 4) if is_3d else (0, 0, 8, 0, 1, 4) + remapping = {1: 2, 2: 1, 3: 20, 4: 25} + one_of_transforms: dict[tio.Transform, float] = { + tio.RandomAffine(): 3, + elastic: 1, + } + transforms: list[tio.Transform] = [ + tio.CropOrPad(cp_args), + tio.EnsureShapeMultiple(2, method='crop'), + tio.Resize(resize_args), + tio.ToCanonical(), + tio.RandomAnisotropy(downsampling=(1.75, 2), axes=axes_downsample), + tio.CopyAffine(channels[0]), + tio.Resample((1, 1.1, 1.25)), + tio.RandomFlip(axes=flip_axes, flip_probability=1), + tio.RandomMotion(), + tio.RandomGhosting(axes=(0, 1, 2)), + tio.RandomSpike(), + tio.RandomNoise(), + tio.RandomBlur(), + tio.RandomSwap(patch_size=swap_patch, num_iterations=5), + tio.Lambda(lambda x: 2 * x, types_to_apply=tio.INTENSITY), + tio.RandomBiasField(), + tio.RescaleIntensity(out_min_max=(0, 1)), + tio.ZNormalization(), + tio.HistogramStandardization(landmarks_dict), + elastic, + tio.RandomAffine(), + affine_elastic, + tio.OneOf(one_of_transforms), + tio.RemapLabels(remapping=remapping, masking_method='Left'), + tio.RemoveLabels([1, 3]), + tio.SequentialLabels(), + tio.Pad(pad_args, padding_mode=3), + tio.Crop(crop_args), + ] + if labels: + transforms.append(tio.RandomLabelsToImage(label_key='label')) + return tio.Compose(transforms) + + def test_transforms_dict(self): + transform = tio.RandomNoise(include=('t1', 't2')) + input_dict: dict[str, object] = { + name: image.data + for name, image in self.sample_subject.get_images_dict( + intensity_only=False + ).items() + } + transformed = transform(input_dict) + assert isinstance(transformed, dict) + + def test_transforms_dict_no_keys(self): + transform = tio.RandomNoise() + input_dict: dict[str, object] = { + name: image.data + for name, image in self.sample_subject.get_images_dict( + intensity_only=False + ).items() + } + with pytest.raises(RuntimeError): + transform(input_dict) + + def test_transforms_image(self): + transform = self.get_transform( + channels=('default_image_name',), + labels=False, + ) + transformed = transform(self.sample_subject.t1) + assert isinstance(transformed, tio.ScalarImage) + + def test_transforms_tensor(self): + tensor = torch.rand(2, 4, 5, 8) + transform = self.get_transform( + channels=('default_image_name',), + labels=False, + ) + transformed = transform(tensor) + assert isinstance(transformed, torch.Tensor) + + def test_transforms_array(self): + tensor = torch.rand(2, 4, 5, 8).numpy() + transform = self.get_transform( + channels=('default_image_name',), + labels=False, + ) + transformed = transform(tensor) + assert isinstance(transformed, np.ndarray) + + def test_transforms_sitk(self): + tensor = torch.rand(2, 4, 5, 8) + affine = np.diag((-1, 2, -3, 1)) + image = nib_to_sitk(tensor, affine) + transform = self.get_transform( + channels=('default_image_name',), + labels=False, + ) + transformed = transform(image) + assert isinstance(transformed, sitk.Image) + + def test_transforms_subject_3d(self): + transform = self.get_transform(channels=('t1', 't2'), is_3d=True) + transformed = transform(self.sample_subject) + assert isinstance(transformed, tio.Subject) + + def test_transforms_subject_2d(self): + transform = self.get_transform(channels=('t1', 't2'), is_3d=False) + subject = self.make_2d(self.sample_subject) + transformed = transform(subject) + assert isinstance(transformed, tio.Subject) + + def test_transforms_subject_4d(self): + composed = self.get_transform(channels=('t1', 't2'), is_3d=True) + subject = self.make_multichannel(self.sample_subject) + subject = self.flip_affine_x(subject) + transformed = None + for transform in composed.transforms: + repr(transform) # cover __repr__ + transformed = transform(subject) + trsf_channels = len(transformed.t1.data) + assert trsf_channels > 1, f'Lost channels in {transform.name}' + exclude = ( + 'RandomLabelsToImage', + 'RemapLabels', + 'RemoveLabels', + 'SequentialLabels', + 'CopyAffine', + ) + if transform.name not in exclude: + assert subject.shape[0] == transformed.shape[0], ( + f'Different number of channels after {transform.name}' + ) + self.assert_tensor_not_equal( + subject.t1.data[1], + transformed.t1.data[1], + msg=f'No changes after {transform.name}', + ) + subject = transformed + assert isinstance(transformed, tio.Subject) + + def test_transform_noop(self): + transform = tio.RandomMotion(p=0) + transformed = transform(self.sample_subject) + assert transformed is self.sample_subject + tensor = torch.rand(2, 4, 5, 8).numpy() + transformed = transform(tensor) + assert transformed is tensor + + def test_original_unchanged(self): + subject = copy.deepcopy(self.sample_subject) + composed = self.get_transform(channels=('t1', 't2'), is_3d=True) + subject = self.flip_affine_x(subject) + for transform in composed.transforms: + original_data = copy.deepcopy(subject.t1.data) + transform(subject) + self.assert_tensor_equal( + subject.t1.data, + original_data, + msg=f'Changes after {transform.name}', + ) + + def test_transforms_use_include(self): + original_subject = copy.deepcopy(self.sample_subject) + transform = tio.RandomNoise(include=['t1']) + transformed = transform(self.sample_subject) + + self.assert_tensor_not_equal( + original_subject.t1.data, + transformed.t1.data, + msg=f'Changes after {transform.name}', + ) + + self.assert_tensor_equal( + original_subject.t2.data, + transformed.t2.data, + msg=f'Changes after {transform.name}', + ) + + def test_transforms_use_exclude(self): + original_subject = copy.deepcopy(self.sample_subject) + transform = tio.RandomNoise(exclude=['t2']) + transformed = transform(self.sample_subject) + + self.assert_tensor_not_equal( + original_subject.t1.data, + transformed.t1.data, + msg=f'Changes after {transform.name}', + ) + + self.assert_tensor_equal( + original_subject.t2.data, + transformed.t2.data, + msg=f'Changes after {transform.name}', + ) + + def test_transforms_use_include_and_exclude(self): + with pytest.raises(ValueError): + tio.RandomNoise(include=['t2'], exclude=['t1']) + + def test_keys_deprecated(self): + with pytest.warns(FutureWarning): + tio.RandomNoise(keys=['t2']) + + def test_keep_original(self): + subject = copy.deepcopy(self.sample_subject) + old, new = 't1', 't1_original' + transformed = tio.RandomAffine(keep={old: new})(subject) + assert old in transformed + assert new in transformed + self.assert_tensor_equal( + transformed[new].data, + subject[old].data, + ) + self.assert_tensor_not_equal( + transformed[new].data, + transformed[old].data, + ) + + +class TestTransform(TorchioTestCase): + def test_abstract_transform(self): + with pytest.raises(TypeError): + tio.Transform() + + def test_arguments_are_not_dict(self): + transform = tio.Noise(0, 1, 0) + assert not transform.arguments_are_dict() + + def test_arguments_are_dict(self): + transform = tio.Noise({'im': 0}, {'im': 1}, {'im': 0}) + assert transform.arguments_are_dict() + + def test_arguments_are_and_are_not_dict(self): + transform = tio.Noise(0, {'im': 1}, {'im': 0}) + with pytest.raises(ValueError): + transform.arguments_are_dict() + + def test_min_constraint(self): + transform = tio.RandomNoise() + assert transform._parse_range(3, 'name', min_constraint=0) == (0, 3) + + def test_bad_over_max(self): + transform = tio.RandomNoise() + with pytest.raises(ValueError): + transform._parse_range(2, 'name', max_constraint=1) + + def test_bad_over_max_range(self): + transform = tio.RandomNoise() + with pytest.raises(ValueError): + transform._parse_range((0, 2), 'name', max_constraint=1) + + def test_bad_type(self): + transform = tio.RandomNoise() + with pytest.raises(ValueError): + transform._parse_range(2.5, 'name', type_constraint=int) + + def test_no_numbers(self): + transform = tio.RandomNoise() + with pytest.raises(ValueError): + _parse_range_with_invalid_value(transform, 'j') + + def test_apply_transform_missing(self): + class T(tio.Transform): + pass + + with pytest.raises(TypeError): + T() + + def test_non_invertible(self): + transform = tio.RandomBlur() + with pytest.raises(RuntimeError): + transform.inverse() + + def test_batch_history(self): + # https://github.com/TorchIO-project/torchio/discussions/743 + subject = self.sample_subject + transform = tio.Compose( + [ + tio.RandomAffine(), + tio.CropOrPad(5), + tio.OneHot(), + ] + ) + dataset = tio.SubjectsDataset([subject], transform=transform) + loader = tio.SubjectsLoader( + dataset, + collate_fn=tio.utils.history_collate, + ) + batch = tio.utils.get_first_item(loader) + transformed: tio.Subject = tio.utils.get_subjects_from_batch(batch)[0] + inverse = transformed.apply_inverse_transform() + images1 = subject.get_images(intensity_only=False) + images2 = inverse.get_images(intensity_only=False) + for image1, image2 in zip(images1, images2, strict=True): + assert image1.shape == image2.shape + + def test_bad_bounds_mask(self): + transform = tio.ZNormalization(masking_method='test') + with pytest.raises(ValueError): + transform(self.sample_subject) + + def test_bounds_mask(self): + transform = tio.ZNormalization() + tensor = torch.rand((1, 2, 2, 2)) + with pytest.raises(ValueError): + transform.get_mask_from_anatomical_label('test', tensor) + + def get_mask(label): + mask = transform.get_mask_from_anatomical_label(label, tensor) + return mask + + left = get_mask('Left') + assert left[:, 0].sum() == 4 and left[:, 1].sum() == 0 + right = get_mask('Right') + assert right[:, 1].sum() == 4 and right[:, 0].sum() == 0 + posterior = get_mask('Posterior') + assert posterior[:, :, 0].sum() == 4 and posterior[:, :, 1].sum() == 0 + anterior = get_mask('Anterior') + assert anterior[:, :, 1].sum() == 4 and anterior[:, :, 0].sum() == 0 + inferior = get_mask('Inferior') + assert inferior[..., 0].sum() == 4 and inferior[..., 1].sum() == 0 + superior = get_mask('Superior') + assert superior[..., 1].sum() == 4 and superior[..., 0].sum() == 0 + + mask = transform.get_mask_from_bounds(3 * (0, 1), tensor) + assert mask[0, 0, 0, 0] == 1 + assert mask.sum() == 1 + + def test_label_keys(self): + # Adapted from the issue in which the feature was requested: + # https://github.com/TorchIO-project/torchio/issues/866#issue-1222255576 + size = 1, 10, 10, 10 + image = torch.rand(size) + num_classes = 2 # excluding background + label = torch.randint(num_classes + 1, size) + + data_dict: dict[str, object] = {'image': image, 'label': label} + + transform = tio.RandomAffine( + include=['image', 'label'], + label_keys=['label'], + ) + transformed_dict = transform(data_dict) + transformed_label = transformed_dict['label'] + assert isinstance(transformed_label, torch.Tensor) + + # If the image is indeed transformed as a label map, nearest neighbor + # interpolation is used by default and therefore no intermediate values + # can exist in the output + num_unique_values = len(torch.unique(transformed_label)) + assert num_unique_values <= num_classes + 1 + + def test_nibabel_input(self): + image = self.sample_subject.t1 + image_nib = Nifti1Image(image.data[0].numpy(), image.affine) + transformed = tio.RandomAffine()(image_nib) + transformed.get_fdata() + _ = transformed.affine + + image = self.subject_4d.t1 + tensor_5d = image.data[np.newaxis].permute(2, 3, 4, 0, 1) + image_nib = Nifti1Image(tensor_5d.numpy(), image.affine) + transformed = tio.RandomAffine()(image_nib) + transformed.get_fdata() + _ = transformed.affine + + def test_bad_shape(self): + tensor = torch.rand(1, 2, 3) + with pytest.raises(ValueError, match='must be a 4D tensor'): + tio.RandomAffine()(tensor) + + def test_bad_keys_type(self): + # From https://github.com/TorchIO-project/torchio/issues/923 + with self.assertRaises(ValueError): + tio.RandomAffine(include='t1') + + def test_init_args(self): + transform = tio.Compose([tio.RandomNoise()]) + base_args = transform._get_base_args() + assert 'parse_input' not in base_args + + transform = tio.OneOf([tio.RandomNoise()]) + base_args = transform._get_base_args() + assert 'parse_input' not in base_args + + transform = tio.RandomNoise() + base_args = transform._get_base_args() + assert all( + arg in base_args + for arg in [ + 'copy', + 'include', + 'exclude', + 'keep', + 'parse_input', + 'label_keys', + ] + ) diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 000000000..6046726bb --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,258 @@ +import copy +import os +import random +import shutil +import tempfile +import unittest +from collections.abc import Sequence +from pathlib import Path +from random import shuffle + +import numpy as np +import pytest +import torch + +import torchio as tio + + +class TorchioTestCase(unittest.TestCase): + def setUp(self): + """Set up test fixtures, if any.""" + self.dir = Path(tempfile.gettempdir()) / os.urandom(24).hex() + self.dir.mkdir(exist_ok=True) + random.seed(42) + np.random.seed(42) + + registration_matrix = np.array( + [ + [1, 0, 0, 10], + [0, 1, 0, 0], + [0, 0, 1.2, 0], + [0, 0, 0, 1], + ] + ) + + subject_a = tio.Subject( + t1=tio.ScalarImage(self.get_image_path('t1_a')), + ) + subject_b = tio.Subject( + t1=tio.ScalarImage(self.get_image_path('t1_b')), + label=tio.LabelMap(self.get_image_path('label_b', binary=True)), + ) + subject_c = tio.Subject( + label=tio.LabelMap(self.get_image_path('label_c', binary=True)), + ) + subject_d = tio.Subject( + t1=tio.ScalarImage( + self.get_image_path('t1_d'), + pre_affine=registration_matrix, + ), + t2=tio.ScalarImage(self.get_image_path('t2_d')), + label=tio.LabelMap(self.get_image_path('label_d', binary=True)), + ) + subject_a4 = tio.Subject( + t1=tio.ScalarImage(self.get_image_path('t1_a'), components=4), + ) + self.subjects_list = [ + subject_a, + subject_a4, + subject_b, + subject_c, + subject_d, + ] + self.dataset = tio.SubjectsDataset(self.subjects_list) + self.sample_subject = self.dataset[-1] # subject_d + self.subject_4d = self.dataset[1] + + def make_2d(self, subject): + subject = copy.deepcopy(subject) + for image in subject.get_images(intensity_only=False): + image.set_data(image.data[..., :1]) + return subject + + def make_multichannel(self, subject): + subject = copy.deepcopy(subject) + for image in subject.get_images(intensity_only=False): + image.set_data(torch.cat(4 * (image.data,))) + return subject + + def flip_affine_x(self, subject): + subject = copy.deepcopy(subject) + for image in subject.get_images(intensity_only=False): + image.affine = np.diag((-1, 1, 1, 1)) @ image.affine + return subject + + def get_inconsistent_shape_subject(self): + """Return a subject containing images of different shape.""" + subject = tio.Subject( + t1=tio.ScalarImage(self.get_image_path('t1_inc')), + t2=tio.ScalarImage( + self.get_image_path('t2_inc', shape=(10, 20, 31)), + ), + label=tio.LabelMap( + self.get_image_path( + 'label_inc', + shape=(8, 17, 25), + binary=True, + ), + ), + label2=tio.LabelMap( + self.get_image_path( + 'label2_inc', + shape=(18, 17, 25), + binary=True, + ), + ), + ) + return subject + + def get_reference_image_and_path(self): + """Return a reference image and its path.""" + path = self.get_image_path( + 'ref', + shape=(10, 20, 31), + spacing=(1, 1, 2), + ) + image = tio.ScalarImage(path) + return image, path + + def get_subject_with_partial_volume_label_map(self, components=1): + """Return a subject with a partial-volume label map.""" + return tio.Subject( + t1=tio.ScalarImage( + self.get_image_path('t1_d'), + ), + label=tio.LabelMap( + self.get_image_path( + 'label_d2', + binary=False, + components=components, + ), + ), + ) + + def get_subject_with_labels(self, labels): + return tio.Subject( + label=tio.LabelMap( + self.get_image_path( + 'label_multi', + labels=labels, + ), + ), + ) + + @staticmethod + def get_unique_labels(data: torch.Tensor) -> set[int]: + labels = data.unique().tolist() + return set(labels) + + @staticmethod + def get_tensor_with_labels(labels: Sequence) -> torch.Tensor: + tensor = torch.as_tensor(list(labels)) + return tensor.repeat_interleave(2).reshape(1, 1, 1, -1) + + def tearDown(self): + """Tear down test fixtures, if any.""" + shutil.rmtree(self.dir) + + def get_ixi_tiny(self): + root_dir = Path(tempfile.gettempdir()) / 'torchio' / 'ixi_tiny' + return tio.datasets.IXITiny(root_dir, download=True) + + def get_image_path( + self, + stem, + binary=False, + labels=None, + shape=(10, 20, 30), + spacing=(1, 1, 1), + components=1, + add_nans=False, + suffix=None, + force_binary_foreground=True, + ): + shape = (*shape, 1) if len(shape) == 2 else shape + data = np.random.rand(components, *shape) + if binary: + data = (data > 0.5).astype(np.uint8) + if not data.sum() and force_binary_foreground: + data[..., 0] = 1 + elif labels is not None: + data = (data * (len(labels) + 1)).astype(np.uint8) + new_data = np.zeros_like(data) + for i, label in enumerate(labels): + new_data[data == (i + 1)] = label + if not (new_data == label).sum(): + new_data[..., i] = label + data = new_data + elif self.flip_coin(): # cast some images + data *= 100 + dtype = np.uint8 if self.flip_coin() else np.uint16 + data = data.astype(dtype) + if add_nans: + data[:] = np.nan + affine = np.diag((*spacing, 1)) + if suffix is None: + extensions = '.nii.gz', '.nii', '.nrrd', '.img', '.mnc' + suffix = random.choice(extensions) + path = self.dir / f'{stem}{suffix}' + if self.flip_coin(): + path = str(path) + image = tio.ScalarImage( + tensor=data, + affine=affine, + check_nans=not add_nans, + ) + image.save(path) + return path + + def flip_coin(self): + return np.random.rand() > 0.5 + + def get_tests_data_dir(self): + return Path(__file__).parent / 'image_data' + + def assert_tensor_not_equal(self, *args, **kwargs): # noqa: N802 + with pytest.raises(AssertionError): + self.assert_tensor_equal(*args, **kwargs) + + @staticmethod + def assert_tensor_equal(*args, **kwargs): # noqa: N802 + torch.testing.assert_close( + *args, + rtol=0, + atol=0, + check_dtype=False, + **kwargs, + ) + + @staticmethod + def assert_tensor_almost_equal(*args, **kwargs): # noqa: N802 + torch.testing.assert_close( + *args, + **kwargs, + check_dtype=False, + ) + + @staticmethod + def assert_tensor_all_zeros(tensor): # noqa: N802 + assert torch.all(tensor == 0) + + def get_large_composed_transform(self): + all_classes = get_all_random_transforms() + shuffle(all_classes) + transforms = [t() for t in all_classes] + # Hack as default patch size for RandomSwap is 15 and sample_subject + # is (10, 20, 30) + for tr in transforms: + if tr.name == 'RandomSwap': + tr.patch_size = np.array((10, 10, 10)) + return tio.Compose(transforms) + + +def get_all_random_transforms(): + transforms_names = [ + name for name in dir(tio.transforms) if name.startswith('Random') + ] + classes = [getattr(tio.transforms, name) for name in transforms_names] + return classes diff --git a/tox.ini b/tox.ini index dd00b3900..e0a662930 100644 --- a/tox.ini +++ b/tox.ini @@ -1,12 +1,12 @@ [tox] -envlist = test, lint, format, types, docs-test, docs-build, prek +envlist = pytest, types, lint, format -[testenv:test] +[testenv:pytest] dependency_groups = test commands = pytest tests \ - --durations=10 \ + --durations=0 \ --capture=no \ {posargs} @@ -37,22 +37,3 @@ dependency_groups = types commands = ty check src --output-format concise - -[testenv:docs-test] -description = Test documentation code snippets -dependency_groups = - test -commands = - pytest --codeblocks docs/ {posargs} - -[testenv:docs-build] -description = Build the documentation -dependency_groups = - docs -commands = - zensical build {posargs} - -[testenv:prek] -description = Pre-commit checks -deps = prek -commands = prek run --all-files diff --git a/tutorials/README.md b/tutorials/README.md new file mode 100644 index 000000000..bb2de41de --- /dev/null +++ b/tutorials/README.md @@ -0,0 +1,32 @@ +# Examples + +All tutorials can be run online on Google Colab, without the need to install +any packages locally. + +## General + +Getting Started tutorial. It includes training and inference of a model for +brain segmentation using a 3D U-Net. + +[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TorchIO-project/torchio-notebooks/blob/main/notebooks/TorchIO_tutorial.ipynb) + +## Transforms + +This tutorial will teach everything about transforms, and how to use them for preprocessing and augmentation. + +[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TorchIO-project/torchio-notebooks/blob/main/notebooks/Data_preprocessing_and_augmentation_using_TorchIO_a_tutorial.ipynb) + +## Inference + +This is an example of using a pre-trained PyTorch model (HighRes3DNet) and TorchIO to perform a full brain parcellation from a T1-weighted MRI. + +[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TorchIO-project/torchio-notebooks/blob/main/notebooks/Brain_parcellation_with_TorchIO_and_HighRes3DNet.ipynb) + +## TorchIO + MONAI + PyTorch Lightning + +In this tutorial, we demonstrate how these three libraries from the +[official PyTorch Ecosystem](https://pytorch.org/ecosystem/) +can be used together to segment the hippocampus on brain MRIs from the +[Medical Segmentation Decathlon](http://medicaldecathlon.com/). + +[![Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TorchIO-project/torchio-notebooks/blob/main/notebooks/TorchIO_MONAI_PyTorch_Lightning.ipynb) diff --git a/tutorials/example_heteromodal.py b/tutorials/example_heteromodal.py new file mode 100644 index 000000000..f3967b504 --- /dev/null +++ b/tutorials/example_heteromodal.py @@ -0,0 +1,77 @@ +"""This is an example of a very particular case in which some modalities might +be missing for some of the subjects, as in. + +Dorent et al. 2019, Hetero-Modal Variational Encoder-Decoder for Joint +Modality Completion and Segmentation +""" + +import logging + +import torch.nn as nn + +import torchio as tio +from torchio import LabelMap +from torchio import Queue +from torchio import ScalarImage +from torchio import Subject +from torchio import SubjectsDataset +from torchio.data import UniformSampler + + +def main(): + # Define training and patches sampling parameters + num_epochs = 20 + patch_size = 128 + queue_length = 100 + patches_per_volume = 5 + batch_size = 2 + + # Populate a list with images + one_subject = Subject( + T1=ScalarImage('../BRATS2018_crop_renamed/LGG75_T1.nii.gz'), + T2=ScalarImage('../BRATS2018_crop_renamed/LGG75_T2.nii.gz'), + label=LabelMap('../BRATS2018_crop_renamed/LGG75_Label.nii.gz'), + ) + + # This subject doesn't have a T2 MRI! + another_subject = Subject( + T1=ScalarImage('../BRATS2018_crop_renamed/LGG74_T1.nii.gz'), + label=LabelMap('../BRATS2018_crop_renamed/LGG74_Label.nii.gz'), + ) + + subjects = [ + one_subject, + another_subject, + ] + + subjects_dataset = SubjectsDataset(subjects) + queue_dataset = Queue( + subjects_dataset, + queue_length, + patches_per_volume, + UniformSampler(patch_size), + ) + + # This collate_fn is needed in the case of missing modalities + # In this case, the batch will be composed by a *list* of samples instead + # of the typical Python dictionary that is collated by default in Pytorch + batch_loader = tio.SubjectsLoader( + queue_dataset, + batch_size=batch_size, + collate_fn=lambda x: x, + ) + + # Mock PyTorch model + model = nn.Identity() + + for epoch_index in range(num_epochs): + logging.info('Epoch %s', epoch_index) + for batch in batch_loader: # batch is a *list* here, not a dictionary + logits = model(batch) + logging.info([batch[idx].keys() for idx in range(batch_size)]) + logging.info(logits.shape) + logging.info('') + + +if __name__ == '__main__': + main() diff --git a/zensical.toml b/zensical.toml index facc9d1d1..938c2587a 100644 --- a/zensical.toml +++ b/zensical.toml @@ -1,6 +1,6 @@ [project] site_name = "TorchIO" -site_url = "https://docs.torchio.org/" +site_url = "https://torchio.org" site_description = "Tools for medical image processing with PyTorch" site_author = "Fernando Pérez-García" copyright = "Copyright © 2026 Fernando Pérez-García" @@ -15,119 +15,99 @@ extra_javascript = [ ] nav = [ - { "Get Started" = [ - "index.md", - "get-started/installation.md", - "get-started/quickstart.md", - "get-started/migration.md", + { Home = "index.md" }, + { "Getting started" = "getting-started.md" }, + { "Data structures" = [ + "data/index.md", + "data/image.md", + "data/subject.md", + "data/dataset.md", + "data/loader.md", ] }, - { Concepts = [ - "concepts/data-model.md", - "concepts/lazy-loading.md", - "concepts/transforms.md", - "concepts/per-instance-augmentation.md", + { "Patch-based pipelines" = [ + "patches/index.md", + "patches/training.md", + "patches/inference.md", ] }, - { Tutorials = [ - "tutorials/first-pipeline.md", - "tutorials/augmentation.md", - "tutorials/large-volumes.md", - "tutorials/annotations.md", - ] }, - { "How-to guides" = [ - "how-to/dataloader.md", - "how-to/monai.md", - "how-to/custom-reader.md", - "how-to/save-nii-zarr.md", - "how-to/remote-nii-zarr.md", - "how-to/annotations.md", - "how-to/visualization.md", - "how-to/patch-inference.md", - "how-to/patch-training.md", - "how-to/tta.md", - ] }, - { "API Reference" = [ - { "Data structures" = [ - "reference/image.md", - "reference/subject.md", - "reference/points.md", - "reference/bboxes.md", - "reference/affine.md", - "reference/axes.md", - "reference/backends.md", - ] }, - { "Data loading" = [ - "reference/loader.md", - ] }, - { Datasets = [ - "datasets.md", - ] }, - { "Transforms" = [ - "reference/transforms.md", - "reference/random_parameters.md", - { "Composition" = [ - "reference/transforms/compose.md", - "reference/transforms/one_of.md", - "reference/transforms/some_of.md", + { Transforms = [ + "transforms/index.md", + { Augmentation = [ + "transforms/augmentation/index.md", + { Composition = [ + "transforms/augmentation/Compose.md", + "transforms/augmentation/OneOf.md", ] }, - { "Spatial" = [ - "reference/transforms/spatial.md", - "reference/transforms/resample.md", - "reference/transforms/affine_transform.md", - "reference/transforms/elastic_deformation.md", - "reference/transforms/anisotropy.md", - "reference/transforms/flip.md", - "reference/transforms/reorient.md", - "reference/transforms/transpose.md", - "reference/transforms/crop.md", - "reference/transforms/pad.md", - "reference/transforms/crop_or_pad.md", - "reference/transforms/resize.md", - "reference/transforms/copy_affine.md", - "reference/transforms/ensure_shape_multiple.md", - "reference/transforms/to_reference_space.md", + { Spatial = [ + "transforms/augmentation/RandomFlip.md", + "transforms/augmentation/RandomAffine.md", + "transforms/augmentation/RandomElasticDeformation.md", + "transforms/augmentation/RandomAffineElasticDeformation.md", + "transforms/augmentation/RandomAnisotropy.md", ] }, - { "Intensity" = [ - "reference/transforms/bias_field.md", - "reference/transforms/blur.md", - "reference/transforms/clamp.md", - "reference/transforms/gamma.md", - "reference/transforms/ghosting.md", - "reference/transforms/histogram_standardization.md", - "reference/transforms/labels_to_image.md", - "reference/transforms/mask.md", - "reference/transforms/motion.md", - "reference/transforms/noise.md", - "reference/transforms/normalize.md", - "reference/transforms/pca.md", - "reference/transforms/spike.md", - "reference/transforms/standardize.md", - "reference/transforms/swap.md", + { Intensity = [ + "transforms/augmentation/RandomMotion.md", + "transforms/augmentation/RandomGhosting.md", + "transforms/augmentation/RandomSpike.md", + "transforms/augmentation/RandomBiasField.md", + "transforms/augmentation/RandomBlur.md", + "transforms/augmentation/RandomNoise.md", + "transforms/augmentation/RandomSwap.md", + "transforms/augmentation/RandomLabelsToImage.md", + "transforms/augmentation/RandomGamma.md", ] }, - { "Label" = [ - "reference/transforms/contour.md", - "reference/transforms/keep_largest.md", - "reference/transforms/one_hot.md", - "reference/transforms/remap_labels.md", - "reference/transforms/remove_labels.md", - "reference/transforms/sequential_labels.md", + ] }, + { Preprocessing = [ + "transforms/preprocessing/index.md", + { Intensity = [ + "transforms/preprocessing/RescaleIntensity.md", + "transforms/preprocessing/ZNormalization.md", + "transforms/preprocessing/HistogramStandardization.md", + "transforms/preprocessing/Mask.md", + "transforms/preprocessing/Clamp.md", + "transforms/preprocessing/PCA.md", + "transforms/preprocessing/To.md", ] }, - { "Other" = [ - "reference/transforms/cornucopia_adapter.md", - "reference/transforms/lambda_transform.md", - "reference/transforms/to.md", - "reference/transforms/monai_adapter.md", + { Spatial = [ + "transforms/preprocessing/CropOrPad.md", + "transforms/preprocessing/Crop.md", + "transforms/preprocessing/Pad.md", + "transforms/preprocessing/Resize.md", + "transforms/preprocessing/Resample.md", + "transforms/preprocessing/ToCanonical.md", + "transforms/preprocessing/ToOrientation.md", + "transforms/preprocessing/ToReferenceSpace.md", + "transforms/preprocessing/Transpose.md", + "transforms/preprocessing/EnsureShapeMultiple.md", + "transforms/preprocessing/CopyAffine.md", + ] }, + { Label = [ + "transforms/preprocessing/RemapLabels.md", + "transforms/preprocessing/RemoveLabels.md", + "transforms/preprocessing/SequentialLabels.md", + "transforms/preprocessing/OneHot.md", + "transforms/preprocessing/Contour.md", + "transforms/preprocessing/KeepLargestComponent.md", ] }, ] }, - { "Visualization" = [ - "reference/visualization.md", - ] }, - { "CLI" = [ - "reference/cli.md", - ] }, - { "Patches" = [ - "reference/patches.md", + { Others = [ + "transforms/others/Lambda.md", + "transforms/others/MonaiAdapter.md", ] }, ] }, + { Datasets = "datasets.md" }, + { Examples = [ + "examples/index.md", + "examples/plot_colin27.md", + "examples/plot_include_exclude.md", + "examples/plot_custom_z_spacing.md", + "examples/plot_3d_to_2d.md", + "examples/plot_history.md", + "examples/plot_video.md", + ] }, + { Interfaces = [ + "interfaces/cli.md", + "interfaces/slicer.md", + ] }, ] [project.theme] @@ -136,8 +116,6 @@ favicon = "images/favicon.ico" language = "en" features = [ - "content.action.edit", - "content.action.view", "content.code.copy", "navigation.footer", "navigation.instant", @@ -172,11 +150,6 @@ accent = "deep purple" toggle.icon = "lucide/moon" toggle.name = "Switch to system preference" -[project.extra.version] -provider = "mike" -default = ["stable", "dev"] -alias = true - [[project.extra.social]] icon = "fontawesome/brands/github" link = "https://github.com/TorchIO-project/torchio" @@ -198,12 +171,11 @@ paths = ["src"] [project.plugins.mkdocstrings.handlers.python.options] docstring_style = "google" inherited_members = true -show_source = true +show_source = false show_bases = true show_root_heading = true show_root_full_path = false members_order = "source" -filters = ["!^_", "^__init__$"] [project.markdown_extensions.admonition] @@ -213,9 +185,6 @@ filters = ["!^_", "^__init__$"] anchor_linenums = true [project.markdown_extensions.pymdownx.superfences] -custom_fences = [ - { name = "mermaid", class = "mermaid", format = "pymdownx.superfences.fence_code_format" }, -] [project.markdown_extensions.pymdownx.arithmatex] generic = true