diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml index 63ac5536d85..a564e2e27f0 100644 --- a/.github/workflows/blossom-ci.yml +++ b/.github/workflows/blossom-ci.yml @@ -51,7 +51,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v7 with: repository: ${{ fromJson(needs.Authorization.outputs.args).repo }} ref: ${{ fromJson(needs.Authorization.outputs.args).ref }} diff --git a/.github/workflows/build_docs.yml b/.github/workflows/build_docs.yml index 18617f55147..aeb2b8a2996 100644 --- a/.github/workflows/build_docs.yml +++ b/.github/workflows/build_docs.yml @@ -23,8 +23,10 @@ jobs: env: # minimum supported version of Python PYTHON_VER1: '3.10' + # force installation of CPU-only PyTorch + PIP_EXTRA_INDEX_URL: 'https://download.pytorch.org/whl/cpu' steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Set up Python ${{ env.PYTHON_VER1 }} uses: actions/setup-python@v6 with: diff --git a/.github/workflows/cicd_tests.yml b/.github/workflows/cicd_tests.yml index ae3694f2769..e9b207951f6 100644 --- a/.github/workflows/cicd_tests.yml +++ b/.github/workflows/cicd_tests.yml @@ -56,7 +56,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - opt: ["codeformat", "mypy"] # "pytype" omitted for being essentially deprecated, see #8865 + opt: ["codeformat", "pyrefly"] steps: - name: Clean unused tools run: | @@ -66,7 +66,7 @@ jobs: sudo rm -rf /opt/ghc /usr/local/.ghcup sudo docker system prune -f - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Set up Python ${{ env.PYTHON_VER1 }} uses: actions/setup-python@v6 with: @@ -75,13 +75,12 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip wheel - python -m pip install --no-build-isolation -r requirements-dev.txt + python -m pip install .[all,testing] - name: Lint and type check run: | # clean up temporary files $(pwd)/runtests.sh --build --clean - # Github actions have multiple cores, so parallelize pytype - $(pwd)/runtests.sh --build --${{ matrix.opt }} -j $(nproc --all) + $(pwd)/runtests.sh --build --${{ matrix.opt }} min-dep: # Test with minumum dependencies installed for different OS, Python, and PyTorch combinations runs-on: ${{ matrix.os }} @@ -129,7 +128,7 @@ jobs: sudo rm -rf /usr/local/lib/android sudo rm -rf /opt/ghc /usr/local/.ghcup sudo docker system prune -f - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v6 with: @@ -138,15 +137,14 @@ jobs: - name: Prepare pip wheel run: | which python - python -m pip install --upgrade pip wheel - python -m pip install --user more-itertools>=8.0 + python -m pip install --upgrade pip wheel tomli - name: Install the minimum dependencies run: | # min. requirements python -m pip install torch==${{ matrix.pytorch-version }} - python -m pip install -r requirements-min.txt + python monai/config/print_dependencies.py build-system | xargs pip install --no-build-isolation + python -m pip install --no-build-isolation .[testing] python -m pip list - BUILD_MONAI=0 python setup.py develop # no compile of extensions shell: bash - if: matrix.os == 'linux-gpu-runner' name: Print GPU Info @@ -171,7 +169,7 @@ jobs: strategy: fail-fast: false matrix: - os: [windows-latest, ubuntu-latest] # macOS-latest omitted for now for being very slow, see #8864 + os: [windows-latest, ubuntu-latest, macOS-latest] # macOS-latest is very slow (#8864), testing install only timeout-minutes: 120 env: QUICKTEST: True @@ -191,7 +189,7 @@ jobs: minimum-size: 8GB maximum-size: 16GB disk-root: "D:" - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Set up Python ${{ env.PYTHON_VER1 }} uses: actions/setup-python@v6 with: @@ -219,28 +217,89 @@ jobs: shell: bash - name: Install the complete dependencies run: | - python -m pip install --user --upgrade pip wheel pybind11 # TODO: pybind11 added for macOS, may not be needed - cat "requirements-dev.txt" - python -m pip install --no-build-isolation -r requirements-dev.txt + python -m pip install --user --upgrade pip wheel tomli + python monai/config/print_dependencies.py build-system | xargs pip install --no-build-isolation + python -m pip install --no-build-isolation .[all,testing] python -m pip list - python -m pip install -e . # test no compile installation shell: bash - name: Run compiled (${{ runner.os }}) run: | python -m pip uninstall -y monai - BUILD_MONAI=1 python -m pip install -e . # compile the cpp extensions + BUILD_MONAI=1 python -m pip install --no-build-isolation -e . # compile the cpp extensions in-place with -e + # ensure extensions were compiled + python -c 'import monai._C' > /dev/null shell: bash - - name: Run quick tests + - if: runner.os != 'macOS' + name: Run full tests run: | python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))' python -c "import monai; monai.config.print_config()" python -m unittest -v shell: bash + - if: runner.os == 'macOS' + name: Run min tests + run: | + python -c 'import torch; print(torch.__version__); print(torch.rand(5,3))' + python -c "import monai; monai.config.print_config()" + # TODO: enable large range of macOS tests which don't take a very long time + ./runtests.sh --min + shell: bash + + hyena-dep: # Optional HyenaND dependency + the no-CUDA Hyena tests. + # nvsubquadratic >= 0.1.1 supports Python >= 3.10 and keeps its CUDA-kernel sdist + # (subquadratic-ops-torch-cu12) plus the megatron / dali / timm packages in opt-in + # extras, so it installs on a CPU runner. We still pass ``--no-deps`` deliberately: + # (1) the HyenaND operators import only torch + einops + omegaconf at runtime, so + # skipping the (still batteries-included: datasets/lightning/wandb) core deps + # keeps this job lean; and + # (2) nvsubquadratic pins torch>=2.10,<2.11, which would otherwise upgrade/clash + # with the torch this job (and MONAI's matrix) installs. + # CUDA-required Hyena tests skip cleanly here; the GPU surface is covered by + # ``.github/workflows/pythonapp-hyena-gpu.yml``. + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - name: Clean unused tools + run: | + find /opt/hostedtoolcache/* -maxdepth 0 ! -name 'Python' -exec rm -rf {} \; + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc /usr/local/.ghcup + sudo docker system prune -f + - uses: actions/checkout@v6 + - name: Set up Python ${{ env.PYTHON_VER1 }} + uses: actions/setup-python@v6 + with: + python-version: ${{ env.PYTHON_VER1 }} + cache: 'pip' + - name: Install dependencies + nvsubquadratic (no-deps) + run: | + python -m pip install --upgrade pip wheel tomli pytest + # need a specific version of torch for nvsubquadratic + python monai/config/print_dependencies.py build-system | \ + xargs -d '\n' pip install --no-build-isolation torch==2.10.0 torchvision==0.25.0 + # # nvsubquadratic runtime imports need only torch + einops + omegaconf; install + # # the package itself without its core dependency tree (see job comment above). + python -m pip install --no-build-isolation omegaconf + python -m pip install --no-build-isolation --no-deps 'nvsubquadratic>=0.1.1' + python -m pip install --no-build-isolation .[hyena,testing] + python -m pip list + shell: bash + - name: Run Hyena tests (CUDA-required cases skip cleanly) + run: | + python -c "from monai.networks.blocks.hyena import is_nvsubquadratic_available; \ + assert is_nvsubquadratic_available(), 'nvsubquadratic must be importable'" + python -m pytest -v \ + tests/networks/blocks/test_hyena_block.py \ + tests/networks/nets/test_hyena_nd_unetr.py \ + tests/networks/nets/test_swin_unetr.py + shell: bash packaging: # Test package generation runs-on: ubuntu-latest env: QUICKTEST: True + INDEX_URL: "https://download.pytorch.org/whl/cpu" steps: - name: Clean unused tools run: | @@ -249,7 +308,7 @@ jobs: sudo rm -rf /usr/local/lib/android sudo rm -rf /opt/ghc /usr/local/.ghcup sudo docker system prune -f - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: fetch-depth: 0 - name: Set up Python ${{ env.PYTHON_VER1 }} @@ -259,11 +318,10 @@ jobs: cache: 'pip' - name: Install dependencies run: | - python -m pip install --user --upgrade pip setuptools wheel twine packaging + python -m pip install --user --upgrade pip setuptools wheel twine packaging tomli # install the latest pytorch for testing - # however, "pip install monai*.tar.gz" will build cpp/cuda with an isolated - # fresh torch installation according to pyproject.toml - python -m pip install torch==${PYTORCH_VER1} torchvision --extra-index-url https://download.pytorch.org/whl/cpu + python monai/config/print_dependencies.py build-system all testing | \ + xargs -d '\n' pip install --no-build-isolation torch==${PYTORCH_VER1} --extra-index-url $INDEX_URL - name: Check packages run: | python -m pip uninstall -y monai @@ -292,7 +350,7 @@ jobs: working-directory: ${{ steps.mktemp.outputs.tmp_dir }} run: | # install from wheel - python -m pip install monai*.whl --extra-index-url https://download.pytorch.org/whl/cpu + python -m pip install --no-build-isolation monai*.whl --extra-index-url $INDEX_URL python -c 'import monai; monai.config.print_config()' 2>&1 | grep -iv "unknown" python -c 'import monai; print(monai.__file__)' python -m pip uninstall -y monai @@ -302,6 +360,14 @@ jobs: run: | for name in *.tar.gz; do break; done echo $name - python -m pip install ${name}[all] --extra-index-url https://download.pytorch.org/whl/cpu + python -m pip install --no-build-isolation ${name}[all] --extra-index-url $INDEX_URL + python -c 'import monai; monai.config.print_config()' 2>&1 | grep -iv "unknown" + python -c 'import monai; print(monai.__file__)' + python -m pip uninstall -y monai + - name: Install using uv + working-directory: ${{ steps.root.outputs.pwd }} + run: | + pip install uv + uv pip install --system --no-build-isolation .[all] --extra-index-url $INDEX_URL python -c 'import monai; monai.config.print_config()' 2>&1 | grep -iv "unknown" python -c 'import monai; print(monai.__file__)' diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml index 4a938e50aba..02d881a9194 100644 --- a/.github/workflows/codeql-analysis.yml +++ b/.github/workflows/codeql-analysis.yml @@ -9,67 +9,110 @@ # the `language` matrix defined below to confirm you have the correct set of # supported CodeQL languages. # -name: "CodeQL" +name: "CodeQL Advanced" on: push: branches: [ dev, main ] pull_request: - # The branches below must be a subset of the branches above branches: [ dev ] schedule: - - cron: '18 1 * * 0' + - cron: '0 2 * * 1' # 2AM Monday + +env: + PYTHON_VER: '3.11' + PYTORCH_VER: '2.8.0' + BUILD_MONAI: 1 + PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cpu" # forces CPU PyTorch installation, should be faster jobs: analyze: - name: Analyze + name: Analyze (${{ matrix.language }}) + # Runner size impacts CodeQL analysis time. To learn more, please see: + # - https://gh.io/recommended-hardware-resources-for-running-codeql + # - https://gh.io/supported-runners-and-hardware-resources + # - https://gh.io/using-larger-runners (GitHub.com only) + # Consider using larger runners or machines with greater resources for possible analysis time improvements. runs-on: ubuntu-latest permissions: + # required for all workflows + security-events: write + + # required to fetch internal or private CodeQL packs + packages: read + + # only required for workflows in private repositories actions: read contents: read - security-events: write strategy: fail-fast: false matrix: - language: [ 'cpp', 'python' ] - # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python', 'ruby' ] - # Learn more about CodeQL language support at https://git.io/codeql-language-support - + include: + - language: actions + build-mode: none + - language: c-cpp + build-mode: none # TODO: get Cpp building working, autobuild doesn't work and manual fails for inexplicable reasons. + - language: python + build-mode: none + # CodeQL supports the following values keywords for 'language': 'actions', 'c-cpp', 'csharp', 'go', 'java-kotlin', 'javascript-typescript', 'python', 'ruby', 'rust', 'swift' + # Use `c-cpp` to analyze code written in C, C++ or both + # Use 'java-kotlin' to analyze code written in Java, Kotlin or both + # Use 'javascript-typescript' to analyze code written in JavaScript, TypeScript or both + # To learn more about changing the languages that are analyzed or customizing the build mode for your analysis, + # see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/customizing-your-advanced-setup-for-code-scanning. + # If you are analyzing a compiled language, you can modify the 'build-mode' for that language to customize how + # your codebase is analyzed, see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/codeql-code-scanning-for-compiled-languages steps: - name: Checkout repository - uses: actions/checkout@v6 + uses: actions/checkout@v7 + + # Add any setup steps before running the `github/codeql-action/init` action. + # This includes steps like installing compilers or runtimes (`actions/setup-node` + # or others). This is typically only required for manual builds. + # - name: Setup runtime (example) + # uses: actions/setup-example@v1 # Initializes the CodeQL tools for scanning. - name: Initialize CodeQL uses: github/codeql-action/init@v4 with: languages: ${{ matrix.language }} + build-mode: ${{ matrix.build-mode }} # If you wish to specify custom queries, you can do so here or in a config file. # By default, queries listed here will override any specified in a config file. # Prefix the list here with "+" to use these queries and those in the config file. - # queries: ./path/to/local/query, your-org/your-repo/queries@main - # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). - # If this step fails, then you should remove it and run the build manually (see below) - # - name: Autobuild - # uses: github/codeql-action/autobuild@v2 + # For more details on CodeQL's query packs, refer to: https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs + # queries: security-extended,security-and-quality + # If the analyze step fails for one of the languages you are analyzing with + # "We were unable to automatically build your code", modify the matrix above + # to set the build mode to "manual" for that language. Then modify this step + # to build your code. # ℹī¸ Command-line programs to run using the OS shell. - # 📚 https://git.io/JvXDl - - # ✏ī¸ If the Autobuild fails above, remove it and uncomment the following three lines - # and modify them (or add more) to build your code if your project - # uses a compiled language + # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun + - name: Set up Python ${{ env.PYTHON_VER }} + if: matrix.language == 'c-cpp' && matrix.build-mode == 'manual' + uses: actions/setup-python@v6 + with: + python-version: ${{ env.PYTHON_VER }} + cache: 'pip' - - name: Build + - name: Run manual build steps + if: matrix.language == 'c-cpp' && matrix.build-mode == 'manual' + shell: bash run: | rm -rf /opt/hostedtoolcache/{node,go,Ruby,Java*} ls -al /opt/hostedtoolcache - rm -rf /usr/share/dotnet/ - python -m pip install -U --no-build-isolation pip wheel wheel-stub - python -m pip install --no-build-isolation -r requirements-dev.txt - BUILD_MONAI=1 ./runtests.sh --build + sudo rm -rf /usr/share/dotnet/ + python -m pip install -U pip wheel wheel-stub + python -m pip install torch==${PYTORCH_VER} torchvision + python -m pip install --user --upgrade pip wheel + python monai/config/print_dependencies.py build-system | xargs pip install --no-build-isolation + python -m pip install --no-build-isolation . - name: Perform CodeQL Analysis uses: github/codeql-action/analyze@v4 + with: + category: "/language:${{matrix.language}}" diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml index 7edd7eb114b..9cb72ab9a47 100644 --- a/.github/workflows/conda.yml +++ b/.github/workflows/conda.yml @@ -32,7 +32,7 @@ jobs: minimum-size: 8GB maximum-size: 16GB disk-root: "D:" - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Clean up disk space run: | find /opt/hostedtoolcache/* -maxdepth 0 ! -name 'Python' -exec rm -rf {} \; diff --git a/.github/workflows/cron-ngc-bundle.yml b/.github/workflows/cron-ngc-bundle.yml index a1618952328..650434eb37e 100644 --- a/.github/workflows/cron-ngc-bundle.yml +++ b/.github/workflows/cron-ngc-bundle.yml @@ -17,7 +17,7 @@ jobs: if: github.repository == 'Project-MONAI/MONAI' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Set up Python 3.10 uses: actions/setup-python@v6 with: @@ -29,8 +29,8 @@ jobs: - name: Install dependencies run: | rm -rf /github/home/.cache/torch/hub/bundle/ - python -m pip install --no-build-isolation --upgrade pip wheel wheel-stub - python -m pip install --no-build-isolation -r requirements-dev.txt + python -m pip install -U pip wheel wheel-stub + python -m pip install .[all,testing] - name: Loading Bundles run: | # clean up temporary files diff --git a/.github/workflows/cron.yml b/.github/workflows/cron.yml index 3bdfe127152..9a01b346e5d 100644 --- a/.github/workflows/cron.yml +++ b/.github/workflows/cron.yml @@ -32,7 +32,7 @@ jobs: options: "--gpus all" runs-on: [self-hosted, linux, x64, common] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: apt install run: | apt-get update @@ -43,7 +43,7 @@ jobs: python -m pip install --upgrade pip wheel python -m pip uninstall -y torch torchvision python -m pip install ${{ matrix.pytorch }} - python -m pip install -r requirements-dev.txt + python -m pip install .[all,testing] python -m pip list - name: Run tests report coverage env: @@ -67,7 +67,7 @@ jobs: if pgrep python; then pkill python; fi shell: bash - name: Upload coverage - uses: codecov/codecov-action@v6 + uses: codecov/codecov-action@v7 with: fail_ci_if_error: false files: ./coverage.xml @@ -82,7 +82,7 @@ jobs: options: "--gpus all" runs-on: [self-hosted, linux, x64, integration] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Install APT dependencies run: | apt-get update @@ -91,7 +91,7 @@ jobs: run: | which python python -m pip install --upgrade pip wheel - python -m pip install -r requirements-dev.txt + python -m pip install .[all,testing] python -m pip list - name: Run tests report coverage env: @@ -115,7 +115,7 @@ jobs: if pgrep python; then pkill python; fi shell: bash - name: Upload coverage - uses: codecov/codecov-action@v6 + uses: codecov/codecov-action@v7 with: fail_ci_if_error: false files: ./coverage.xml @@ -131,13 +131,14 @@ jobs: options: "--gpus all" runs-on: [self-hosted, linux, x64, integration] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: fetch-depth: 0 - name: Install the dependencies run: | which python - python -m pip install --upgrade pip wheel twine + python -m pip install --upgrade pip wheel twine tomli + python monai/config/print_dependencies.py build-system all testing | xargs -d '\n' pip install --no-build-isolation python -m pip list - name: Run tests report coverage shell: bash @@ -171,7 +172,6 @@ jobs: python -c 'import monai; print(monai.__file__)' # run tests - cp $root_dir/requirements*.txt "$tmp_dir" cp -r $root_dir/tests "$tmp_dir" pwd ls -al @@ -186,7 +186,6 @@ jobs: python -c $'import torch\na,b=torch.zeros(1,device="cuda:0"),torch.zeros(1,device="cuda:1");\nwhile True:print(a,b)' > /dev/null & python -c "import torch; print(torch.__version__); print('{} of GPUs available'.format(torch.cuda.device_count()))" - python -m pip install -r requirements-dev.txt PYTHONPATH="$tmp_dir":$PYTHONPATH BUILD_MONAI=1 python ./tests/runner.py -p 'test_((?!integration).)' # unit tests if pgrep python; then pkill python; fi @@ -214,13 +213,13 @@ jobs: python -c "import torch; print(torch.__version__); print('{} of GPUs available'.format(torch.cuda.device_count()))" python -c 'import torch; print(torch.rand(5,3, device=torch.device("cuda:0")))' ngc --version - BUILD_MONAI=1 ./runtests.sh --build --coverage --unittests --disttests # unit tests with pytype checks, coverage report + BUILD_MONAI=1 ./runtests.sh --build --coverage --pyrefly --unittests --disttests # unit tests with pyrefly checks, coverage report BUILD_MONAI=1 ./runtests.sh --build --coverage --net # integration tests with coverage report coverage xml --ignore-errors if pgrep python; then pkill python; fi shell: bash - name: Upload coverage - uses: codecov/codecov-action@v6 + uses: codecov/codecov-action@v7 with: fail_ci_if_error: false files: ./coverage.xml @@ -233,13 +232,13 @@ jobs: options: "--gpus all --ipc=host" runs-on: [self-hosted, linux, x64, integration] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Install MONAI id: monai-install run: | which python - python -m pip install --upgrade pip wheel - python -m pip install -r requirements-dev.txt + python -m pip install --upgrade pip wheel tomli + python monai/config/print_dependencies.py build-system | xargs -d '\n' pip install --no-cache-dir --no-build-isolation BUILD_MONAI=1 python setup.py develop # install monai nvidia-smi export CUDA_VISIBLE_DEVICES=$(python -m tests.utils | tail -n 1) diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index c3e5bab7a1f..921d8f5c89a 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -21,7 +21,7 @@ jobs: if: ${{ false }} # disable docker build job project-monai/monai#7450 runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 # full history so that we can git describe with: ref: dev @@ -53,7 +53,7 @@ jobs: needs: versioning_dev runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: ref: dev - name: Download version @@ -69,8 +69,6 @@ jobs: cat _version.py mv _version.py monai/ - # build "latest": remove flake package as it is not needed on hub.docker.com - sed -i '/flake/d' requirements-dev.txt docker build -t projectmonai/monai:latest -f Dockerfile . # distribute as always w/ tag "latest" to hub.docker.com diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 56d015e190b..d3c1e473c72 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -13,7 +13,7 @@ jobs: runs-on: [self-hosted, linux, x64, command] steps: # checkout the pull request branch - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: token: ${{ secrets.PR_MAINTAIN }} repository: ${{ github.event.client_payload.pull_request.head.repo.full_name }} @@ -22,7 +22,7 @@ jobs: id: pip-cache run: echo "datew=$(date '+%Y-%V')" >> $GITHUB_OUTPUT - name: cache for pip - uses: actions/cache@v5 + uses: actions/cache@v6 id: cache with: path: | @@ -37,8 +37,8 @@ jobs: pip uninstall -y monai pip uninstall -y monai-weekly pip uninstall -y monai-weekly - python -m pip install --upgrade torch torchvision torchaudio torchtext - python -m pip install -r requirements-dev.txt + python monai/config/print_dependencies.py build-system | xargs -d '\n' pip install --no-cache-dir --no-build-isolation + python -m pip install --no-build-isolation .[all,testing] rm -rf /github/home/.cache/torch/hub/mmars/ - name: Clean directory run: | @@ -89,7 +89,7 @@ jobs: runs-on: [self-hosted, linux, x64, command1] steps: # checkout the pull request branch - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: token: ${{ secrets.PR_MAINTAIN }} repository: ${{ github.event.client_payload.pull_request.head.repo.full_name }} @@ -98,7 +98,7 @@ jobs: id: pip-cache run: echo "datew=$(date '+%Y-%V')" >> $GITHUB_OUTPUT - name: cache for pip - uses: actions/cache@v5 + uses: actions/cache@v6 id: cache with: path: | @@ -113,8 +113,7 @@ jobs: pip uninstall -y monai pip uninstall -y monai-weekly pip uninstall -y monai-weekly - python -m pip install --upgrade torch torchvision torchaudio torchtext - python -m pip install -r requirements-dev.txt + python -m pip install .[all,testing] rm -rf /github/home/.cache/torch/hub/mmars/ - name: Clean directory run: | @@ -124,7 +123,7 @@ jobs: nvidia-smi export CUDA_VISIBLE_DEVICES=$(python -m tests.utils -c 1 | tail -n 1) echo $CUDA_VISIBLE_DEVICES - python -c "import torch; print(torch.__version__); print('{} of GPUs available'.format(torch.cuda.device_count()))" + python -c "import torch; print(torch.__version__); print(f'{torch.cuda.device_count()} of GPUs available')" python -c 'import torch; print(torch.rand(5,3, device=torch.device("cuda:0")))' - name: Auto3dseg latest algo diff --git a/.github/workflows/pythonapp-gpu.yml b/.github/workflows/pythonapp-gpu.yml index f851966e01e..8378d2742ab 100644 --- a/.github/workflows/pythonapp-gpu.yml +++ b/.github/workflows/pythonapp-gpu.yml @@ -46,7 +46,7 @@ jobs: options: --gpus all --env NVIDIA_DISABLE_REQUIRE=true # workaround for unsatisfied condition: cuda>=11.6 runs-on: [self-hosted, linux, x64, common] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: apt install if: github.event.pull_request.merged != true run: | @@ -96,7 +96,7 @@ jobs: rm -rf $(python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())")/ruamel* rm -rf $(python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())")/llvmlite* #6377 python -m pip install ${{ matrix.pytorch }} - python -m pip install -r requirements-dev.txt + python -m pip install .[all,testing] python -m pip list - name: Run quick tests (GPU) if: github.event.pull_request.merged != true @@ -127,6 +127,6 @@ jobs: shell: bash - name: Upload coverage if: ${{ github.head_ref != 'dev' && github.event.pull_request.merged != true }} - uses: codecov/codecov-action@v6 + uses: codecov/codecov-action@v7 with: files: ./coverage.xml diff --git a/.github/workflows/pythonapp-hyena-gpu.yml b/.github/workflows/pythonapp-hyena-gpu.yml new file mode 100644 index 00000000000..b080942bf3c --- /dev/null +++ b/.github/workflows/pythonapp-hyena-gpu.yml @@ -0,0 +1,73 @@ +# Optional self-hosted GPU CI for the HyenaND test surface. +# +# This workflow exercises the CUDA-required Hyena tests +# (tests/networks/blocks/test_hyena_block.py CUDA cases, the four-paper-variant +# forward and gradient cases in tests/networks/nets/test_swin_unetr.py and +# tests/networks/nets/test_hyena_nd_unetr.py, the SwinUNETR(use_hyena=False) +# golden-hash backward-compat regression, and sliding-window inference). +# +# Disabled by default (``if: false``). To enable: +# 1. Ensure a self-hosted runner with the labels below is available, AND +# 2. Ensure the runner has CUDA-capable hardware visible (the existing +# ``pythonapp-gpu.yml`` uses ``--gpus all`` against ``[self-hosted, linux, +# x64, common]``). Reuse that pool if possible. +# 3. Flip ``if: false`` to ``if: github.event.pull_request.merged != true`` +# (mirroring ``pythonapp-gpu.yml``'s gating pattern). +# +# nvsubquadratic (Hyena's optional dep) requires Python >= 3.10; any NGC base with +# Python >= 3.10 works. The accelerated [cuda] kernels build against the container nvcc. + +name: hyena-gpu + +on: + workflow_dispatch: + +concurrency: + group: hyena-gpu-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + GPU-Hyena: + if: ${{ false }} # See header for enable instructions. + strategy: + matrix: + environment: + # NGC PyTorch 25.05 ships Python 3.12 and CUDA 12.5. Bump as needed. + - "NGC25.05+PY312" + include: + - environment: NGC25.05+PY312 + base: "nvcr.io/nvidia/pytorch:25.05-py3" + container: + image: ${{ matrix.base }} + options: --gpus all --env NVIDIA_DISABLE_REQUIRE=true + runs-on: [self-hosted, linux, x64, common] + steps: + - uses: actions/checkout@v6 + - name: Install dependencies + run: | + python -m pip install --upgrade pip wheel + python -c "import sys; assert sys.version_info >= (3, 10), f'Python >= 3.10 required for nvsubquadratic, got {sys.version}'" + python -m pip install -e .[all,testing] + # Install nvsubquadratic with --no-deps: the default torch_fft path needs only + # torch + einops + omegaconf, and nvsubquadratic pins torch>=2.10,<2.11 which can + # clash with the container's torch. To exercise the accelerated fused CUDA + # kernels instead, install the [cuda] extra (subquadratic-ops-torch-cu12, builds + # against the container's nvcc) and set fft_backend="subq_ops" in the tests. + python -m pip install omegaconf + python -m pip install --no-deps 'nvsubquadratic>=0.1.1' + python -m pip list + shell: bash + - name: Verify CUDA + nvsubquadratic + run: | + nvidia-smi + python -c "import torch; assert torch.cuda.is_available(); print('CUDA OK:', torch.cuda.get_device_name(0))" + python -c "from monai.networks.blocks.hyena import is_nvsubquadratic_available; \ + assert is_nvsubquadratic_available(), 'nvsubquadratic must be importable'" + shell: bash + - name: Run Hyena test suite (CUDA + no-CUDA) + run: | + python -m pytest -v \ + tests/networks/blocks/test_hyena_block.py \ + tests/networks/nets/test_hyena_nd_unetr.py \ + tests/networks/nets/test_swin_unetr.py + shell: bash diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 0a7087c93a6..d4c55e1f08b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -15,16 +15,17 @@ jobs: matrix: python-version: ['3.10', '3.11', '3.12', '3.13'] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: fetch-depth: 0 - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install setuptools + - name: Install Dependencies run: | - python -m pip install --user --upgrade setuptools wheel packaging + python -m pip install --user --upgrade setuptools wheel packaging tomli + python monai/config/print_dependencies.py build-system all testing | xargs -d '\n' pip install --no-build-isolation - name: Build and test source archive and wheel file run: | find /opt/hostedtoolcache/* -maxdepth 0 ! -name 'Python' -exec rm -rf {} \; @@ -55,11 +56,10 @@ jobs: # clean up cd "$root_dir" rm -r "$tmp_dir" - rm -rf monai/ ls -al . - name: Quick test installed run: | - python -m pip install -r requirements-min.txt + python -m pip install -e .[testing] python -m tests.min_tests env: QUICKTEST: True @@ -93,7 +93,7 @@ jobs: needs: packaging runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 # full history so that we can git describe with: fetch-depth: 0 @@ -105,7 +105,8 @@ jobs: run: | find /opt/hostedtoolcache/* -maxdepth 0 ! -name 'Python' -exec rm -rf {} \; git describe - python -m pip install --user --upgrade setuptools wheel packaging + python -m pip install --user --upgrade setuptools wheel packaging tomli + python monai/config/print_dependencies.py build-system all testing | xargs -d '\n' pip install --no-build-isolation python setup.py build cat build/lib/monai/_version.py - name: Upload version @@ -125,7 +126,7 @@ jobs: needs: versioning runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - name: Download version uses: actions/download-artifact@v8 with: @@ -159,8 +160,6 @@ jobs: echo "unmatched version string, please check the tagging branch." exit 1 fi - # remove flake package as it is not needed on hub.docker.com - sed -i '/flake/d' requirements-dev.txt docker build -t projectmonai/monai:"$RELEASE_VERSION" -f Dockerfile . # distribute with a tag to hub.docker.com echo "${{ secrets.DOCKER_PW }}" | docker login -u projectmonai --password-stdin diff --git a/.github/workflows/setupapp.yml b/.github/workflows/setupapp.yml index 26db41f6c48..d7bf1dc6dd0 100644 --- a/.github/workflows/setupapp.yml +++ b/.github/workflows/setupapp.yml @@ -34,7 +34,7 @@ jobs: # options: --gpus all # runs-on: [self-hosted, linux, x64, integration] # steps: - # - uses: actions/checkout@v6 + # - uses: actions/checkout@v7 # - name: cache weekly timestamp # id: pip-cache # run: | @@ -90,7 +90,7 @@ jobs: # matrix: # python-version: ['3.10', '3.11', '3.12'] # steps: - # - uses: actions/checkout@v6 + # - uses: actions/checkout@v7 # with: # fetch-depth: 0 # - name: Set up Python ${{ matrix.python-version }} @@ -156,16 +156,15 @@ jobs: python -c 'import monai; monai.config.print_config()' - name: Get the test cases (dev branch only) if: github.ref == 'refs/heads/dev' - uses: actions/checkout@v6 + uses: actions/checkout@v7 with: ref: dev - name: Quick test installed (dev branch only) if: github.ref == 'refs/heads/dev' run: | cd $GITHUB_WORKSPACE - rm -rf monai/ ls -al . - python -m pip install -r requirements-min.txt + python -m pip install -e .[testing] python -m tests.min_tests env: QUICKTEST: True diff --git a/.github/workflows/weekly-preview.yml b/.github/workflows/weekly-preview.yml index 6a2d07386fa..7b5dd442aaf 100644 --- a/.github/workflows/weekly-preview.yml +++ b/.github/workflows/weekly-preview.yml @@ -6,13 +6,22 @@ permissions: on: schedule: - cron: "0 2 * * 0" # 02:00 of every Sunday + pull_request: + branches: + - dev + +env: + PYTHON_VER: '3.10' + PYTORCH_VER: '2.8.0' + PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cpu" # forces CPU PyTorch installation, should be faster jobs: static-checks: + if: github.event_name == 'schedule' # only check on cron run, these checks are redundant in a PR runs-on: ubuntu-latest strategy: matrix: - opt: ["codeformat", "mypy"] + opt: ["codeformat", "pyrefly"] steps: - name: Clean unused tools run: | @@ -22,59 +31,67 @@ jobs: sudo rm -rf /opt/ghc /usr/local/.ghcup sudo docker system prune -f - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: persist-credentials: false - - name: Set up Python 3.10 + - name: Set up Python ${{ env.PYTHON_VER }} uses: actions/setup-python@v6 with: - python-version: '3.10' + python-version: ${{ env.PYTHON_VER }} cache: 'pip' - name: Install dependencies run: | - python -m pip install --upgrade pip wheel - python -m pip install --no-build-isolation -r requirements-dev.txt + python -m pip install -U pip wheel + python -m pip install .[all,testing] - name: Lint and type check run: | # clean up temporary files $(pwd)/runtests.sh --build --clean $(pwd)/runtests.sh --build --${{ matrix.opt }} - packaging: + publish: if: github.repository == 'Project-MONAI/MONAI' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: - ref: dev + # get the ref for the PR branch or dev if this is a cron job + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || 'dev' }} fetch-depth: 0 persist-credentials: false - - name: Set up Python 3.10 + - name: Set up Python ${{ env.PYTHON_VER }} uses: actions/setup-python@v6 with: - python-version: '3.10' - - name: Install setuptools + python-version: ${{ env.PYTHON_VER }} + cache: 'pip' + - name: Install tools run: | - python -m pip install --user --upgrade setuptools wheel packaging + python -m pip install -U pip build - name: Build distribution run: | export HEAD_COMMIT_ID=$(git rev-parse HEAD) - sed -i 's/name\ =\ monai$/name\ =\ monai-weekly/g' setup.cfg + sed -i 's/name\ =\ "monai"$/name\ =\ "monai-weekly"/g' pyproject.toml echo "__commit_id__ = \"$HEAD_COMMIT_ID\"" >> monai/__init__.py - git diff setup.cfg monai/__init__.py + git diff pyproject.toml monai/__init__.py git config user.name "CI Builder" git config user.email "monai.contact@gmail.com" - git add setup.cfg monai/__init__.py + git add pyproject.toml monai/__init__.py git commit -m "Weekly build at $HEAD_COMMIT_ID" export YEAR_WEEK=$(date +'%y%U') echo "Year week for tag is ${YEAR_WEEK}" if ! [[ $YEAR_WEEK =~ ^[0-9]{4}$ ]] ; then echo "Wrong 'year week' format. Should be 4 digits."; exit 1 ; fi - git tag "1.6.dev${YEAR_WEEK}" + git tag "1.7.dev${YEAR_WEEK}" git log -1 git tag --list - python setup.py sdist bdist_wheel - + python -m build + ls -lh dist + - name: Test Installation + run: | + pip install dist/*.whl + pip list + (cd "$(mktemp -d)" && python -c 'import monai; print(monai.__version__)') - name: Publish to PyPI + if: github.event_name == 'schedule' # only publish on cron run uses: pypa/gh-action-pypi-publish@release/v1 with: user: __token__ diff --git a/.gitignore b/.gitignore index 76c6ab0d124..d0bdc54018d 100644 --- a/.gitignore +++ b/.gitignore @@ -110,12 +110,17 @@ venv.bak/ # pytype cache .pytype/ +# pyrefly cache +.pyrefly_cache/ + # mypy .mypy_cache/ +.dmypy.json + examples/scd_lvsegs.npz temp/ .idea/ -.dmypy.json +.plans/ *~ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b6c1a3c1128..ae03b5bae9a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -13,11 +13,13 @@ repos: hooks: - id: end-of-file-fixer - id: trailing-whitespace + - id: check-ast - id: check-yaml - id: check-docstring-first - id: check-executables-have-shebangs - id: check-toml - id: check-case-conflict + - id: check-illegal-windows-names - id: check-added-large-files args: ['--maxkb=1024'] - id: detect-private-key @@ -26,8 +28,9 @@ repos: args: ['--autofix', '--no-sort-keys', '--indent=4'] - id: end-of-file-fixer - id: mixed-line-ending + - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.11 + rev: v0.16.5 hooks: - id: ruff-check args: ["--fix"] @@ -37,8 +40,25 @@ repos: ^monai/_version.py ) - - repo: https://github.com/hadialqattan/pycln - rev: v2.6.0 + - repo: https://github.com/psf/black-pre-commit-mirror + rev: 26.5.1 # Black version, keep synced with MONAI requirements hooks: - - id: pycln - args: [--config=pyproject.toml] + - id: black + language_version: python3 + # black will be given individual file names and so will ignore the excludes in pyproject.toml + exclude: | + (?x)( + ^versioneer.py| + ^monai/_version.py + ) + + - repo: https://github.com/pycqa/isort + rev: 9.0.1 # isort version, keep synced with MONAI requirements + hooks: + - id: isort + name: isort (python) + exclude: | + (?x)( + ^versioneer.py| + ^monai/_version.py + ) diff --git a/CHANGELOG.md b/CHANGELOG.md index 419210a9030..987a1d3a2e4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ All notable changes to MONAI are documented in this file. The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ## [Unreleased] +### Added +* `NaViT` (`monai.networks.nets.NaViT`): Native Resolution Vision Transformer with Patch n' Pack, supporting variable-resolution 2D and 3D inputs. Implements factorized positional embeddings, token dropout, attention pooling, and QK normalization, based on ["Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution"](https://arxiv.org/abs/2307.06304). +* `HyenaMixer`, `HyenaTransformerBlock`, and `DepthwiseFFTConv{2,3}d` in `monai.networks.blocks`: subquadratic O(N log N) alternatives to windowed self-attention, backed by the HyenaND operator from the optional `nvsubquadratic` package. +* `HyenaNDUNETR` (`monai.networks.nets.HyenaNDUNETR`): thin `SwinUNETR` subclass with a `get_variant(name)` classmethod for the three Hyena variants (`HHHH`, `HAHA`, `HHAA`) from the NeurIPS 2026 paper "Native Multi-Dimensional Subquadratic Operators via Input Dependent Long Convolutions" (paper id 26539). +* `SwinUNETR.use_hyena` and `SwinUNETR.hyena_stages` kwargs to thread HyenaND blocks through any subset of Swin stages. Default `use_hyena=False` preserves bit-identical forward behavior of the existing code path. +* New `[hyena]` extras_require in setup.cfg (`pip install monai[hyena]`). ## [1.6.0] - 2026-06-12 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2ad171abb1b..1fb68b4b583 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -38,7 +38,7 @@ Please note that, as per PyTorch, MONAI uses American English spelling. This mea ### Preparing pull requests To ensure the code quality, MONAI relies on several linting tools ([black](https://github.com/psf/black), [isort](https://github.com/timothycrosley/isort), [ruff](https://github.com/astral-sh/ruff)), -static type analysis tools ([mypy](https://github.com/python/mypy), [pytype](https://github.com/google/pytype)), as well as a set of unit/integration tests. +static type analysis tools ([pyrefly](https://github.com/facebook/pyrefly)), as well as a set of unit/integration tests. This section highlights all the necessary preparation steps required before sending a pull request. To collaborate efficiently, please read through this section and follow them. @@ -57,10 +57,6 @@ Before submitting a pull request, we recommend that all linting should pass, by ```bash # optionally update the dependencies and dev tools python -m pip install -U pip -python -m pip install -U -r requirements-dev.txt - -# run the linting and type checking tools -./runtests.sh --codeformat # try to fix the coding style errors automatically ./runtests.sh --autofix @@ -132,7 +128,7 @@ It is recommended that the new test `test_[module_name].py` is constructed by us python 3.9+ build-in functions, `torch`, `numpy`, `coverage` (for reporting code coverages) and `parameterized` (for organising test cases) packages. If it requires any other external packages, please make sure: -- the packages are listed in [`requirements-dev.txt`](requirements-dev.txt) +- the packages are listed in [`pyproject.toml`](pyproject.toml) - the new test `test_[module_name].py` is added to the `exclude_cases` in [`./tests/min_tests.py`](./tests/min_tests.py) so that the minimal CI runner will not execute it. @@ -219,10 +215,8 @@ Integration tests with minimal requirements are deployed to ensure this strategy To add new optional dependencies, please communicate with the core team during pull request reviews, and add the necessary information (at least) to the following files: -- [setup.cfg](https://github.com/Project-MONAI/MONAI/blob/dev/setup.cfg) (for package's `[options.extras_require]` config) -- [requirements-dev.txt](https://github.com/Project-MONAI/MONAI/blob/dev/requirements-dev.txt) (pip requirements file) +- [pyproject.toml](https://github.com/Project-MONAI/MONAI/blob/dev/pyproject.toml) (for package's `[project.optional-dependencies]` config) - [docs/requirements.txt](https://github.com/Project-MONAI/MONAI/blob/dev/docs/requirements.txt) (docs pip requirements file) -- [environment-dev.yml](https://github.com/Project-MONAI/MONAI/blob/dev/environment-dev.yml) (conda environment file) - [installation.md](https://github.com/Project-MONAI/MONAI/blob/dev/docs/source/installation.md) (documentation) When writing unit tests that use 3rd-party packages, it is a good practice to always consider @@ -303,6 +297,13 @@ By making a contribution to this project, I certify that: this project or the open source license(s) involved. ``` +> **Tip:** If you need to add a DCO remediation commit (e.g., after a force-push +> or rebase), include `[skip ci]` in the commit message so the remediation +> does not trigger unnecessary CI pipelines: +> ```bash +> git commit -s --allow-empty -m 'DCO Remediation Commit for... [skip ci]' +> ``` + #### Utility functions MONAI provides a set of generic utility functions and frequently used routines. @@ -358,6 +359,90 @@ Ideally, the new branch should be based on the latest `dev` branch. 1. Reviewer and contributor may have discussions back and forth until all comments addressed. 1. Wait for the pull request to be merged. +## Skipping CI + +MONAI's CI pipelines run automatically on every push and pull request. +These pipelines can be resource-intensive, especially the full premerge matrix +which spans multiple OSes, Python versions, and PyTorch versions. + +To reduce unnecessary resource consumption and speed up iteration, you can +skip CI on commits that don't need automated validation — for example, +documentation-only changes, README updates, workflow YAML changes, or WIP +commits during development. + +### Mechanism + +GitHub Actions natively supports skipping `push` and `pull_request` workflows +when the commit message contains any of the following strings: + +- `[skip ci]` +- `[ci skip]` +- `[no ci]` +- `[skip actions]` +- `[actions skip]` + +These are case-insensitive. `[skip ci]` is the recommended convention for +this repository. + +Alternatively, you can add a `skip-checks: true` trailer at the end of the +commit message, preceded by two blank lines: + +``` +commit message + +skip-checks: true +``` + +### Usage + +Add the keyword anywhere in the commit message when committing: + +```bash +git commit -s -m 'update docs [skip ci]' +``` + +If the HEAD commit of a pull request contains the skip instruction, +the entire PR's pull_request-triggered workflows are skipped. + +### Which workflows are affected + +The skip instruction applies only to workflows triggered by `on: push` or +`on: pull_request` events. All other workflows — those using `issue_comment`, +`repository_dispatch`, `schedule`, or `workflow_dispatch` — use different +event types and are **not** affected by `[skip ci]`. + +### Important caveat + +If a workflow is skipped via `[skip ci]`, its associated checks remain in +"Pending" state. If your pull request requires those checks to pass before +merging, you will need to push a new commit **without** the skip instruction +to trigger the CI pipelines. + +### When to use + +Use `[skip ci]` for commits that are safe to skip CI: + +- Documentation-only changes (`docs/`, `README.md`, docstrings) +- Workflow configuration changes (`.github/`) +- Repository metadata (`.gitignore`, `CONTRIBUTING.md`, `LICENSE`) +- WIP or draft commits during local development + +Do **not** use `[skip ci]` for commits that change: + +- Source code in `monai/` +- Test files in `tests/` +- Dependencies (`pyproject.toml`, `setup.py`, `docs/requirements.txt`) +- Anything that could affect correctness or compatibility + +### Quick example + +```bash +git commit -s -m 'fix typo in README [skip ci]' +``` + +This commit will be recorded in the repository history but will not +consume CI minutes. + ## The code reviewing process ### Reviewing pull requests diff --git a/Dockerfile b/Dockerfile index e9f005e75ef..79d9d7677e1 100644 --- a/Dockerfile +++ b/Dockerfile @@ -16,6 +16,8 @@ FROM ${PYTORCH_IMAGE} LABEL maintainer="monai.contact@gmail.com" +ENV BUILD_MONAI=1 + # TODO: remark for issue [revise the dockerfile](https://github.com/zarr-developers/numcodecs/issues/431) RUN if [[ $(uname -m) =~ "aarch64" ]]; then \ export CFLAGS="-O3" && \ @@ -24,27 +26,6 @@ RUN if [[ $(uname -m) =~ "aarch64" ]]; then \ pip install numcodecs; \ fi -WORKDIR /opt/monai - -# install full deps -COPY requirements.txt requirements-min.txt requirements-dev.txt /tmp/ -RUN cp /tmp/requirements.txt /tmp/req.bak \ - && awk '!/torch/' /tmp/requirements.txt > /tmp/tmp && mv /tmp/tmp /tmp/requirements.txt \ - && python -m pip install --upgrade --no-cache-dir --no-build-isolation pip wheel wheel-stub \ - && python -m pip install --no-cache-dir --no-build-isolation -r /tmp/requirements-dev.txt - -# compile ext and remove temp files -# TODO: remark for issue [revise the dockerfile #1276](https://github.com/Project-MONAI/MONAI/issues/1276) -# please specify exact files and folders to be copied -- else, basically always, the Docker build process cannot cache -# this or anything below it and always will build from at most here; one file change leads to no caching from here on... - -COPY LICENSE CHANGELOG.md CODE_OF_CONDUCT.md CONTRIBUTING.md README.md versioneer.py setup.py setup.cfg runtests.sh MANIFEST.in ./ -COPY tests ./tests -COPY monai ./monai - -RUN BUILD_MONAI=1 FORCE_CUDA=1 python setup.py develop \ - && rm -rf build __pycache__ - # NGC Client WORKDIR /opt/tools ARG NGC_CLI_URI="https://ngc.nvidia.com/downloads/ngccli_linux.zip" @@ -59,5 +40,17 @@ RUN apt-get update \ ENV PATH=${PATH}:/opt/tools ENV POLYGRAPHY_AUTOINSTALL_DEPS=1 - WORKDIR /opt/monai + +# TODO: remark for issue [revise the dockerfile #1276](https://github.com/Project-MONAI/MONAI/issues/1276) +# please specify exact files and folders to be copied -- else, basically always, the Docker build process cannot cache +# this or anything below it and always will build from at most here; one file change leads to no caching from here on... +COPY LICENSE CHANGELOG.md CODE_OF_CONDUCT.md CONTRIBUTING.md README.md versioneer.py setup.py pyproject.toml runtests.sh MANIFEST.in ./ +COPY tests ./tests +COPY monai ./monai + +# Need to install build requirements explicitly so that no-build-isolation can be used. This is needed to make pip build +# against the included version of PyTorch, rather than install a new version in the isolated environment. Constraint +# files will not work for this image which installed things like PyTorch through files which are no longer present. +RUN python monai/config/print_dependencies.py build-system | xargs -d '\n' pip install --no-cache-dir --no-build-isolation \ + && FORCE_CUDA=1 pip install --no-cache-dir --no-build-isolation -e .[all,testing] diff --git a/Dockerfile.slim b/Dockerfile.slim index 1b3cf3bcd4c..bfaf897967f 100644 --- a/Dockerfile.slim +++ b/Dockerfile.slim @@ -28,9 +28,9 @@ RUN apt update && apt upgrade -y && \ wget https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/cuda-keyring_1.1-1_all.deb && \ dpkg -i cuda-keyring_1.1-1_all.deb && \ apt update && \ - ${APT_INSTALL} cuda-toolkit-12-9 && \ + ${APT_INSTALL} cuda-toolkit-13-3 && \ rm -rf /usr/lib/python*/EXTERNALLY-MANAGED /var/lib/apt/lists/* && \ - python -m pip install --upgrade --no-cache-dir --no-build-isolation pip + python -m pip install --upgrade --no-cache-dir pip # TODO: remark for issue [revise the dockerfile](https://github.com/zarr-developers/numcodecs/issues/431) RUN if [[ $(uname -m) =~ "aarch64" ]]; then \ @@ -46,18 +46,14 @@ RUN wget -q ${NGC_CLI_URI} && unzip ngccli_linux.zip && chmod u+x ngc-cli/ngc && WORKDIR /opt/monai -# copy relevant parts of repo -COPY requirements.txt requirements-min.txt requirements-dev.txt versioneer.py setup.py setup.cfg pyproject.toml ./ -COPY LICENSE CHANGELOG.md CODE_OF_CONDUCT.md CONTRIBUTING.md README.md MANIFEST.in runtests.sh ./ +# TODO: remark for issue [revise the dockerfile #1276](https://github.com/Project-MONAI/MONAI/issues/1276) +# please specify exact files and folders to be copied -- else, basically always, the Docker build process cannot cache +# this or anything below it and always will build from at most here; one file change leads to no caching from here on... +COPY LICENSE CHANGELOG.md CODE_OF_CONDUCT.md CONTRIBUTING.md README.md versioneer.py setup.py pyproject.toml runtests.sh MANIFEST.in ./ COPY tests ./tests COPY monai ./monai -# install full deps -RUN python -m pip install --no-cache-dir --no-build-isolation -U wheel wheel-stub -RUN python -m pip install --no-cache-dir --no-build-isolation "torch>=2.8.0,<2.11" -r requirements-dev.txt - -# compile ext -RUN CUDA_HOME=/usr/local/cuda FORCE_CUDA=1 USE_COMPILED=1 BUILD_MONAI=1 python setup.py develop +RUN BUILD_MONAI=1 FORCE_CUDA=1 pip install --no-cache-dir -e .[all,testing] # recreate the image without the installed CUDA packages then copy the installed MONAI and Python directories FROM ${IMAGE} AS build2 @@ -66,10 +62,9 @@ ENV DEBIAN_FRONTEND=noninteractive ENV APT_INSTALL="apt install -y --no-install-recommends" RUN apt update && apt upgrade -y && \ - ${APT_INSTALL} ca-certificates python3-pip python-is-python3 git libopenslide0 && \ + ${APT_INSTALL} ca-certificates python-is-python3 git libopenslide0 && \ apt clean && \ - rm -rf /usr/lib/python*/EXTERNALLY-MANAGED /var/lib/apt/lists/* && \ - python -m pip install --upgrade --no-cache-dir --no-build-isolation pip + rm -rf /usr/lib/python*/EXTERNALLY-MANAGED /var/lib/apt/lists/* COPY --from=build /opt/monai /opt/monai COPY --from=build /opt/tools /opt/tools diff --git a/README.md b/README.md index d0927ad8c3a..92e4ca7eba8 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ MONAI works with the [currently supported versions of Python](https://devguide.p * Major releases of MONAI will have dependency versions stated for them. The current state of the `dev` branch in this repository is the unreleased development version of MONAI which typically will support current versions of dependencies and include updates and bug fixes to do so. * PyTorch support covers [the current version](https://github.com/pytorch/pytorch/releases) plus three previous minor versions. If compatibility issues with a PyTorch version and other dependencies arise, support for a version may be delayed until a major release. * Our support policy for other dependencies adheres for the most part to [SPEC0](https://scientific-python.org/specs/spec-0000), where dependency versions are supported where possible for up to two years. Discovered vulnerabilities or defects may require certain versions to be explicitly not supported. -* See the `requirements*.txt` files for dependency version information. +* See the `pyproject.toml` file for dependency version information. ## Installation diff --git a/docs/requirements.txt b/docs/requirements.txt index 3027d401646..6598722dcab 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -20,8 +20,8 @@ sphinxcontrib-serializinghtml sphinx-autodoc-typehints==1.11.1 pandas einops -transformers>=4.53.0 -mlflow>=2.12.2,<3.13 +transformers>=5.5.0 +mlflow>=3.15.2 clearml>=1.10.0rc0 tensorboardX imagecodecs; platform_system == "Linux" or platform_system == "Darwin" @@ -33,6 +33,7 @@ pynrrd pydicom h5py nni==2.10.1; platform_system == "Linux" and "arm" not in platform_machine and "aarch" not in platform_machine +filelock<3.12.0 optuna opencv-python-headless onnx>=1.13.0 diff --git a/docs/source/installation.md b/docs/source/installation.md index 5123bc3e6bb..006ac23cda7 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -19,13 +19,13 @@ --- -MONAI's core functionality is written in Python 3 (>= 3.10) and only requires [Numpy](https://numpy.org/) and [Pytorch](https://pytorch.org/). +MONAI's core functionality is written in Python 3 (>= 3.10) and only requires [Numpy](https://numpy.org/) and [PyTorch](https://pytorch.org/). The package is currently distributed via Github as the primary source code repository, and the Python package index (PyPI). The pre-built Docker images are made available on DockerHub. To install optional features such as handling the NIfTI files using -[Nibabel](https://nipy.org/nibabel/), or building workflows using [Pytorch +[Nibabel](https://nipy.org/nibabel/), or building workflows using [PyTorch Ignite](https://pytorch.org/ignite/), please follow the instructions: - [Installing the recommended dependencies](#installing-the-recommended-dependencies) @@ -49,6 +49,26 @@ To install the [current milestone release](https://pypi.org/project/monai/): pip install monai ``` +MONAI supports the extras syntax such as `pip install 'monai[nibabel]'`. The options are + +```text +clearml, cucim, cupy, einops, fire, gdown, h5py, huggingface_hub, hyena, ignite, imagecodecs, itk, jsonschema, lmdb, lpips, matplotlib, metrics_reloaded, mlflow, nibabel, nni, onnx, openslide, optuna, pandas, pillow, polygraphy, psutil, pyamg, pybind11, pydicom, pynrrd, pynvml, pyyaml, requests, segment_anything, scipy, skimage, tensorboard, tensorboardX, tifffile, torchio, torchvision, tqdm, transformers, zarr +``` + +which correspond to the packages: `clearml`, `cucim` (`cucim-cu12` or `cucim-cu13`), `cupy-cuda13x`, `einops`, `fire`, `gdown`, `h5py`, `huggingface_hub`, `nvsubquadratic`, `omegaconf`, `pytorch-ignite`, `imagecodecs`, `itk`, `jsonschema`, `lmdb`, `lpips`, `matplotlib`, `MetricsReloaded`, `mlflow`, `nibabel`, `nni`, `filelock`, `onnx`, `onnxruntime`, `onnx_graphsurgeon`, `onnxscript`, `openslide-python`, `openslide-bin`, `optuna`, `pandas`, `pillow`, `polygraphy`, `psutil`, `pyamg`, `pybind11`, `pydicom`, `pynrrd`, `nvidia-ml-py`, `pyyaml`, `requests`, `segment_anything`, `scipy`, `scikit-image`, `tensorboard`, `tensorboardX`, `tifffile`, `torchio`, `torchvision`, `tqdm`, `transformers`, `zarr`. + +Almost all of these can be installed together with the `all` option. For development on MONAI, this should be accompanied by `testing` which will install the testing static checking packages. Cupy is omitted from `all` since the choice between +Cuda 12 and 13 versions of the library can't be resolved when installing and must be manually installed. + +The `hyena` extra pulls in [`nvsubquadratic`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic), +required by `HyenaNDUNETR` / `HyenaMixer` / `HyenaTransformerBlock` (subquadratic +O(N log N) alternatives to windowed self-attention). Install with +`pip install 'monai[hyena]'`. + +The command `pip install 'monai[all,hyena,testing]'` installs almost all the optional dependencies. + +When installing MONAI, the compiled extensions are not compiled by default. Set the environment variable `BUILD_MONAI` to `1` when invoking `pip` to compile these, see below for details. + ### Weekly preview release To install the [weekly preview release](https://pypi.org/project/monai-weekly/): @@ -110,13 +130,41 @@ or, to build with MONAI C++/CUDA extensions: BUILD_MONAI=1 pip install git+https://github.com/Project-MONAI/MONAI ``` -To build the extensions, if the system environment already has a version of Pytorch installed, -`--no-build-isolation` might be preferred: +On Windows the inline `BUILD_MONAI=1 pip install ...` form is not supported by +`cmd.exe` or PowerShell. Set the environment variable first, then run either +install command shown above: + +```bat +:: cmd.exe +set BUILD_MONAI=1 +pip install git+https://github.com/Project-MONAI/MONAI +``` + +```powershell +# PowerShell +$env:BUILD_MONAI="1" +pip install git+https://github.com/Project-MONAI/MONAI +``` + +To build the extensions, if the system environment already has a version of PyTorch installed, `--no-build-isolation` might be preferred: ```bash BUILD_MONAI=1 pip install --no-build-isolation git+https://github.com/Project-MONAI/MONAI ``` +When using build isolation (pip's default behaviour), a version of PyTorch must be installed which may not be the same as an existing install. This can cause the compiled libraries to be built against an ABI-incompatible PyTorch and thus not function at runtime. Building without isolation requires the current environment to have the necessary building libraries already installed. See the `build-system` section of `pyproject.toml` for these libraries, or use the following to install them in a bash environment: + +```bash +python monai/config/print_dependencies.py build-system | xargs -d '\n' pip install --no-build-isolation +``` + +An alternative solution is to use built constraints during installation: + +```bash +pip freeze | grep torch > constraints.txt +pip install --build-constraint constraints.txt git+https://github.com/Project-MONAI/MONAI +``` + this command will download and install the current `dev` branch of [MONAI from GitHub](https://github.com/Project-MONAI/MONAI). @@ -136,6 +184,7 @@ You can install it by running: ```bash cd MONAI/ pip install -e . +# or pip install -e .[all,testing] to include most of the dependencies ``` or, to build with MONAI C++/CUDA extensions and install: @@ -147,6 +196,24 @@ BUILD_MONAI=1 pip install -e . BUILD_MONAI=1 CC=clang CXX=clang++ pip install -e . ``` +On Windows set the environment variable before running `pip install -e .`: + +```bat +:: cmd.exe +cd MONAI/ +set BUILD_MONAI=1 +pip install -e . +``` + +```powershell +# PowerShell +cd MONAI/ +$env:BUILD_MONAI="1" +pip install -e . +``` + +If the compiled extensions were built by pip against a different version of PyTorch than the one in your environment, you may need to run the above with the `--no-build-isoloation` flag to force the use of that version, or use the `--build-constraint` method. + To uninstall the package please run: ```bash @@ -231,12 +298,14 @@ cd MONAI/ pip install -e ".[all]" ``` -To install all optional dependencies with `pip` based on MONAI development environment settings: +To install all optional dependencies with `pip` based on MONAI development environment settings without installing +MONAI itself: ```bash git clone https://github.com/Project-MONAI/MONAI.git cd MONAI/ -pip install -r requirements-dev.txt +python monai/config/print_dependencies.py \* > requirements.txt +pip install -r requirements.txt ``` To install all optional dependencies with `conda` based on MONAI development environment settings (`environment-dev.yml`; @@ -248,16 +317,3 @@ cd MONAI/ conda create -n python= # eg 3.10 conda env update -n -f environment-dev.yml ``` - -Since MONAI v0.2.0, the extras syntax such as `pip install 'monai[nibabel]'` is available via PyPI. - -- The options are - -``` -[nibabel, skimage, scipy, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, cucim, openslide, pandas, einops, transformers, mlflow, clearml, matplotlib, tensorboardX, tifffile, imagecodecs, pyyaml, fire, jsonschema, ninja, pynrrd, pydicom, h5py, nni, optuna, onnx, onnxruntime, zarr, lpips, pynvml, huggingface_hub] -``` - -which correspond to `nibabel`, `scikit-image`,`scipy`, `pillow`, `tensorboard`, -`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `clearml`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `ninja`, `pynrrd`, `pydicom`, `h5py`, `nni`, `optuna`, `onnx`, `onnxruntime`, `zarr`, `lpips`, `nvidia-ml-py`, `huggingface_hub` and `pyamg` respectively. - -- `pip install 'monai[all]'` installs all the optional dependencies. diff --git a/docs/source/lazy_resampling.rst b/docs/source/lazy_resampling.rst index 7b809965f33..9776cd94385 100644 --- a/docs/source/lazy_resampling.rst +++ b/docs/source/lazy_resampling.rst @@ -253,7 +253,7 @@ so the user must set lazy=True on the transforms that they still wish to execute .. figure:: ../images/lazy_resampling_none_example.svg - Figure shwoing the effect of using ``lazy=False`` when ``Compose`` is being executed with ``lazy=None``. Note that + Figure showing the effect of using ``lazy=False`` when ``Compose`` is being executed with ``lazy=None``. Note that the additional resamples that occur due to ``RandRotate90d`` being executed in a non-lazy fashion. @@ -270,4 +270,4 @@ the following transform is a lazy transform, or is configured to execute lazily. .. figure:: ../images/lazy_resampling_apply_pending_example.svg Figure showing the use of :class:`ApplyPendingd` to cause - resampling to occur in the midele of a chain of lazy transforms. + resampling to occur in the middle of a chain of lazy transforms. diff --git a/docs/source/losses.rst b/docs/source/losses.rst index baeebbbe9c0..a5be560324b 100644 --- a/docs/source/losses.rst +++ b/docs/source/losses.rst @@ -78,6 +78,11 @@ Segmentation Losses .. autoclass:: BarlowTwinsLoss :members: +`BoundaryLoss` +~~~~~~~~~~~~~~ +.. autoclass:: BoundaryLoss + :members: + `HausdorffDTLoss` ~~~~~~~~~~~~~~~~~ .. autoclass:: HausdorffDTLoss diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst index 654958bbbf4..ebd745e795f 100644 --- a/docs/source/metrics.rst +++ b/docs/source/metrics.rst @@ -116,6 +116,13 @@ Metrics .. autoclass:: SurfaceDiceMetric :members: +`Absolute volume difference` +---------------------------- +.. autofunction:: compute_absolute_volume_difference + +.. autoclass:: AbsoluteVolumeDifferenceMetric + :members: + `PanopticQualityMetric` ----------------------- .. autofunction:: compute_panoptic_quality diff --git a/docs/source/modules.md b/docs/source/modules.md index b2e95658bf7..a0f24b64a88 100644 --- a/docs/source/modules.md +++ b/docs/source/modules.md @@ -205,7 +205,7 @@ The workflow and some of MONAI event handlers are shown as below [[Workflow exam ### EnsembleEvaluator -A typical ensemble procoess is implemented as a ready-to-use workflow [[Cross validation and model ensemble tutorial]](https://github.com/Project-MONAI/tutorials/blob/main/modules/cross_validation_models_ensemble.ipynb): +A typical ensemble process is implemented as a ready-to-use workflow [[Cross validation and model ensemble tutorial]](https://github.com/Project-MONAI/tutorials/blob/main/modules/cross_validation_models_ensemble.ipynb): 1. Split all the training dataset into K folds. 2. Train K models with every K-1 folds data. 3. Execute inference on the test data with all the K models. diff --git a/docs/source/networks.rst b/docs/source/networks.rst index de0aece3f7a..b4857c20a3f 100644 --- a/docs/source/networks.rst +++ b/docs/source/networks.rst @@ -129,6 +129,23 @@ Blocks .. autoclass:: TransformerBlock :members: +`Hyena Mixer` +~~~~~~~~~~~~~ +.. autoclass:: HyenaMixer + :members: + +`Hyena Transformer Block` +~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autoclass:: HyenaTransformerBlock + :members: + +`Depthwise FFT Convolution` +~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autoclass:: DepthwiseFFTConv2d + :members: +.. autoclass:: DepthwiseFFTConv3d + :members: + `UNETR Block` ~~~~~~~~~~~~~ .. autoclass:: UnetrBasicBlock @@ -591,6 +608,11 @@ Nets .. autoclass:: SwinUNETR :members: +`HyenaNDUNETR` +~~~~~~~~~~~~~~ +.. autoclass:: HyenaNDUNETR + :members: + `BasicUNet` ~~~~~~~~~~~ .. autoclass:: BasicUNet @@ -640,6 +662,11 @@ Nets .. autoclass:: VarAutoEncoder :members: +`NaViT` +~~~~~~~ +.. autoclass:: NaViT + :members: + `ViT` ~~~~~ .. autoclass:: ViT diff --git a/docs/source/utils.rst b/docs/source/utils.rst index ae3b476c3e1..958d27337e6 100644 --- a/docs/source/utils.rst +++ b/docs/source/utils.rst @@ -80,3 +80,8 @@ Ordering -------- .. automodule:: monai.utils.ordering :members: + +Safe Evaluation +--------------- +.. automodule:: monai.utils.safeeval + :members: diff --git a/docs/source/whatsnew_1_5.md b/docs/source/whatsnew_1_5.md index 8b68d716866..5fd7d779234 100644 --- a/docs/source/whatsnew_1_5.md +++ b/docs/source/whatsnew_1_5.md @@ -3,7 +3,7 @@ - Support numpy 2.x and Pytorch 2.6 - MAISI inference accelerate -- Bundles storage changed to huggingface and correspoinding api updated in core +- Bundles storage changed to huggingface and corresponding api updated in core - Ported remaining generative tutorials and bundles - New tutorials: - [2d_regression/image_restoration.ipynb](https://github.com/Project-MONAI/tutorials/blob/main/2d_regression/image_restoration.ipynb) diff --git a/environment-dev.yml b/environment-dev.yml index b2457006c89..7d6d95a306e 100644 --- a/environment-dev.yml +++ b/environment-dev.yml @@ -1,15 +1,8 @@ name: monai channels: - - pytorch - defaults - - nvidia - - conda-forge dependencies: - - numpy>=1.24,<3.0 - - pytorch>=2.8.0 - - torchio - - torchvision - - pytorch-cuda>=11.6 + - python>=3.10 - pip - pip: - - -r requirements-dev.txt + - -e .[all,testing] diff --git a/monai/__init__.py b/monai/__init__.py index 45e15ddcd9a..2d6f5dfc37b 100644 --- a/monai/__init__.py +++ b/monai/__init__.py @@ -62,8 +62,8 @@ def filter(self, record): PY_REQUIRED_MINOR = 9 version_dict = get_versions() -__version__: str = version_dict.get("version", "0+unknown") -__revision_id__: str = version_dict.get("full-revisionid") +__version__: str = str(version_dict.get("version", "0+unknown")) +__revision_id__: str = str(version_dict.get("full-revisionid") or "") del get_versions, version_dict __copyright__ = "(c) MONAI Consortium" diff --git a/monai/_version.py b/monai/_version.py index f2342271048..d14412be66c 100644 --- a/monai/_version.py +++ b/monai/_version.py @@ -5,8 +5,9 @@ # directories (produced by setup.py build) will contain a much shorter file # that just contains the computed version number. -# This file is released into the public domain. Generated by -# versioneer-0.23 (https://github.com/python-versioneer/python-versioneer) +# This file is released into the public domain. +# Generated by versioneer-0.29 +# https://github.com/python-versioneer/python-versioneer """Git implementation of _version.py.""" @@ -15,11 +16,11 @@ import re import subprocess import sys -from collections.abc import Callable +from typing import Any, Callable, Dict, List, Optional, Tuple import functools -def get_keywords(): +def get_keywords() -> Dict[str, str]: """Get the keywords needed to look up the version information.""" # these strings will be replaced by git during git-archive. # setup.py/versioneer.py will grep for the variable names, so they must @@ -35,8 +36,15 @@ def get_keywords(): class VersioneerConfig: """Container for Versioneer configuration parameters.""" + VCS: str + style: str + tag_prefix: str + parentdir_prefix: str + versionfile_source: str + verbose: bool -def get_config(): + +def get_config() -> VersioneerConfig: """Create, populate and return the VersioneerConfig() object.""" # these strings are filled in when 'setup.py versioneer' creates # _version.py @@ -54,13 +62,13 @@ class NotThisMethod(Exception): """Exception raised if a method is not valid for the current scenario.""" -LONG_VERSION_PY: dict[str, str] = {} -HANDLERS: dict[str, dict[str, Callable]] = {} +LONG_VERSION_PY: Dict[str, str] = {} +HANDLERS: Dict[str, Dict[str, Callable]] = {} -def register_vcs_handler(vcs, method): # decorator +def register_vcs_handler(vcs: str, method: str) -> Callable: # decorator """Create decorator to mark a method as the handler of a VCS.""" - def decorate(f): + def decorate(f: Callable) -> Callable: """Store f in HANDLERS[vcs][method].""" if vcs not in HANDLERS: HANDLERS[vcs] = {} @@ -69,13 +77,19 @@ def decorate(f): return decorate -def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, - env=None): +def run_command( + commands: List[str], + args: List[str], + cwd: Optional[str] = None, + verbose: bool = False, + hide_stderr: bool = False, + env: Optional[Dict[str, str]] = None, +) -> Tuple[Optional[str], Optional[int]]: """Call the given command(s).""" assert isinstance(commands, list) process = None - popen_kwargs = {} + popen_kwargs: Dict[str, Any] = {} if sys.platform == "win32": # This hides the console window if pythonw.exe is used startupinfo = subprocess.STARTUPINFO() @@ -91,8 +105,7 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, stderr=(subprocess.PIPE if hide_stderr else None), **popen_kwargs) break - except OSError: - e = sys.exc_info()[1] + except OSError as e: if e.errno == errno.ENOENT: continue if verbose: @@ -112,7 +125,11 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, return stdout, process.returncode -def versions_from_parentdir(parentdir_prefix, root, verbose): +def versions_from_parentdir( + parentdir_prefix: str, + root: str, + verbose: bool, +) -> Dict[str, Any]: """Try to determine the version from the parent directory name. Source tarballs conventionally unpack into a directory that includes both @@ -137,15 +154,15 @@ def versions_from_parentdir(parentdir_prefix, root, verbose): @register_vcs_handler("git", "get_keywords") -def git_get_keywords(versionfile_abs): +def git_get_keywords(versionfile_abs: str) -> Dict[str, str]: """Extract version information from the given file.""" # the code embedded in _version.py can just fetch the value of these # keywords. When used from setup.py, we don't want to import _version.py, # so we do it with a regexp instead. This function is not used from # _version.py. - keywords = {} + keywords: Dict[str, str] = {} try: - with open(versionfile_abs) as fobj: + with open(versionfile_abs, "r") as fobj: for line in fobj: if line.strip().startswith("git_refnames ="): mo = re.search(r'=\s*"(.*)"', line) @@ -165,7 +182,11 @@ def git_get_keywords(versionfile_abs): @register_vcs_handler("git", "keywords") -def git_versions_from_keywords(keywords, tag_prefix, verbose): +def git_versions_from_keywords( + keywords: Dict[str, str], + tag_prefix: str, + verbose: bool, +) -> Dict[str, Any]: """Get version information from git keywords.""" if "refnames" not in keywords: raise NotThisMethod("Short version file found") @@ -229,7 +250,12 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose): @register_vcs_handler("git", "pieces_from_vcs") -def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): +def git_pieces_from_vcs( + tag_prefix: str, + root: str, + verbose: bool, + runner: Callable = run_command +) -> Dict[str, Any]: """Get version from 'git describe' in the root of the source tree. This only gets called if the git-archive 'subst' keywords were *not* @@ -248,7 +274,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): runner = functools.partial(runner, env=env) _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, - hide_stderr=True) + hide_stderr=not verbose) if rc != 0: if verbose: print("Directory %s not under git control" % root) @@ -259,7 +285,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): describe_out, rc = runner(GITS, [ "describe", "--tags", "--dirty", "--always", "--long", "--match", f"{tag_prefix}[[:digit:]]*" - ], cwd=root) + ], cwd=root) # --long was added in git-1.5.5 if describe_out is None: raise NotThisMethod("'git describe' failed") @@ -269,7 +295,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): raise NotThisMethod("'git rev-parse' failed") full_out = full_out.strip() - pieces = {} + pieces: Dict[str, Any] = {} pieces["long"] = full_out pieces["short"] = full_out[:7] # maybe improved later pieces["error"] = None @@ -361,14 +387,14 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): return pieces -def plus_or_dot(pieces): +def plus_or_dot(pieces: Dict[str, Any]) -> str: """Return a + if we don't already have one, else return a .""" if "+" in pieces.get("closest-tag", ""): return "." return "+" -def render_pep440(pieces): +def render_pep440(pieces: Dict[str, Any]) -> str: """Build up version string, with post-release "local version identifier". Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you @@ -393,7 +419,7 @@ def render_pep440(pieces): return rendered -def render_pep440_branch(pieces): +def render_pep440_branch(pieces: Dict[str, Any]) -> str: """TAG[[.dev0]+DISTANCE.gHEX[.dirty]] . The ".dev0" means not master branch. Note that .dev0 sorts backwards @@ -423,7 +449,7 @@ def render_pep440_branch(pieces): return rendered -def pep440_split_post(ver): +def pep440_split_post(ver: str) -> Tuple[str, Optional[int]]: """Split pep440 version string at the post-release segment. Returns the release segments before the post-release and the @@ -433,7 +459,7 @@ def pep440_split_post(ver): return vc[0], int(vc[1] or 0) if len(vc) == 2 else None -def render_pep440_pre(pieces): +def render_pep440_pre(pieces: Dict[str, Any]) -> str: """TAG[.postN.devDISTANCE] -- No -dirty. Exceptions: @@ -457,7 +483,7 @@ def render_pep440_pre(pieces): return rendered -def render_pep440_post(pieces): +def render_pep440_post(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]+gHEX] . The ".dev0" means dirty. Note that .dev0 sorts backwards @@ -484,7 +510,7 @@ def render_pep440_post(pieces): return rendered -def render_pep440_post_branch(pieces): +def render_pep440_post_branch(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] . The ".dev0" means not master branch. @@ -513,7 +539,7 @@ def render_pep440_post_branch(pieces): return rendered -def render_pep440_old(pieces): +def render_pep440_old(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]] . The ".dev0" means dirty. @@ -535,7 +561,7 @@ def render_pep440_old(pieces): return rendered -def render_git_describe(pieces): +def render_git_describe(pieces: Dict[str, Any]) -> str: """TAG[-DISTANCE-gHEX][-dirty]. Like 'git describe --tags --dirty --always'. @@ -555,7 +581,7 @@ def render_git_describe(pieces): return rendered -def render_git_describe_long(pieces): +def render_git_describe_long(pieces: Dict[str, Any]) -> str: """TAG-DISTANCE-gHEX[-dirty]. Like 'git describe --tags --dirty --always -long'. @@ -575,7 +601,7 @@ def render_git_describe_long(pieces): return rendered -def render(pieces, style): +def render(pieces: Dict[str, Any], style: str) -> Dict[str, Any]: """Render the given version pieces into the requested style.""" if pieces["error"]: return {"version": "unknown", @@ -611,7 +637,7 @@ def render(pieces, style): "date": pieces.get("date")} -def get_versions(): +def get_versions() -> Dict[str, Any]: """Get version information or return default if unable to do so.""" # I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have # __file__, we can work backwards from there to the root. Some diff --git a/monai/apps/__init__.py b/monai/apps/__init__.py index 9cc7aeb8e05..a9f99599377 100644 --- a/monai/apps/__init__.py +++ b/monai/apps/__init__.py @@ -13,4 +13,13 @@ from .datasets import CrossValidation, DecathlonDataset, MedNISTDataset, TciaDataset from .mmars import MODEL_DESC, RemoteMMARKeys, download_mmar, get_model_spec, load_from_mmar -from .utils import SUPPORTED_HASH_TYPES, check_hash, download_and_extract, download_url, extractall, get_logger, logger +from .utils import ( + SUPPORTED_HASH_TYPES, + HashCheckError, + check_hash, + download_and_extract, + download_url, + extractall, + get_logger, + logger, +) diff --git a/monai/apps/auto3dseg/auto_runner.py b/monai/apps/auto3dseg/auto_runner.py index 8421893f514..37ccf7f19d6 100644 --- a/monai/apps/auto3dseg/auto_runner.py +++ b/monai/apps/auto3dseg/auto_runner.py @@ -790,6 +790,7 @@ def _train_algo_in_nni(self, history: list[dict[str, Any]]) -> None: nni_config_filename = os.path.abspath(os.path.join(self.work_dir, f"{name}_nni_config.yaml")) ConfigParser.export_config_file(nni_config, nni_config_filename, fmt="yaml", default_flow_style=None) + # pyrefly: ignore [redundant-cast] max_trial = min(self.hpo_tasks, cast(int, default_nni_config["maxTrialNumber"])) cmd = "nnictl create --config " + nni_config_filename + " --port 8088" @@ -805,6 +806,7 @@ def _train_algo_in_nni(self, history: list[dict[str, Any]]) -> None: n_trainings = len(import_bundle_algo_history(self.work_dir, only_trained=True)) cmd = "nnictl stop --all" + # pyrefly: ignore [bad-argument-type] run_cmd(cmd.split(), check=True) logger.info(f"NNI completes HPO on {name}") last_total_tasks = n_trainings diff --git a/monai/apps/deepedit/transforms.py b/monai/apps/deepedit/transforms.py index d2f89d2eea7..d15b2bec3cf 100644 --- a/monai/apps/deepedit/transforms.py +++ b/monai/apps/deepedit/transforms.py @@ -434,6 +434,7 @@ def _randomize(self, d, key_label): else: logger.info(f"Not slice IDs for label: {key_label}") sid = None + # pyrefly: ignore [unsupported-operation] self.sid[key_label] = sid def __call__(self, data: Mapping[Hashable, np.ndarray]) -> dict[Hashable, np.ndarray]: @@ -561,6 +562,7 @@ def __init__( self.guidance: dict[str, list[list[int]]] = {} def randomize(self, data=None): + # pyrefly: ignore [unsupported-operation] probability = data[self.probability] self._will_interact = self.R.choice([True, False], p=[probability, 1.0 - probability]) @@ -885,6 +887,7 @@ def _randomize(self, d, key_label): else: logger.info(f"Not slice IDs for label: {key_label}") sid = None + # pyrefly: ignore [unsupported-operation] self.sid[key_label] = sid def __call__(self, data: Mapping[Hashable, np.ndarray]) -> dict[Hashable, np.ndarray]: diff --git a/monai/apps/deepgrow/dataset.py b/monai/apps/deepgrow/dataset.py index e597188e745..0d1c11b1193 100644 --- a/monai/apps/deepgrow/dataset.py +++ b/monai/apps/deepgrow/dataset.py @@ -175,6 +175,7 @@ def _save_data_2d(vol_idx, vol_image, vol_label, dataset_dir, relative_path): continue # For all Labels + # pyrefly: ignore [missing-attribute] unique_labels = np.unique(label.flatten()) unique_labels = unique_labels[unique_labels != 0] unique_labels_count = max(unique_labels_count, len(unique_labels)) diff --git a/monai/apps/deepgrow/transforms.py b/monai/apps/deepgrow/transforms.py index d92a79a16a5..624eed342da 100644 --- a/monai/apps/deepgrow/transforms.py +++ b/monai/apps/deepgrow/transforms.py @@ -288,6 +288,7 @@ def __init__(self, guidance: str = "guidance", discrepancy: str = "discrepancy", self._will_interact = None def randomize(self, data=None): + # pyrefly: ignore [unsupported-operation] probability = data[self.probability] self._will_interact = self.R.choice([True, False], p=[probability, 1.0 - probability]) diff --git a/monai/apps/detection/networks/retinanet_detector.py b/monai/apps/detection/networks/retinanet_detector.py index 95b29b8285c..9b9bf269113 100644 --- a/monai/apps/detection/networks/retinanet_detector.py +++ b/monai/apps/detection/networks/retinanet_detector.py @@ -525,6 +525,7 @@ def forward( ) # 4. Generate anchors and store it in self.anchors: List[Tensor] + # pyrefly: ignore [bad-argument-type] self.generate_anchors(images, head_outputs) # num_anchor_locs_per_level: List[int], list of HW or HWD for each level num_anchor_locs_per_level = [x.shape[2:].numel() for x in head_outputs[self.cls_key]] @@ -535,6 +536,7 @@ def forward( # reshape to Tensor sized(B, sum(HWA), self.num_classes) for self.cls_key # or (B, sum(HWA), 2* self.spatial_dims) for self.box_reg_key # A = self.num_anchors_per_loc + # pyrefly: ignore [bad-argument-type] head_outputs[key] = self._reshape_maps(head_outputs[key]) # 6(1). If during training, return losses diff --git a/monai/apps/detection/transforms/array.py b/monai/apps/detection/transforms/array.py index 301a636b6cc..635506c08ab 100644 --- a/monai/apps/detection/transforms/array.py +++ b/monai/apps/detection/transforms/array.py @@ -257,10 +257,14 @@ def __call__(self, boxes: NdarrayTensor, src_spatial_size: Sequence[int] | int | diff = od - zd half = abs(diff) // 2 if diff > 0: # need padding (half, diff - half) + # pyrefly: ignore [bad-index, unsupported-operation] zoomed_boxes[:, axis] = zoomed_boxes[:, axis] + half + # pyrefly: ignore [bad-index, unsupported-operation] zoomed_boxes[:, axis + spatial_dims] = zoomed_boxes[:, axis + spatial_dims] + half elif diff < 0: # need slicing (half, half + od) + # pyrefly: ignore [bad-index, unsupported-operation] zoomed_boxes[:, axis] = zoomed_boxes[:, axis] - half + # pyrefly: ignore [bad-index, unsupported-operation] zoomed_boxes[:, axis + spatial_dims] = zoomed_boxes[:, axis + spatial_dims] - half return zoomed_boxes diff --git a/monai/apps/detection/transforms/box_ops.py b/monai/apps/detection/transforms/box_ops.py index fa714daad12..54bbc8fd31b 100644 --- a/monai/apps/detection/transforms/box_ops.py +++ b/monai/apps/detection/transforms/box_ops.py @@ -186,7 +186,9 @@ def flip_boxes( _flip_boxes: NdarrayTensor = boxes.clone() if isinstance(boxes, torch.Tensor) else deepcopy(boxes) # type: ignore[assignment] for axis in flip_axes: + # pyrefly: ignore [bad-index, unsupported-operation] _flip_boxes[:, axis + spatial_dims] = spatial_size[axis] - boxes[:, axis] - TO_REMOVE + # pyrefly: ignore [bad-index, unsupported-operation] _flip_boxes[:, axis] = spatial_size[axis] - boxes[:, axis + spatial_dims] - TO_REMOVE return _flip_boxes diff --git a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py index 39f84592246..90ffff6c340 100644 --- a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py +++ b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py @@ -246,7 +246,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # update padding length if necessary padding = 3 + # pyrefly: ignore [unsupported-operation] if padding % self.stride > 0: + # pyrefly: ignore [unsupported-operation] padding = (padding // self.stride + 1) * self.stride if self.print_info: logger.info(f"Padding size: {padding}") diff --git a/monai/apps/nnunet/nnunetv2_runner.py b/monai/apps/nnunet/nnunetv2_runner.py index 547e73332fe..c18e7bcd055 100644 --- a/monai/apps/nnunet/nnunetv2_runner.py +++ b/monai/apps/nnunet/nnunetv2_runner.py @@ -274,6 +274,7 @@ def convert_dataset(self): modality = [modality] create_new_dataset_json( + # pyrefly: ignore [bad-argument-type] modality=modality, num_foreground_classes=num_foreground_classes, num_input_channels=num_input_channels, @@ -598,7 +599,11 @@ def train_single_model_command( for _key, _value in kwargs.items(): prefix = "-" if _key in {"p", "pretrained_weights"} else "--" - cmd += [f"{prefix}{_key}", str(_value)] + if isinstance(_value, bool): + if _value: + cmd.append(f"{prefix}{_key}") + else: + cmd += [f"{prefix}{_key}", str(_value)] cmd_str: list[str] = [str(c) for c in cmd] @@ -758,7 +763,7 @@ def validate_single_model(self, config: str, fold: int, **kwargs: Any) -> None: kwargs: this optional parameter allows you to specify additional arguments defined in the ``train_single_model`` method. """ - self.train_single_model(config=config, fold=fold, only_run_validation=True, **kwargs) + self.train_single_model(config=config, fold=fold, val=True, **kwargs) def validate( self, configs: tuple = (M.N_3D_FULLRES, M.N_2D, M.N_3D_LOWRES, M.N_3D_CASCADE_FULLRES), **kwargs: Any diff --git a/monai/apps/nnunet/utils.py b/monai/apps/nnunet/utils.py index 1278eacd56b..c5102357f99 100644 --- a/monai/apps/nnunet/utils.py +++ b/monai/apps/nnunet/utils.py @@ -149,7 +149,6 @@ def create_new_dataset_json( """ new_json_data: dict = {} - # modality = self.input_info.pop("modality") modality = ensure_tuple(modality) # type: ignore new_json_data["channel_names"] = {} @@ -161,18 +160,11 @@ def create_new_dataset_json( for _j in range(num_foreground_classes): new_json_data["labels"][f"class{_j + 1}"] = _j + 1 - # new_json_data["numTraining"] = len(datalist_json["training"]) new_json_data["numTraining"] = num_training_data new_json_data["file_ending"] = ".nii.gz" ConfigParser.export_config_file( - config=new_json_data, - # filepath=os.path.join(raw_data_foldername, "dataset.json"), - filepath=output_filepath, - fmt="json", - sort_keys=True, - indent=4, - ensure_ascii=False, + config=new_json_data, filepath=output_filepath, fmt="json", sort_keys=True, indent=4, ensure_ascii=False ) return diff --git a/monai/apps/pathology/transforms/post/dictionary.py b/monai/apps/pathology/transforms/post/dictionary.py index 1e2540daee0..b1dde44ec1f 100644 --- a/monai/apps/pathology/transforms/post/dictionary.py +++ b/monai/apps/pathology/transforms/post/dictionary.py @@ -403,6 +403,7 @@ def __call__(self, data): d = dict(data) for key in self.key_iterator(d): offset = d[self.offset_key] if self.offset_key else None + # pyrefly: ignore [bad-argument-type] centroid = self.converter(d[key], offset) key_to_add = f"{key}_{self.centroid_key_postfix}" if key_to_add in d: diff --git a/monai/apps/utils.py b/monai/apps/utils.py index 856bc64c9ea..57cc7b18a4a 100644 --- a/monai/apps/utils.py +++ b/monai/apps/utils.py @@ -42,12 +42,24 @@ else: tqdm, has_tqdm = optional_import("tqdm", "4.47.0", min_version, "tqdm") -__all__ = ["check_hash", "download_url", "extractall", "download_and_extract", "get_logger", "SUPPORTED_HASH_TYPES"] +__all__ = [ + "HashCheckError", + "check_hash", + "download_url", + "extractall", + "download_and_extract", + "get_logger", + "SUPPORTED_HASH_TYPES", +] DEFAULT_FMT = "%(asctime)s - %(levelname)s - %(message)s" SUPPORTED_HASH_TYPES = {"md5": hashlib.md5, "sha1": hashlib.sha1, "sha256": hashlib.sha256, "sha512": hashlib.sha512} +class HashCheckError(ValueError): + pass + + def get_logger( module_name: str = "monai.apps", fmt: str = DEFAULT_FMT, @@ -154,20 +166,20 @@ def safe_extract_member(member, extract_to): return full_path -def check_hash(filepath: PathLike, val: str | None = None, hash_type: str = "md5") -> bool: +def check_hash(filepath: PathLike, val: str | None = None, hash_type: str = "sha256") -> bool: """ Verify hash signature of specified file. Args: filepath: path of source file to verify hash value. val: expected hash value of the file. - hash_type: type of hash algorithm to use, default is `"md5"`. + hash_type: type of hash algorithm to use, default is `"sha256"`. The supported hash types are `"md5"`, `"sha1"`, `"sha256"`, `"sha512"`. See also: :py:data:`monai.apps.utils.SUPPORTED_HASH_TYPES`. """ if val is None: - logger.info(f"Expected {hash_type} is None, skip {hash_type} check for file {filepath}.") + warnings.warn(f"No hash value provided for {filepath}; file integrity is NOT verified.", stacklevel=2) return True actual_hash_func = look_up_option(hash_type.lower(), SUPPORTED_HASH_TYPES) @@ -192,7 +204,7 @@ def download_url( url: str, filepath: PathLike = "", hash_val: str | None = None, - hash_type: str = "md5", + hash_type: str = "sha256", progress: bool = True, **gdown_kwargs: Any, ) -> None: @@ -205,7 +217,8 @@ def download_url( If undefined, `os.path.basename(url)` will be used. hash_val: expected hash value to validate the downloaded file. if None, skip hash validation. - hash_type: 'md5' or 'sha1', defaults to 'md5'. + hash_type: type of hash algorithm to use, default is `"sha256"`. + The supported hash types are `"md5"`, `"sha1"`, `"sha256"`, `"sha512"`. progress: whether to display a progress bar. gdown_kwargs: other args for `gdown` except for the `url`, `output` and `quiet`. these args will only be used if download from google drive. @@ -220,8 +233,7 @@ def download_url( HTTPError: See urllib.request.urlretrieve. ContentTooShortError: See urllib.request.urlretrieve. IOError: See urllib.request.urlretrieve. - RuntimeError: When the hash validation of the ``url`` downloaded file fails. - + HashCheckError: When the hash validation of the ``url`` downloaded file fails. """ if not filepath: filepath = Path(".", _basename(url)).resolve() @@ -229,9 +241,7 @@ def download_url( filepath = Path(filepath) if filepath.exists(): if not check_hash(filepath, hash_val, hash_type): - raise RuntimeError( - f"{hash_type} check of existing file failed: filepath={filepath}, expected {hash_type}={hash_val}." - ) + raise HashCheckError(f"{hash_type} hash check of existing file failed: {filepath=}, expected {hash_type=}.") logger.info(f"File exists: {filepath}, skipped downloading.") return try: @@ -240,7 +250,7 @@ def download_url( if urlparse(url).netloc == "drive.google.com": if not has_gdown: raise RuntimeError("To download files from Google Drive, please install the gdown dependency.") - if "fuzzy" not in gdown_kwargs: + if "fuzzy" not in gdown_kwargs and not min_version(gdown, "6.0.0"): # "fuzzy" dropped in gdown 6.0.0 gdown_kwargs["fuzzy"] = True # default to true for flexible url gdown.download(url, f"{tmp_name}", quiet=not progress, **gdown_kwargs) elif urlparse(url).netloc == "cloud-api.yandex.net": @@ -260,6 +270,13 @@ def download_url( raise RuntimeError( f"Download of file from {url} to {filepath} failed due to network issue or denied permission." ) + if not check_hash(tmp_name, hash_val, hash_type): + raise HashCheckError( + f"{hash_type} hash check of downloaded file failed: {url=}, " + f"{filepath=}, expected {hash_type}={hash_val}, " + f"The file may be corrupted or tampered with. " + "Please retry the download or verify the source." + ) file_dir = filepath.parent if file_dir: os.makedirs(file_dir, exist_ok=True) @@ -267,11 +284,6 @@ def download_url( except (PermissionError, NotADirectoryError): # project-monai/monai issue #3613 #3757 for windows pass logger.info(f"Downloaded: {filepath}") - if not check_hash(filepath, hash_val, hash_type): - raise RuntimeError( - f"{hash_type} check of downloaded file failed: URL={url}, " - f"filepath={filepath}, expected {hash_type}={hash_val}." - ) def _extract_zip(filepath, output_dir): @@ -304,7 +316,7 @@ def extractall( filepath: PathLike, output_dir: PathLike = ".", hash_val: str | None = None, - hash_type: str = "md5", + hash_type: str = "sha256", file_type: str = "", has_base: bool = True, ) -> None: @@ -317,7 +329,7 @@ def extractall( output_dir: target directory to save extracted files. hash_val: expected hash value to validate the compressed file. if None, skip hash validation. - hash_type: 'md5' or 'sha1', defaults to 'md5'. + hash_type: type of hash algorithm to use, default is `"sha256"`. file_type: string of file type for decompressing. Leave it empty to infer the type from the filepath basename. has_base: whether the extracted files have a base folder. This flag is used when checking if the existing folder is a result of `extractall`, if it is, the extraction is skipped. For example, if A.zip is unzipped @@ -325,10 +337,15 @@ def extractall( be False. Raises: - RuntimeError: When the hash validation of the ``filepath`` compressed file fails. + HashCheckError: When the hash validation of the ``filepath`` compressed file fails. NotImplementedError: When the ``filepath`` file extension is not one of [zip", "tar.gz", "tar"]. """ + filepath = Path(filepath) + if hash_val and not check_hash(filepath, hash_val, hash_type): + raise HashCheckError( + f"{hash_type} hash check of compressed file failed: " f"{filepath=}, expected {hash_type}={hash_val}." + ) if has_base: # the extracted files will be in this folder cache_dir = Path(output_dir, _basename(filepath).split(".")[0]) @@ -337,11 +354,6 @@ def extractall( if cache_dir.exists() and next(cache_dir.iterdir(), None) is not None: logger.info(f"Non-empty folder exists in {cache_dir}, skipped extracting.") return - filepath = Path(filepath) - if hash_val and not check_hash(filepath, hash_val, hash_type): - raise RuntimeError( - f"{hash_type} check of compressed file failed: " f"filepath={filepath}, expected {hash_type}={hash_val}." - ) logger.info(f"Writing into directory: {output_dir}.") _file_type = file_type.lower().strip() if filepath.name.endswith("zip") or _file_type == "zip": @@ -383,7 +395,7 @@ def download_and_extract( filepath: PathLike = "", output_dir: PathLike = ".", hash_val: str | None = None, - hash_type: str = "md5", + hash_type: str = "sha256", file_type: str = "", has_base: bool = True, progress: bool = True, @@ -399,7 +411,7 @@ def download_and_extract( default is the current directory. hash_val: expected hash value to validate the downloaded file. if None, skip hash validation. - hash_type: 'md5' or 'sha1', defaults to 'md5'. + hash_type: type of hash algorithm to use, default is `"sha256"`. file_type: string of file type for decompressing. Leave it empty to infer the type from url's base file name. has_base: whether the extracted files have a base folder. This flag is used when checking if the existing folder is a result of `extractall`, if it is, the extraction is skipped. For example, if A.zip is unzipped diff --git a/monai/apps/vista3d/transforms.py b/monai/apps/vista3d/transforms.py index bd7fb194934..7860e3db406 100644 --- a/monai/apps/vista3d/transforms.py +++ b/monai/apps/vista3d/transforms.py @@ -160,8 +160,6 @@ def __call__(self, data): pred = pred.argmax(0).unsqueeze(0).float() + 1.0 pred[is_bk] = 0.0 else: - # AsDiscrete will remove NaN - # pred = monai.transforms.AsDiscrete(threshold=0.5)(pred) pred[pred > 0] = 1.0 if "label_prompt" in data and data["label_prompt"] is not None: pred += 0.5 # inplace mapping to avoid cloning pred diff --git a/monai/auto3dseg/operations.py b/monai/auto3dseg/operations.py index 404a6d326ee..7b64a04a6cd 100644 --- a/monai/auto3dseg/operations.py +++ b/monai/auto3dseg/operations.py @@ -149,4 +149,5 @@ def evaluate(self, data: Any, **kwargs: Any) -> dict: Args: data: input data """ + # pyrefly: ignore [missing-attribute] return {k: v(data[k], **kwargs).tolist() for k, v in self.data.items() if (callable(v) and k in data)} diff --git a/monai/auto3dseg/seg_summarizer.py b/monai/auto3dseg/seg_summarizer.py index 14a10635df2..8fdd9652455 100644 --- a/monai/auto3dseg/seg_summarizer.py +++ b/monai/auto3dseg/seg_summarizer.py @@ -208,6 +208,7 @@ def summarize(self, data: list[dict]) -> dict[str, dict]: for analyzer in self.summary_analyzers: if callable(analyzer): + # pyrefly: ignore [missing-attribute] report.update({analyzer.stats_name: analyzer(data)}) return report diff --git a/monai/auto3dseg/utils.py b/monai/auto3dseg/utils.py index d6fb561242a..518c91919da 100644 --- a/monai/auto3dseg/utils.py +++ b/monai/auto3dseg/utils.py @@ -284,7 +284,7 @@ def verify_report_format(report: dict, report_format: dict) -> bool: if isinstance(v_fmt, list) and isinstance(v, list): if len(v_fmt) != 1: - raise UserWarning("list length in report_format is not 1") + warnings.warn("list length in report_format is not 1", stacklevel=2) if len(v_fmt) > 0 and len(v) > 0: return verify_report_format(v[0], v_fmt[0]) else: @@ -493,6 +493,14 @@ def algo_from_json(filename: str, template_path: PathLike | None = None, **kwarg if state_template_path: algo_config["template_path"] = state_template_path + warnings.warn( + f"Loading {filename}: the file's `_target_` value is resolved to an imported callable and " + "invoked, and template directories from the file may be added to `sys.path`; only load " + "algo_object.json files from a source you trust " + "(see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-2wx3-8x3w-r8qv).", + stacklevel=2, + ) + parser = ConfigParser(algo_config) algo = parser.get_parsed_content() used_template_path = path diff --git a/monai/bundle/config_parser.py b/monai/bundle/config_parser.py index 08c6fc296c3..0e1ea4e77a3 100644 --- a/monai/bundle/config_parser.py +++ b/monai/bundle/config_parser.py @@ -162,7 +162,14 @@ def __getattr__(self, key: str) -> Any: try: return self._chain(key) except KeyError: - return getattr(self._value, key) + pass + if isinstance(self._value, dict) and key in self._value: + # the chained id is absent from the resolver (for example when this proxy is + # backed by a `$@ref`, whose children have no ids of their own), but the key + # does exist in the container: resolve it like `__getitem__` does, so dot- and + # bracket-notation agree and config keys keep precedence over dict methods. + return self._value[key] + return getattr(self._value, key) def __getitem__(self, key: str | int) -> Any: try: diff --git a/monai/bundle/reference_resolver.py b/monai/bundle/reference_resolver.py index b55c62174b7..27f34ebd5eb 100644 --- a/monai/bundle/reference_resolver.py +++ b/monai/bundle/reference_resolver.py @@ -254,6 +254,7 @@ def iter_subconfigs(cls, id: str, config: Any) -> Iterator[tuple[str, str, Any]] """ for k, v in config.items() if isinstance(config, dict) else enumerate(config): sub_id = f"{id}{cls.sep}{k}" if id != "" else f"{k}" + # pyrefly: ignore [invalid-yield] yield k, sub_id, v @classmethod diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index ab02cd552e3..b2809191281 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -25,6 +25,7 @@ from textwrap import dedent from typing import Any +import numpy as np import torch from torch.cuda import is_available @@ -51,13 +52,13 @@ min_version, optional_import, pprint_edges, + safe_eval, ) validate, _ = optional_import("jsonschema", name="validate") ValidationError, _ = optional_import("jsonschema.exceptions", name="ValidationError") Checkpoint, has_ignite = optional_import("ignite.handlers", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Checkpoint") requests, has_requests = optional_import("requests") -onnx, _ = optional_import("onnx") huggingface_hub, _ = optional_import("huggingface_hub") logger = get_logger(module_name=__name__) @@ -158,10 +159,12 @@ def _get_fake_spatial_shape(shape: Sequence[str | int], p: int = 1, n: int = 1, if i == "*": ret.append(any) else: - for c in _get_var_names(i): - if c not in ["p", "n"]: - raise ValueError(f"only support variables 'p' and 'n' so far, but got: {c}.") - ret.append(eval(i, {"p": p, "n": n})) + bad_names = set(c for c in _get_var_names(i) if c not in {"p", "n"}) + if bad_names: + raise ValueError(f"Only variables `p` and `n` currently supported. Invalid names: {bad_names}") + + # evaluate using Numpy types to prevent slow Python DoS attacks + ret.append(int(safe_eval(i, {"p": np.int32(p), "n": np.int32(n)}, rewrite_np=True))) else: raise ValueError(f"spatial shape items must be int or string, but got: {type(i)} {i}.") return tuple(ret) @@ -648,6 +651,14 @@ def load( """ Load model weights or TorchScript module of a bundle. + Security note: if `model` is `None`, building `network_def` requires parsing the bundle's own + "{workflow_type}.json" config, which can define `"_target_"` components resolved to any importable + callable and `"$"`-prefixed expressions evaluated with Python `eval()`. Only call `load()` this way + for bundles from a source you trust; a warning is printed every time this happens + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-873f-pvrv-4x83). To skip parsing + the bundle's config entirely, pass an explicit `model=` — only the weights are then loaded, via + `torch.load(..., weights_only=True)`. + Args: name: bundle name. If `None` and `url` is `None`, it must be provided in `args_file`. for example: @@ -935,6 +946,12 @@ def run( """ Specify `config_file` to run monai bundle components and workflows. + Security note: parsing `config_file` can run arbitrary code. Any `"_target_"` value is resolved to an + importable callable and invoked with no allow list, and any `"$"`-prefixed value is passed to Python + `eval()`. Only point this at config files you wrote or otherwise fully trust; never at a config + downloaded from, or otherwise sourced from, an untrusted party. A warning is printed every time this + happens (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-873f-pvrv-4x83). + Typical usage examples: .. code-block:: bash @@ -980,7 +997,8 @@ def run( common parameters shown below will be added and can be passed through the `override` parameter of this method. - ``"output_dir"``: the path to save mlflow tracking outputs locally, default to "/eval". - - ``"tracking_uri"``: uri to save mlflow tracking outputs, default to "/output_dir/mlruns". + - ``"tracking_uri"``: uri to save mlflow tracking outputs, default to a local SQLite database + at "/mlruns.db" with run artifacts kept under "/mlruns". - ``"experiment_name"``: experiment name for this run, default to "monai_experiment". - ``"run_name"``: the name of current run. - ``"save_execute_config"``: whether to save the executed config files. It can be `False`, `/path/to/artifacts` @@ -1419,6 +1437,7 @@ def onnx_export( converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_}) def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str, **kwargs: Any) -> None: + onnx, _ = optional_import("onnx") onnx.save(onnx_obj, filename_prefix_or_stream) _export( @@ -1929,6 +1948,12 @@ def create_workflow( The workflow should be subclass of `BundleWorkflow` and be available to import. It can be MONAI existing bundle workflows or user customized workflows. + Security note: parsing `config_file` can run arbitrary code. Any `"_target_"` value is resolved to an + importable callable and invoked with no allow list, and any `"$"`-prefixed value is passed to Python + `eval()`. Only point this at config files you wrote or otherwise fully trust; never at a config + downloaded from, or otherwise sourced from, an untrusted party. A warning is printed every time this + happens (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-873f-pvrv-4x83). + Typical usage examples: .. code-block:: python @@ -1966,6 +1991,14 @@ def create_workflow( ) if config_file is not None: + warnings.warn( + f'parsing config_file {config_file}: any `"_target_"` value in it is resolved to an importable ' + 'callable and invoked with no allow list, and any `"$"`-prefixed value is passed to Python ' + "`eval()`. Only proceed if this config is from a source you trust " + "(see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-873f-pvrv-4x83).", + stacklevel=2, + ) + # pyrefly: ignore [unexpected-keyword] workflow_ = workflow_class(config_file=config_file, **_args) else: workflow_ = workflow_class(**_args) @@ -1975,6 +2008,17 @@ def create_workflow( return workflow_ +def _safe_large_file_path(bundle_path: PathLike, filepath: str) -> str: + """Securely resolve a large-file target path to prevent traversal outside the bundle directory.""" + bundle_root = os.path.realpath(bundle_path) + target = os.path.normpath(os.path.join(bundle_path, filepath)) + target_real = os.path.realpath(target) + # Ensure the resolved path stays within the bundle root + if os.path.commonpath([bundle_root, target_real]) != bundle_root: + raise ValueError(f"Unsafe path: path traversal {filepath} for bundle_path {bundle_path}") + return target + + def download_large_files(bundle_path: str | None = None, large_file_name: str | None = None) -> None: """ This utility allows you to download large files from a bundle. It supports file suffixes like ".yml", ".yaml", and ".json". @@ -2005,11 +2049,10 @@ def download_large_files(bundle_path: str | None = None, large_file_name: str | parser.read_config(large_file_path) large_files_list = parser.get()["large_files"] for lf_data in large_files_list: - lf_data["fuzzy"] = True if "hash_val" in lf_data and lf_data.get("hash_val", "") == "": lf_data.pop("hash_val") if "hash_type" in lf_data and lf_data.get("hash_type", "") == "": lf_data.pop("hash_type") - lf_data["filepath"] = os.path.join(bundle_path, lf_data["path"]) + lf_data["filepath"] = _safe_large_file_path(bundle_path, lf_data["path"]) lf_data.pop("path") download_url(**lf_data) diff --git a/monai/bundle/utils.py b/monai/bundle/utils.py index d37d7f1c05c..ebd521a93f1 100644 --- a/monai/bundle/utils.py +++ b/monai/bundle/utils.py @@ -21,7 +21,6 @@ from monai.utils import optional_import yaml, _ = optional_import("yaml") - __all__ = [ "ID_REF_KEY", "ID_SEP_KEY", @@ -39,7 +38,6 @@ MERGE_KEY = "+" # prefix indicating merge instead of override in case of multiple configs. _conf_values = get_config_values() - DEFAULT_METADATA = { "version": "0.0.1", "changelog": {"0.0.1": "Initial version"}, @@ -118,8 +116,10 @@ "configs": { # if no "output_dir" in the bundle config, default to "/eval" "output_dir": "$@bundle_root + '/eval'", - # use URI to support linux, mac and windows os - "tracking_uri": "$monai.utils.path_to_uri(@output_dir) + '/mlruns'", + # MLflow 3.13+ rejects the filesystem (file store) tracking backend, so default tracking + # to a local SQLite database. The handler keeps run artifacts under "/mlruns" + # (next to the db). A URI is used so the path is valid on linux, mac and windows os. + "tracking_uri": "$monai.utils.path_to_sqlite_uri(@output_dir + '/mlruns.db')", "experiment_name": "monai_experiment", "run_name": None, # may fill it at runtime @@ -211,20 +211,15 @@ def load_bundle_config(bundle_path: str, *config_names: str, **load_kw_args: Any name, _ = os.path.splitext(os.path.basename(bundle_path)) archive = zipfile.ZipFile(bundle_path, "r") - all_files = archive.namelist() - zip_meta_name = f"{name}/configs/metadata.json" - if zip_meta_name in all_files: prefix = f"{name}/configs/" # zipped directory location for files else: zip_meta_name = f"{name}/extra/metadata.json" prefix = f"{name}/extra/" # Torchscript location for files - meta_json = json.loads(archive.read(zip_meta_name)) parser.read_meta(f=meta_json) - for cname in config_names: full_cname = prefix + cname if full_cname not in all_files: diff --git a/monai/bundle/workflows.py b/monai/bundle/workflows.py index 5b95441d51f..3d8637cb343 100644 --- a/monai/bundle/workflows.py +++ b/monai/bundle/workflows.py @@ -15,6 +15,7 @@ import os import sys import time +import warnings from abc import ABC, abstractmethod from collections.abc import Sequence from copy import copy @@ -34,6 +35,23 @@ logger = get_logger(module_name=__name__) +def _warn_logging_file_execution(logging_file: str) -> None: + """ + Warn that ``logging_file`` is about to be executed by `logging.config.fileConfig`. + + Called immediately before every `fileConfig` invocation in this module, so the warning is only + raised when the file is really executed -- not when it is missing or logging is disabled. + """ + warnings.warn( + f"applying logging config {logging_file}: `logging.config.fileConfig` passes the `class=` and " + "`args=` fields of the INI's handler and formatter sections to Python `eval()`, so this file " + "runs as code. A bundle ships its own `configs/logging.conf` and it is applied by default, " + "before any of the bundle's config is parsed. Only proceed if this file is from a source you " + "trust (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-wvpx-5qmp-46g3).", + stacklevel=3, + ) + + class BundleWorkflow(ABC): """ Base class for the workflow specification in bundle, it can be a training, evaluation or inference workflow. @@ -55,6 +73,10 @@ class BundleWorkflow(ABC): meta_file: filepath of the metadata file, if this is a list of file paths, their contents will be merged in order. logging_file: config file for `logging` module in the program. for more details: https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig. + Security note: `fileConfig` passes the INI's `class=` and `args=` fields to Python + `eval()`, so this file runs as code and applying it raises a warning -- once per call + site, as Python's default warning filter suppresses repeats + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-wvpx-5qmp-46g3). """ @@ -72,6 +94,7 @@ def __init__( if not os.path.isfile(logging_file): raise FileNotFoundError(f"Cannot find the logging config file: {logging_file}.") logger.info(f"Setting logging properties based on config: {logging_file}.") + _warn_logging_file_execution(logging_file) fileConfig(logging_file, disable_existing_loggers=False) if meta_file is not None: @@ -224,6 +247,7 @@ def add_property(self, name: str, required: str, desc: str | None = None) -> Non desc: descriptions for the property. """ if self.properties is None: + # pyrefly: ignore [bad-assignment] self.properties = {} if name in self.properties: logger.warning(f"property '{name}' already exists in the properties list, overriding it.") @@ -272,6 +296,10 @@ class PythonicWorkflow(BundleWorkflow): meta_file: filepath of the metadata file, if this is a list of file paths, their contents will be merged in order. logging_file: config file for `logging` module in the program. for more details: https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig. + Security note: `fileConfig` passes the INI's `class=` and `args=` fields to Python + `eval()`, so this file runs as code and applying it raises a warning -- once per call + site, as Python's default warning filter suppresses repeats + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-wvpx-5qmp-46g3). """ @@ -329,6 +357,7 @@ def _get_property(self, name: str, property: dict) -> Any: elif name in self._props_vals: value = self._props_vals[name] elif name in self.parser.config[self.parser.meta_key]: # type: ignore[index] + # pyrefly: ignore [missing-attribute] id = self.properties.get(name, None).get(BundlePropertyConfig.ID, None) value = self.parser[id] else: @@ -373,6 +402,10 @@ class ConfigWorkflow(BundleWorkflow): https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig. If None, default to "configs/logging.conf", which is commonly used for bundles in MONAI model zoo. If False, the logging logic for the bundle will not be modified. + Security note: `fileConfig` passes the INI's `class=` and `args=` fields to Python + `eval()`, so this file runs as code and applying it raises a warning -- once per call + site, as Python's default warning filter suppresses repeats + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-wvpx-5qmp-46g3). init_id: ID name of the expected config expression to initialize before running, default to "initialize". allow a config to have no `initialize` logic and the ID. run_id: ID name of the expected config expression to run, default to "run". @@ -442,6 +475,7 @@ def __init__( else: raise FileNotFoundError(f"Cannot find the logging config file: {logging_file}.") else: + _warn_logging_file_execution(str(logging_file)) fileConfig(str(logging_file), disable_existing_loggers=False) logger.info(f"Setting logging properties based on config: {logging_file}.") @@ -621,6 +655,7 @@ def _check_optional_id(self, name: str, property: dict) -> bool: else: ref = self.parser.get(ref_id, None) # for reference IDs that not refer to a property directly but using expressions, skip the check + # pyrefly: ignore [unsupported-operation] if ref is not None and not ref.startswith(EXPR_KEY) and ref != ID_REF_KEY + id: return False return True diff --git a/monai/config/print_dependencies.py b/monai/config/print_dependencies.py new file mode 100644 index 00000000000..a099949ecac --- /dev/null +++ b/monai/config/print_dependencies.py @@ -0,0 +1,85 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This program prints the MONAI dependencies for the optional names given on the command line. The printed values can +be piped to a requirements file to work with pip. All required dependencies are always printed, those for builing are +included in "build-system" is given as an argument, and all optional requirements are included if "*" is given. This +assumes the pyproject.toml file is in the current working directory. +""" + +from __future__ import annotations + +import sys +from collections.abc import Collection + +BUILD_SYSTEM_KEY = "build-system" +PROJ_KEY = "project" +OPTS_KEY = "optional-dependencies" +DEP_KEY = "dependencies" +REQ_KEY = "requires" +TOML_FILE = "pyproject.toml" + + +def parse_dependencies(filename: str | None = None, sections: Collection[str] | None = None) -> list[str]: + """ + Parse the toml file given by `filename` and return the dependency sections selected by `sections`. + + Args: + filename: TOML file to parse, if None this defaults to TOML_FILE. + sections: "optional-dependencies" sections to print in addition to the required dependencies. If + "build-system" is included, the build requirements will be included in the output. If "*" is included, all + of the optional dependencies will be included in the output. + + Returns: + List of requirements in alphabetical order. + """ + # these imports should be here to avoid attempting to import when MONAI is imported and both packages are missing + # isort: off + if sys.version_info.minor >= 11: + from tomllib import loads + else: + from tomli import loads + # isort: on + + with open(filename or TOML_FILE) as o: + data = loads(o.read()) + + proj = data[PROJ_KEY] + opts = proj[OPTS_KEY] + dependencies = list(proj[DEP_KEY]) + sections = set(sections or []) + + if BUILD_SYSTEM_KEY in sections: + sections.remove(BUILD_SYSTEM_KEY) + dependencies += data[BUILD_SYSTEM_KEY][REQ_KEY] + + if "*" in sections: + dependencies += sum(opts.values(), []) + else: + for s in sections: + dependencies += opts[s] + + return sorted(set(dependencies)) + + +def print_dependencies_argv(): + """ + Print dependencies specified through argv. + """ + dependencies = parse_dependencies(sections=set(sys.argv[1:])) + + for d in dependencies: + print(d) + + +if __name__ == "__main__": + print_dependencies_argv() diff --git a/monai/data/__init__.py b/monai/data/__init__.py index 971d5121f71..ef04160425a 100644 --- a/monai/data/__init__.py +++ b/monai/data/__init__.py @@ -71,7 +71,7 @@ monai_to_itk_ddf, ) from .meta_obj import MetaObj, get_track_meta, set_track_meta -from .meta_tensor import MetaTensor +from .meta_tensor import MetaTensor, get_spatial_ndim from .samplers import DistributedSampler, DistributedWeightedRandomSampler from .synthetic import create_test_image_2d, create_test_image_3d from .test_time_augmentation import TestTimeAugmentation diff --git a/monai/data/box_utils.py b/monai/data/box_utils.py index b0c41b5d7d4..b344434715a 100644 --- a/monai/data/box_utils.py +++ b/monai/data/box_utils.py @@ -1035,8 +1035,8 @@ def spatial_crop_boxes( # convert to float32 since torch.clamp_ does not support float16 boxes_t = boxes_t.to(dtype=COMPUTE_DTYPE) - roi_start_t = convert_to_dst_type(src=roi_start, dst=boxes_t, wrap_sequence=True)[0].to(torch.int16) - roi_end_t = convert_to_dst_type(src=roi_end, dst=boxes_t, wrap_sequence=True)[0].to(torch.int16) + roi_start_t = convert_to_dst_type(src=roi_start, dst=boxes_t, wrap_sequence=True)[0] + roi_end_t = convert_to_dst_type(src=roi_end, dst=boxes_t, wrap_sequence=True)[0] roi_end_t = torch.maximum(roi_end_t, roi_start_t) # makes sure the bounding boxes are within the patch @@ -1200,7 +1200,7 @@ def batched_nms( # from different classes do not overlap max_coordinate = boxes_t.max() offsets = labels_t.to(boxes_t) * (max_coordinate + 1) - boxes_for_nms = boxes + offsets[:, None] + boxes_for_nms = boxes_t + offsets[:, None] keep = non_max_suppression(boxes_for_nms, scores_t, nms_thresh, max_proposals, box_overlap_metric) # convert tensor back to numpy if needed diff --git a/monai/data/dataset.py b/monai/data/dataset.py index 2511ce22191..62aced937dd 100644 --- a/monai/data/dataset.py +++ b/monai/data/dataset.py @@ -210,8 +210,12 @@ class PersistentDataset(Dataset): Cached data is expected to be tensors, primitives, or dictionaries keying to these values. Numpy arrays will be converted to tensors, however any other object type returned by transforms will not be loadable since - `torch.load` will be used with `weights_only=True` to prevent loading of potentially malicious objects. - Legacy cache files may not be loadable and may need to be recomputed. + `torch.load` will be used with `weights_only=True` by default to prevent loading of potentially malicious + objects. Legacy cache files may not be loadable and may need to be recomputed. MetaTensor objects can be saved + and loaded with their metadata preserved if `track_meta` is True, however the objects stored in the metadata + must be acceptable as serialisable by `torch.load` by default or if they have been white-listed with + `torch.serialization.add_safe_globals`. Any other object type may be stored but will fail to load and force + a cache recompute. Lazy Resampling: If you make use of the lazy resampling feature of `monai.transforms.Compose`, please refer to @@ -245,8 +249,8 @@ def __init__( may share a common cache dir provided that the transforms pre-processing is consistent. If `cache_dir` doesn't exist, will automatically create it. If `cache_dir` is `None`, there is effectively no caching. - hash_func: a callable to compute hash from data items to be cached. - defaults to `monai.data.utils.pickle_hashing`. + hash_func: a callable to compute hash from data items to be cached, defaults to + `monai.data.utils.pickle_hashing` which uses sha256 (previously md5 so old caches will not work). pickle_module: string representing the module used for pickling metadata and objects, default to `"pickle"`. due to the pickle limitation in multi-processing of Dataloader, we can't use `pickle` as arg directly, so here we use a string name instead. @@ -266,17 +270,12 @@ def __init__( When this is enabled, the traced transform instance IDs will be removed from the cached MetaTensors. This is useful for skipping the transform instance checks when inverting applied operations using the cached content and with re-created transform instances. - track_meta: whether to track the meta information, if `True`, will convert to `MetaTensor`. - default to `False`. Cannot be used with `weights_only=True`. + track_meta: whether to track the meta information, defaults to False. If `True`, converts to `MetaTensor`. weights_only: keyword argument passed to `torch.load` when reading cached files. - default to `True`. When set to `True`, `torch.load` restricts loading to tensors and - other safe objects. Setting this to `False` is required for loading `MetaTensor` - objects saved with `track_meta=True`, however this creates the possibility of remote - code execution through `torch.load` so be aware of the security implications of doing so. - - Raises: - ValueError: When both `track_meta=True` and `weights_only=True`, since this combination - prevents cached MetaTensors from being reloaded and causes perpetual cache regeneration. + default to `True`. When `True`, `torch.load` restricts loading to tensors and other safe objects. + Setting to `False` should only be done if it's absolutely necessary to load unsafe pickled data, + eg. MetaTensor objects with unsafe objects in their metadata. Users must verify the safety of the data + they intend to load before doing so. """ super().__init__(data=data, transform=transform) self.cache_dir = Path(cache_dir) if cache_dir is not None else None @@ -292,11 +291,6 @@ def __init__( if hash_transform is not None: self.set_transform_hash(hash_transform) self.reset_ops_id = reset_ops_id - if track_meta and weights_only: - raise ValueError( - "Invalid argument combination: `track_meta=True` cannot be used with `weights_only=True`. " - "To cache and reload MetaTensors, set `track_meta=True` and `weights_only=False`." - ) self.track_meta = track_meta self.weights_only = weights_only @@ -390,9 +384,9 @@ def _cachecheck(self, item_transformed): """ hashfile = None if self.cache_dir is not None: - data_item_md5 = self.hash_func(item_transformed).decode("utf-8") - data_item_md5 += self.transform_hash - hashfile = self.cache_dir / f"{data_item_md5}.pt" + data_item_hash = self.hash_func(item_transformed).decode("utf-8") + data_item_hash += self.transform_hash + hashfile = self.cache_dir / f"{data_item_hash}.pt" if hashfile is not None and hashfile.is_file(): # cache hit try: @@ -1624,9 +1618,9 @@ def _cachecheck(self, item_transformed): hashfile = None # compute a cache id if self.cache_dir is not None: - data_item_md5 = self.hash_func(item_transformed).decode("utf-8") - data_item_md5 += self.transform_hash - hashfile = self.cache_dir / f"{data_item_md5}.pt" + data_item_hash = self.hash_func(item_transformed).decode("utf-8") + data_item_hash += self.transform_hash + hashfile = self.cache_dir / f"{data_item_hash}.pt" if hashfile is not None and hashfile.is_file(): # cache hit with cp.cuda.Device(self.device): @@ -1654,6 +1648,7 @@ def _cachecheck(self, item_transformed): item_k = kvikio_numpy.fromfile( f"{hashfile}-{k}-{i}", dtype=meta_i_k["dtype"], like=cp.empty(()) ) + # pyrefly: ignore [missing-attribute] item_k = convert_to_tensor(item[i].reshape(meta_i_k["shape"]), device=f"cuda:{self.device}") item[i].update({k: item_k, f"{k}_meta_dict": meta_i_k}) return item diff --git a/monai/data/grid_dataset.py b/monai/data/grid_dataset.py index 689138179ae..a4860cfcad5 100644 --- a/monai/data/grid_dataset.py +++ b/monai/data/grid_dataset.py @@ -142,6 +142,7 @@ def __call__( self, data: Mapping[Hashable, NdarrayTensor] ) -> Generator[tuple[Mapping[Hashable, NdarrayTensor], np.ndarray], None, None]: d = dict(data) + # pyrefly: ignore [missing-attribute] original_spatial_shape = d[first(self.keys)].shape[1:] for patch in zip(*[self.patch_iter(d[key]) for key in self.keys]): diff --git a/monai/data/image_reader.py b/monai/data/image_reader.py index a85eb95c208..53fcaa95451 100644 --- a/monai/data/image_reader.py +++ b/monai/data/image_reader.py @@ -22,7 +22,7 @@ from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any, TypeAlias +from typing import TYPE_CHECKING, Any, TypeAlias # pyrefly: ignore [missing-module-attribute] import numpy as np from torch.utils.data._utils.collate import np_str_obj_array_pattern @@ -347,6 +347,7 @@ def _get_affine(self, img, lps_to_ras: bool = True): affine: np.ndarray = np.eye(sr + 1) affine[:sr, :sr] = direction[:sr, :sr] @ np.diag(spacing[:sr]) affine[:sr, -1] = origin[:sr] + if lps_to_ras: affine = orientation_ras_lps(affine) return affine @@ -735,17 +736,42 @@ def _get_affine(self, metadata: dict, lps_to_ras: bool = True): metadata: metadata with dict type. lps_to_ras: whether to convert the affine matrix from "LPS" to "RAS". Defaults to True. + Warns: + UserWarning: when ImageOrientationPatient (00200037) or ImagePositionPatient + (00200032) is missing from metadata. The affine matrix is set to identity, + which may be incorrect. Common with multiframe DICOM files. + """ affine: np.ndarray = np.eye(4) if not ("00200037" in metadata and "00200032" in metadata): + warnings.warn( + "PydicomReader: ImageOrientationPatient (0020,0037) and/or " + "ImagePositionPatient (0020,0032) tags are missing, so the affine " + "matrix cannot be derived and defaults to the identity. The image " + "orientation and spacing may be incorrect (e.g. for multi-frame " + "Enhanced DICOM); consider using ITKReader for such files.", + stacklevel=2, + ) return affine + + def _raise_if_not_finite(values: Sequence[Any], tag: str) -> None: + if not np.isfinite(tuple(values)).all(): + raise ValueError( + f"PydicomReader: cannot derive affine matrix because DICOM tag {tag} " + f"has a non-finite value: {values}." + ) + # "00200037" is the tag of `ImageOrientationPatient` rx, ry, rz, cx, cy, cz = metadata["00200037"]["Value"] + _raise_if_not_finite((rx, ry, rz, cx, cy, cz), "ImageOrientationPatient (0020,0037)") # "00200032" is the tag of `ImagePositionPatient` sx, sy, sz = metadata["00200032"]["Value"] + _raise_if_not_finite((sx, sy, sz), "ImagePositionPatient (0020,0032)") # "00280030" is the tag of `PixelSpacing` spacing = metadata["00280030"]["Value"] if "00280030" in metadata else (1.0, 1.0) + _raise_if_not_finite(tuple(spacing), "PixelSpacing (0028,0030)") dr, dc = metadata.get("spacing", spacing)[:2] + _raise_if_not_finite((dr, dc), "spacing") affine[0, 0] = cx * dr affine[0, 1] = rx * dc affine[0, 3] = sx @@ -760,11 +786,15 @@ def _get_affine(self, metadata: dict, lps_to_ras: bool = True): # 3d if "lastImagePositionPatient" in metadata: t1n, t2n, t3n = metadata["lastImagePositionPatient"] + _raise_if_not_finite((t1n, t2n, t3n), "lastImagePositionPatient") n = metadata[MetaKeys.SPATIAL_SHAPE][-1] - k1, k2, k3 = (t1n - sx) / (n - 1), (t2n - sy) / (n - 1), (t3n - sz) / (n - 1) - affine[0, 2] = k1 - affine[1, 2] = k2 - affine[2, 2] = k3 + if n > 1: + affine[0, 2] = (t1n - sx) / (n - 1) + affine[1, 2] = (t2n - sy) / (n - 1) + affine[2, 2] = (t3n - sz) / (n - 1) + + if not np.isfinite(affine).all(): + raise ValueError("PydicomReader: affine matrix not finite after composition.") if lps_to_ras: affine = orientation_ras_lps(affine) diff --git a/monai/data/meta_obj.py b/monai/data/meta_obj.py index 15e6e8be15a..df1bc713344 100644 --- a/monai/data/meta_obj.py +++ b/monai/data/meta_obj.py @@ -24,6 +24,9 @@ _TRACK_META = True +# Default number of spatial dimensions for medical imaging (3D volumetric data) +_DEFAULT_SPATIAL_NDIM = 3 + __all__ = ["get_track_meta", "set_track_meta", "MetaObj"] @@ -84,6 +87,7 @@ def __init__(self) -> None: self._applied_operations: list = MetaObj.get_default_applied_operations() self._pending_operations: list = MetaObj.get_default_applied_operations() # the same default as applied_ops self._is_batch: bool = False + self._spatial_ndim: int = 3 # default: 3 spatial dimensions @staticmethod def flatten_meta_objs(*args: Iterable): diff --git a/monai/data/meta_tensor.py b/monai/data/meta_tensor.py index 12bd76ba605..0dc5abb1b20 100644 --- a/monai/data/meta_tensor.py +++ b/monai/data/meta_tensor.py @@ -13,22 +13,60 @@ import functools import warnings -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from copy import deepcopy +from numbers import Integral from typing import Any import numpy as np import torch import monai -from monai.config.type_definitions import NdarrayTensor -from monai.data.meta_obj import MetaObj, get_track_meta -from monai.data.utils import affine_to_spacing, decollate_batch, list_data_collate, remove_extra_metadata +from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor +from monai.data.meta_obj import _DEFAULT_SPATIAL_NDIM, MetaObj, get_track_meta +from monai.data.utils import affine_to_spacing, decollate_batch, is_no_channel, list_data_collate, remove_extra_metadata from monai.utils import look_up_option from monai.utils.enums import LazyAttr, MetaKeys, PostFix, SpaceKeys from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_numpy, convert_to_tensor -__all__ = ["MetaTensor"] +__all__ = ["MetaTensor", "get_spatial_ndim"] + + +def _normalize_spatial_ndim(spatial_ndim: int, tensor_ndim: int, no_channel: bool = False) -> int: + """Clamp spatial dims to a valid range for the current tensor shape.""" + limit = max(int(tensor_ndim), 1) if no_channel else max(int(tensor_ndim) - 1, 1) + return max(1, min(int(spatial_ndim), limit)) + + +def _has_explicit_no_channel(meta: Mapping | None) -> bool: + return ( + isinstance(meta, Mapping) + and MetaKeys.ORIGINAL_CHANNEL_DIM in meta + and is_no_channel(meta[MetaKeys.ORIGINAL_CHANNEL_DIM]) + ) + + +def get_spatial_ndim(img: NdarrayOrTensor) -> int: + """Return the number of spatial dimensions assuming channel-first layout. + + Uses ``MetaTensor.spatial_ndim`` when available, otherwise falls back to + ``img.ndim - 1``. Always assumes channel-first (``no_channel=False``) + because callers run after ``EnsureChannelFirst`` has already added one. + """ + if isinstance(img, MetaTensor): + return _normalize_spatial_ndim(img.spatial_ndim, img.ndim) + return img.ndim - 1 + + +def _is_batch_only_index(index: Any) -> bool: + """True when indexing pattern selects only the batch axis (e.g., ``x[0]`` or ``x[0, ...]``).""" + if isinstance(index, (int, np.integer)): + return True + if not isinstance(index, Sequence) or not index: + return False + if not isinstance(index[0], (int, np.integer)): + return False + return all(i in (slice(None, None, None), Ellipsis, None) for i in index[1:]) @functools.lru_cache(None) @@ -111,6 +149,7 @@ def __new__( meta: dict | None = None, applied_operations: list | None = None, *args, + spatial_ndim: int | None = None, **kwargs, ) -> MetaTensor: _kwargs = {"device": kwargs.pop("device", None), "dtype": kwargs.pop("dtype", None)} if kwargs else {} @@ -123,6 +162,7 @@ def __init__( meta: dict | None = None, applied_operations: list | None = None, *_args, + spatial_ndim: int | None = None, **_kwargs, ) -> None: """ @@ -134,6 +174,8 @@ def __init__( the list is typically maintained by `monai.transforms.TraceableTransform`. See also: :py:class:`monai.transforms.TraceableTransform` _args: additional args (currently not in use in this constructor). + spatial_ndim: optional number of spatial dimensions. If ``None``, derived + from the affine matrix clamped by the tensor shape. _kwargs: additional kwargs (currently not in use in this constructor). Note: @@ -158,6 +200,14 @@ def __init__( self.affine = self.meta[MetaKeys.AFFINE] else: self.affine = self.get_default_affine() + # Initialize spatial_ndim from affine matrix (source of truth), clamped by tensor shape. + # This cached value is kept in sync via the affine setter for hot-path performance. + no_channel = _has_explicit_no_channel(self.meta) + if spatial_ndim is not None: + self.spatial_ndim = _normalize_spatial_ndim(spatial_ndim, self.ndim, no_channel=no_channel) + elif self.affine.ndim == 2: + self.spatial_ndim = _normalize_spatial_ndim(self.affine.shape[-1] - 1, self.ndim, no_channel=no_channel) + # applied_operations if applied_operations is not None: self.applied_operations = applied_operations @@ -237,6 +287,7 @@ def _handle_batched(cls, ret, idx, metas, func, args, kwargs): if func == torch.Tensor.__getitem__: if idx > 0 or len(args) < 2 or len(args[0]) < 1: return ret + full_idx = args[1] batch_idx = args[1][0] if isinstance(args[1], Sequence) else args[1] # if using e.g., `batch[:, -1]` or `batch[..., -1]`, then the # first element will be `slice(None, None, None)` and `Ellipsis`, @@ -258,6 +309,8 @@ def _handle_batched(cls, ret, idx, metas, func, args, kwargs): ret_meta.is_batch = False if hasattr(ret_meta, "__dict__"): ret.__dict__ = ret_meta.__dict__.copy() + if _is_batch_only_index(full_idx): + ret.spatial_ndim = _normalize_spatial_ndim(ret.spatial_ndim, ret.ndim) # `unbind` is used for `next(iter(batch))`. Also for `decollate_batch`. # But we only want to split the batch if the `unbind` is along the 0th dimension. elif func == torch.Tensor.unbind: @@ -442,7 +495,8 @@ def astype(self, dtype, device=None, *_args, **_kwargs): _kwargs: additional kwargs (currently unused). Returns: - data array instance + ``MetaTensor`` when a torch dtype is given (metadata is preserved), + or ``np.ndarray`` when a numpy dtype is given. """ if isinstance(dtype, str): mod_str, *dtype = dtype.split(".", 1) @@ -453,7 +507,7 @@ def astype(self, dtype, device=None, *_args, **_kwargs): out_type: type[torch.Tensor] | type[np.ndarray] | None if mod_str == "torch": - out_type = torch.Tensor + out_type = type(self) elif mod_str in ("numpy", "np"): out_type = np.ndarray else: @@ -467,15 +521,42 @@ def affine(self) -> torch.Tensor: @affine.setter def affine(self, d: NdarrayTensor) -> None: - """Set the affine.""" - self.meta[MetaKeys.AFFINE] = torch.as_tensor(d, device=torch.device("cpu"), dtype=torch.float64) + """Set the affine. + + When setting a non-batched affine matrix, automatically synchronizes the cached + spatial_ndim attribute to maintain consistency between the affine matrix (source of truth) + and the cached spatial dimension count. + """ + a = torch.as_tensor(d, device=torch.device("cpu"), dtype=torch.float64) + self.meta[MetaKeys.AFFINE] = a + if a.ndim == 2: # non-batched: sync spatial_ndim from affine (source of truth) + no_channel = _has_explicit_no_channel(self.meta) + self.spatial_ndim = _normalize_spatial_ndim(a.shape[-1] - 1, self.ndim, no_channel=no_channel) + + @property + def spatial_ndim(self) -> int: + """Get the number of spatial dimensions. + + This value is cached for hot-path performance and is kept in sync with the affine matrix + via the affine setter. The affine matrix is the source of truth for spatial dimensions. + """ + return getattr(self, "_spatial_ndim", _DEFAULT_SPATIAL_NDIM) + + @spatial_ndim.setter + def spatial_ndim(self, val: int) -> None: + """Set the number of spatial dimensions.""" + if not isinstance(val, Integral): + raise TypeError(f"'val' must be an numbers.Integral type; got {type(val)}.") + if val < 1: + raise ValueError(f"spatial_ndim must be >= 1, got {val}") + self._spatial_ndim = int(val) @property def pixdim(self): """Get the spacing""" if self.is_batch: - return [affine_to_spacing(a) for a in self.affine] - return affine_to_spacing(self.affine) + return [affine_to_spacing(a, r=self.spatial_ndim) for a in self.affine] + return affine_to_spacing(self.affine, r=self.spatial_ndim) def peek_pending_shape(self): """ @@ -490,7 +571,7 @@ def peek_pending_shape(self): def peek_pending_affine(self): res = self.affine - r = len(res) - 1 + r = res.shape[-1] - 1 if res.ndim >= 2 else self.spatial_ndim if r not in (2, 3): warnings.warn(f"Only 2d and 3d affine are supported, got {r}d input.") for p in self.pending_operations: @@ -503,8 +584,10 @@ def peek_pending_affine(self): return res def peek_pending_rank(self): - a = self.pending_operations[-1].get(LazyAttr.AFFINE, None) if self.pending_operations else self.affine - return 1 if a is None else int(max(1, len(a) - 1)) + if self.pending_operations: + a = self.pending_operations[-1].get(LazyAttr.AFFINE, None) + return 1 if a is None else int(max(1, len(a) - 1)) + return self.spatial_ndim def new_empty(self, size, dtype=None, device=None, requires_grad=False): # type: ignore[override] """ diff --git a/monai/data/utils.py b/monai/data/utils.py index d548ed72486..64bd79c7128 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -17,7 +17,6 @@ import math import os import pickle -import sys from collections import abc, defaultdict from collections.abc import Generator, Iterable, Mapping, Sequence, Sized from copy import deepcopy @@ -31,7 +30,7 @@ from torch.utils.data._utils.collate import default_collate from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor, PathLike -from monai.data.meta_obj import MetaObj +from monai.data.meta_obj import _DEFAULT_SPATIAL_NDIM, MetaObj from monai.utils import ( MAX_SEED, BlendMode, @@ -432,6 +431,9 @@ def collate_meta_tensor_fn(batch, *, collate_fn_map=None): collated.meta = default_collate(meta_dicts) collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch] collated.is_batch = True + collated.spatial_ndim = min( + min(getattr(t, "spatial_ndim", _DEFAULT_SPATIAL_NDIM) for t in batch), max(collated.ndim - 1, 1) + ) return collated @@ -1367,13 +1369,8 @@ def json_hashing(item) -> bytes: """ # TODO: Find way to hash transforms content as part of the cache - cache_key = "" - if sys.version_info.minor < 9: - cache_key = hashlib.md5(json.dumps(item, sort_keys=True).encode("utf-8")).hexdigest() - else: - cache_key = hashlib.md5( - json.dumps(item, sort_keys=True).encode("utf-8"), usedforsecurity=False # type: ignore - ).hexdigest() + dump = json.dumps(item, sort_keys=True).encode("utf-8") + cache_key = hashlib.sha256(dump, usedforsecurity=False).hexdigest() # type: ignore return f"{cache_key}".encode() @@ -1388,13 +1385,8 @@ def pickle_hashing(item, protocol=pickle.HIGHEST_PROTOCOL) -> bytes: Returns: the corresponding hash key """ - cache_key = "" - if sys.version_info.minor < 9: - cache_key = hashlib.md5(pickle.dumps(sorted_dict(item), protocol=protocol)).hexdigest() - else: - cache_key = hashlib.md5( - pickle.dumps(sorted_dict(item), protocol=protocol), usedforsecurity=False # type: ignore - ).hexdigest() + dump = pickle.dumps(sorted_dict(item), protocol=protocol) + cache_key = hashlib.sha256(dump, usedforsecurity=False).hexdigest() # type: ignore return f"{cache_key}".encode() diff --git a/monai/data/wsi_datasets.py b/monai/data/wsi_datasets.py index 2ee8c9d3634..b1830018358 100644 --- a/monai/data/wsi_datasets.py +++ b/monai/data/wsi_datasets.py @@ -250,8 +250,10 @@ def __init__( self.offset_limits = None elif isinstance(offset_limits, tuple): if isinstance(offset_limits[0], int): + # pyrefly: ignore [bad-assignment] self.offset_limits = (offset_limits, offset_limits) elif isinstance(offset_limits[0], tuple): + # pyrefly: ignore [bad-assignment] self.offset_limits = offset_limits else: raise ValueError( diff --git a/monai/data/wsi_reader.py b/monai/data/wsi_reader.py index b377234d10d..8a465c11979 100644 --- a/monai/data/wsi_reader.py +++ b/monai/data/wsi_reader.py @@ -319,6 +319,7 @@ def _get_metadata( } return metadata + # pyrefly: ignore [bad-override] def get_data( self, wsi, diff --git a/monai/engines/evaluator.py b/monai/engines/evaluator.py index 62d5f838477..2748ca34508 100644 --- a/monai/engines/evaluator.py +++ b/monai/engines/evaluator.py @@ -489,11 +489,13 @@ def _iteration(self, engine: EnsembleEvaluator, batchdata: dict[str, torch.Tenso if engine.amp: with torch.autocast("cuda", **engine.amp_kwargs): if isinstance(engine.state.output, dict): + # pyrefly: ignore [no-matching-overload] engine.state.output.update( {engine.pred_keys[idx]: engine.inferer(inputs, network, *args, **kwargs)} ) else: if isinstance(engine.state.output, dict): + # pyrefly: ignore [no-matching-overload] engine.state.output.update( {engine.pred_keys[idx]: engine.inferer(inputs, network, *args, **kwargs)} ) diff --git a/monai/engines/trainer.py b/monai/engines/trainer.py index 921d54a59cb..1f0c75620fa 100644 --- a/monai/engines/trainer.py +++ b/monai/engines/trainer.py @@ -774,4 +774,5 @@ def _compute_discriminator_loss() -> None: engine.state.output[AdversarialKeys.DISCRIMINATOR_LOSS].backward() engine.state.d_optimizer.step() + # pyrefly: ignore [bad-return] return engine.state.output diff --git a/monai/fl/client/monai_algo.py b/monai/fl/client/monai_algo.py index 6e9a6fd1fe2..4bf8de7c77a 100644 --- a/monai/fl/client/monai_algo.py +++ b/monai/fl/client/monai_algo.py @@ -13,6 +13,7 @@ import os import time +import warnings from collections.abc import Mapping, MutableMapping from typing import Any, cast @@ -34,6 +35,26 @@ logger = get_logger(__name__) +def _warn_provisioned_config_execution(bundle_root: str) -> None: + """ + Warn that the bundle under ``bundle_root`` is about to be executed. + + In federated learning the whole app directory -- configs included -- is provisioned by the FL + system, and the aggregation server dispatches `initialize`/`train` tasks that the client runs + on its own, so there is no per-round human interaction to catch a poisoned config. + """ + warnings.warn( + f"executing the bundle config under {bundle_root}, which is provisioned by the FL system: " + 'any `"_target_"` value in it is resolved to an importable callable and invoked with no ' + 'allow list, and any `"$"`-prefixed value is passed to Python `eval()`. A malicious or ' + "compromised aggregation server therefore gets code execution on this client, without any " + "per-round human interaction. Only join a federation whose server and app-provisioning " + "channel you trust (see " + "https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-x6pr-233j-x5cw).", + stacklevel=3, + ) + + def convert_global_weights(global_weights: Mapping, local_var_dict: MutableMapping) -> tuple[MutableMapping, int]: """Helper function to convert global weights to local weights format""" # Before loading weights, tensors might need to be reshaped to support HE for secure aggregation. @@ -86,6 +107,15 @@ class MonaiAlgoStats(ClientAlgoStats): """ Implementation of ``ClientAlgoStats`` to allow federated learning with MONAI bundle configurations. + Security note: the bundle under `bundle_root` is provisioned by the FL system -- `initialize()` + resolves it against `extra[ExtraItems.APP_ROOT]`, which the aggregation server supplies -- and + executing it runs whatever its config contains: any `"_target_"` value is resolved to an + importable callable and invoked with no allow list, and any `"$"`-prefixed value is passed to + Python `eval()`. A malicious or compromised server therefore gets code execution on this client, + with no per-round human interaction. Executing a config raises a warning -- once per call site, + as Python's default warning filter suppresses repeats + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-x6pr-233j-x5cw). + Args: bundle_root: directory path of the bundle. config_train_filename: bundle training config path relative to bundle_root. Can be a list of files; @@ -135,18 +165,29 @@ def initialize(self, extra=None): Args: extra: Dict with additional information that should be provided by FL system, i.e., `ExtraItems.CLIENT_NAME`, `ExtraItems.APP_ROOT` and `ExtraItems.LOGGING_FILE`. - You can diable the logging logic in the monai bundle by setting {ExtraItems.LOGGING_FILE} to False. + `{ExtraItems.LOGGING_FILE}` defaults to False here, and an explicit `None` is + treated the same way, so the bundle's own "configs/logging.conf" is not applied: + it is provisioned by the FL system and `logging.config.fileConfig` runs the INI's + `class=`/`args=` fields through `eval()` + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-wvpx-5qmp-46g3). + Set it to a logging config file path to opt back in to configuring logging. """ if extra is None: extra = {} self.client_name = extra.get(ExtraItems.CLIENT_NAME, "noname") - logging_file = extra.get(ExtraItems.LOGGING_FILE, None) + logging_file = extra.get(ExtraItems.LOGGING_FILE, False) + if logging_file is None: + # `ConfigWorkflow` reads `None` as "fall back to the bundle's own configs/logging.conf", + # the FL-provisioned file this default exists to keep away from `fileConfig`. Passing the + # key explicitly as `None` has to mean the same as leaving it out. + logging_file = False self.logger.info(f"Initializing {self.client_name} ...") # FL platform needs to provide filepath to configuration files self.app_root = extra.get(ExtraItems.APP_ROOT, "") self.bundle_root = os.path.join(self.app_root, self.bundle_root) + _warn_provisioned_config_execution(self.bundle_root) if self.workflow is None: config_train_files = self._add_config_files(self.config_train_filename) @@ -251,6 +292,7 @@ def _get_data_key_stats(self, data, data_key, hist_bins, hist_range, output_path dataroot=self.workflow.dataset_dir, # type: ignore hist_bins=hist_bins, hist_range=hist_range, + # pyrefly: ignore [bad-argument-type] output_path=output_path, histogram_only=self.histogram_only, ) @@ -312,6 +354,15 @@ class MonaiAlgo(ClientAlgo, MonaiAlgoStats): """ Implementation of ``ClientAlgo`` to allow federated learning with MONAI bundle configurations. + Security note: the bundle under `bundle_root` is provisioned by the FL system -- `initialize()` + resolves it against `extra[ExtraItems.APP_ROOT]`, which the aggregation server supplies -- and + executing it runs whatever its config contains: any `"_target_"` value is resolved to an + importable callable and invoked with no allow list, and any `"$"`-prefixed value is passed to + Python `eval()`. A malicious or compromised server therefore gets code execution on this client, + with no per-round human interaction. Executing a config raises a warning -- once per call site, + as Python's default warning filter suppresses repeats + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-x6pr-233j-x5cw). + Args: bundle_root: directory path of the bundle. local_epochs: number of local epochs to execute during each round of local training; defaults to 1. @@ -415,19 +466,30 @@ def initialize(self, extra=None): Args: extra: Dict with additional information that should be provided by FL system, i.e., `ExtraItems.CLIENT_NAME`, `ExtraItems.APP_ROOT` and `ExtraItems.LOGGING_FILE`. - You can diable the logging logic in the monai bundle by setting {ExtraItems.LOGGING_FILE} to False. + `{ExtraItems.LOGGING_FILE}` defaults to False here, and an explicit `None` is + treated the same way, so the bundle's own "configs/logging.conf" is not applied: + it is provisioned by the FL system and `logging.config.fileConfig` runs the INI's + `class=`/`args=` fields through `eval()` + (see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-wvpx-5qmp-46g3). + Set it to a logging config file path to opt back in to configuring logging. """ self._set_cuda_device() if extra is None: extra = {} self.client_name = extra.get(ExtraItems.CLIENT_NAME, "noname") - logging_file = extra.get(ExtraItems.LOGGING_FILE, None) + logging_file = extra.get(ExtraItems.LOGGING_FILE, False) + if logging_file is None: + # `ConfigWorkflow` reads `None` as "fall back to the bundle's own configs/logging.conf", + # the FL-provisioned file this default exists to keep away from `fileConfig`. Passing the + # key explicitly as `None` has to mean the same as leaving it out. + logging_file = False timestamp = time.strftime("%Y%m%d_%H%M%S") self.logger.info(f"Initializing {self.client_name} ...") # FL platform needs to provide filepath to configuration files self.app_root = extra.get(ExtraItems.APP_ROOT, "") self.bundle_root = os.path.join(self.app_root, self.bundle_root) + _warn_provisioned_config_execution(self.bundle_root) if self.train_workflow is None and self.config_train_filename is not None: config_train_files = self._add_config_files(self.config_train_filename) diff --git a/monai/handlers/mlflow_handler.py b/monai/handlers/mlflow_handler.py index 3078d89f97c..1cd26d5287a 100644 --- a/monai/handlers/mlflow_handler.py +++ b/monai/handlers/mlflow_handler.py @@ -22,7 +22,16 @@ from torch.utils.data import Dataset from monai.apps.utils import get_logger -from monai.utils import CommonKeys, IgniteInfo, ensure_tuple, flatten_dict, min_version, optional_import +from monai.utils import ( + CommonKeys, + IgniteInfo, + ensure_tuple, + flatten_dict, + min_version, + optional_import, + path_to_sqlite_uri, + path_to_uri, +) Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events") mlflow, _ = optional_import("mlflow", descriptor="Please install mlflow before using MLFlowHandler.") @@ -68,7 +77,16 @@ class MLFlowHandler: tracking_uri: connects to a tracking URI. can also set the `MLFLOW_TRACKING_URI` environment variable to have MLflow find a URI from there. in both cases, the URI can either be an HTTP/HTTPS URI for a remote server, a database connection string, or a local path - to log data to a directory. The URI defaults to path `mlruns`. + to log data to a directory. When no ``tracking_uri`` is provided and the + ``MLFLOW_TRACKING_URI`` environment variable is unset, the handler now + defaults to a local SQLite database backend at ``sqlite:////mlruns.db`` with + artifacts stored under ``/mlruns``. The default was changed from the filesystem + (file store) backend because MLflow 3.13+ raises an exception for the file store unless + ``MLFLOW_ALLOW_FILE_STORE=true`` is set; SQLite is the backend MLflow recommends and it + does not raise. Any explicitly provided ``tracking_uri`` is passed through unchanged + unless ``MLFLOW_TRACKING_URI`` is set (which takes precedence); local file paths and + ``file://`` URIs are rejected because MLflow no longer supports the filesystem (file + store) tracking backend. for more details: https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.set_tracking_uri. iteration_log: whether to log data to MLFlow when iteration completed, default to `True`. ``iteration_log`` can be also a function and it will be interpreted as an event filter @@ -113,6 +131,11 @@ class MLFlowHandler: optimizer_param_names: parameter names in the optimizer that need to be recorded during running the workflow, default to `'lr'`. close_on_complete: whether to close the mlflow run in `complete` phase in workflow, default to False. + artifact_location: the location to store run artifacts in, passed to MLflow when the experiment is + created. When ``None`` and a local SQLite backend is used (from the ``tracking_uri`` argument + or the ``MLFLOW_TRACKING_URI`` environment variable), it defaults to an ``mlruns`` directory + next to the database file; for other backends ``None`` lets MLflow decide based on the + ``tracking_uri``. Has no effect if the experiment already exists. For more details of MLFlow usage, please refer to: https://mlflow.org/docs/latest/index.html. @@ -141,6 +164,7 @@ def __init__( artifacts: str | Sequence[Path] | None = None, optimizer_param_names: str | Sequence[str] = "lr", close_on_complete: bool = False, + artifact_location: str | None = None, ) -> None: self.iteration_log = iteration_log self.epoch_log = epoch_log @@ -156,7 +180,39 @@ def __init__( self.experiment_param = experiment_param self.artifacts = ensure_tuple(artifacts) self.optimizer_param_names = ensure_tuple(optimizer_param_names) - self.client = mlflow.MlflowClient(tracking_uri=tracking_uri if tracking_uri else None) + # When no tracking_uri is provided, default to a local SQLite backend instead of the + # filesystem (file store) backend. MLflow 3.13+ raises for the file store unless + # `MLFLOW_ALLOW_FILE_STORE=true` is set, while SQLite is the recommended backend and does + # not raise. Artifacts cannot live inside a database, so by default they are stored under + # the `./mlruns` directory (where the previous file store default kept them) via the + # experiment `artifact_location`. Any explicitly provided tracking_uri is left unchanged. + self.artifact_location = artifact_location + # Resolve the effective tracking URI. The `MLFLOW_TRACKING_URI` environment variable takes + # priority so it can override a hard-coded `tracking_uri` argument; both configure the + # artifact location the same way. + env_tracking_uri = os.environ.get("MLFLOW_TRACKING_URI") + effective_tracking_uri = env_tracking_uri or tracking_uri + # When neither is set, fall back to the local SQLite default described above. + if not effective_tracking_uri: + tracking_uri = effective_tracking_uri = path_to_sqlite_uri(os.path.join(os.getcwd(), "mlruns.db")) + # For a local SQLite backend, keep run artifacts in an `mlruns` directory next to the + # database file (mirroring the previous file-store layout) unless the caller set + # `artifact_location`. Other backends (e.g. a remote server) are left to MLflow to decide. + if self.artifact_location is None and effective_tracking_uri.startswith("sqlite:///"): + db_path = Path(effective_tracking_uri[len("sqlite:///") :]) + self.artifact_location = path_to_uri(db_path.parent / "mlruns") + # MLflow 3.13+ refuses the filesystem (file store) tracking backend, and 3.14+ resolves + # the store eagerly at client construction, so a local path or ``file://`` URI would raise + # an opaque MlflowException. Reject those here with an actionable message instead. + if effective_tracking_uri.startswith("file://") or "://" not in effective_tracking_uri: + raise ValueError( + "MLflow no longer supports the filesystem (file store) tracking backend; got " + f"tracking_uri={effective_tracking_uri!r}. Use a SQLite URI " + "(sqlite:////mlruns.db) or a remote tracking URI instead." + ) + # Only the argument is passed to the client; when `MLFLOW_TRACKING_URI` took priority it + # is left None so MLflow resolves the environment variable itself. + self.client = mlflow.MlflowClient(tracking_uri=None if env_tracking_uri else tracking_uri) self.run_finish_status = mlflow.entities.RunStatus.to_string(mlflow.entities.RunStatus.FINISHED) self.close_on_complete = close_on_complete self.experiment = None @@ -234,6 +290,7 @@ def start(self, engine: Engine) -> None: self._log_params(attrs) if self.dataset_logger: + # pyrefly: ignore [bad-argument-type] self.dataset_logger(self.dataset_dict) else: self._default_dataset_log(self.dataset_dict) @@ -245,7 +302,12 @@ def _set_experiment(self): try: experiment = self.client.get_experiment_by_name(self.experiment_name) if not experiment: - experiment_id = self.client.create_experiment(self.experiment_name) + # pass an explicit artifact_location (set for the default SQLite backend, or + # by the caller) so artifacts land in the intended directory; when it is + # None MLflow decides based on the tracking_uri. + experiment_id = self.client.create_experiment( + self.experiment_name, artifact_location=self.artifact_location + ) experiment = self.client.get_experiment(experiment_id) break except MlflowException as e: @@ -257,6 +319,7 @@ def _set_experiment(self): else: raise e + # pyrefly: ignore [missing-attribute] if experiment.lifecycle_stage != mlflow.entities.LifecycleStage.ACTIVE: raise ValueError(f"Cannot set a deleted experiment '{self.experiment_name}' as the active experiment") self.experiment = experiment @@ -336,14 +399,43 @@ def complete(self) -> None: for artifact in artifact_list: self.client.log_artifact(self.cur_run.info.run_id, artifact) + def _dispose_sqlite_store(self) -> None: + """ + Release MLflow's SQLAlchemy engine when a local SQLite tracking backend is used. + + MLflow keeps the SQLite connection open for the lifetime of the client, which on + Windows prevents the database file from being deleted. MLflow exposes no public + client close/dispose API, so this reaches into its internals defensively to release + the engine. It is a no-op for non-SQLite backends. + """ + tracking_uri = getattr(self.client, "tracking_uri", "") + if not isinstance(tracking_uri, str) or not tracking_uri.startswith("sqlite:"): + return + store = getattr(getattr(self.client, "_tracking_client", None), "store", None) + if store is None: + return + dispose = getattr(store, "_dispose_engine", None) + if callable(dispose): + dispose() + else: + engine = getattr(store, "engine", None) + if engine is not None: + engine.dispose() + read_engine = getattr(store, "read_engine", None) + if read_engine is not None: + read_engine.dispose() + def close(self) -> None: """ - Stop current running logger of MLFlow. + Stop current running logger of MLFlow and release local SQLite resources. """ - if self.cur_run: - self.client.set_terminated(self.cur_run.info.run_id, self.run_finish_status) - self.cur_run = None + try: + if self.cur_run: + self.client.set_terminated(self.cur_run.info.run_id, self.run_finish_status) + self.cur_run = None + finally: + self._dispose_sqlite_store() def epoch_completed(self, engine: Engine) -> None: """ diff --git a/monai/handlers/utils.py b/monai/handlers/utils.py index 02975039b38..47cc94838a8 100644 --- a/monai/handlers/utils.py +++ b/monai/handlers/utils.py @@ -122,15 +122,15 @@ class mean median max 5percentile 95percentile notnans # add the average value of all classes to v if class_labels is None: - class_labels = ["class" + str(i) for i in range(v.shape[1])] + labels = ["class" + str(i) for i in range(v.shape[1])] else: - class_labels = [str(i) for i in class_labels] # ensure to have a list of str + labels = [str(i) for i in class_labels] # ensure to have a list of str - class_labels += ["mean"] + labels += ["mean"] v = np.concatenate([v, np.nanmean(v, axis=1, keepdims=True)], axis=1) with open(os.path.join(save_dir, f"{k}_raw.csv"), "w") as f: - f.write(f"filename{deli}{deli.join(class_labels)}\n") + f.write(f"filename{deli}{deli.join(labels)}\n") for i, b in enumerate(v): f.write( f"{images[i] if images is not None else str(i)}{deli}" @@ -164,7 +164,7 @@ def _compute_op(op: str, d: np.ndarray) -> Any: with open(os.path.join(save_dir, f"{k}_summary.csv"), "w") as f: f.write(f"class{deli}{deli.join(ops)}\n") for i, c in enumerate(np.transpose(v)): - f.write(f"{class_labels[i]}{deli}{deli.join([f'{_compute_op(k, c):.4f}' for k in ops])}\n") + f.write(f"{labels[i]}{deli}{deli.join([f'{_compute_op(k, c):.4f}' for k in ops])}\n") def from_engine(keys: KeysCollection, first: bool = False) -> Callable: diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py index ee94b1ebdbe..bc55cef15ca 100644 --- a/monai/inferers/inferer.py +++ b/monai/inferers/inferer.py @@ -841,6 +841,7 @@ def network_wrapper( if isinstance(out, Mapping): for k in out.keys(): + # pyrefly: ignore [unsupported-operation] out[k] = out[k].unsqueeze(dim=self.spatial_dim + 2) return out diff --git a/monai/inferers/merger.py b/monai/inferers/merger.py index 31e9b5d6322..3d07925e457 100644 --- a/monai/inferers/merger.py +++ b/monai/inferers/merger.py @@ -309,9 +309,6 @@ def __init__( self.chunks = chunks - # Handle compressor/codecs based on zarr version - is_zarr_v3 = version_geq(get_package_version("zarr"), "3.0.0") - # Initialize codecs/compressor attributes with proper types self.codecs: list | None = None self.value_codecs: list | None = None diff --git a/monai/losses/__init__.py b/monai/losses/__init__.py index 087a24f9d76..9f35e5f0750 100644 --- a/monai/losses/__init__.py +++ b/monai/losses/__init__.py @@ -14,6 +14,7 @@ from .adversarial_loss import PatchAdversarialLoss from .aucm_loss import AUCMLoss from .barlow_twins import BarlowTwinsLoss +from .boundary_loss import BoundaryLoss from .cldice import SoftclDiceLoss, SoftDiceclDiceLoss from .contrastive import ContrastiveLoss from .deform import BendingEnergyLoss, DiffusionLoss diff --git a/monai/losses/adversarial_loss.py b/monai/losses/adversarial_loss.py index b2c27a41eed..8be05bab89e 100644 --- a/monai/losses/adversarial_loss.py +++ b/monai/losses/adversarial_loss.py @@ -129,7 +129,8 @@ def forward( target_is_real = True # With generator, we always want this to be true! warnings.warn( "Variable target_is_real has been set to False, but for_discriminator is set" - "to False. To optimise a generator, target_is_real must be set to True." + "to False. To optimise a generator, target_is_real must be set to True.", + stacklevel=2, ) if not isinstance(input, list): diff --git a/monai/losses/boundary_loss.py b/monai/losses/boundary_loss.py new file mode 100644 index 00000000000..169763e6222 --- /dev/null +++ b/monai/losses/boundary_loss.py @@ -0,0 +1,233 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import warnings +from collections.abc import Callable + +import torch +from torch.nn.modules.loss import _Loss + +from monai.networks import one_hot +from monai.transforms.utils import distance_transform_edt +from monai.utils import LossReduction + +__all__ = ["BoundaryLoss"] + + +class BoundaryLoss(_Loss): + """ + Compute the boundary loss for highly unbalanced segmentation. + + The boundary loss is a distance-based loss that operates on the interface between segmentation + regions rather than on the regions themselves. This makes it particularly effective for + highly imbalanced segmentation tasks (e.g., small lesions, thin structures), where standard + Dice or Cross-Entropy losses struggle due to foreground-background imbalance. + + The loss is formulated as a pixel-wise weighted sum of predicted probabilities and a + signed distance map derived from the ground truth. The signed distance map is negative + inside the foreground region and positive outside, with zero on the boundary. + + The data `input` (BNHW[D] where N is number of classes) is compared with ground truth `target` + (BNHW[D]). + Note that axis N of `input` is expected to be logits or probabilities for each class, if passing logits as input, + must set `sigmoid=True` or `softmax=True`, or specifying `other_act`. And the same axis of `target` + can be 1 or N (one-hot format). + + The original paper: + Kervadec, H. et al. (2019) Boundary loss for highly unbalanced segmentation. MIDL 2019. + https://arxiv.org/abs/1812.07032 + + Example: + >>> import torch + >>> from monai.losses import BoundaryLoss + >>> B, C, H, W = 2, 3, 5, 5 + >>> input = torch.rand(B, C, H, W) + >>> target = torch.randint(0, C, size=(B, H, W)) + >>> bl = BoundaryLoss(softmax=True, to_onehot_y=True) + >>> loss = bl(input, target) + """ + + def __init__( + self, + include_background: bool = True, + to_onehot_y: bool = False, + sigmoid: bool = False, + softmax: bool = False, + other_act: Callable | None = None, + reduction: LossReduction | str = LossReduction.MEAN, + batch: bool = False, + ) -> None: + """ + Args: + include_background: if False, channel index 0 (background category) is excluded from the calculation. + if the non-background segmentations are small compared to the total image size they can get overwhelmed + by the signal from the background so excluding it in such cases helps convergence. + to_onehot_y: whether to convert the ``target`` into the one-hot format, + using the number of classes inferred from `input` (``input.shape[1]``). Defaults to False. + sigmoid: if True, apply a sigmoid function to the prediction. + softmax: if True, apply a softmax function to the prediction. + other_act: callable function to execute other activation layers, Defaults to ``None``. for example: + ``other_act = torch.tanh``. + reduction: {``"none"``, ``"mean"``, ``"sum"``} + Specifies the reduction to apply to the output. Defaults to ``"mean"``. + + - ``"none"``: no reduction will be applied. + - ``"mean"``: the sum of the output will be divided by the number of elements in the output. + - ``"sum"``: the output will be summed. + batch: whether to compute the distance map and loss over the batch dimension before the dividing. + Defaults to False, a boundary loss value is computed independently from each item in the batch + before any `reduction`. + + Raises: + TypeError: When ``other_act`` is not an ``Optional[Callable]``. + ValueError: When more than 1 of [``sigmoid=True``, ``softmax=True``, ``other_act is not None``]. + Incompatible values. + """ + super().__init__(reduction=LossReduction(reduction).value) + if other_act is not None and not callable(other_act): + raise TypeError(f"other_act must be None or callable but is {type(other_act).__name__}.") + if int(sigmoid) + int(softmax) + int(other_act is not None) > 1: + raise ValueError("Incompatible values: more than 1 of [sigmoid=True, softmax=True, other_act is not None].") + + self.include_background = include_background + self.to_onehot_y = to_onehot_y + self.sigmoid = sigmoid + self.softmax = softmax + self.other_act = other_act + self.batch = batch + + @torch.no_grad() + def compute_distance_map(self, target: torch.Tensor) -> torch.Tensor: + """ + Compute the signed distance map for each class in the target. + + The signed distance map is negative inside the foreground region and positive outside, + with zero on the boundary. + + Args: + target: target tensor of shape BNHW[D], with values in {0, 1} (one-hot encoded). + + Returns: + Signed distance map of the same shape as target. + """ + if target.dim() not in (4, 5): + raise ValueError("Only 2D (BNHW) and 3D (BNHWD) supported") + + distance_map = torch.zeros_like(target, dtype=torch.float32) + + for batch_idx in range(target.shape[0]): + for channel_idx in range(target.shape[1]): + mask = target[batch_idx, channel_idx : channel_idx + 1] > 0.5 + + # Empty or full masks do not have a foreground/background interface. + if not mask.any() or mask.all(): + continue + + fg_dist: torch.Tensor = distance_transform_edt(mask) # type: ignore + bg_dist: torch.Tensor = distance_transform_edt(~mask) # type: ignore + + signed = torch.zeros_like(mask, dtype=torch.float32) + signed[mask] = -(fg_dist[mask].to(torch.float32) - 1) + signed[~mask] = bg_dist[~mask].to(torch.float32) + + distance_map[batch_idx, channel_idx] = signed[0] + + return distance_map + + def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """ + Args: + input: the shape should be BNHW[D], where N is the number of classes. + target: the shape should be BNHW[D] or B1HW[D], where N is the number of classes. + + Raises: + ValueError: If the input is not 2D (BNHW) or 3D (BNHWD). + AssertionError: When input and target (after one hot transform if set) + have different shapes. + ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. + + Example: + >>> import torch + >>> from monai.losses import BoundaryLoss + >>> B, C, H, W = 2, 3, 5, 5 + >>> input = torch.rand(B, C, H, W) + >>> target_idx = torch.randint(0, C, size=(B, H, W)).long() + >>> target = one_hot(target_idx[:, None, ...], num_classes=C) + >>> bl = BoundaryLoss(softmax=True) + >>> loss = bl(input, target) + """ + if input.dim() not in (4, 5): + raise ValueError("Only 2D (BNHW) and 3D (BNHWD) supported") + + n_pred_ch = input.shape[1] + + # Apply activation to input + if self.sigmoid: + input = torch.sigmoid(input) + + if self.softmax: + if n_pred_ch == 1: + warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2) + else: + input = torch.softmax(input, dim=1) + + if self.other_act is not None: + input = self.other_act(input) + + # Convert target to one-hot if needed + if self.to_onehot_y: + if n_pred_ch == 1: + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) + else: + if target.dim() == input.dim() - 1: + target = target.unsqueeze(dim=1) + target = one_hot(target, num_classes=n_pred_ch) + + # Validate shapes match + if input.shape != target.shape: + raise AssertionError(f"input and target shapes do not match: {input.shape} vs {target.shape}") + + # Exclude background if requested + if not self.include_background: + if n_pred_ch == 1: + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2) + else: + input = input[:, 1:] + target = target[:, 1:] + + # Compute signed distance maps from target + distance_map = self.compute_distance_map(target) + + # Compute boundary loss: sum over spatial dimensions of (probabilities * distance_map) + # Then average over classes and batch + spatial_axes = list(range(2, input.dim())) + + loss = torch.sum(input * distance_map, dim=spatial_axes) + + # Normalize by number of pixels per class per batch element + num_pixels = torch.prod(torch.as_tensor(input.shape[2:], device=input.device)) + loss = loss / num_pixels + if self.batch: + loss = loss.mean(dim=0) + + if self.reduction == LossReduction.MEAN.value: + loss = loss.mean() + elif self.reduction == LossReduction.SUM.value: + loss = loss.sum() + elif self.reduction == LossReduction.NONE.value: + # Return shape (B, C') unless batch=True reduces the batch dimension first. + pass + else: + raise ValueError(f"Unsupported reduction: {self.reduction}") + + return loss diff --git a/monai/losses/cldice.py b/monai/losses/cldice.py index 7d7e447c549..60e86014590 100644 --- a/monai/losses/cldice.py +++ b/monai/losses/cldice.py @@ -134,7 +134,8 @@ def __init__( Args: iter_: Number of iterations for skeletonization. Must be a non-negative integer. Defaults to 3. smooth_nr: a small constant added to the numerator to avoid zero. Defaults to 1.0. - smooth_dr: a small constant added to the denominator to avoid nan. Defaults to 1.0. + smooth_dr: a small constant added to the denominator of the individual precision / + sensitivity ratios and the internal Dice denominator to avoid nan. Defaults to 1.0. smooth: a small constant added to the denominator of the harmonic mean to avoid nan. Defaults to 1e-4. include_background: if False, channel index 0 (background category) is excluded from the calculation. if the non-background segmentations are small compared to the total image size they can get overwhelmed diff --git a/monai/losses/deform.py b/monai/losses/deform.py index 37e4468d4b4..80da8bafcc8 100644 --- a/monai/losses/deform.py +++ b/monai/losses/deform.py @@ -44,9 +44,58 @@ def spatial_gradient(x: torch.Tensor, dim: int) -> torch.Tensor: return (x[slicing_s] - x[slicing_e]) / 2.0 +def spatial_gradient_squared(x: torch.Tensor, dim_1: int, dim_2: int) -> torch.Tensor: + """ + Calculate the second-order partial derivative of ``x`` with respect to spatial dims + ``dim_1`` and ``dim_2`` using compact central finite differences. + + For ``dim_1 == dim_2`` the pure second derivative uses the ``[1, -2, 1]`` stencil: + ``d2x[i] = x[i+1] - 2 * x[i] + x[i-1]``. + + For ``dim_1 != dim_2`` the mixed partial uses the compact 4-point stencil: + ``d2x[i, j] = (x[i+1, j+1] - x[i+1, j-1] - x[i-1, j+1] + x[i-1, j-1]) / 4``. + + Every spatial dimension is sliced to ``[1:-1]`` so the output shape is independent of + ``(dim_1, dim_2)``; this lets terms be summed together. Requires ``x.shape[d] > 2`` + for every spatial dim ``d``. + + Args: + x: the shape should be BCH(WD). + dim_1: first spatial dimension index. + dim_2: second spatial dimension index. + + Returns: + Tensor with batch and channel axes preserved and every spatial axis sliced to + ``[1:-1]``. + """ + slice_inner = slice(1, -1) + slice_plus = slice(2, None) + slice_minus = slice(None, -2) + slice_all = slice(None) + + def _idx(overrides: dict) -> list: + out: list = [slice_all, slice_all] + for d in range(2, x.ndim): + out.append(overrides.get(d, slice_inner)) + return out + + if dim_1 == dim_2: + return x[_idx({dim_1: slice_plus})] - 2 * x[_idx({})] + x[_idx({dim_1: slice_minus})] + return ( + x[_idx({dim_1: slice_plus, dim_2: slice_plus})] + - x[_idx({dim_1: slice_plus, dim_2: slice_minus})] + - x[_idx({dim_1: slice_minus, dim_2: slice_plus})] + + x[_idx({dim_1: slice_minus, dim_2: slice_minus})] + ) / 4.0 + + class BendingEnergyLoss(_Loss): """ - Calculate the bending energy based on second-order differentiation of ``pred`` using central finite difference. + Calculate the bending energy based on second-order differentiation of ``pred``. + + Pure second derivatives use the compact ``[1, -2, 1]`` stencil; mixed partials use a + compact 4-point central scheme. Both span three voxels per axis, so each spatial + dimension of ``pred`` only needs to be greater than 2. For more information, see https://github.com/Project-MONAI/tutorials/blob/main/modules/bending_energy_diffusion_loss_notes.ipynb. @@ -79,41 +128,41 @@ def forward(self, pred: torch.Tensor) -> torch.Tensor: Raises: ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. ValueError: When ``pred`` is not 3-d, 4-d or 5-d. - ValueError: When any spatial dimension of ``pred`` has size less than or equal to 4. + ValueError: When any spatial dimension of ``pred`` has size less than or equal to 2. ValueError: When the number of channels of ``pred`` does not match the number of spatial dimensions. """ if pred.ndim not in [3, 4, 5]: raise ValueError(f"Expecting 3-d, 4-d or 5-d pred, instead got pred of shape {pred.shape}") for i in range(pred.ndim - 2): - if pred.shape[-i - 1] <= 4: - raise ValueError(f"All spatial dimensions must be > 4, got spatial dimensions {pred.shape[2:]}") + if pred.shape[-i - 1] <= 2: + raise ValueError(f"All spatial dimensions must be > 2, got spatial dimensions {pred.shape[2:]}") if pred.shape[1] != pred.ndim - 2: raise ValueError( f"Number of vector components, i.e. number of channels of the input DDF, {pred.shape[1]}, " f"does not match number of spatial dimensions, {pred.ndim - 2}" ) - # first order gradient - first_order_gradient = [spatial_gradient(pred, dim) for dim in range(2, pred.ndim)] - # spatial dimensions in a shape suited for broadcasting below if self.normalize: spatial_dims = torch.tensor(pred.shape, device=pred.device)[2:].reshape((1, -1) + (pred.ndim - 2) * (1,)) - energy = torch.tensor(0) - for dim_1, g in enumerate(first_order_gradient): - dim_1 += 2 + # Initialize on pred.device so a GPU `pred` does not get added to a CPU + # accumulator, and as a float so an integer-dtype `pred` still produces a + # floating-point energy (the compact pure-derivative stencil has no + # division, so a Long input would otherwise propagate as Long and fail + # `torch.mean` at the reduction step). + energy = torch.tensor(0.0, device=pred.device) + for dim_1 in range(2, pred.ndim): + d2 = spatial_gradient_squared(pred, dim_1, dim_1) if self.normalize: - g *= pred.shape[dim_1] / spatial_dims - energy = energy + (spatial_gradient(g, dim_1) * pred.shape[dim_1]) ** 2 - else: - energy = energy + spatial_gradient(g, dim_1) ** 2 + d2 = d2 * (pred.shape[dim_1] ** 2 / spatial_dims) + energy = energy + d2**2 for dim_2 in range(dim_1 + 1, pred.ndim): + d2_mixed = spatial_gradient_squared(pred, dim_1, dim_2) if self.normalize: - energy = energy + 2 * (spatial_gradient(g, dim_2) * pred.shape[dim_2]) ** 2 - else: - energy = energy + 2 * spatial_gradient(g, dim_2) ** 2 + d2_mixed = d2_mixed * (pred.shape[dim_1] * pred.shape[dim_2] / spatial_dims) + energy = energy + 2 * d2_mixed**2 if self.reduction == LossReduction.MEAN.value: energy = torch.mean(energy) # the batch and channel average diff --git a/monai/losses/dice.py b/monai/losses/dice.py index b4558f930c3..2c4010176a4 100644 --- a/monai/losses/dice.py +++ b/monai/losses/dice.py @@ -156,7 +156,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: n_pred_ch = input.shape[1] if self.softmax: if n_pred_ch == 1: - warnings.warn("single channel prediction, `softmax=True` ignored.") + warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2) else: input = torch.softmax(input, 1) @@ -165,13 +165,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: target = one_hot(target, num_classes=n_pred_ch) if not self.include_background: if n_pred_ch == 1: - warnings.warn("single channel prediction, `include_background=False` ignored.") + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2) else: # if skipping background, removing first channel target = target[:, 1:] @@ -405,7 +405,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: n_pred_ch = input.shape[1] if self.softmax: if n_pred_ch == 1: - warnings.warn("single channel prediction, `softmax=True` ignored.") + warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2) else: input = torch.softmax(input, 1) @@ -414,13 +414,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: target = one_hot(target, num_classes=n_pred_ch) if not self.include_background: if n_pred_ch == 1: - warnings.warn("single channel prediction, `include_background=False` ignored.") + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2) else: # if skipping background, removing first channel target = target[:, 1:] @@ -987,7 +987,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: n_pred_ch = input.shape[1] if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: target = one_hot(target, num_classes=n_pred_ch) dice_loss = self.dice(input, target) diff --git a/monai/losses/focal_loss.py b/monai/losses/focal_loss.py index 7773cbdc9a1..b6d10c711c0 100644 --- a/monai/losses/focal_loss.py +++ b/monai/losses/focal_loss.py @@ -146,13 +146,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: target = one_hot(target, num_classes=n_pred_ch) if not self.include_background: if n_pred_ch == 1: - warnings.warn("single channel prediction, `include_background=False` ignored.") + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2) else: # if skipping background, removing first channel target = target[:, 1:] diff --git a/monai/losses/hausdorff_loss.py b/monai/losses/hausdorff_loss.py index 680ff7bc82a..d10fdb9fd5b 100644 --- a/monai/losses/hausdorff_loss.py +++ b/monai/losses/hausdorff_loss.py @@ -154,7 +154,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: n_pred_ch = input.shape[1] if self.softmax: if n_pred_ch == 1: - warnings.warn("single channel prediction, `softmax=True` ignored.") + warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2) else: input = torch.softmax(input, 1) @@ -163,13 +163,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: target = one_hot(target, num_classes=n_pred_ch) if not self.include_background: if n_pred_ch == 1: - warnings.warn("single channel prediction, `include_background=False` ignored.") + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2) else: # If skipping background, removing first channel target = target[:, 1:] diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 37c78fae60f..2b39ab57634 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -232,12 +232,16 @@ def __init__( sigma = torch.mean(bin_centers[1:] - bin_centers[:-1]) * sigma_ratio self.kernel_type = look_up_option(kernel_type, ["gaussian", "b-spline"]) self.num_bins = num_bins - self.kernel_type = kernel_type + # declared as buffers so they move with the module (e.g. ``.to(device)``); only populated for the + # gaussian kernel, hence the ``Tensor`` annotation reflects the type at the use sites in that path. + self.preterm: torch.Tensor | None self.bin_centers: torch.Tensor | None + self.register_buffer("preterm", None, persistent=False) self.register_buffer("bin_centers", None, persistent=False) if self.kernel_type == "gaussian": - self.preterm = 1 / (2 * sigma**2) + self.register_buffer("preterm", 1 / (2 * sigma**2), persistent=False) self.register_buffer("bin_centers", bin_centers[None, None, ...], persistent=False) + self.smooth_nr = float(smooth_nr) self.smooth_dr = float(smooth_dr) @@ -316,8 +320,8 @@ def parzen_windowing_gaussian(self, img: torch.Tensor) -> tuple[torch.Tensor, to """ img = torch.clamp(img, 0, 1) img = img.reshape(img.shape[0], -1, 1) # (batch, num_sample, 1) - if self.bin_centers is None: - raise ValueError("bin_centers must be defined for gaussian parzen windowing.") + if self.bin_centers is None or self.preterm is None: + raise ValueError("bin_centers and preterm must be defined for gaussian parzen windowing.") weight = torch.exp( -self.preterm.to(img) * (img - self.bin_centers.to(img)) ** 2 ) # (batch, num_sample, num_bin) diff --git a/monai/losses/mcc_loss.py b/monai/losses/mcc_loss.py index ac2877e5f7e..17323f89417 100644 --- a/monai/losses/mcc_loss.py +++ b/monai/losses/mcc_loss.py @@ -133,7 +133,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: n_pred_ch = input.shape[1] if self.softmax: if n_pred_ch == 1: - warnings.warn("single channel prediction, `softmax=True` ignored.") + warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2) else: input = torch.softmax(input, 1) @@ -142,13 +142,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: target = one_hot(target, num_classes=n_pred_ch) if not self.include_background: if n_pred_ch == 1: - warnings.warn("single channel prediction, `include_background=False` ignored.") + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2) else: target = target[:, 1:] input = input[:, 1:] diff --git a/monai/losses/nacl_loss.py b/monai/losses/nacl_loss.py index 7447478dadc..792cc372c46 100644 --- a/monai/losses/nacl_loss.py +++ b/monai/losses/nacl_loss.py @@ -138,7 +138,7 @@ def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: if self.distance_type == "l1": loss_conf = utargets.sub(inputs).abs_().mean() elif self.distance_type == "l2": - loss_conf = utargets.sub(inputs).pow_(2).abs_().mean() + loss_conf = utargets.sub(inputs).pow_(2).mean() loss: torch.Tensor = loss_ce + self.alpha * loss_conf diff --git a/monai/losses/perceptual.py b/monai/losses/perceptual.py index 635a3e75ced..6c0608f605a 100644 --- a/monai/losses/perceptual.py +++ b/monai/losses/perceptual.py @@ -112,7 +112,7 @@ def __init__( ) if not channel_wise: warnings.warn( - "MedicalNet networks supp, ort channel-wise loss. Consider setting channel_wise=True.", stacklevel=2 + "MedicalNet networks support channel-wise loss. Consider setting channel_wise=True.", stacklevel=2 ) # Channel-wise only for MedicalNet @@ -127,7 +127,8 @@ def __init__( torch.hub.set_dir(cache_dir) # raise a warning that this may change the default cache dir for all torch.hub calls warnings.warn( - f"Setting cache_dir to {cache_dir}, this may change the default cache dir for all torch.hub calls." + f"Setting cache_dir to {cache_dir}, this may change the default cache dir for all torch.hub calls.", + stacklevel=2, ) self.spatial_dims = spatial_dims @@ -311,8 +312,10 @@ def spatial_average_3d(x: torch.Tensor, keepdim: bool = True) -> torch.Tensor: def normalize_tensor(x: torch.Tensor, eps: float = 1e-10) -> torch.Tensor: - norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True)) - return x / (norm_factor + eps) + # Add eps inside the sqrt so the gradient stays finite when the norm is zero + # (e.g. identical input/target features), avoiding NaNs from SqrtBackward. See issue #8412. + norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True) + eps) + return x / norm_factor def medicalnet_intensity_normalisation(volume): diff --git a/monai/losses/spatial_mask.py b/monai/losses/spatial_mask.py index 0f823410dde..ba91c22fdef 100644 --- a/monai/losses/spatial_mask.py +++ b/monai/losses/spatial_mask.py @@ -55,16 +55,18 @@ def forward(self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor mask: the shape should be B1H[WD] or 11H[WD]. """ if mask is None: - warnings.warn("No mask value specified for the MaskedLoss.") + warnings.warn("No mask value specified for the MaskedLoss.", stacklevel=2) return self.loss(input, target) if input.dim() != mask.dim(): - warnings.warn(f"Dim of input ({input.shape}) is different from mask ({mask.shape}).") + warnings.warn(f"Dim of input ({input.shape}) is different from mask ({mask.shape}).", stacklevel=2) if input.shape[0] != mask.shape[0] and mask.shape[0] != 1: raise ValueError(f"Batch size of mask ({mask.shape}) must be one or equal to input ({input.shape}).") if target.dim() > 1: if mask.shape[1] != 1: raise ValueError(f"Mask ({mask.shape}) must have only one channel.") if input.shape[2:] != mask.shape[2:]: - warnings.warn(f"Spatial size of input ({input.shape}) is different from mask ({mask.shape}).") + warnings.warn( + f"Spatial size of input ({input.shape}) is different from mask ({mask.shape}).", stacklevel=2 + ) return self.loss(input * mask, target * mask) diff --git a/monai/losses/tversky.py b/monai/losses/tversky.py index 154f34c5261..5db4025be0f 100644 --- a/monai/losses/tversky.py +++ b/monai/losses/tversky.py @@ -118,7 +118,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: n_pred_ch = input.shape[1] if self.softmax: if n_pred_ch == 1: - warnings.warn("single channel prediction, `softmax=True` ignored.") + warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2) else: input = torch.softmax(input, 1) @@ -127,13 +127,13 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: target = one_hot(target, num_classes=n_pred_ch) if not self.include_background: if n_pred_ch == 1: - warnings.warn("single channel prediction, `include_background=False` ignored.") + warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2) else: # if skipping background, removing first channel target = target[:, 1:] diff --git a/monai/losses/unified_focal_loss.py b/monai/losses/unified_focal_loss.py index 745513fec03..98dbf124c6a 100644 --- a/monai/losses/unified_focal_loss.py +++ b/monai/losses/unified_focal_loss.py @@ -58,7 +58,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: y_true = one_hot(y_true, num_classes=n_pred_ch) @@ -122,7 +122,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor: if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: y_true = one_hot(y_true, num_classes=n_pred_ch) @@ -223,7 +223,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor: n_pred_ch = y_pred.shape[1] if self.to_onehot_y: if n_pred_ch == 1: - warnings.warn("single channel prediction, `to_onehot_y=True` ignored.") + warnings.warn("single channel prediction, `to_onehot_y=True` ignored.", stacklevel=2) else: y_true = one_hot(y_true, num_classes=n_pred_ch) diff --git a/monai/metrics/__init__.py b/monai/metrics/__init__.py index 2265dd3a3f2..f55f92db1b8 100644 --- a/monai/metrics/__init__.py +++ b/monai/metrics/__init__.py @@ -11,6 +11,7 @@ from __future__ import annotations +from .absolute_volume_difference import AbsoluteVolumeDifferenceMetric, compute_absolute_volume_difference from .active_learning_metrics import LabelQualityScore, VarianceMetric, compute_variance, label_quality_score from .average_precision import AveragePrecisionMetric, compute_average_precision from .calibration import CalibrationErrorMetric, CalibrationReduction, calibration_binning diff --git a/monai/metrics/absolute_volume_difference.py b/monai/metrics/absolute_volume_difference.py new file mode 100644 index 00000000000..92e18acea2a --- /dev/null +++ b/monai/metrics/absolute_volume_difference.py @@ -0,0 +1,180 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import torch + +from monai.metrics.utils import do_metric_reduction, ignore_background +from monai.utils import MetricReduction + +from .metric import CumulativeIterationMetric + +__all__ = ["AbsoluteVolumeDifferenceMetric", "compute_absolute_volume_difference"] + + +class AbsoluteVolumeDifferenceMetric(CumulativeIterationMetric): + """ + Compute the Absolute Volume Difference (AVD) between predicted and ground-truth + segmentation masks. + + AVD measures the absolute difference in the number of foreground voxels between + prediction and ground truth, per class. It is particularly useful for small-object + segmentation (e.g. retinal fluid in OCT volumes) where Dice score is known to be + overly sensitive to volume size and does not directly reflect volume discrepancies. + + .. note:: + For 2D inputs this computes the difference in foreground **areas** rather than + volumes. In all cases the returned values are raw voxel/pixel counts and are + **not** scaled by the voxel/pixel spacing, so they are not expressed in the + physical units of the original image. + + Reference: + Bogunovic et al. (2019). RETOUCH: The Retinal OCT Fluid Detection and + Segmentation Benchmark and Challenge. + IEEE Transactions on Medical Imaging, 38(8), 1858-1874. + https://ieeexplore.ieee.org/document/8653407 + + The inputs ``y_pred`` and ``y`` are expected to be binarized one-hot tensors with + shape BCHW[D]. If they contain continuous values (e.g. sigmoid outputs), binarize + them first with a suitable threshold transform. + + The typical execution steps of this metric class follow + :py:class:`monai.metrics.metric.Cumulative`. + + Example: + + .. code-block:: python + + import torch + from monai.metrics import AbsoluteVolumeDifferenceMetric + + batch_size, n_classes = 4, 3 + y_pred = torch.randint(0, 2, (batch_size, n_classes, 64, 64, 32)).float() + y = torch.randint(0, 2, (batch_size, n_classes, 64, 64, 32)).float() + + metric = AbsoluteVolumeDifferenceMetric(include_background=False) + metric(y_pred, y) # accumulate + result = metric.aggregate() # shape: (n_classes - 1,) after mean reduction + metric.reset() + + Args: + include_background: whether to include AVD computation on the first channel + (index 0), which is by convention assumed to be background. Defaults to + ``True``. Set to ``False`` when the background class dominates and you only + care about foreground classes (e.g. fluid sub-types in OCT). + reduction: defines how to aggregate per-batch-per-class results. Available + modes are enumerated in :py:class:`monai.utils.enums.MetricReduction`. + Defaults to ``"mean"``. + get_not_nans: if ``True``, :meth:`aggregate` returns ``(metric, not_nans)`` + where ``not_nans`` counts the number of valid (non-NaN) values. + Defaults to ``False``. + ignore_empty: if ``True``, cases where the ground-truth channel is entirely + empty (zero voxels) are excluded from aggregation by setting their value + to ``NaN``. If ``False``, the raw absolute difference (equal to the + predicted volume for that class) is returned. Defaults to ``True``. + """ + + def __init__( + self, + include_background: bool = True, + reduction: MetricReduction | str = MetricReduction.MEAN, + get_not_nans: bool = False, + ignore_empty: bool = True, + ) -> None: + super().__init__() + self.include_background = include_background + self.reduction = reduction + self.get_not_nans = get_not_nans + self.ignore_empty = ignore_empty + + def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # type: ignore[override] + """ + Args: + y_pred: binarized prediction tensor, shape BCHW[D]. + y: binarized ground-truth tensor, shape BCHW[D]. + + Raises: + ValueError: when ``y_pred`` has fewer than three dimensions. + """ + if y_pred.ndimension() < 3: + raise ValueError( + f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {y_pred.ndimension()}." + ) + return compute_absolute_volume_difference( + y_pred=y_pred, y=y, include_background=self.include_background, ignore_empty=self.ignore_empty + ) + + def aggregate( + self, reduction: MetricReduction | str | None = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Execute reduction logic for the accumulated AVD values. + + Args: + reduction: optional override for the reduction mode set at construction. + """ + data = self.get_buffer() + if not isinstance(data, torch.Tensor): + raise ValueError("the data to aggregate must be a PyTorch Tensor.") + + f, not_nans = do_metric_reduction(data, reduction or self.reduction) + return (f, not_nans) if self.get_not_nans else f + + +def compute_absolute_volume_difference( + y_pred: torch.Tensor, y: torch.Tensor, include_background: bool = True, ignore_empty: bool = True +) -> torch.Tensor: + """ + Compute the Absolute Volume Difference (AVD) for a batch of segmentation predictions. + + AVD is defined per class as:: + + AVD_c = | sum_{spatial}(y_pred_c) - sum_{spatial}(y_c) | + + where the sum counts the number of foreground voxels in each channel. + + Args: + y_pred: binarized prediction tensor with shape BCHW[D]. + y: binarized ground-truth tensor with shape BCHW[D]. + include_background: whether to include the first channel (background). + Defaults to ``True``. + ignore_empty: if ``True``, entries where the ground-truth channel contains no + foreground voxels are set to ``NaN`` so they are excluded during reduction. + Defaults to ``True``. + + Returns: + AVD per batch item and per class, shape ``[batch_size, num_classes]``. + + Raises: + ValueError: when ``y_pred`` and ``y`` have different shapes. + """ + if y_pred.ndim < 3: + raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {y_pred.ndim}.") + + if not include_background: + y_pred, y = ignore_background(y_pred=y_pred, y=y) + + if y_pred.shape != y.shape: + raise ValueError(f"y_pred and y should have the same shape, got {y_pred.shape} and {y.shape}.") + + # sum over all spatial dimensions; keep batch (dim 0) and channel (dim 1) + reduce_axis = list(range(2, y_pred.ndim)) + vol_pred = torch.sum(y_pred, dim=reduce_axis) # [B, C] + vol_true = torch.sum(y, dim=reduce_axis) # [B, C] + + avd = torch.abs(vol_pred - vol_true) # [B, C] + + if ignore_empty: + # mark cases with no ground-truth foreground as NaN + avd = torch.where(vol_true > 0, avd, torch.tensor(float("nan"), device=avd.device)) + + return avd diff --git a/monai/metrics/active_learning_metrics.py b/monai/metrics/active_learning_metrics.py index 5c51d262ed2..a31d2fc1500 100644 --- a/monai/metrics/active_learning_metrics.py +++ b/monai/metrics/active_learning_metrics.py @@ -137,7 +137,7 @@ def compute_variance( n_len = len(y_pred.shape) if n_len < 4 and spatial_map: - warnings.warn("Spatial map requires a 2D/3D image with N-repeats and C-channels") + warnings.warn("Spatial map requires a 2D/3D image with N-repeats and C-channels", stacklevel=2) return None # Create new shape list @@ -190,7 +190,10 @@ def label_quality_score( n_len = len(y_pred.shape) if n_len < 4 and scalar_reduction == "none": - warnings.warn("Reduction set to None, Spatial map return requires a 2D/3D image of B-Batchsize and C-channels") + warnings.warn( + "Reduction set to None, Spatial map return requires a 2D/3D image of B-Batchsize and C-channels", + stacklevel=2, + ) return None abs_diff_map = torch.abs(y_pred - y) diff --git a/monai/metrics/average_precision.py b/monai/metrics/average_precision.py index 7dd277bde61..3bb7b0bcb4c 100644 --- a/monai/metrics/average_precision.py +++ b/monai/metrics/average_precision.py @@ -88,10 +88,12 @@ def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float: raise AssertionError("y and y_pred must be 1 dimension data with same length.") y_unique = y.unique() if len(y_unique) == 1: - warnings.warn(f"y values can not be all {y_unique.item()}, skip AP computation and return `Nan`.") + warnings.warn(f"y values can not be all {y_unique.item()}, skip AP computation and return `Nan`.", stacklevel=2) return float("nan") if not y_unique.equal(torch.tensor([0, 1], dtype=y.dtype, device=y.device)): - warnings.warn(f"y values must be 0 or 1, but in {y_unique.tolist()}, skip AP computation and return `Nan`.") + warnings.warn( + f"y values must be 0 or 1, but in {y_unique.tolist()}, skip AP computation and return `Nan`.", stacklevel=2 + ) return float("nan") n = len(y) diff --git a/monai/metrics/confusion_matrix.py b/monai/metrics/confusion_matrix.py index 26ec823081b..3152b023c0e 100644 --- a/monai/metrics/confusion_matrix.py +++ b/monai/metrics/confusion_matrix.py @@ -93,7 +93,7 @@ def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor) -> torch.Tensor raise ValueError("y_pred should have at least two dimensions.") if dims == 2 or (dims == 3 and y_pred.shape[-1] == 1): if self.compute_sample: - warnings.warn("As for classification task, compute_sample should be False.") + warnings.warn("As for classification task, compute_sample should be False.", stacklevel=2) self.compute_sample = False return get_confusion_matrix(y_pred=y_pred, y=y, include_background=self.include_background) diff --git a/monai/metrics/cumulative_average.py b/monai/metrics/cumulative_average.py index dccf7b094b8..8fbc4700580 100644 --- a/monai/metrics/cumulative_average.py +++ b/monai/metrics/cumulative_average.py @@ -154,7 +154,7 @@ def append(self, val: Any, count: Any | None = 1) -> None: # account for possible non-finite numbers in val and replace them with 0s nfin = torch.isfinite(val) if not torch.all(nfin): - warnings.warn(f"non-finite inputs received: val: {val}, count: {count}") + warnings.warn(f"non-finite inputs received: val: {val}, count: {count}", stacklevel=2) count = torch.where(nfin, count, torch.zeros_like(count)) val = torch.where(nfin, val, torch.zeros_like(val)) diff --git a/monai/metrics/froc.py b/monai/metrics/froc.py index 81a890aa681..3faef849170 100644 --- a/monai/metrics/froc.py +++ b/monai/metrics/froc.py @@ -11,7 +11,7 @@ from __future__ import annotations -from typing import Any, cast +from typing import Any import numpy as np import torch @@ -67,12 +67,14 @@ def compute_fp_tp_probs_nd( hittedlabel = evaluation_mask[tuple(coords.T)] fp_probs = probs[np.where(hittedlabel == 0)] + num_targets = 0 for i in range(1, max_label + 1): - if i not in labels_to_exclude and i in hittedlabel: - tp_probs[i - 1] = probs[np.where(hittedlabel == i)].max() + if i not in labels_to_exclude: + num_targets += 1 + if i in hittedlabel: + tp_probs[i - 1] = probs[np.where(hittedlabel == i)].max() - num_targets = max_label - len(labels_to_exclude) - return fp_probs, tp_probs, cast(int, num_targets) + return fp_probs, tp_probs, num_targets def compute_fp_tp_probs( diff --git a/monai/metrics/generalized_dice.py b/monai/metrics/generalized_dice.py index 05eb94af48f..4651322a585 100644 --- a/monai/metrics/generalized_dice.py +++ b/monai/metrics/generalized_dice.py @@ -181,7 +181,6 @@ def compute_generalized_dice( else: numer = 2.0 * (intersection * w) denom = denominator * w - y_pred_o = y_pred_o # Compute the score generalized_dice_score = numer / denom diff --git a/monai/metrics/rocauc.py b/monai/metrics/rocauc.py index 72f0b3730ca..6b4ac368dac 100644 --- a/monai/metrics/rocauc.py +++ b/monai/metrics/rocauc.py @@ -77,10 +77,14 @@ def _calculate(y_pred: torch.Tensor, y: torch.Tensor) -> float: raise AssertionError("y and y_pred must be 1 dimension data with same length.") y_unique = y.unique() if len(y_unique) == 1: - warnings.warn(f"y values can not be all {y_unique.item()}, skip AUC computation and return `Nan`.") + warnings.warn( + f"y values can not be all {y_unique.item()}, skip AUC computation and return `Nan`.", stacklevel=2 + ) return float("nan") if not y_unique.equal(torch.tensor([0, 1], dtype=y.dtype, device=y.device)): - warnings.warn(f"y values must be 0 or 1, but in {y_unique.tolist()}, skip AUC computation and return `Nan`.") + warnings.warn( + f"y values must be 0 or 1, but in {y_unique.tolist()}, skip AUC computation and return `Nan`.", stacklevel=2 + ) return float("nan") n = len(y) diff --git a/monai/metrics/utils.py b/monai/metrics/utils.py index 3070764e066..a5927a0a5be 100644 --- a/monai/metrics/utils.py +++ b/monai/metrics/utils.py @@ -38,6 +38,7 @@ binary_erosion, _ = optional_import("scipy.ndimage", name="binary_erosion") distance_transform_edt, _ = optional_import("scipy.ndimage", name="distance_transform_edt") distance_transform_cdt, _ = optional_import("scipy.ndimage", name="distance_transform_cdt") +KDTree, has_scipy_kdtree = optional_import("scipy.spatial", name="KDTree") scipy_ndimage, has_scipy_ndimage = optional_import("scipy.ndimage") cupy, has_cupy = optional_import("cupy") @@ -216,6 +217,7 @@ def get_mask_edges( or_vol = seg_pred | seg_gt if not or_vol.any(): pred, gt = lib.zeros(seg_pred.shape, dtype=bool), lib.zeros(seg_gt.shape, dtype=bool) + # pyrefly: ignore [bad-return] return (pred, gt) if spacing is None else (pred, gt, pred, gt) channel_first = [seg_pred[None], seg_gt[None], or_vol[None]] if spacing is None and not use_cucim: # cpu only erosion @@ -269,7 +271,8 @@ def get_surface_distance( distance_metric: : [``"euclidean"``, ``"chessboard"``, ``"taxicab"``] the metric used to compute surface distance. Defaults to ``"euclidean"``. - - ``"euclidean"``, uses Exact Euclidean distance transform. + - ``"euclidean"``, the exact Euclidean distance (a KD-tree over the edge voxels on + CPU, or the cuCIM distance transform when the inputs are on a CUDA device). - ``"chessboard"``, uses `chessboard` metric in chamfer type of transform. - ``"taxicab"``, uses `taxicab` metric in chamfer type of transform. spacing: spacing of pixel (or voxel). This parameter is relevant only if ``distance_metric`` is set to ``"euclidean"``. @@ -291,6 +294,27 @@ def get_surface_distance( dis = dis[seg_gt] return convert_to_dst_type(dis, seg_pred, dtype=dis.dtype)[0] if distance_metric == "euclidean": + # The euclidean surface distance only needs the distance from each `seg_pred` + # edge voxel to the nearest `seg_gt` edge voxel. CPU and GPU favour different + # algorithms for this: + # * On CPU, a KD-tree over the (sparse) edge-voxel coordinates avoids the dense + # full-volume distance transform, and handles outlier points that expand the + # bounding box. + # * On GPU, the dense EDT is embarrassingly parallel and significantly faster than + # cupy's KDTree (as of this writing anyway) + # When scipy's KDTree is unavailable we fall back to the dense distance transform. + on_gpu = isinstance(seg_gt, torch.Tensor) and seg_gt.device.type == "cuda" + if not on_gpu and has_scipy_kdtree: + gt_coords = np.argwhere(convert_to_numpy(seg_gt)).astype(np.float64) + pred_coords = np.argwhere(convert_to_numpy(seg_pred)).astype(np.float64) + if spacing is not None: + scale = np.asarray(spacing, dtype=np.float64) + gt_coords *= scale + pred_coords *= scale + # leafsize larger than the default (16) is faster here: we build the tree + # for a single batched query rather than amortizing it over many queries. + surface_distance = KDTree(gt_coords, leafsize=32).query(pred_coords, k=1)[0] + return convert_to_dst_type(surface_distance, seg_pred, dtype=lib.float32)[0] dis = monai_distance_transform_edt((~seg_gt)[None, ...], sampling=spacing)[0] # type: ignore elif distance_metric in {"chessboard", "taxicab"}: dis = distance_transform_cdt(convert_to_numpy(~seg_gt), metric=distance_metric) @@ -341,12 +365,14 @@ def get_edge_surface_distance( if not edges_gt.any(): warnings.warn( f"the ground truth of class {class_index if class_index != -1 else 'Unknown'} is all 0," - " this may result in nan/inf distance." + " this may result in nan/inf distance.", + stacklevel=2, ) if not edges_pred.any(): warnings.warn( f"the prediction of class {class_index if class_index != -1 else 'Unknown'} is all 0," - " this may result in nan/inf distance." + " this may result in nan/inf distance.", + stacklevel=2, ) distances: tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor] if symmetric: @@ -375,7 +401,7 @@ def is_binary_tensor(input: torch.Tensor, name: str) -> None: if not isinstance(input, torch.Tensor): raise ValueError(f"{name} must be of type PyTorch Tensor.") if not torch.all(input.byte() == input) or input.max() > 1 or input.min() < 0: - warnings.warn(f"{name} should be a binarized tensor.") + warnings.warn(f"{name} should be a binarized tensor.", stacklevel=2) def remap_instance_id(pred: torch.Tensor, by_size: bool = False) -> torch.Tensor: @@ -510,7 +536,8 @@ def compute_voronoi_regions_fast(labels: np.ndarray | torch.Tensor) -> torch.Ten if isinstance(labels, torch.Tensor): warnings.warn( "Voronoi computation is running on CPU. " - "To accelerate, move the input tensor to GPU and ensure 'cupy' with 'cupyx.scipy.ndimage' is installed." + "To accelerate, move the input tensor to GPU and ensure 'cupy' with 'cupyx.scipy.ndimage' is installed.", + stacklevel=2, ) x = labels.cpu().numpy() else: diff --git a/monai/networks/blocks/__init__.py b/monai/networks/blocks/__init__.py index 22af82d3161..5932aba7fe2 100644 --- a/monai/networks/blocks/__init__.py +++ b/monai/networks/blocks/__init__.py @@ -26,6 +26,13 @@ from .encoder import BaseEncoder from .fcn import FCN, GCN, MCFCN, Refine from .feature_pyramid_network import ExtraFPNBlock, FeaturePyramidNetwork, LastLevelMaxPool, LastLevelP6P7 +from .hyena import ( + DepthwiseFFTConv2d, + DepthwiseFFTConv3d, + HyenaMixer, + HyenaTransformerBlock, + is_nvsubquadratic_available, +) from .localnet_block import LocalNetDownSampleBlock, LocalNetFeatureExtractorBlock, LocalNetUpSampleBlock from .mednext_block import MedNeXtBlock, MedNeXtDownBlock, MedNeXtOutBlock, MedNeXtUpBlock from .mlp import MLPBlock diff --git a/monai/networks/blocks/hyena.py b/monai/networks/blocks/hyena.py new file mode 100644 index 00000000000..51277cb32f8 --- /dev/null +++ b/monai/networks/blocks/hyena.py @@ -0,0 +1,555 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +HyenaND-based building blocks for MONAI networks. + +These blocks provide a subquadratic O(N log N) alternative to windowed self-attention +in transformer-style segmentation networks. The operator is HyenaND from the +``nvsubquadratic`` package, gated through a thin :class:`HyenaMixer` and wrapped in +the conventional pre-norm / MLP residual pattern by :class:`HyenaTransformerBlock`. + +The supporting :class:`DepthwiseFFTConv2d` / :class:`DepthwiseFFTConv3d` classes are +drop-in depthwise convolutions implemented via FFT. They preserve the +``nn.Conv{2,3}d`` weight layout and ``isinstance`` relationship but route the forward +pass through ``torch.fft.rfftn`` to avoid PyTorch's INT32 unfold limit, which caps +``F.conv3d`` at ROI ~128 for typical medical-imaging channel counts. + +``nvsubquadratic`` is an optional dependency. The Hyena classes raise ``ImportError`` +with an install hint at construction time if the library is unavailable; the +FFT-conv classes have no such dependency and always work. +""" + +from __future__ import annotations + +from typing import cast + +import torch +import torch.nn as nn +from torch.nn import LayerNorm + +from monai.networks.blocks.mlp import MLPBlock +from monai.networks.layers.drop_path import DropPath +from monai.utils import optional_import + +__all__ = [ + "DepthwiseFFTConv2d", + "DepthwiseFFTConv3d", + "HyenaMixer", + "HyenaTransformerBlock", + "is_nvsubquadratic_available", +] + +# Optional ``nvsubquadratic`` symbols. Resolved at module-import time; the missing-dep +# error is raised lazily inside the consuming class ``__init__``. Every symbol the Hyena +# classes use carries its own availability flag so a partial / broken install (e.g. +# ``lazy_config`` present but ``modules.hyena_nd`` missing) reports unavailable rather than +# failing later with an opaque ``AttributeError``. +_LazyConfig, _has_lazyconfig = optional_import("nvsubquadratic.lazy_config", name="LazyConfig") +_instantiate, _has_instantiate = optional_import("nvsubquadratic.lazy_config", name="instantiate") +_Hyena, _has_hyena = optional_import("nvsubquadratic.modules.hyena_nd", name="Hyena") +_CKConvND, _has_ckconv = optional_import("nvsubquadratic.modules.ckconv_nd", name="CKConvND") +_SIRENKernelND, _has_siren = optional_import("nvsubquadratic.modules.kernels_nd", name="SIRENKernelND") +_GaussianModulationND, _has_gaussian = optional_import("nvsubquadratic.modules.masks_nd", name="GaussianModulationND") +_has_nvsubq = all((_has_lazyconfig, _has_instantiate, _has_hyena, _has_ckconv, _has_siren, _has_gaussian)) + +_NVSUBQ_INSTALL_HINT = ( + "HyenaND operators require the optional 'nvsubquadratic' package (Python >= 3.10). " + "Install it with: pip install 'monai[hyena]' " + "(equivalently: pip install 'nvsubquadratic>=0.1.1'). " + "See https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies" +) + + +def is_nvsubquadratic_available() -> bool: + """Return ``True`` if the optional ``nvsubquadratic`` package is importable.""" + return bool(_has_nvsubq) + + +# --------------------------------------------------------------------------- +# Depthwise FFT convolutions — no nvsubquadratic dependency +# --------------------------------------------------------------------------- + + +class _DepthwiseFFTForward: + """Mixin providing FFT-based forward for depthwise ``nn.Conv{2,3}d`` subclasses. + + No ``nn.Module`` parent: module machinery comes from ``nn.Conv2d`` / ``nn.Conv3d`` + in the concrete subclasses. Placed first in the MRO so ``forward`` resolves here + (FFT) rather than to ``nn.Conv{2,3}d.forward`` (im2col / unfold). + + Avoids PyTorch's im2col INT32 overflow, which caps ``F.conv3d`` at ROI ~128 for + typical medical-imaging channel counts. There is no spatial-size restriction. + + ``fft_chunk_size > 0`` enables channel-chunked FFT to cap peak memory: + + peak ≈ (B × chunk × spatial × 4 + B × chunk × rfft_spatial × 8) bytes + + instead of the full ``(B × C × ...)`` allocation. + """ + + _spatial_dims: int # set by subclasses + fft_chunk_size: int = 0 # 0 = no chunking; set in subclass __init__ + + def forward(self, input: torch.Tensor) -> torch.Tensor: + spatial = input.shape[2:] + kernel_shape = self.weight.shape[2:] # type: ignore[attr-defined] + fft_dims = tuple(range(-self._spatial_dims, 0)) + fft_size = [s + k - 1 for s, k in zip(spatial, kernel_shape)] + in_dtype = input.dtype + + slices = (slice(None), slice(None)) + tuple(slice(k // 2, k // 2 + s) for s, k in zip(spatial, kernel_shape)) + + chunk = getattr(self, "fft_chunk_size", 0) + if chunk > 0 and input.shape[1] > chunk: + parts = [] + for c0 in range(0, input.shape[1], chunk): + c1 = min(c0 + chunk, input.shape[1]) + xc = input[:, c0:c1].float() + kc = self.weight[c0:c1].squeeze(1).float() # type: ignore[attr-defined] + kc = kc.flip(list(range(1, self._spatial_dims + 1))) + xc_fft = torch.fft.rfftn(xc, s=fft_size, dim=fft_dims) + kc_fft = torch.fft.rfftn(kc, s=fft_size, dim=fft_dims) + out_fft = xc_fft * kc_fft.unsqueeze(0) + del xc_fft, kc_fft + out_c = torch.fft.irfftn(out_fft, s=fft_size, dim=fft_dims) + del out_fft + parts.append(out_c[slices].to(in_dtype)) + del out_c + return torch.cat(parts, dim=1) + + x_f32 = input.float() + k_f32 = self.weight.squeeze(1).float() # type: ignore[attr-defined] + # PyTorch ``F.conv*`` computes cross-correlation; FFT computes convolution. + # Flip the kernel so the FFT output matches ``Conv{2,3}d`` exactly. + k_f32 = k_f32.flip(list(range(1, self._spatial_dims + 1))) + + x_fft = torch.fft.rfftn(x_f32, s=fft_size, dim=fft_dims) + k_fft = torch.fft.rfftn(k_f32, s=fft_size, dim=fft_dims) + + out_fft = x_fft * k_fft.unsqueeze(0) + out = torch.fft.irfftn(out_fft, s=fft_size, dim=fft_dims) + return cast(torch.Tensor, out[slices].to(in_dtype)) + + +def _validate_depthwise_fft_args( + in_channels: int, out_channels: int, kernel_size: int, groups: int, padding: int, bias: bool +) -> None: + """Validate the constructor arguments shared by ``DepthwiseFFTConv{2,3}d``. + + The FFT forward only implements depthwise, bias-free, ``"same"``-style convolution: it + crops the full convolution back to the input spatial size assuming ``padding == + kernel_size // 2`` with an odd kernel. Reject anything else up front rather than return a + silently wrong shape. + """ + if not (in_channels == out_channels == groups): + raise ValueError( + "DepthwiseFFTConv only supports depthwise (groups == in_channels == out_channels); " + f"got in_channels={in_channels}, out_channels={out_channels}, groups={groups}" + ) + if bias: + raise ValueError("bias is not supported in DepthwiseFFTConv") + if kernel_size % 2 == 0 or padding != kernel_size // 2: + raise ValueError( + "DepthwiseFFTConv only supports 'same'-style padding: kernel_size must be odd and " + f"padding must equal kernel_size // 2; got kernel_size={kernel_size}, padding={padding}. " + "The FFT forward crops to the input spatial size and does not implement general padding." + ) + + +class DepthwiseFFTConv2d(_DepthwiseFFTForward, nn.Conv2d): + """2-D depthwise FFT convolution. ``isinstance(x, nn.Conv2d)`` remains ``True``. + + Drop-in replacement for an ``nn.Conv2d`` with ``groups == in_channels == out_channels`` + and ``bias=False``. Useful as the short-conv inside :class:`HyenaMixer` at large 2-D + inputs, where ``F.conv2d`` would not yet hit the INT32 limit but the unified + Conv2d/Conv3d API is convenient. + + Only ``"same"``-style padding is supported: ``padding`` must equal ``kernel_size // 2`` + (and ``kernel_size`` must be odd). The FFT forward crops its output back to the input + spatial size and does not implement general (e.g. ``"valid"``) padding. + """ + + _spatial_dims = 2 + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + groups: int, + padding: int, + bias: bool = False, + fft_chunk_size: int = 0, + ) -> None: + _validate_depthwise_fft_args(in_channels, out_channels, kernel_size, groups, padding, bias) + nn.Conv2d.__init__( + self, in_channels, out_channels, kernel_size, stride=1, padding=padding, groups=groups, bias=False + ) + self.fft_chunk_size = fft_chunk_size + + +class DepthwiseFFTConv3d(_DepthwiseFFTForward, nn.Conv3d): + """3-D depthwise FFT convolution. ``isinstance(x, nn.Conv3d)`` remains ``True``. + + Drop-in replacement for an ``nn.Conv3d`` with ``groups == in_channels == out_channels`` + and ``bias=False``. Avoids the INT32 unfold limit that prevents ``F.conv3d`` from + running at ROI > ~128 for typical medical-imaging channel counts. + + Only ``"same"``-style padding is supported: ``padding`` must equal ``kernel_size // 2`` + (and ``kernel_size`` must be odd). The FFT forward crops its output back to the input + spatial size and does not implement general (e.g. ``"valid"``) padding. + """ + + _spatial_dims = 3 + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + groups: int, + padding: int, + bias: bool = False, + fft_chunk_size: int = 0, + ) -> None: + _validate_depthwise_fft_args(in_channels, out_channels, kernel_size, groups, padding, bias) + nn.Conv3d.__init__( + self, in_channels, out_channels, kernel_size, stride=1, padding=padding, groups=groups, bias=False + ) + self.fft_chunk_size = fft_chunk_size + + +# --------------------------------------------------------------------------- +# HyenaMixer — QKV-projected gated long-conv mixer +# --------------------------------------------------------------------------- + + +class HyenaMixer(nn.Module): + """QKV-projected gated long-convolution mixer using HyenaND. + + Replaces self-attention with the HyenaND operator from ``nvsubquadratic``, providing + a global receptive field at O(N log N) cost via FFT. ``HyenaMixer`` matches the + channels-last layout ``[B, *spatial, C]`` expected by Swin-style transformer + blocks; the underlying HyenaND operator handles the 2-D / 3-D distinction. + + The HyenaND FFT path requires float32 precision, so the inner mixer call is + wrapped in ``torch.amp.autocast("cuda", enabled=False)``. Inputs are cast to + float32 before the mixer call and back to the original dtype after, so the block + is transparent under ``torch.autocast``. + + Args: + dim: hidden dimension. + spatial_dims: 2 or 3. + use_rope: kept for forward-compatibility with older configurations. + ``nvsubquadratic`` removed RoPE from HyenaND on 2026-04-27 and this kwarg + is now a silent no-op. + apply_qk_norm: whether to apply per-channel LayerNorm to Q (and to K when the + first gate is the identity, as is the default here). + short_conv_kernel_size: depthwise short-convolution kernel size on the + concatenated ``[Q; K; V]`` tensor. + kernel_mlp_hidden_dim: hidden dim of the SIREN implicit kernel MLP. + kernel_num_layers: depth of the SIREN implicit kernel. + kernel_omega_0: SIREN frequency. Default 10.0 (stable). Higher values allow + higher-frequency kernels but reduce training stability. + kernel_l_cache: SIREN coordinate-grid cache size per spatial dim. Memory + scales as ``(2L-1) ** D × D × 4`` bytes; set to ``>= max(spatial_dims)`` + to pre-allocate. + mask_max_attenuation: Gaussian-modulation attenuation at the grid boundary + for the widest channel (0–1). Default 0.95. + fft_padding: ``"circular"`` or ``"zero"``. + grid_type: ``"single"`` (kernel size = input size; required for circular + padding) or ``"double"`` (kernel size = 2× input size; only valid with + ``"zero"`` padding). + use_chunked_fftconv: chunk the FFT convolution by channel to reduce peak + memory ~26% with ~11% compute overhead. Requires ``fft_padding="zero"``. + use_fft_short_conv: replace the depthwise short conv with + :class:`DepthwiseFFTConv{2,3}d`, eliminating the INT32 unfold limit and + enabling unlimited ROI sizes. Adds ~11% compute overhead. + short_conv_fft_chunk_size: channel chunk size for the FFT short conv + (0 = no chunking). + + Raises: + ImportError: if ``nvsubquadratic`` is not installed. + ValueError: on invalid ``spatial_dims`` / ``fft_padding`` / ``grid_type`` + combinations. + """ + + def __init__( + self, + dim: int, + spatial_dims: int = 3, + use_rope: bool = True, + apply_qk_norm: bool = True, + short_conv_kernel_size: int = 3, + kernel_mlp_hidden_dim: int = 32, + kernel_num_layers: int = 3, + kernel_omega_0: float = 10.0, + kernel_l_cache: int = 32, + mask_max_attenuation: float = 0.95, + fft_padding: str = "circular", + grid_type: str = "single", + use_chunked_fftconv: bool = False, + use_fft_short_conv: bool = False, + short_conv_fft_chunk_size: int = 0, + ) -> None: + super().__init__() + + if not _has_nvsubq: + raise ImportError(_NVSUBQ_INSTALL_HINT) + + if fft_padding not in ("circular", "zero"): + raise ValueError(f"fft_padding must be 'circular' or 'zero', got '{fft_padding}'") + if grid_type not in ("single", "double"): + raise ValueError(f"grid_type must be 'single' or 'double', got '{grid_type}'") + if fft_padding == "circular" and grid_type != "single": + raise ValueError( + "fft_padding='circular' requires grid_type='single' " + f"(kernel size must match input size for periodic convolution); got grid_type='{grid_type}'" + ) + if use_chunked_fftconv and fft_padding != "zero": + raise ValueError( + "use_chunked_fftconv=True requires fft_padding='zero'; " f"got fft_padding='{fft_padding}'" + ) + + self.dim = dim + self.spatial_dims = spatial_dims + + conv_class: type[nn.Module] + if use_fft_short_conv: + if spatial_dims == 2: + conv_class = DepthwiseFFTConv2d + elif spatial_dims == 3: + conv_class = DepthwiseFFTConv3d + else: + raise ValueError(f"spatial_dims must be 2 or 3, got {spatial_dims}") + else: + if spatial_dims == 2: + conv_class = nn.Conv2d + elif spatial_dims == 3: + conv_class = nn.Conv3d + else: + raise ValueError(f"spatial_dims must be 2 or 3, got {spatial_dims}") + + global_conv_cfg = _LazyConfig(_CKConvND)( + data_dim=spatial_dims, + hidden_dim=dim, + kernel_cfg=_LazyConfig(_SIRENKernelND)( + data_dim=spatial_dims, + out_dim=dim, + mlp_hidden_dim=kernel_mlp_hidden_dim, + num_layers=kernel_num_layers, + embedding_dim=kernel_mlp_hidden_dim, + omega_0=kernel_omega_0, + L_cache=kernel_l_cache, + use_bias=True, + hidden_omega_0=1.0, + ), + mask_cfg=_LazyConfig(_GaussianModulationND)( + data_dim=spatial_dims, + num_channels=dim, + min_attenuation_at_step=0.1, + max_attenuation_at_limit=mask_max_attenuation, + init_extent=1.0, + parametrization="direct", + ), + grid_type=grid_type, + fft_padding=fft_padding, + use_chunked_fftconv=use_chunked_fftconv, + ) + + short_conv_kwargs: dict = dict( + in_channels=3 * dim, + out_channels=3 * dim, + kernel_size=short_conv_kernel_size, + groups=3 * dim, + padding=short_conv_kernel_size // 2, + bias=False, + ) + if use_fft_short_conv and short_conv_fft_chunk_size > 0: + short_conv_kwargs["fft_chunk_size"] = short_conv_fft_chunk_size + short_conv_cfg = _LazyConfig(conv_class)(**short_conv_kwargs) + + # ``use_rope`` retained on the API only; ``nvsubquadratic`` removed RoPE + # from ``Hyena.__init__`` on 2026-04-27. Saving the flag here keeps caller + # introspection intact ("did the user ask for RoPE?") while not affecting + # the constructed operator. + self._use_rope_requested = use_rope + + self.mixer = _instantiate( + _LazyConfig(_Hyena)( + global_conv_cfg=global_conv_cfg, + short_conv_cfg=short_conv_cfg, + gate_nonlinear_cfg=_LazyConfig(nn.Identity)(), + pixelhyena_norm_cfg=_LazyConfig(nn.GroupNorm)(num_groups=1, num_channels=dim), + qk_norm_cfg=_LazyConfig(nn.LayerNorm)(normalized_shape=dim) if apply_qk_norm else None, + ) + ) + + self.qkv_proj = nn.Linear(dim, 3 * dim, bias=False) + self.out_proj = nn.Linear(dim, dim, bias=False) + self._init_weights() + + def _init_weights(self) -> None: + nn.init.normal_(self.qkv_proj.weight, std=0.02) + nn.init.normal_(self.out_proj.weight, std=0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass. + + Args: + x: tensor of shape ``[batch, *spatial, dim]``. + + Returns: + Tensor of the same shape and dtype as the input. + """ + qkv = self.qkv_proj(x) + q, k, v = torch.chunk(qkv, 3, dim=-1) + + # HyenaND requires float32 internally; disable autocast for the mixer only, + # then restore the original (autocast) dtype afterwards. + with torch.amp.autocast("cuda", enabled=False): + q = q.float() + k = k.float() + v = v.float() + x = self.mixer(q, k, v) + x = x.to(qkv.dtype) + return cast(torch.Tensor, self.out_proj(x)) + + +# --------------------------------------------------------------------------- +# HyenaTransformerBlock — pre-norm Hyena + MLP residual +# --------------------------------------------------------------------------- + + +class HyenaTransformerBlock(nn.Module): + """Pre-norm transformer block with HyenaND in place of self-attention. + + Sandwiches a :class:`HyenaMixer` and an MLP between :class:`~torch.nn.LayerNorm` + layers in the standard transformer residual pattern, with optional gradient + checkpointing on each half. + + Args: + dim: number of feature channels. + spatial_dims: 2 or 3. + mlp_ratio: hidden / input ratio for the MLP. + drop: dropout rate inside the MLP. + drop_path: stochastic-depth rate for the MLP residual. + act_layer: activation name passed to :class:`monai.networks.blocks.MLPBlock`. + norm_layer: normalization class (default :class:`~torch.nn.LayerNorm`). + use_checkpoint: enable gradient checkpointing on the mixer and MLP halves. + use_rope: forward-compatibility flag for older configs; no-op after the + ``nvsubquadratic`` 2026-04-27 RoPE removal. + apply_qk_norm: per-channel LayerNorm on Q (and K when the first gate is + identity, as is the default). + hyena_kernel_size: short-convolution kernel size on the QKV tensor. + hyena_kernel_mlp_dim: SIREN kernel MLP hidden dimension. + hyena_kernel_layers: SIREN kernel depth. + hyena_mask_max_attenuation: Gaussian-modulation boundary attenuation (0–1). + hyena_fft_padding: ``"circular"`` or ``"zero"``. + hyena_grid_type: ``"single"`` or ``"double"``. + hyena_use_chunked_fft: enable chunked FFT (requires zero padding). + hyena_use_fft_short_conv: use FFT for the short conv (no INT32 limit). + hyena_omega_0: SIREN ``omega_0``. Default 10.0. + hyena_l_cache: SIREN coordinate-grid cache size per dim. + hyena_short_conv_fft_chunks: channel chunk size for the FFT short conv. + + Raises: + ImportError: if ``nvsubquadratic`` is not installed. + """ + + def __init__( + self, + dim: int, + spatial_dims: int = 3, + mlp_ratio: float = 4.0, + drop: float = 0.0, + drop_path: float = 0.0, + act_layer: str = "GELU", + norm_layer: type[LayerNorm] = nn.LayerNorm, # type: ignore[assignment] + use_checkpoint: bool = False, + use_rope: bool = True, + apply_qk_norm: bool = True, + hyena_kernel_size: int = 3, + hyena_kernel_mlp_dim: int = 32, + hyena_kernel_layers: int = 3, + hyena_mask_max_attenuation: float = 0.95, + hyena_fft_padding: str = "circular", + hyena_grid_type: str = "single", + hyena_use_chunked_fft: bool = False, + hyena_use_fft_short_conv: bool = False, + hyena_omega_0: float = 10.0, + hyena_l_cache: int = 32, + hyena_short_conv_fft_chunks: int = 0, + ) -> None: + super().__init__() + self.dim = dim + self.spatial_dims = spatial_dims + self.mlp_ratio = mlp_ratio + self.use_checkpoint = use_checkpoint + + self.norm1 = norm_layer(dim) + self.mixer = HyenaMixer( + dim=dim, + spatial_dims=spatial_dims, + use_rope=use_rope, + apply_qk_norm=apply_qk_norm, + short_conv_kernel_size=hyena_kernel_size, + kernel_mlp_hidden_dim=hyena_kernel_mlp_dim, + kernel_num_layers=hyena_kernel_layers, + kernel_omega_0=hyena_omega_0, + kernel_l_cache=hyena_l_cache, + mask_max_attenuation=hyena_mask_max_attenuation, + fft_padding=hyena_fft_padding, + grid_type=hyena_grid_type, + use_chunked_fftconv=hyena_use_chunked_fft, + use_fft_short_conv=hyena_use_fft_short_conv, + short_conv_fft_chunk_size=hyena_short_conv_fft_chunks, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = MLPBlock( + hidden_size=dim, mlp_dim=mlp_hidden_dim, act=act_layer, dropout_rate=drop, dropout_mode="swin" + ) + + def forward_part1(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm1(x) + return cast(torch.Tensor, self.mixer(x)) + + def forward_part2(self, x: torch.Tensor) -> torch.Tensor: + return cast(torch.Tensor, self.drop_path(self.mlp(self.norm2(x)))) + + def forward(self, x: torch.Tensor, mask_matrix: torch.Tensor | None = None) -> torch.Tensor: + """Forward pass. + + Args: + x: input tensor of shape ``[batch, *spatial, dim]``. + mask_matrix: unused; accepted for signature parity with Swin's + ``WindowAttention``-based block so the two can be swapped at the + ``BasicLayer`` level without per-call branching. + + Returns: + Tensor of the same shape as the input. + """ + del mask_matrix + shortcut = x + if self.use_checkpoint: + x = torch.utils.checkpoint.checkpoint(self.forward_part1, x, use_reentrant=False) + else: + x = self.forward_part1(x) + x = shortcut + self.drop_path(x) + + if self.use_checkpoint: + x = x + torch.utils.checkpoint.checkpoint(self.forward_part2, x, use_reentrant=False) + else: + x = x + self.forward_part2(x) + return x diff --git a/monai/networks/blocks/text_embedding.py b/monai/networks/blocks/text_embedding.py index 6f2990e35c1..cae64d92c7c 100644 --- a/monai/networks/blocks/text_embedding.py +++ b/monai/networks/blocks/text_embedding.py @@ -67,7 +67,7 @@ def __init__( if pretrained: model_url = url_map[self.encoding] - pretrain_state_dict = model_zoo.load_url(model_url, map_location="cpu") + pretrain_state_dict = model_zoo.load_url(model_url, map_location="cpu", weights_only=True) self.text_embedding.data = pretrain_state_dict.float() # type: ignore else: print(f"{self.encoding} is not implemented, and can not be downloaded, please load your own") @@ -79,7 +79,6 @@ def forward(self): # text embedding as random initialized 'rand_embedding' text_embedding = self.text_embedding.weight else: - print(self.text_embedding) text_embedding = nn.functional.relu(self.text_to_vision(self.text_embedding)) if self.spatial_dims == 3: diff --git a/monai/networks/blocks/warp.py b/monai/networks/blocks/warp.py index ddd3a350d56..1878662916f 100644 --- a/monai/networks/blocks/warp.py +++ b/monai/networks/blocks/warp.py @@ -121,12 +121,13 @@ def get_reference_grid(self, ddf: torch.Tensor, jitter: bool = False, seed: int mesh_points = [torch.arange(0, dim) for dim in ddf.shape[2:]] grid = torch.stack(meshgrid_ij(*mesh_points), dim=0) # (spatial_dims, ...) grid = torch.stack([grid] * ddf.shape[0], dim=0) # (batch, spatial_dims, ...) - self.ref_grid = grid.to(ddf) + grid = grid.to(ddf) if jitter: # Define reference grid on non-integer values - with torch.random.fork_rng(enabled=seed): + with torch.random.fork_rng(): torch.random.manual_seed(seed) grid += torch.rand_like(grid) + self.ref_grid = grid self.ref_grid.requires_grad = False return self.ref_grid @@ -155,7 +156,10 @@ def forward(self, image: torch.Tensor, ddf: torch.Tensor): if not _use_compiled: # pytorch native grid_sample for i, dim in enumerate(grid.shape[1:-1]): - grid[..., i] = grid[..., i] * 2 / (dim - 1) - 1 + # guard against a singleton spatial dim (e.g. a single-slice volume), where + # ``dim - 1 == 0`` would divide by zero; clamp the denominator to 1 so the lone + # voxel maps to -1, matching ``monai.networks.utils.normalize_transform``. + grid[..., i] = grid[..., i] * 2 / max(dim - 1, 1) - 1 index_ordering: list[int] = list(range(spatial_dims - 1, -1, -1)) grid = grid[..., index_ordering] # z, y, x -> x, y, z return F.grid_sample( diff --git a/monai/networks/layers/filtering.py b/monai/networks/layers/filtering.py index 5c6647f6211..db86cff04d8 100644 --- a/monai/networks/layers/filtering.py +++ b/monai/networks/layers/filtering.py @@ -96,9 +96,6 @@ def forward(ctx, input, features, sigmas=None): @staticmethod def backward(ctx, grad_output): raise NotImplementedError("PHLFilter does not currently support Backpropagation") - # scaled_features, = ctx.saved_variables - # grad_input = _C.phl_filter(grad_output, scaled_features) - # return grad_input class TrainableBilateralFilterFunction(torch.autograd.Function): diff --git a/monai/networks/layers/simplelayers.py b/monai/networks/layers/simplelayers.py index 56f7192e4df..f044b6f3a79 100644 --- a/monai/networks/layers/simplelayers.py +++ b/monai/networks/layers/simplelayers.py @@ -671,7 +671,6 @@ def __init__(self, spatial_dims: int, size: int) -> None: size: edge length of the filter """ filter = torch.ones([size] * spatial_dims) - filter = filter super().__init__(filter=filter) diff --git a/monai/networks/nets/__init__.py b/monai/networks/nets/__init__.py index c1917e5293a..d33f4dbf9cd 100644 --- a/monai/networks/nets/__init__.py +++ b/monai/networks/nets/__init__.py @@ -53,6 +53,7 @@ from .generator import Generator from .highresnet import HighResBlock, HighResNet from .hovernet import Hovernet, HoVernet, HoVerNet, HoverNet +from .hyena_nd_unetr import HyenaNDUNETR from .masked_autoencoder_vit import MaskedAutoEncoderViT from .mednext import ( MedNeXt, @@ -74,6 +75,7 @@ MedNextSmall, ) from .milmodel import MILModel +from .navit import NaViT from .netadapter import NetAdapter from .patchgan_discriminator import MultiScalePatchDiscriminator, PatchDiscriminator from .quicknat import Quicknat diff --git a/monai/networks/nets/basic_unet.py b/monai/networks/nets/basic_unet.py index d2a655f981a..3b47fa0b030 100644 --- a/monai/networks/nets/basic_unet.py +++ b/monai/networks/nets/basic_unet.py @@ -235,7 +235,6 @@ def __init__( """ super().__init__() fea = ensure_tuple_rep(features, 6) - print(f"BasicUNet features: {fea}.") self.conv_0 = TwoConv(spatial_dims, in_channels, features[0], act, norm, bias, dropout) self.down_1 = Down(spatial_dims, fea[0], fea[1], act, norm, bias, dropout) diff --git a/monai/networks/nets/basic_unetplusplus.py b/monai/networks/nets/basic_unetplusplus.py index f7ae7685137..dc5711b0bd6 100644 --- a/monai/networks/nets/basic_unetplusplus.py +++ b/monai/networks/nets/basic_unetplusplus.py @@ -94,7 +94,6 @@ def __init__( self.deep_supervision = deep_supervision fea = ensure_tuple_rep(features, 6) - print(f"BasicUNetPlusPlus features: {fea}.") self.conv_0_0 = TwoConv(spatial_dims, in_channels, fea[0], act, norm, bias, dropout) self.conv_1_0 = Down(spatial_dims, fea[0], fea[1], act, norm, bias, dropout) diff --git a/monai/networks/nets/densenet.py b/monai/networks/nets/densenet.py index 42463b2493c..7e9c7ab5a85 100644 --- a/monai/networks/nets/densenet.py +++ b/monai/networks/nets/densenet.py @@ -277,7 +277,7 @@ def _load_state_dict(model: nn.Module, arch: str, progress: bool): r"^(.*denselayer\d+)(\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$" ) - state_dict = load_state_dict_from_url(model_url, progress=progress) + state_dict = load_state_dict_from_url(model_url, progress=progress, weights_only=True) for key in list(state_dict.keys()): res = pattern.match(key) if res: diff --git a/monai/networks/nets/dints.py b/monai/networks/nets/dints.py index 98f8de57722..ad4b350d30e 100644 --- a/monai/networks/nets/dints.py +++ b/monai/networks/nets/dints.py @@ -36,20 +36,23 @@ __all__ = ["DiNTS", "TopologyConstruction", "TopologyInstance", "TopologySearch"] -@torch.jit.interface -class CellInterface(torch.nn.Module): - """interface for torchscriptable Cell""" +# TODO: added temporarily for PyTorch 2.14 warnings, remove when factoring out deprecated Torchscript components +with warnings.catch_warnings(): + warnings.simplefilter("ignore") - def forward(self, x: torch.Tensor, weight: torch.Tensor | None) -> torch.Tensor: # type: ignore - pass + @torch.jit.interface + class CellInterface(torch.nn.Module): + """interface for torchscriptable Cell""" + def forward(self, x: torch.Tensor, weight: torch.Tensor | None) -> torch.Tensor: # type: ignore + pass -@torch.jit.interface -class StemInterface(torch.nn.Module): - """interface for torchscriptable Stem""" + @torch.jit.interface + class StemInterface(torch.nn.Module): + """interface for torchscriptable Stem""" - def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore - pass + def forward(self, x: torch.Tensor) -> torch.Tensor: # type: ignore + pass class StemTS(StemInterface): diff --git a/monai/networks/nets/efficientnet.py b/monai/networks/nets/efficientnet.py index e9b7675144c..e8b510e47d4 100644 --- a/monai/networks/nets/efficientnet.py +++ b/monai/networks/nets/efficientnet.py @@ -793,7 +793,7 @@ def _load_state_dict(model: nn.Module, arch: str, progress: bool, adv_prop: bool else: # load state dict from url model_url = url_map[arch] - pretrain_state_dict = model_zoo.load_url(model_url, progress=progress) + pretrain_state_dict = model_zoo.load_url(model_url, progress=progress, weights_only=True) model_state_dict = model.state_dict() pattern = re.compile(r"(.+)\.\d+(\.\d+\..+)") diff --git a/monai/networks/nets/flexible_unet.py b/monai/networks/nets/flexible_unet.py index c27b0fc17b9..bc244ba6e7a 100644 --- a/monai/networks/nets/flexible_unet.py +++ b/monai/networks/nets/flexible_unet.py @@ -61,9 +61,13 @@ def register_class(self, name: type[Any] | str): "or implement all interfaces specified by it." ) + # pyrefly: ignore [missing-attribute] name_string_list = name.get_encoder_names() + # pyrefly: ignore [missing-attribute] feature_number_list = name.num_outputs() + # pyrefly: ignore [missing-attribute] feature_channel_list = name.num_channels_per_output() + # pyrefly: ignore [missing-attribute] parameter_list = name.get_encoder_parameters() assert len(name_string_list) == len(feature_number_list) == len(feature_channel_list) == len(parameter_list) diff --git a/monai/networks/nets/hovernet.py b/monai/networks/nets/hovernet.py index f0cb5ab74dc..1d652782faf 100644 --- a/monai/networks/nets/hovernet.py +++ b/monai/networks/nets/hovernet.py @@ -632,7 +632,7 @@ def _remap_preact_resnet_model(model_url: str): pattern_bna = re.compile(r"^(.+\.d\d+)\.blk_bna\.(.+)") # download the pretrained weights into torch hub's default dir weights_dir = os.path.join(torch.hub.get_dir(), "preact-resnet50.pth") - download_url(model_url, fuzzy=True, filepath=weights_dir, progress=False) + download_url(model_url, filepath=weights_dir, progress=False) map_location = None if torch.cuda.is_available() else torch.device("cpu") state_dict = torch.load(weights_dir, map_location=map_location, weights_only=True)["desc"] @@ -667,7 +667,7 @@ def _remap_standard_resnet_model(model_url: str, state_dict_key: str | None = No pattern_downsample1 = re.compile(r"^(res_blocks.d\d+).+\.downsample\.1\.(.+)") # download the pretrained weights into torch hub's default dir weights_dir = os.path.join(torch.hub.get_dir(), "resnet50.pth") - download_url(model_url, fuzzy=True, filepath=weights_dir, progress=False) + download_url(model_url, filepath=weights_dir, progress=False) map_location = None if torch.cuda.is_available() else torch.device("cpu") state_dict = torch.load(weights_dir, map_location=map_location, weights_only=True) if state_dict_key is not None: diff --git a/monai/networks/nets/hyena_nd_unetr.py b/monai/networks/nets/hyena_nd_unetr.py new file mode 100644 index 00000000000..b6d8c120acc --- /dev/null +++ b/monai/networks/nets/hyena_nd_unetr.py @@ -0,0 +1,151 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +HyenaNDUNETR: SwinUNETR with the HyenaND subquadratic operator in place of (or +mixed with) windowed self-attention. + +This is a thin convenience subclass of :class:`monai.networks.nets.SwinUNETR` whose +defaults make Hyena placement explicit: + +* ``use_hyena`` is forced to ``True``. +* ``hyena_stages`` is **required** -- callers must explicitly declare which of the + four Swin stages run HyenaND vs windowed attention. + +The classmethod :meth:`HyenaNDUNETR.get_variant` provides the three Hyena +variants from Table 4 of the NeurIPS 2026 paper "Native Multi-Dimensional Subquadratic +Operators via Input Dependent Long Convolutions" (paper id 26539): + +========== ================================= ============================== +Variant ``hyena_stages`` Notes +========== ================================= ============================== +``HHHH`` ``(True, True, True, True)`` Hyena at every Swin stage +``HAHA`` ``(True, False, True, False)`` striped/interleaved +``HHAA`` ``(True, True, False, False)`` paper-best (outer Hyena, inner attention) +========== ================================= ============================== + +``AAAA`` (pure attention) is intentionally not exposed here -- it is plain +:class:`SwinUNETR` and constructing a "HyenaNDUNETR" with no Hyena stages would be a +contradiction. + +Requires the optional ``nvsubquadratic`` package; install with +``pip install monai[hyena]``. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from monai.networks.nets.swin_unetr import SwinUNETR + +__all__ = ["HyenaNDUNETR"] + + +# Per-stage Hyena patterns for the three paper variants. +PAPER_VARIANTS: dict[str, tuple[bool, ...]] = { + "HHHH": (True, True, True, True), + "HAHA": (True, False, True, False), + "HHAA": (True, True, False, False), +} + + +class HyenaNDUNETR(SwinUNETR): + """SwinUNETR with HyenaND replacing windowed self-attention at selected stages. + + See the module docstring for the paper-variant table and + :meth:`get_variant` for a convenience constructor matching the NeurIPS + 2026 paper. All other kwargs are forwarded to :class:`SwinUNETR`. + + Args: + in_channels: dimension of input channels. + out_channels: dimension of output channels. + hyena_stages: required 4-tuple of bools, one per Swin stage. At least one + element must be ``True`` (otherwise use :class:`SwinUNETR` directly). + feature_size: dimension of network feature size. Must be a multiple of 12 + (inherited from :class:`SwinUNETR`). + **kwargs: forwarded to :class:`SwinUNETR`. + + Raises: + ValueError: if ``hyena_stages`` is missing, has the wrong length, or has no + ``True`` element. + ImportError: if the optional ``nvsubquadratic`` package is not installed. + """ + + def __init__( + self, in_channels: int, out_channels: int, hyena_stages: Sequence[bool], feature_size: int = 48, **kwargs + ) -> None: + if hyena_stages is None: + raise ValueError( + "HyenaNDUNETR requires `hyena_stages` (a 4-tuple of bools); " + "use SwinUNETR directly for pure attention." + ) + stages_tuple = tuple(bool(s) for s in hyena_stages) + if len(stages_tuple) != 4: + raise ValueError( + f"hyena_stages must have length 4 (one bool per Swin stage); got length {len(stages_tuple)}." + ) + if not any(stages_tuple): + raise ValueError( + "hyena_stages must enable HyenaND at at least one stage; " "use SwinUNETR directly for pure attention." + ) + + # ``use_hyena`` is forced True here; reject it in kwargs rather than silently + # override -- the subclass exists to make Hyena placement explicit. (``hyena_stages`` + # is an explicit parameter above, so it can never reach ``kwargs``.) + if "use_hyena" in kwargs: + raise TypeError("HyenaNDUNETR forces use_hyena=True; do not pass use_hyena via kwargs.") + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + feature_size=feature_size, + use_hyena=True, + hyena_stages=stages_tuple, + **kwargs, + ) + + @classmethod + def get_variant(cls, variant: str, **kwargs) -> HyenaNDUNETR: + """Build a :class:`HyenaNDUNETR` matching one of the NeurIPS 2026 paper variants. + + Args: + variant: one of ``"HHHH"``, ``"HAHA"``, ``"HHAA"`` (case-insensitive). + **kwargs: forwarded to :class:`HyenaNDUNETR.__init__`. Must include at + least ``in_channels`` and ``out_channels``. Must NOT include + ``hyena_stages`` (set by the variant). + + Returns: + A :class:`HyenaNDUNETR` with ``hyena_stages`` set per the variant. + + Raises: + ValueError: if ``variant`` is not one of the three known names, or if + ``hyena_stages`` is also passed via kwargs. + + Example:: + + >>> net = HyenaNDUNETR.get_variant( + ... "HHAA", + ... in_channels=1, + ... out_channels=29, + ... feature_size=48, + ... ) + """ + key = variant.upper() + if key not in PAPER_VARIANTS: + raise ValueError( + f"Unknown paper variant '{variant}'. " + f"Known variants: {sorted(PAPER_VARIANTS)}. " + "(AAAA is plain SwinUNETR; use that class directly.)" + ) + if "hyena_stages" in kwargs: + raise ValueError( + "get_variant sets hyena_stages from the variant name; " "do not also pass hyena_stages via kwargs." + ) + return cls(hyena_stages=PAPER_VARIANTS[key], **kwargs) diff --git a/monai/networks/nets/milmodel.py b/monai/networks/nets/milmodel.py index a31f1051106..a0f27008da0 100644 --- a/monai/networks/nets/milmodel.py +++ b/monai/networks/nets/milmodel.py @@ -160,6 +160,7 @@ def hook(module, input, output): ] ) self.transformer = transformer_list + # pyrefly: ignore [unsupported-operation] nfc = nfc + 256 self.attention = nn.Sequential(nn.Linear(nfc, 2048), nn.Tanh(), nn.Linear(2048, 1)) diff --git a/monai/networks/nets/navit.py b/monai/networks/nets/navit.py new file mode 100644 index 00000000000..6342d96e79e --- /dev/null +++ b/monai/networks/nets/navit.py @@ -0,0 +1,585 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from functools import partial +from typing import cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.nn.utils.rnn import pad_sequence as orig_pad_sequence + +from monai.networks.blocks.mlp import MLPBlock +from monai.utils import ensure_tuple_rep, optional_import + +rearrange, has_einops = optional_import("einops", name="rearrange") +repeat, _ = optional_import("einops", name="repeat") + +__all__ = ["NaViT"] + + +def _group_images_by_max_seq_len( + images: list[Tensor], patch_size: int, calc_token_dropout: Callable | None = None, max_seq_len: int = 2048 +) -> list[list[Tensor]]: + """Group a flat list of variable-size images into packed batches that each fit within ``max_seq_len`` tokens. + + Args: + images: flat list of image tensors, each of shape ``(C, *spatial)``. + patch_size: isotropic patch size used to compute the number of tokens per image. + calc_token_dropout: optional callable ``(*spatial_dims) -> float`` returning the fraction of tokens + to drop for an image of the given spatial size. If ``None``, no dropout is assumed. + max_seq_len: maximum number of tokens allowed per packed group. + + Returns: + List of groups, where each group is a list of image tensors that together fit within ``max_seq_len``. + """ + groups: list[list[Tensor]] = [] + group: list[Tensor] = [] + seq_len = 0 + + for image in images: + assert isinstance(image, Tensor) + spatial_dims = image.shape[1:] + num_patches = 1 + for d in spatial_dims: + num_patches *= d // patch_size + + image_seq_len = num_patches + if calc_token_dropout is not None: + image_seq_len = max(1, int(image_seq_len * (1.0 - calc_token_dropout(*spatial_dims)))) + + if image_seq_len > max_seq_len: + raise ValueError( + f"Image with spatial dimensions {spatial_dims} produces {image_seq_len} tokens, " + f"which exceeds max_seq_len={max_seq_len}." + ) + + if (seq_len + image_seq_len) > max_seq_len: + groups.append(group) + group = [] + seq_len = 0 + + group.append(image) + seq_len += image_seq_len + + if group: + groups.append(group) + + return groups + + +class _RMSNorm(nn.Module): + """Per-head RMS normalization applied to query and key tensors. + + Equivalent to the QK-norm introduced in ViT-22B + (Dehghani et al., https://arxiv.org/abs/2302.05442). + + Args: + num_heads: number of attention heads. + dim_head: dimension of each head. + """ + + def __init__(self, num_heads: int, dim_head: int) -> None: + super().__init__() + self.scale = dim_head**0.5 + self.gamma = nn.Parameter(torch.ones(num_heads, 1, dim_head)) + + def forward(self, x: Tensor) -> Tensor: + return cast(Tensor, F.normalize(x, dim=-1) * self.scale * self.gamma) + + +class _NaViTAttention(nn.Module): + """Multi-head attention with QK-normalization and support for packed-sequence attention masks. + + This block is used both for the main transformer layers (self-attention) and for the final + attention-pooling step (cross-attention between learned queries and patch tokens). + + Args: + hidden_size: dimension of the token embeddings. + num_heads: number of attention heads. + dim_head: dimension of each head. Defaults to ``hidden_size // num_heads``. + dropout_rate: dropout probability applied to attention weights and output projection. + qkv_bias: whether to add a bias term to the QKV linear projections. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + dim_head: int | None = None, + dropout_rate: float = 0.0, + qkv_bias: bool = False, + ) -> None: + super().__init__() + + if not (0 <= dropout_rate <= 1): + raise ValueError("dropout_rate should be between 0 and 1.") + + self.num_heads = num_heads + self.dim_head = dim_head if dim_head is not None else hidden_size // num_heads + inner_dim = self.num_heads * self.dim_head + + self.norm = nn.LayerNorm(hidden_size) + self.q_norm = _RMSNorm(num_heads, self.dim_head) + self.k_norm = _RMSNorm(num_heads, self.dim_head) + + self.to_q = nn.Linear(hidden_size, inner_dim, bias=qkv_bias) + self.to_k = nn.Linear(hidden_size, inner_dim, bias=qkv_bias) + self.to_v = nn.Linear(hidden_size, inner_dim, bias=qkv_bias) + self.out_proj = nn.Linear(inner_dim, hidden_size, bias=False) + + self.drop_weights = nn.Dropout(dropout_rate) + self.drop_output = nn.Dropout(dropout_rate) + self.scale = self.dim_head**-0.5 + + def forward(self, x: Tensor, context: Tensor | None = None, attn_mask: Tensor | None = None) -> Tensor: + """ + Args: + x: query tensor of shape ``(B, N, C)``. + context: key/value source tensor of shape ``(B, M, C)``. When ``None``, self-attention is performed. + attn_mask: boolean mask of shape ``(B, 1, N, M)`` where ``True`` indicates positions that + **should** be attended to. Positions with ``False`` are masked out (set to ``-inf``). + + Returns: + Tensor of shape ``(B, N, C)``. + """ + x = self.norm(x) + kv_src = context if context is not None else x + + # project and reshape to (B, heads, seq, dim_head) + q = self.to_q(x).unflatten(-1, (self.num_heads, self.dim_head)).transpose(1, 2) + k = self.to_k(kv_src).unflatten(-1, (self.num_heads, self.dim_head)).transpose(1, 2) + v = self.to_v(kv_src).unflatten(-1, (self.num_heads, self.dim_head)).transpose(1, 2) + + # QK normalization for training stability + q = self.q_norm(q) + k = self.k_norm(k) + + dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale + + if attn_mask is not None: + dots = dots.masked_fill(~attn_mask, -torch.finfo(dots.dtype).max) + + attn = self.drop_weights(dots.softmax(dim=-1)) + out = torch.matmul(attn, v) # (B, heads, N, dim_head) + out = out.transpose(1, 2).flatten(-2) # (B, N, inner_dim) + return cast(Tensor, self.drop_output(self.out_proj(out))) + + +class _NaViTTransformerBlock(nn.Module): + """Single NaViT transformer block: pre-norm self-attention followed by pre-norm MLP. + + Args: + hidden_size: token embedding dimension. + mlp_dim: hidden dimension of the feed-forward network. + num_heads: number of attention heads. + dim_head: per-head dimension. Defaults to ``hidden_size // num_heads``. + dropout_rate: dropout probability. + qkv_bias: whether to add bias to QKV projections. + """ + + def __init__( + self, + hidden_size: int, + mlp_dim: int, + num_heads: int, + dim_head: int | None = None, + dropout_rate: float = 0.0, + qkv_bias: bool = False, + ) -> None: + super().__init__() + self.attn = _NaViTAttention(hidden_size, num_heads, dim_head, dropout_rate, qkv_bias) + self.norm = nn.LayerNorm(hidden_size) + self.mlp = MLPBlock(hidden_size, mlp_dim, dropout_rate) + + def forward(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor: + """ + Args: + x: input tensor of shape ``(B, N, C)``. + attn_mask: packed-sequence attention mask of shape ``(B, 1, N, N)``. + + Returns: + Tensor of shape ``(B, N, C)``. + """ + x = self.attn(x, attn_mask=attn_mask) + x + x = self.mlp(self.norm(x)) + x + return x + + +class NaViT(nn.Module): + """NaViT: Native Resolution Vision Transformer with Patch n' Pack, extended to 2D and 3D. + + Based on: "Dehghani et al., Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution + " + + NaViT removes the fixed-resolution constraint of standard ViT by packing multiple variable-size images + into a single sequence. Key features: + + - **Patch n' Pack**: multiple images (possibly different resolutions) are concatenated into one sequence + per batch element, separated by a per-image attention mask. + - **Factorized positional embeddings**: separate learnable embeddings for each spatial axis are summed, + allowing generalization to unseen resolutions. + - **Token dropout**: a configurable fraction of patch tokens can be randomly dropped during training, + acting as a form of masked-image modelling. + - **Attention pooling**: a learned query vector attends over each image's tokens to produce a fixed-size + per-image representation, cleanly handling variable numbers of images per packed sequence. + - **QK normalization**: RMS normalization on queries and keys for training stability (ViT-22B). + - **spatial_dims support**: works for 2D images ``(C, H, W)`` and 3D volumes ``(C, H, W, D)``. + + Args: + image_size (Union[Sequence[int], int]): reference spatial size used to derive the positional embedding + tables. Each axis gets a table of size ``image_size[i] // patch_size``. Images at inference time + may differ from this size as long as each dimension is divisible by ``patch_size``. + patch_size (int): isotropic patch size. All spatial dimensions of every input image must be divisible + by this value. + num_classes (int): number of output classes for the classification head. + hidden_size (int): token embedding dimension. + mlp_dim (int): hidden dimension of the feed-forward network inside each transformer block. + num_layers (int): number of transformer blocks. + num_heads (int): number of attention heads. + in_channels (int): number of input image channels. Defaults to 1 (grayscale medical images). + dim_head (int, optional): per-head dimension. Defaults to ``hidden_size // num_heads``. + dropout_rate (float): dropout probability applied inside transformer blocks. Defaults to 0.0. + emb_dropout_rate (float): dropout probability applied to patch embeddings. Defaults to 0.0. + token_dropout_prob (Union[float, Callable, None]): fraction of patch tokens to randomly drop during + training. Accepts a float in ``(0, 1)``, or a callable ``(*spatial_dims) -> float`` for + resolution-dependent dropout. ``None`` disables token dropout. Defaults to ``None``. + spatial_dims (int): number of spatial dimensions, either 2 or 3. Defaults to 3. + qkv_bias (bool): whether to add bias to QKV projections. Defaults to False. + + Raises: + ValueError: When ``spatial_dims`` is not 2 or 3. + ValueError: When ``dropout_rate`` or ``emb_dropout_rate`` is outside ``[0, 1]``. + ValueError: When ``hidden_size`` is not divisible by ``num_heads``. + ValueError: When ``token_dropout_prob`` is a float outside ``(0, 1)``. + + Examples:: + + # 3D single-channel (e.g. CT) classification with variable-resolution volumes + >>> net = NaViT( + ... image_size=96, patch_size=16, num_classes=2, + ... hidden_size=768, mlp_dim=3072, num_layers=12, num_heads=12, + ... in_channels=1, spatial_dims=3, + ... ) + >>> volumes = [ + ... [torch.randn(1, 96, 96, 96), torch.randn(1, 64, 64, 64)], + ... [torch.randn(1, 80, 96, 80)], + ... ] + >>> logits = net(volumes) # shape: (3, 2) + + # 2D RGB classification (e.g. pathology patches) with token dropout + >>> net2d = NaViT( + ... image_size=256, patch_size=32, num_classes=10, + ... hidden_size=512, mlp_dim=2048, num_layers=6, num_heads=8, + ... in_channels=3, spatial_dims=2, token_dropout_prob=0.1, + ... ) + >>> images = [ + ... [torch.randn(3, 256, 256), torch.randn(3, 128, 128)], + ... [torch.randn(3, 192, 256)], + ... ] + >>> logits2d = net2d(images) # shape: (3, 10) + """ + + def __init__( + self, + image_size: Sequence[int] | int, + patch_size: int, + num_classes: int, + hidden_size: int, + mlp_dim: int, + num_layers: int, + num_heads: int, + in_channels: int = 1, + dim_head: int | None = None, + dropout_rate: float = 0.0, + emb_dropout_rate: float = 0.0, + token_dropout_prob: float | Callable | None = None, + spatial_dims: int = 3, + qkv_bias: bool = False, + ) -> None: + super().__init__() + + if spatial_dims not in (2, 3): + raise ValueError("spatial_dims must be 2 or 3.") + if not (0 <= dropout_rate <= 1): + raise ValueError("dropout_rate should be between 0 and 1.") + if not (0 <= emb_dropout_rate <= 1): + raise ValueError("emb_dropout_rate should be between 0 and 1.") + if num_heads <= 0: + raise ValueError("num_heads must be a positive integer.") + if hidden_size % num_heads != 0: + raise ValueError("hidden_size should be divisible by num_heads.") + + self.spatial_dims = spatial_dims + self.patch_size = patch_size + self.in_channels = in_channels + + # --- token dropout --- + self.calc_token_dropout: Callable | None = None + if callable(token_dropout_prob): + self.calc_token_dropout = token_dropout_prob + elif isinstance(token_dropout_prob, (float, int)): + if not (0.0 < float(token_dropout_prob) < 1.0): + raise ValueError("token_dropout_prob must be in (0, 1) when given as a float.") + _prob = float(token_dropout_prob) + self.calc_token_dropout = lambda *_dims: _prob + + # --- patch embedding --- + # patch_dim = channels * patch_size^spatial_dims + patch_dim = in_channels * (patch_size**spatial_dims) + self.to_patch_embedding = nn.Sequential( + nn.LayerNorm(patch_dim), nn.Linear(patch_dim, hidden_size), nn.LayerNorm(hidden_size) + ) + + # --- factorized positional embeddings (one table per spatial axis) --- + image_size_t = ensure_tuple_rep(image_size, spatial_dims) + for i, img_d in enumerate(image_size_t): + if img_d % patch_size != 0: + raise ValueError(f"image_size dimension {i} ({img_d}) must be divisible by patch_size ({patch_size}).") + self.pos_embed_axes = nn.ParameterList( + [nn.Parameter(torch.randn(img_d // patch_size, hidden_size)) for img_d in image_size_t] + ) + + self.emb_dropout = nn.Dropout(emb_dropout_rate) + + # --- transformer --- + self.blocks = nn.ModuleList( + [ + _NaViTTransformerBlock(hidden_size, mlp_dim, num_heads, dim_head, dropout_rate, qkv_bias) + for _ in range(num_layers) + ] + ) + self.norm = nn.LayerNorm(hidden_size) + + # --- attention pooling --- + self.attn_pool_query = nn.Parameter(torch.randn(hidden_size)) + self.attn_pool = _NaViTAttention(hidden_size, num_heads, dim_head, dropout_rate, qkv_bias) + + # --- classification head --- + self.mlp_head = nn.Sequential(nn.LayerNorm(hidden_size), nn.Linear(hidden_size, num_classes, bias=False)) + + self._init_weights() + + def _init_weights(self) -> None: + """Initialise weights following standard ViT practice.""" + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + nn.init.ones_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + for param in self.pos_embed_axes: + nn.init.trunc_normal_(param, std=0.02) + nn.init.trunc_normal_(self.attn_pool_query, std=0.02) + + @property + def device(self) -> torch.device: + """Return the device on which the model parameters reside.""" + return next(self.parameters()).device + + def _image_to_patches(self, image: Tensor) -> tuple[Tensor, Tensor]: + """Rearrange a single image into a sequence of flattened patch tokens and their grid positions. + + Args: + image: tensor of shape ``(C, *spatial)`` where ``spatial`` has ``self.spatial_dims`` dimensions. + + Returns: + seq: patch token tensor of shape ``(num_patches, patch_dim)``. + pos: integer grid-coordinate tensor of shape ``(num_patches, spatial_dims)``. + """ + p = self.patch_size + spatial = image.shape[1:] # (H, W) or (H, W, D) + + if self.spatial_dims == 2: + h, w = spatial + # (C, H, W) -> (H/p * W/p, C * p * p) + seq = image.unfold(1, p, p).unfold(2, p, p) # (C, H/p, W/p, p, p) + seq = seq.permute(1, 2, 0, 3, 4).reshape(-1, self.in_channels * p * p) + gh, gw = h // p, w // p + pos_h = torch.arange(gh, device=image.device) + pos_w = torch.arange(gw, device=image.device) + grid = torch.stack(torch.meshgrid(pos_h, pos_w, indexing="ij"), dim=-1) # (gh, gw, 2) + pos = grid.reshape(-1, 2) + else: + h, w, d = spatial + # (C, H, W, D) -> (H/p * W/p * D/p, C * p * p * p) + seq = image.unfold(1, p, p).unfold(2, p, p).unfold(3, p, p) # (C, H/p, W/p, D/p, p, p, p) + seq = seq.permute(1, 2, 3, 0, 4, 5, 6).reshape(-1, self.in_channels * p * p * p) + gh, gw, gd = h // p, w // p, d // p + pos_h = torch.arange(gh, device=image.device) + pos_w = torch.arange(gw, device=image.device) + pos_d = torch.arange(gd, device=image.device) + grid = torch.stack(torch.meshgrid(pos_h, pos_w, pos_d, indexing="ij"), dim=-1) # (gh, gw, gd, 3) + pos = grid.reshape(-1, 3) + + return seq, pos + + def forward( + self, batched_images: list[list[Tensor]], group_images: bool = False, group_max_seq_len: int = 2048 + ) -> Tensor: + """Run NaViT on a batch of packed image groups. + + Args: + batched_images: a list of groups, where each group is a list of image tensors. + Each image must have shape ``(C, *spatial)`` with ``C == in_channels`` and every + spatial dimension divisible by ``patch_size``. The outer list corresponds to the + batch dimension; images within the same group are packed into a single sequence. + If ``group_images=True``, a flat ``list[Tensor]`` may be passed instead and the + packing is performed automatically. + group_images: when ``True``, treat ``batched_images`` as a flat list of tensors and + automatically pack them into groups of at most ``group_max_seq_len`` tokens. + group_max_seq_len: maximum sequence length used when ``group_images=True``. + + Returns: + Tensor of shape ``(total_images, num_classes)`` containing one logit vector per image + across all groups in the batch. + + Raises: + ValueError: If an image does not have the expected number of dimensions + (``spatial_dims + 1``). + ValueError: If an image's channel count does not match ``in_channels``. + ValueError: If any spatial dimension of an image is not divisible by ``patch_size``. + """ + device = self.device + pad_sequence = partial(orig_pad_sequence, batch_first=True) + + # optional auto-packing + if group_images: + batched_images = _group_images_by_max_seq_len( + batched_images, # type: ignore[arg-type] + patch_size=self.patch_size, + calc_token_dropout=self.calc_token_dropout, + max_seq_len=group_max_seq_len, + ) + + # ------------------------------------------------------------------ # + # 1. Convert each image to patch tokens + grid positions # + # ------------------------------------------------------------------ # + num_images_per_group: list[int] = [] + batched_sequences: list[Tensor] = [] + batched_positions: list[Tensor] = [] + batched_image_ids: list[Tensor] = [] + + for images in batched_images: + num_images_per_group.append(len(images)) + sequences: list[Tensor] = [] + positions: list[Tensor] = [] + image_ids = torch.empty((0,), device=device, dtype=torch.long) + + for image_id, image in enumerate(images): + if image.ndim != self.spatial_dims + 1: + raise ValueError( + f"Expected image with {self.spatial_dims + 1} dimensions (C, *spatial), " + f"got shape {tuple(image.shape)}." + ) + if image.shape[0] != self.in_channels: + raise ValueError(f"Expected {self.in_channels} input channels, got {image.shape[0]}.") + spatial = image.shape[1:] + for dim_size in spatial: + if dim_size % self.patch_size != 0: + raise ValueError( + f"All spatial dimensions must be divisible by patch_size={self.patch_size}, " + f"got spatial shape {spatial}." + ) + + seq, pos = self._image_to_patches(image) # (N, patch_dim), (N, spatial_dims) + + # optional token dropout (training only) + if self.calc_token_dropout is not None and self.training: + dropout_frac = self.calc_token_dropout(*spatial) + num_keep = max(1, int(seq.shape[0] * (1.0 - dropout_frac))) + keep_idx = torch.randn(seq.shape[0], device=device).topk(num_keep).indices + seq = seq[keep_idx] + pos = pos[keep_idx] + + image_ids = F.pad(image_ids, (0, seq.shape[0]), value=image_id) + sequences.append(seq) + positions.append(pos) + + batched_image_ids.append(image_ids) + batched_sequences.append(torch.cat(sequences, dim=0)) + batched_positions.append(torch.cat(positions, dim=0)) + + # ------------------------------------------------------------------ # + # 2. Pad sequences to the same length and build attention masks # + # ------------------------------------------------------------------ # + lengths = torch.tensor([s.shape[0] for s in batched_sequences], device=device, dtype=torch.long) + max_len = int(lengths.amax().item()) + len_range = torch.arange(max_len, device=device) + + # key-padding mask: True for valid (non-padded) positions + key_pad_mask = len_range.unsqueeze(0) < lengths.unsqueeze(1) # (B, max_len) + + # per-image attention mask: tokens from different images must not attend to each other + batched_image_ids_padded = pad_sequence(batched_image_ids) # (B, max_len) + same_image = batched_image_ids_padded.unsqueeze(2) == batched_image_ids_padded.unsqueeze( + 1 + ) # (B, max_len, max_len) + attn_mask = same_image & key_pad_mask.unsqueeze(1) # (B, max_len, max_len) + attn_mask = attn_mask.unsqueeze(1) # (B, 1, max_len, max_len) + + # ------------------------------------------------------------------ # + # 3. Patch embedding + factorized positional encoding ## + # ------------------------------------------------------------------ # + patches = pad_sequence(batched_sequences) # (B, max_len, patch_dim) + patch_positions = pad_sequence(batched_positions) # (B, max_len, spatial_dims) + + x = self.to_patch_embedding(patches) # (B, max_len, hidden_size) + + # sum positional embeddings from each axis + for axis_idx, pos_embed in enumerate(self.pos_embed_axes): + axis_indices = patch_positions[..., axis_idx] # (B, max_len) + # clamp to handle positions beyond the reference image_size table + axis_indices = axis_indices.clamp(max=pos_embed.shape[0] - 1) + x = x + pos_embed[axis_indices] + + x = self.emb_dropout(x) + + # ------------------------------------------------------------------ # + # 4. Transformer # + # ------------------------------------------------------------------ # + for block in self.blocks: + x = block(x, attn_mask=attn_mask) + x = self.norm(x) + + # ------------------------------------------------------------------ # + # 5. Attention pooling: one query per image in the group # + # ------------------------------------------------------------------ # + num_images_t = torch.tensor(num_images_per_group, device=device, dtype=torch.long) + max_queries = int(num_images_t.amax().item()) + + # expand the shared query vector to (B, max_queries, hidden_size) + queries = self.attn_pool_query.unsqueeze(0).unsqueeze(0).expand(x.shape[0], max_queries, -1) + + # build cross-attention mask: query i attends only to tokens belonging to image i + image_id_range = torch.arange(max_queries, device=device) + pool_mask = image_id_range.unsqueeze(1) == batched_image_ids_padded.unsqueeze(1) # (B, max_queries, max_len) + pool_mask = pool_mask & key_pad_mask.unsqueeze(1) # (B, max_queries, max_len) + pool_mask = pool_mask.unsqueeze(1) # (B, 1, max_queries, max_len) + + pooled = self.attn_pool(queries, context=x, attn_mask=pool_mask) + queries # (B, max_queries, hidden_size) + + # ------------------------------------------------------------------ # + # 6. Flatten, filter padding queries, and classify # + # ------------------------------------------------------------------ # + pooled = pooled.reshape(-1, pooled.shape[-1]) # (B * max_queries, hidden_size) + + is_valid = (image_id_range.unsqueeze(0) < num_images_t.unsqueeze(1)).reshape(-1) # (B * max_queries,) + pooled = pooled[is_valid] # (total_images, hidden_size) + + return cast(Tensor, self.mlp_head(pooled)) # (total_images, num_classes) diff --git a/monai/networks/nets/senet.py b/monai/networks/nets/senet.py index 4c7dd0f0c24..668125b1428 100644 --- a/monai/networks/nets/senet.py +++ b/monai/networks/nets/senet.py @@ -304,7 +304,7 @@ def _load_state_dict(model: nn.Module, arch: str, progress: bool): download_url(model_url["url"], filepath=model_url["filename"]) state_dict = torch.load(model_url["filename"], map_location=None, weights_only=True) else: - state_dict = load_state_dict_from_url(model_url, progress=progress) + state_dict = load_state_dict_from_url(model_url, progress=progress, weights_only=True) for key in list(state_dict.keys()): new_key = None if pattern_conv.match(key): diff --git a/monai/networks/nets/swin_unetr.py b/monai/networks/nets/swin_unetr.py index 0db2d50d26a..7ab55ab3960 100644 --- a/monai/networks/nets/swin_unetr.py +++ b/monai/networks/nets/swin_unetr.py @@ -23,6 +23,7 @@ from monai.networks.blocks import MLPBlock as Mlp from monai.networks.blocks import PatchEmbed, UnetOutBlock, UnetrBasicBlock, UnetrUpBlock +from monai.networks.blocks.hyena import HyenaTransformerBlock from monai.networks.layers import DropPath, trunc_normal_ from monai.utils import ensure_tuple_rep, look_up_option, optional_import @@ -84,6 +85,20 @@ def __init__( spatial_dims: int = 3, downsample: str | nn.Module = "merging", use_v2: bool = False, + use_hyena: bool = False, + hyena_stages: Sequence[bool] | None = None, + hyena_kernel_size: int = 3, + hyena_kernel_mlp_dim: int = 32, + hyena_kernel_layers: int = 3, + hyena_mask_max_attenuation: float = 0.95, + hyena_fft_padding: str = "circular", + hyena_grid_type: str = "single", + hyena_use_chunked_fft: bool = False, + hyena_use_fft_short_conv: bool = False, + hyena_omega_0: float = 10.0, + hyena_l_cache: int = 32, + hyena_short_conv_fft_chunks: int = 0, + use_flash_attention: bool = False, ) -> None: """ Args: @@ -110,6 +125,32 @@ def __init__( user-specified `nn.Module` following the API defined in :py:class:`monai.networks.nets.PatchMerging`. The default is currently `"merging"` (the original version defined in v0.9.0). use_v2: using swinunetr_v2, which adds a residual convolution block at the beggining of each swin stage. + use_hyena: replace windowed self-attention with the HyenaND operator (subquadratic O(N log N) + global convolution) in every Swin stage. Default ``False`` keeps the model bit-identical + to the pre-HyenaND code path. Requires the optional ``nvsubquadratic`` package + (``pip install monai[hyena]``). When combined with ``hyena_stages``, the per-stage flag + overrides this master switch on a per-stage basis. + hyena_stages: optional 4-tuple of bools selecting which Swin stages use HyenaND vs windowed + attention. ``True`` at index ``i`` builds :class:`HyenaTransformerBlock` at stage ``i``, + ``False`` builds the conventional :class:`SwinTransformerBlock`. The four NeurIPS 2026 + paper variants are: ``None`` (AAAA, all attention, requires ``use_hyena=False``); + ``(True, True, True, True)`` (HHHH, equivalent to ``use_hyena=True``); ``(True, False, + True, False)`` (HAHA); ``(True, True, False, False)`` (HHAA, paper-best). + hyena_kernel_size: HyenaND short-convolution kernel size (depthwise on QKV). + hyena_kernel_mlp_dim: SIREN implicit-kernel MLP hidden dimension. + hyena_kernel_layers: SIREN implicit-kernel depth. + hyena_mask_max_attenuation: Gaussian-modulation boundary attenuation (0-1). + hyena_fft_padding: ``"circular"`` or ``"zero"``. ``"circular"`` was the paper-best setting. + hyena_grid_type: ``"single"`` (kernel = input size, required for circular) or ``"double"`` + (kernel = 2x input size, requires zero padding). + hyena_use_chunked_fft: enable chunked FFT for ~26 percent memory savings; requires + ``hyena_fft_padding="zero"``. + hyena_use_fft_short_conv: replace the short conv with :class:`DepthwiseFFTConv{2,3}d` to + eliminate the INT32 unfold limit and enable ROI > 128. + hyena_omega_0: SIREN frequency. Default 10.0 (stable). + hyena_l_cache: SIREN coordinate-grid cache size per spatial dim. + hyena_short_conv_fft_chunks: channel chunk size for the FFT short conv (0 = no chunking). + use_flash_attention: use flash attention (scaled dot product attention) at inference. Examples:: @@ -151,6 +192,8 @@ def __init__( raise ValueError("feature_size should be divisible by 12.") self.normalize = normalize + self.use_hyena = use_hyena + self.hyena_stages = tuple(bool(s) for s in hyena_stages) if hyena_stages is not None else None self.swinViT = SwinTransformer( in_chans=in_channels, @@ -170,6 +213,20 @@ def __init__( spatial_dims=spatial_dims, downsample=look_up_option(downsample, MERGING_MODE) if isinstance(downsample, str) else downsample, use_v2=use_v2, + use_hyena=use_hyena, + hyena_stages=self.hyena_stages, + hyena_kernel_size=hyena_kernel_size, + hyena_kernel_mlp_dim=hyena_kernel_mlp_dim, + hyena_kernel_layers=hyena_kernel_layers, + hyena_mask_max_attenuation=hyena_mask_max_attenuation, + hyena_fft_padding=hyena_fft_padding, + hyena_grid_type=hyena_grid_type, + hyena_use_chunked_fft=hyena_use_chunked_fft, + hyena_use_fft_short_conv=hyena_use_fft_short_conv, + hyena_omega_0=hyena_omega_0, + hyena_l_cache=hyena_l_cache, + hyena_short_conv_fft_chunks=hyena_short_conv_fft_chunks, + use_flash_attention=use_flash_attention, ) self.encoder1 = UnetrBasicBlock( @@ -274,50 +331,52 @@ def __init__( self.out = UnetOutBlock(spatial_dims=spatial_dims, in_channels=feature_size, out_channels=out_channels) def load_from(self, weights): + """Load pretrained Swin weights into the matching submodules. + + When a stage uses :class:`HyenaTransformerBlock` instead of + :class:`SwinTransformerBlock`, the per-block ``load_from`` call is skipped for + that stage and a warning is issued -- HyenaND has a different parameter layout + and there are no compatible attention weights to copy. PatchMerging + downsample weights are still loaded for all stages (the downsample layer is + the same in both code paths). + """ + import warnings + layers1_0: BasicLayer = self.swinViT.layers1[0] # type: ignore[assignment] layers2_0: BasicLayer = self.swinViT.layers2[0] # type: ignore[assignment] layers3_0: BasicLayer = self.swinViT.layers3[0] # type: ignore[assignment] layers4_0: BasicLayer = self.swinViT.layers4[0] # type: ignore[assignment] wstate = weights["state_dict"] + def _stage_is_hyena(stage_layer: BasicLayer) -> bool: + first_block = next(iter(stage_layer.blocks.children())) + return isinstance(first_block, HyenaTransformerBlock) + with torch.no_grad(): self.swinViT.patch_embed.proj.weight.copy_(wstate["module.patch_embed.proj.weight"]) self.swinViT.patch_embed.proj.bias.copy_(wstate["module.patch_embed.proj.bias"]) - for bname, block in layers1_0.blocks.named_children(): - block.load_from(weights, n_block=bname, layer="layers1") # type: ignore[operator] - - if layers1_0.downsample is not None: - d = layers1_0.downsample - d.reduction.weight.copy_(wstate["module.layers1.0.downsample.reduction.weight"]) # type: ignore - d.norm.weight.copy_(wstate["module.layers1.0.downsample.norm.weight"]) # type: ignore - d.norm.bias.copy_(wstate["module.layers1.0.downsample.norm.bias"]) # type: ignore - - for bname, block in layers2_0.blocks.named_children(): - block.load_from(weights, n_block=bname, layer="layers2") # type: ignore[operator] - - if layers2_0.downsample is not None: - d = layers2_0.downsample - d.reduction.weight.copy_(wstate["module.layers2.0.downsample.reduction.weight"]) # type: ignore - d.norm.weight.copy_(wstate["module.layers2.0.downsample.norm.weight"]) # type: ignore - d.norm.bias.copy_(wstate["module.layers2.0.downsample.norm.bias"]) # type: ignore - - for bname, block in layers3_0.blocks.named_children(): - block.load_from(weights, n_block=bname, layer="layers3") # type: ignore[operator] - - if layers3_0.downsample is not None: - d = layers3_0.downsample - d.reduction.weight.copy_(wstate["module.layers3.0.downsample.reduction.weight"]) # type: ignore - d.norm.weight.copy_(wstate["module.layers3.0.downsample.norm.weight"]) # type: ignore - d.norm.bias.copy_(wstate["module.layers3.0.downsample.norm.bias"]) # type: ignore - - for bname, block in layers4_0.blocks.named_children(): - block.load_from(weights, n_block=bname, layer="layers4") # type: ignore[operator] - - if layers4_0.downsample is not None: - d = layers4_0.downsample - d.reduction.weight.copy_(wstate["module.layers4.0.downsample.reduction.weight"]) # type: ignore - d.norm.weight.copy_(wstate["module.layers4.0.downsample.norm.weight"]) # type: ignore - d.norm.bias.copy_(wstate["module.layers4.0.downsample.norm.bias"]) # type: ignore + + for layer_name, stage in [ + ("layers1", layers1_0), + ("layers2", layers2_0), + ("layers3", layers3_0), + ("layers4", layers4_0), + ]: + if _stage_is_hyena(stage): + warnings.warn( + f"Skipping {layer_name} block weights: stage uses HyenaTransformerBlock, " + "which has no compatible Swin attention weights. Blocks remain at their " + "random initialization.", + stacklevel=2, + ) + else: + for bname, block in stage.blocks.named_children(): + block.load_from(weights, n_block=bname, layer=layer_name) # type: ignore[operator] + if stage.downsample is not None: + d = stage.downsample + d.reduction.weight.copy_(wstate[f"module.{layer_name}.0.downsample.reduction.weight"]) # type: ignore + d.norm.weight.copy_(wstate[f"module.{layer_name}.0.downsample.norm.weight"]) # type: ignore + d.norm.bias.copy_(wstate[f"module.{layer_name}.0.downsample.norm.bias"]) # type: ignore @torch.jit.unused def _check_input_size(self, spatial_shape): @@ -457,6 +516,7 @@ def __init__( qkv_bias: bool = False, attn_drop: float = 0.0, proj_drop: float = 0.0, + use_flash_attention: bool = False, ) -> None: """ Args: @@ -466,12 +526,17 @@ def __init__( qkv_bias: add a learnable bias to query, key, value. attn_drop: attention dropout rate. proj_drop: dropout rate of output. + use_flash_attention: if True, use ``torch.nn.functional.scaled_dot_product_attention`` for the + windowed attention. Equivalent to the default path but faster at inference; only used when + autograd is disabled (e.g. under ``torch.no_grad()`` or ``torch.inference_mode()``, not + ``eval()`` alone) and the module is not scripted. """ super().__init__() self.dim = dim self.window_size = window_size self.num_heads = num_heads + self.use_flash_attention = use_flash_attention head_dim = dim // num_heads self.scale = head_dim**-0.5 mesh_args = torch.meshgrid.__kwdefaults__ @@ -528,12 +593,26 @@ def forward(self, x, mask): b, n, c = x.shape qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] - q = q * self.scale - attn = q @ k.transpose(-2, -1) relative_position_bias = self.relative_position_bias_table[ self.relative_position_index.clone()[:n, :n].reshape(-1) # type: ignore[operator] ].reshape(n, n, -1) relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + if self.use_flash_attention and not torch.jit.is_scripting() and not torch.is_grad_enabled(): + # additive bias combines the relative position bias and, for shifted windows, the attention mask + if mask is not None: + nw = mask.shape[0] + bias = relative_position_bias.view(1, 1, self.num_heads, n, n) + mask.reshape(1, nw, 1, n, n) + bias = bias.expand(b // nw, nw, self.num_heads, n, n).reshape(b, self.num_heads, n, n) + else: + bias = relative_position_bias.unsqueeze(0) + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=bias.to(q.dtype), dropout_p=0.0, scale=self.scale + ) + x = x.transpose(1, 2).reshape(b, n, c) + return self.proj_drop(self.proj(x)) + + q = q * self.scale + attn = q @ k.transpose(-2, -1) attn = attn + relative_position_bias.unsqueeze(0) if mask is not None: nw = mask.shape[0] @@ -572,6 +651,7 @@ def __init__( act_layer: str = "GELU", norm_layer: type[LayerNorm] = nn.LayerNorm, use_checkpoint: bool = False, + use_flash_attention: bool = False, ) -> None: """ Args: @@ -587,6 +667,7 @@ def __init__( act_layer: activation layer. norm_layer: normalization layer. use_checkpoint: use gradient checkpointing for reduced memory usage. + use_flash_attention: use flash attention (scaled dot product attention) at inference. """ super().__init__() @@ -604,6 +685,7 @@ def __init__( qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, + use_flash_attention=use_flash_attention, ) self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() @@ -856,6 +938,19 @@ def __init__( norm_layer: type[LayerNorm] = nn.LayerNorm, downsample: nn.Module | None = None, use_checkpoint: bool = False, + use_hyena: bool = False, + hyena_kernel_size: int = 3, + hyena_kernel_mlp_dim: int = 32, + hyena_kernel_layers: int = 3, + hyena_mask_max_attenuation: float = 0.95, + hyena_fft_padding: str = "circular", + hyena_grid_type: str = "single", + hyena_use_chunked_fft: bool = False, + hyena_use_fft_short_conv: bool = False, + hyena_omega_0: float = 10.0, + hyena_l_cache: int = 32, + hyena_short_conv_fft_chunks: int = 0, + use_flash_attention: bool = False, ) -> None: """ Args: @@ -871,6 +966,14 @@ def __init__( norm_layer: normalization layer. downsample: an optional downsampling layer at the end of the layer. use_checkpoint: use gradient checkpointing for reduced memory usage. + use_hyena: replace :class:`SwinTransformerBlock` with :class:`HyenaTransformerBlock` + in this stage. See :class:`SwinUNETR` for the per-stage selection mechanism. + hyena_kernel_size, hyena_kernel_mlp_dim, hyena_kernel_layers, + hyena_mask_max_attenuation, hyena_fft_padding, hyena_grid_type, + hyena_use_chunked_fft, hyena_use_fft_short_conv, hyena_omega_0, hyena_l_cache, + hyena_short_conv_fft_chunks: forwarded to :class:`HyenaTransformerBlock`. See its + docstring for semantics. + use_flash_attention: use flash attention (scaled dot product attention) at inference. """ super().__init__() @@ -879,24 +982,53 @@ def __init__( self.no_shift = tuple(0 for i in window_size) self.depth = depth self.use_checkpoint = use_checkpoint - self.blocks = nn.ModuleList( - [ - SwinTransformerBlock( - dim=dim, - num_heads=num_heads, - window_size=self.window_size, - shift_size=self.no_shift if (i % 2 == 0) else self.shift_size, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - drop=drop, - attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, - norm_layer=norm_layer, - use_checkpoint=use_checkpoint, - ) - for i in range(depth) - ] - ) + self.use_hyena = use_hyena + if use_hyena: + self.blocks = nn.ModuleList( + [ + HyenaTransformerBlock( + dim=dim, + spatial_dims=len(self.window_size), + mlp_ratio=mlp_ratio, + drop=drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + hyena_kernel_size=hyena_kernel_size, + hyena_kernel_mlp_dim=hyena_kernel_mlp_dim, + hyena_kernel_layers=hyena_kernel_layers, + hyena_mask_max_attenuation=hyena_mask_max_attenuation, + hyena_fft_padding=hyena_fft_padding, + hyena_grid_type=hyena_grid_type, + hyena_use_chunked_fft=hyena_use_chunked_fft, + hyena_use_fft_short_conv=hyena_use_fft_short_conv, + hyena_omega_0=hyena_omega_0, + hyena_l_cache=hyena_l_cache, + hyena_short_conv_fft_chunks=hyena_short_conv_fft_chunks, + ) + for i in range(depth) + ] + ) + else: + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=self.window_size, + shift_size=self.no_shift if (i % 2 == 0) else self.shift_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + use_flash_attention=use_flash_attention, + ) + for i in range(depth) + ] + ) self.downsample = downsample if callable(self.downsample): self.downsample = downsample(dim=dim, norm_layer=norm_layer, spatial_dims=len(self.window_size)) @@ -910,7 +1042,8 @@ def forward(self, x): dp = int(np.ceil(d / window_size[0])) * window_size[0] hp = int(np.ceil(h / window_size[1])) * window_size[1] wp = int(np.ceil(w / window_size[2])) * window_size[2] - attn_mask = compute_mask([dp, hp, wp], window_size, shift_size, x.device) + # HyenaTransformerBlock ignores the attention mask; skip building it for Hyena stages. + attn_mask = None if self.use_hyena else compute_mask([dp, hp, wp], window_size, shift_size, x.device) for blk in self.blocks: x = blk(x, attn_mask) x = x.view(b, d, h, w, -1) @@ -924,7 +1057,8 @@ def forward(self, x): x = rearrange(x, "b c h w -> b h w c") hp = int(np.ceil(h / window_size[0])) * window_size[0] wp = int(np.ceil(w / window_size[1])) * window_size[1] - attn_mask = compute_mask([hp, wp], window_size, shift_size, x.device) + # HyenaTransformerBlock ignores the attention mask; skip building it for Hyena stages. + attn_mask = None if self.use_hyena else compute_mask([hp, wp], window_size, shift_size, x.device) for blk in self.blocks: x = blk(x, attn_mask) x = x.view(b, h, w, -1) @@ -961,6 +1095,20 @@ def __init__( spatial_dims: int = 3, downsample="merging", use_v2=False, + use_hyena: bool = False, + hyena_stages: Sequence[bool] | None = None, + hyena_kernel_size: int = 3, + hyena_kernel_mlp_dim: int = 32, + hyena_kernel_layers: int = 3, + hyena_mask_max_attenuation: float = 0.95, + hyena_fft_padding: str = "circular", + hyena_grid_type: str = "single", + hyena_use_chunked_fft: bool = False, + hyena_use_fft_short_conv: bool = False, + hyena_omega_0: float = 10.0, + hyena_l_cache: int = 32, + hyena_short_conv_fft_chunks: int = 0, + use_flash_attention: bool = False, ) -> None: """ Args: @@ -983,6 +1131,18 @@ def __init__( user-specified `nn.Module` following the API defined in :py:class:`monai.networks.nets.PatchMerging`. The default is currently `"merging"` (the original version defined in v0.9.0). use_v2: using swinunetr_v2, which adds a residual convolution block at the beginning of each swin stage. + use_hyena: build :class:`HyenaTransformerBlock` instead of :class:`SwinTransformerBlock` + in every stage. See :class:`SwinUNETR` for paper-variant patterns. + hyena_stages: optional per-stage override (4-tuple of bools); a stage flagged ``True`` + builds a Hyena block regardless of ``use_hyena``, and a stage flagged ``False`` + builds a Swin block regardless of ``use_hyena``. ``None`` falls back to ``use_hyena`` + for all stages. + hyena_kernel_size, hyena_kernel_mlp_dim, hyena_kernel_layers, + hyena_mask_max_attenuation, hyena_fft_padding, hyena_grid_type, + hyena_use_chunked_fft, hyena_use_fft_short_conv, hyena_omega_0, hyena_l_cache, + hyena_short_conv_fft_chunks: HyenaND configuration. See + :class:`monai.networks.blocks.HyenaTransformerBlock` for semantics. + use_flash_attention: use flash attention (scaled dot product attention) at inference. """ super().__init__() @@ -991,6 +1151,33 @@ def __init__( self.patch_norm = patch_norm self.window_size = window_size self.patch_size = patch_size + + # Per-stage Hyena selection: explicit ``hyena_stages`` overrides the master flag. + self._per_stage_hyena: list[bool] = ( + [bool(s) for s in hyena_stages] if hyena_stages is not None else [bool(use_hyena)] * self.num_layers + ) + if len(self._per_stage_hyena) != self.num_layers: + raise ValueError( + f"hyena_stages must have length {self.num_layers} (one bool per Swin stage); " + f"got length {len(self._per_stage_hyena)}." + ) + + # Legacy RoPE-divisibility guard: kept as defensive validation for callers that bypass the + # SwinUNETR-level ``feature_size % 12 == 0`` check. ``nvsubquadratic`` removed RoPE from + # the HyenaND operator on 2026-04-27, so this check is now slightly conservative; it does + # not affect any valid SwinUNETR configuration. + if any(self._per_stage_hyena): + div = 6 if spatial_dims == 3 else 4 + for i, use_h in enumerate(self._per_stage_hyena): + if use_h: + dim_at_layer = int(embed_dim * 2**i) + if dim_at_layer % div != 0: + raise ValueError( + f"For {spatial_dims}D Hyena, embed_dim * 2^layer must be divisible by {div}. " + f"At layer {i}, dim={dim_at_layer} is not. " + "Use embed_dim that is a multiple of 12 (the SwinUNETR default check)." + ) + self.patch_embed = PatchEmbed( patch_size=self.patch_size, in_chans=in_chans, @@ -1025,6 +1212,19 @@ def __init__( norm_layer=norm_layer, downsample=down_sample_mod, use_checkpoint=use_checkpoint, + use_hyena=self._per_stage_hyena[i_layer], + hyena_kernel_size=hyena_kernel_size, + hyena_kernel_mlp_dim=hyena_kernel_mlp_dim, + hyena_kernel_layers=hyena_kernel_layers, + hyena_mask_max_attenuation=hyena_mask_max_attenuation, + hyena_fft_padding=hyena_fft_padding, + hyena_grid_type=hyena_grid_type, + hyena_use_chunked_fft=hyena_use_chunked_fft, + hyena_use_fft_short_conv=hyena_use_fft_short_conv, + hyena_omega_0=hyena_omega_0, + hyena_l_cache=hyena_l_cache, + hyena_short_conv_fft_chunks=hyena_short_conv_fft_chunks, + use_flash_attention=use_flash_attention, ) if i_layer == 0: self.layers1.append(layer) diff --git a/monai/networks/nets/transchex.py b/monai/networks/nets/transchex.py index 6c40cae2aa3..cd71b307422 100644 --- a/monai/networks/nets/transchex.py +++ b/monai/networks/nets/transchex.py @@ -23,6 +23,7 @@ transformers = optional_import("transformers") load_tf_weights_in_bert = optional_import("transformers", name="load_tf_weights_in_bert")[0] cached_file = optional_import("transformers.utils", name="cached_file")[0] +BertConfig = optional_import("transformers", name="BertConfig")[0] BertEmbeddings = optional_import("transformers.models.bert.modeling_bert", name="BertEmbeddings")[0] BertLayer = optional_import("transformers.models.bert.modeling_bert", name="BertLayer")[0] @@ -74,6 +75,7 @@ def from_pretrained( return load_tf_weights_in_bert(model, weights_path) old_keys = [] new_keys = [] + # pyrefly: ignore [missing-attribute] for key in state_dict.keys(): new_key = None if "gamma" in key: @@ -84,11 +86,13 @@ def from_pretrained( old_keys.append(key) new_keys.append(new_key) for old_key, new_key in zip(old_keys, new_keys): + # pyrefly: ignore [missing-attribute, unsupported-operation] state_dict[new_key] = state_dict.pop(old_key) missing_keys: list = [] unexpected_keys: list = [] error_msgs: list = [] metadata = getattr(state_dict, "_metadata", None) + # pyrefly: ignore [missing-attribute] state_dict = state_dict.copy() if metadata is not None: state_dict._metadata = metadata @@ -219,7 +223,11 @@ def __init__( """ super().__init__() - self.config = type("obj", (object,), bert_config) + self.config = BertConfig(**bert_config) + # explicitly select the eager attention path: transformers>=4.48 dispatches attention + # implementations via `config._attn_implementation`, which is otherwise left unset since + # `bert_config` above does not come from a `from_pretrained` call. + self.config._attn_implementation = "eager" self.embeddings = BertEmbeddings(self.config) self.language_encoder = nn.ModuleList([BertLayer(self.config) for _ in range(num_language_layers)]) self.vision_encoder = nn.ModuleList([BertLayer(self.config) for _ in range(num_vision_layers)]) diff --git a/monai/networks/nets/unetr.py b/monai/networks/nets/unetr.py index 79ea0e23f7a..2ea5c0e1587 100644 --- a/monai/networks/nets/unetr.py +++ b/monai/networks/nets/unetr.py @@ -25,6 +25,16 @@ class UNETR(nn.Module): """ UNETR based on: "Hatamizadeh et al., UNETR: Transformers for 3D Medical Image Segmentation " + + Spatial Shape Constraints: + Each spatial dimension of ``img_size`` must be divisible by ``patch_size``. + UNETR uses a fixed patch size of 16, so each spatial dimension must be + divisible by **16**. This is required by the ViT patch embedding step. + + Valid 3D input sizes: ``(16, 16, 16)``, ``(32, 32, 32)``, ``(64, 64, 64)``, + ``(96, 96, 96)``, ``(128, 128, 128)``, ``(96, 64, 128)``. + + A ``ValueError`` is raised in ``__init__`` if ``img_size`` is not divisible by 16. """ def __init__( @@ -81,6 +91,15 @@ def __init__( if not (0 <= dropout_rate <= 1): raise ValueError("dropout_rate should be between 0 and 1.") + img_size = ensure_tuple_rep(img_size, spatial_dims) + patch_size = ensure_tuple_rep(16, spatial_dims) + for i, (img_d, p_d) in enumerate(zip(img_size, patch_size)): + if img_d % p_d != 0: + raise ValueError( + f"img_size[{i}]={img_d} is not divisible by patch_size={p_d}. " + f"Each spatial dimension of img_size must be divisible by 16." + ) + if hidden_size % num_heads != 0: raise ValueError("hidden_size should be divisible by num_heads.") diff --git a/monai/networks/nets/vqvae.py b/monai/networks/nets/vqvae.py index 43ba48585c2..690bf48de08 100644 --- a/monai/networks/nets/vqvae.py +++ b/monai/networks/nets/vqvae.py @@ -361,18 +361,22 @@ def __init__( else: downsample_parameters_tuple = downsample_parameters + # pyrefly: ignore [not-iterable] if not all(all(isinstance(value, int) for value in sub_item) for sub_item in downsample_parameters_tuple): raise ValueError("`downsample_parameters` should be a single tuple of integer or a tuple of tuples.") # check if downsample_parameters is a tuple of ints or a tuple of tuples of ints + # pyrefly: ignore [not-iterable] if not all(all(isinstance(value, int) for value in sub_item) for sub_item in upsample_parameters_tuple): raise ValueError("`upsample_parameters` should be a single tuple of integer or a tuple of tuples.") for parameter in downsample_parameters_tuple: + # pyrefly: ignore [bad-argument-type] if len(parameter) != 4: raise ValueError("`downsample_parameters` should be a tuple of tuples with 4 integers.") for parameter in upsample_parameters_tuple: + # pyrefly: ignore [bad-argument-type] if len(parameter) != 5: raise ValueError("`upsample_parameters` should be a tuple of tuples with 5 integers.") @@ -396,6 +400,7 @@ def __init__( channels=channels, num_res_layers=num_res_layers, num_res_channels=num_res_channels, + # pyrefly: ignore [bad-argument-type] downsample_parameters=downsample_parameters_tuple, dropout=dropout, act=act, @@ -408,6 +413,7 @@ def __init__( channels=channels, num_res_layers=num_res_layers, num_res_channels=num_res_channels, + # pyrefly: ignore [bad-argument-type] upsample_parameters=upsample_parameters_tuple, dropout=dropout, act=act, diff --git a/monai/networks/utils.py b/monai/networks/utils.py index f56c39dcd18..0c40e5318b6 100644 --- a/monai/networks/utils.py +++ b/monai/networks/utils.py @@ -34,9 +34,6 @@ from monai.utils.module import look_up_option, optional_import from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor -onnx, _ = optional_import("onnx") -onnxreference, _ = optional_import("onnx.reference") -onnxruntime, _ = optional_import("onnxruntime") polygraphy, polygraphy_imported = optional_import("polygraphy") torch_tensorrt, _ = optional_import("torch_tensorrt", "1.4.0") @@ -601,6 +598,7 @@ def copy_model_state( dst_dict[dst_key] = val updated_keys.append(dst_key) for s in mapping if mapping else {}: + # pyrefly: ignore [unsupported-operation] dst_key = f"{dst_prefix}{mapping[s]}" if dst_key in dst_dict and dst_key not in to_skip: if dst_dict[dst_key].shape != src_dict[s].shape: @@ -708,6 +706,8 @@ def convert_to_onnx( https://pytorch.org/docs/master/generated/torch.jit.script.html. """ + onnx, _ = optional_import("onnx") + model.eval() with torch.no_grad(): torch_versioned_kwargs = {} @@ -777,11 +777,13 @@ def convert_to_onnx( model_input_names = [i.name for i in onnx_model.graph.input] input_dict = dict(zip(model_input_names, [i.cpu().numpy() for i in inputs])) if use_ort: + onnxruntime, _ = optional_import("onnxruntime") ort_sess = onnxruntime.InferenceSession( onnx_model.SerializeToString(), providers=ort_provider if ort_provider else ["CPUExecutionProvider"] ) onnx_out = ort_sess.run(None, input_dict) else: + onnxreference, _ = optional_import("onnx.reference") sess = onnxreference.ReferenceEvaluator(onnx_model) onnx_out = sess.run(None, input_dict) set_determinism(seed=None) diff --git a/monai/optimizers/novograd.py b/monai/optimizers/novograd.py index 9ca612fc564..5d3c5504e17 100644 --- a/monai/optimizers/novograd.py +++ b/monai/optimizers/novograd.py @@ -112,6 +112,7 @@ def step(self, closure: Callable[[], T] | None = None) -> T | None: # type: ign norm = torch.sum(torch.pow(grad, 2)) if exp_avg_sq == 0: + # pyrefly: ignore [missing-attribute] exp_avg_sq.copy_(norm) else: exp_avg_sq.mul_(beta2).add_(norm, alpha=1 - beta2) diff --git a/monai/transforms/adaptors.py b/monai/transforms/adaptors.py index b3c8b34f5c5..3d22f3fb746 100644 --- a/monai/transforms/adaptors.py +++ b/monai/transforms/adaptors.py @@ -140,7 +140,7 @@ def must_be_types(variable_name, variable, types): raise TypeError(f"'{variable_name}' must be one of {types} but is {type(variable)}") def map_names(ditems, input_map): - return {input_map(k, k): v for k, v in ditems.items()} + return {input_map.get(k, k): v for k, v in ditems.items()} def map_only_names(ditems, input_map): return {v: ditems[k] for k, v in input_map.items()} diff --git a/monai/transforms/croppad/array.py b/monai/transforms/croppad/array.py index b23fbac7d9d..fc913fa767d 100644 --- a/monai/transforms/croppad/array.py +++ b/monai/transforms/croppad/array.py @@ -342,6 +342,24 @@ def compute_pad_width(self, spatial_shape: Sequence[int]) -> tuple[tuple[int, in return spatial_pad.compute_pad_width(spatial_shape) +def _to_int_list(data: Sequence[int] | int | NdarrayOrTensor) -> list[int]: + """Coerce an ROI spec (scalar, sequence, tensor or ndarray) to a list of Python ints.""" + if isinstance(data, (str, bytes)): + raise TypeError("ROI specs must be integers or sequences of integers, not strings.") + return [int(i) for i in ensure_tuple(data)] + + +def _broadcast_int_pair( + a: Sequence[int] | int | NdarrayOrTensor, b: Sequence[int] | int | NdarrayOrTensor +) -> tuple[list[int], list[int]]: + """Coerce a pair of ROI specs to two equal-length int lists, broadcasting a scalar to match.""" + list_a, list_b = _to_int_list(a), _to_int_list(b) + n = max(len(list_a), len(list_b)) + if len(list_a) not in (1, n) or len(list_b) not in (1, n): + raise ValueError(f"ROI specs must have matching lengths or be scalar, got {len(list_a)} and {len(list_b)}.") + return (list_a * n if len(list_a) == 1 else list_a), (list_b * n if len(list_b) == 1 else list_b) + + class Crop(InvertibleTransform, LazyTransform): """ Perform crop operations on the input image. @@ -379,31 +397,22 @@ def compute_slices( roi_slices: list of slices for each of the spatial dimensions. """ - roi_start_t: torch.Tensor - if roi_slices: if not all(s.step is None or s.step == 1 for s in roi_slices): raise ValueError(f"only slice steps of 1/None are currently supported, got {roi_slices}.") return ensure_tuple(roi_slices) else: if roi_center is not None and roi_size is not None: - roi_center_t = convert_to_tensor(data=roi_center, dtype=torch.int16, wrap_sequence=True, device="cpu") - roi_size_t = convert_to_tensor(data=roi_size, dtype=torch.int16, wrap_sequence=True, device="cpu") - _zeros = torch.zeros_like(roi_center_t) - half = torch.divide(roi_size_t, 2, rounding_mode="floor") - roi_start_t = torch.maximum(roi_center_t - half, _zeros) - roi_end_t = torch.maximum(roi_start_t + roi_size_t, roi_start_t) + centers, sizes = _broadcast_int_pair(roi_center, roi_size) + starts = [max(c - s // 2, 0) for c, s in zip(centers, sizes)] + ends = [st + s for st, s in zip(starts, sizes)] else: if roi_start is None or roi_end is None: raise ValueError("please specify either roi_center, roi_size or roi_start, roi_end.") - roi_start_t = convert_to_tensor(data=roi_start, dtype=torch.int16, wrap_sequence=True) - roi_start_t = torch.maximum(roi_start_t, torch.zeros_like(roi_start_t)) - roi_end_t = convert_to_tensor(data=roi_end, dtype=torch.int16, wrap_sequence=True) - roi_end_t = torch.maximum(roi_end_t, roi_start_t) - # convert to slices (accounting for 1d) - if roi_start_t.numel() == 1: - return ensure_tuple([slice(int(roi_start_t.item()), int(roi_end_t.item()))]) - return ensure_tuple([slice(int(s), int(e)) for s, e in zip(roi_start_t.tolist(), roi_end_t.tolist())]) + starts, ends = _broadcast_int_pair(roi_start, roi_end) + starts = [max(s, 0) for s in starts] + # clamp each end to its own start so no slice has negative width + return ensure_tuple(slice(s, max(e, s)) for s, e in zip(starts, ends)) def __call__( # type: ignore[override] self, img: torch.Tensor, slices: tuple[slice, ...], lazy: bool | None = None @@ -965,6 +974,7 @@ class RandWeightedCrop(Randomizable, TraceableTransform, LazyTransform, MultiSam weight_map: weight map used to generate patch samples. The weights must be non-negative. Each element denotes a sampling weight of the spatial location. 0 indicates no sampling. It should be a single-channel array in shape, for example, `(1, spatial_dim_0, spatial_dim_1, ...)`. + The weight map is only used to compute the patch sample locations; it is not cropped itself. lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False. """ @@ -1129,6 +1139,7 @@ def __init__( self.bg_indices = bg_indices self.allow_smaller = allow_smaller + # pyrefly: ignore [bad-override] def randomize( self, label: torch.Tensor | None = None, @@ -1318,6 +1329,7 @@ def __init__( self.warn = warn self.max_samples_per_class = max_samples_per_class + # pyrefly: ignore [bad-override] def randomize( self, label: torch.Tensor | None = None, diff --git a/monai/transforms/croppad/dictionary.py b/monai/transforms/croppad/dictionary.py index 510ff72938c..d089cea457b 100644 --- a/monai/transforms/croppad/dictionary.py +++ b/monai/transforms/croppad/dictionary.py @@ -19,7 +19,7 @@ from collections.abc import Callable, Hashable, Mapping, Sequence from copy import deepcopy -from typing import Any, TypeAlias, cast +from typing import Any, TypeAlias, cast # pyrefly: ignore [missing-module-attribute] import numpy as np import torch @@ -943,7 +943,11 @@ class RandWeightedCropd(Randomizable, MapTransform, LazyTransform, MultiSampleTr keys: keys of the corresponding items to be transformed. See also: :py:class:`monai.transforms.compose.MapTransform` w_key: key for the weight map. The corresponding value will be used as the sampling weights, - it should be a single-channel array in size, for example, `(1, spatial_dim_0, spatial_dim_1, ...)` + it should be a single-channel array with shape, for example, `(1, spatial_dim_0, spatial_dim_1, ...)`. + The weight map is only used to compute the patch sample locations; it is not cropped itself. + To obtain the cropped weight map (e.g. to batch it alongside the image), include ``w_key`` in + ``keys`` so it is cropped with the same sample centers; otherwise it is passed through unchanged + at its original spatial size. spatial_size: the spatial size of the image patch e.g. [224, 224, 128]. If its components have non-positive values, the corresponding size of `img` will be used. num_samples: number of samples (image patches) to take in the returned list. @@ -1100,6 +1104,7 @@ def set_random_state( self.cropper.set_random_state(seed, state) return self + # pyrefly: ignore [bad-override] def randomize( self, label: torch.Tensor | None = None, @@ -1262,6 +1267,7 @@ def set_random_state( self.cropper.set_random_state(seed, state) return self + # pyrefly: ignore [bad-override] def randomize( self, label: torch.Tensor, indices: list[NdarrayOrTensor] | None = None, image: torch.Tensor | None = None ) -> None: diff --git a/monai/transforms/croppad/functional.py b/monai/transforms/croppad/functional.py index acf42849d37..378f1cf688c 100644 --- a/monai/transforms/croppad/functional.py +++ b/monai/transforms/croppad/functional.py @@ -22,7 +22,7 @@ from monai.config.type_definitions import NdarrayTensor from monai.data.meta_obj import get_track_meta -from monai.data.meta_tensor import MetaTensor +from monai.data.meta_tensor import MetaTensor, get_spatial_ndim from monai.data.utils import to_affine_nd from monai.transforms.inverse import TraceableTransform from monai.transforms.utils import convert_pad_mode, create_translate @@ -132,7 +132,7 @@ def crop_or_pad_nd(img: torch.Tensor, translation_mat, spatial_size: tuple[int, mode: the padding mode. kwargs: other arguments for the `np.pad` or `torch.pad` function. """ - ndim = len(img.shape) - 1 + ndim = get_spatial_ndim(img) matrix_np = np.round(to_affine_nd(ndim, convert_to_numpy(translation_mat, wrap_sequence=True).copy())) matrix_np = to_affine_nd(len(spatial_size), matrix_np) cc = np.asarray(np.meshgrid(*[[0.5, x - 0.5] for x in spatial_size], indexing="ij")) diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py index 23a57ae9fbe..a23c867d60a 100644 --- a/monai/transforms/intensity/array.py +++ b/monai/transforms/intensity/array.py @@ -26,6 +26,7 @@ from monai.config import DtypeLike from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor from monai.data.meta_obj import get_track_meta +from monai.data.meta_tensor import get_spatial_ndim from monai.data.ultrasound_confidence_map import UltrasoundConfidenceMap from monai.data.utils import get_random_patch, get_valid_patch_size from monai.networks.layers import GaussianFilter, HilbertTransform, MedianFilter, SavitzkyGolayFilter @@ -353,17 +354,14 @@ def __init__( self.dtype = dtype def _stdshift(self, img: NdarrayOrTensor) -> NdarrayOrTensor: - ones: Callable std: Callable if isinstance(img, torch.Tensor): - ones = torch.ones std = partial(torch.std, unbiased=False) else: - ones = np.ones std = np.std - slices = (img != 0) if self.nonzero else ones(img.shape, dtype=bool) - if slices.any(): + slices = (img != 0) if self.nonzero else () + if not self.nonzero or (isinstance(slices, (np.ndarray, torch.Tensor)) and slices.any()): offset = self.factor * std(img[slices]) img[slices] = img[slices] + offset return img @@ -1130,12 +1128,12 @@ def __init__( self.upper = upper self.sharpness_factor = sharpness_factor self.channel_wise = channel_wise - if return_clipping_values: - self.clipping_values: list[tuple[float | None, float | None]] = [] self.return_clipping_values = return_clipping_values self.dtype = dtype - def _clip(self, img: NdarrayOrTensor) -> NdarrayOrTensor: + def _clip( + self, img: NdarrayOrTensor, clipping_values: list[tuple[float | None, float | None]] | None = None + ) -> NdarrayOrTensor: if self.sharpness_factor is not None: lower_percentile = percentile(img, self.lower) if self.lower is not None else None upper_percentile = percentile(img, self.upper) if self.upper is not None else None @@ -1145,8 +1143,8 @@ def _clip(self, img: NdarrayOrTensor) -> NdarrayOrTensor: upper_percentile = percentile(img, self.upper) if self.upper is not None else percentile(img, 100) img = clip(img, lower_percentile, upper_percentile) - if self.return_clipping_values: - self.clipping_values.append( + if clipping_values is not None: + clipping_values.append( ( ( lower_percentile @@ -1167,16 +1165,17 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ Apply the transform to `img`. """ + clipping_values: list[tuple[float | None, float | None]] | None = [] if self.return_clipping_values else None img = convert_to_tensor(img, track_meta=get_track_meta()) img_t = convert_to_tensor(img, track_meta=False) if self.channel_wise: - img_t = torch.stack([self._clip(img=d) for d in img_t]) # type: ignore + img_t = torch.stack([self._clip(img=d, clipping_values=clipping_values) for d in img_t]) # type: ignore else: - img_t = self._clip(img=img_t) + img_t = self._clip(img=img_t, clipping_values=clipping_values) img = convert_to_dst_type(img_t, dst=img)[0] - if self.return_clipping_values: - img.meta["clipping_values"] = self.clipping_values # type: ignore + if clipping_values is not None: + img.meta["clipping_values"] = clipping_values # type: ignore return img @@ -1603,7 +1602,7 @@ def __init__(self, radius: Sequence[int] | int = 1) -> None: def __call__(self, img: NdarrayTensor) -> NdarrayTensor: img = convert_to_tensor(img, track_meta=get_track_meta()) img_t, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float) - spatial_dims = img_t.ndim - 1 + spatial_dims = get_spatial_ndim(img) r = ensure_tuple_rep(self.radius, spatial_dims) median_filter_instance = MedianFilter(r, spatial_dims=spatial_dims) out_t: torch.Tensor = median_filter_instance(img_t) @@ -1639,7 +1638,7 @@ def __call__(self, img: NdarrayTensor) -> NdarrayTensor: sigma = [torch.as_tensor(s, device=img_t.device) for s in self.sigma] else: sigma = torch.as_tensor(self.sigma, device=img_t.device) - gaussian_filter = GaussianFilter(img_t.ndim - 1, sigma, approx=self.approx) + gaussian_filter = GaussianFilter(get_spatial_ndim(img), sigma, approx=self.approx) out_t: torch.Tensor = gaussian_filter(img_t.unsqueeze(0)).squeeze(0) out, *_ = convert_to_dst_type(out_t, dst=img, dtype=out_t.dtype) @@ -1696,7 +1695,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen if not self._do_transform: return img - sigma = ensure_tuple_size(vals=(self.x, self.y, self.z), dim=img.ndim - 1) + sigma = ensure_tuple_size(vals=(self.x, self.y, self.z), dim=get_spatial_ndim(img)) return GaussianSmooth(sigma=sigma, approx=self.approx)(img) @@ -1746,7 +1745,7 @@ def __call__(self, img: NdarrayTensor) -> NdarrayTensor: img_t, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float32) gf1, gf2 = ( - GaussianFilter(img_t.ndim - 1, sigma, approx=self.approx).to(img_t.device) + GaussianFilter(get_spatial_ndim(img), sigma, approx=self.approx).to(img_t.device) for sigma in (self.sigma1, self.sigma2) ) blurred_f = gf1(img_t.unsqueeze(0)) @@ -1834,8 +1833,9 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen if self.x2 is None or self.y2 is None or self.z2 is None or self.a is None: raise RuntimeError("please call the `randomize()` function first.") - sigma1 = ensure_tuple_size(vals=(self.x1, self.y1, self.z1), dim=img.ndim - 1) - sigma2 = ensure_tuple_size(vals=(self.x2, self.y2, self.z2), dim=img.ndim - 1) + _sp = get_spatial_ndim(img) + sigma1 = ensure_tuple_size(vals=(self.x1, self.y1, self.z1), dim=_sp) + sigma2 = ensure_tuple_size(vals=(self.x2, self.y2, self.z2), dim=_sp) return GaussianSharpen(sigma1=sigma1, sigma2=sigma2, alpha=self.a, approx=self.approx)(img) diff --git a/monai/transforms/inverse.py b/monai/transforms/inverse.py index 154fa07647d..f250fdfaf65 100644 --- a/monai/transforms/inverse.py +++ b/monai/transforms/inverse.py @@ -215,7 +215,7 @@ def track_transform_meta( orig_affine = data_t.peek_pending_affine() orig_affine = convert_to_dst_type(orig_affine, affine, dtype=torch.float64)[0] try: - affine = orig_affine @ to_affine_nd(len(orig_affine) - 1, affine, dtype=torch.float64) + affine = orig_affine @ to_affine_nd(orig_affine.shape[-1] - 1, affine, dtype=torch.float64) except RuntimeError as e: if orig_affine.ndim > 2: if data_t.is_batch: diff --git a/monai/transforms/io/array.py b/monai/transforms/io/array.py index aadd96763d9..e5d127c4b90 100644 --- a/monai/transforms/io/array.py +++ b/monai/transforms/io/array.py @@ -210,9 +210,7 @@ def __init__( try: self.register(the_reader(*args, **kwargs)) except OptionalImportError: - warnings.warn( - f"required package for reader {_r} is not installed, or the version doesn't match requirement." - ) + raise except TypeError: # the reader doesn't have the corresponding args/kwargs warnings.warn(f"{_r} is not supported with the given parameters {args} {kwargs}.") self.register(the_reader()) diff --git a/monai/transforms/lazy/functional.py b/monai/transforms/lazy/functional.py index 55fd7ef031e..4120dece388 100644 --- a/monai/transforms/lazy/functional.py +++ b/monai/transforms/lazy/functional.py @@ -257,9 +257,11 @@ def apply_pending(data: torch.Tensor | MetaTensor, pending: list | None = None, if not pending: return data, [] + _rank = data.spatial_ndim if isinstance(data, MetaTensor) else 3 + cumulative_xform = affine_from_pending(pending[0]) - if cumulative_xform.shape[0] == 3: - cumulative_xform = to_affine_nd(3, cumulative_xform) + if cumulative_xform.shape[0] < _rank + 1: + cumulative_xform = to_affine_nd(_rank, cumulative_xform) cur_kwargs = kwargs_from_pending(pending[0]) override_kwargs: dict[str, Any] = {} @@ -284,8 +286,8 @@ def apply_pending(data: torch.Tensor | MetaTensor, pending: list | None = None, data = resample(data.to(device), cumulative_xform, _cur_kwargs) next_matrix = affine_from_pending(p) - if next_matrix.shape[0] == 3: - next_matrix = to_affine_nd(3, next_matrix) + if next_matrix.shape[0] < _rank + 1: + next_matrix = to_affine_nd(_rank, next_matrix) cumulative_xform = combine_transforms(cumulative_xform, next_matrix) cur_kwargs.update(new_kwargs) diff --git a/monai/transforms/lazy/utils.py b/monai/transforms/lazy/utils.py index 75f1e3529d0..1f8506031a9 100644 --- a/monai/transforms/lazy/utils.py +++ b/monai/transforms/lazy/utils.py @@ -228,11 +228,13 @@ def resample(data: torch.Tensor, matrix: NdarrayOrTensor, kwargs: dict | None = img.affine = call_kwargs["dst_affine"] img = img.to(torch.float32) # consistent with monai.transforms.spatial.functional.spatial_resample return img + # pyrefly: ignore [bad-argument-type, implicit-import] img = monai.transforms.crop_or_pad_nd(img, matrix_np, out_spatial_size, mode=call_kwargs["padding_mode"]) img = img.to(torch.float32) # consistent with monai.transforms.spatial.functional.spatial_resample img.affine = call_kwargs["dst_affine"] return img + # pyrefly: ignore [bad-argument-type, implicit-import] resampler = monai.transforms.SpatialResample(**init_kwargs) resampler.lazy = False # resampler is a lazytransform with resampler.trace_transform(False): # don't track this transform in `img` diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index 47623b748d3..3b5d38cf52d 100644 --- a/monai/transforms/post/array.py +++ b/monai/transforms/post/array.py @@ -23,7 +23,7 @@ from monai.config.type_definitions import NdarrayOrTensor from monai.data.meta_obj import get_track_meta -from monai.data.meta_tensor import MetaTensor +from monai.data.meta_tensor import MetaTensor, get_spatial_ndim from monai.networks import one_hot from monai.networks.layers import GaussianFilter, apply_filter, separable_filtering from monai.transforms.inverse import InvertibleTransform @@ -624,7 +624,11 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ img = convert_to_tensor(img, track_meta=get_track_meta()) img_: torch.Tensor = convert_to_tensor(img, track_meta=False) - spatial_dims = len(img_.shape) - 1 + spatial_dims = get_spatial_ndim(img) + # Validate actual tensor shape against tracked spatial_ndim + actual_spatial = img_.ndim - 1 # channel-first layout + if actual_spatial != spatial_dims: + spatial_dims = actual_spatial img_ = img_.unsqueeze(0) # adds a batch dim if spatial_dims == 2: kernel = torch.tensor([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype=torch.float32) @@ -1104,7 +1108,7 @@ def __call__(self, image: NdarrayOrTensor) -> torch.Tensor: image_tensor = convert_to_tensor(image, track_meta=get_track_meta()) # Check/set spatial axes - n_spatial_dims = image_tensor.ndim - 1 # excluding the channel dimension + n_spatial_dims = get_spatial_ndim(image_tensor) valid_spatial_axes = list(range(n_spatial_dims)) + list(range(-n_spatial_dims, 0)) # Check gradient axes to be valid diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py index 420c8c8d8e8..2d3edaeed57 100644 --- a/monai/transforms/spatial/array.py +++ b/monai/transforms/spatial/array.py @@ -27,7 +27,7 @@ from monai.config.type_definitions import NdarrayOrTensor from monai.data.box_utils import BoxMode, StandardMode from monai.data.meta_obj import get_track_meta -from monai.data.meta_tensor import MetaTensor +from monai.data.meta_tensor import MetaTensor, get_spatial_ndim from monai.data.utils import AFFINE_TOL, affine_to_spacing, compute_shape_offset, iter_patch, to_affine_nd, zoom_affine from monai.networks.layers import AffineTransform, GaussianFilter, grid_pull from monai.networks.utils import meshgrid_ij @@ -850,12 +850,14 @@ def __call__( anti_aliasing = self.anti_aliasing if anti_aliasing is None else anti_aliasing anti_aliasing_sigma = self.anti_aliasing_sigma if anti_aliasing_sigma is None else anti_aliasing_sigma - input_ndim = img.ndim - 1 # spatial ndim + input_ndim = get_spatial_ndim(img) if self.size_mode == "all": output_ndim = len(ensure_tuple(self.spatial_size)) if output_ndim > input_ndim: input_shape = ensure_tuple_size(img.shape, output_ndim + 1, 1) img = img.reshape(input_shape) + if isinstance(img, MetaTensor): + img.spatial_ndim = output_ndim elif output_ndim < input_ndim: raise ValueError( "len(spatial_size) must be greater or equal to img spatial dimensions, " @@ -934,6 +936,9 @@ class Rotate(InvertibleTransform, LazyTransform): the output data type is always ``float32``. lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False + rotate_order: for 3D inputs, the order in which the axes are rotated about, following the convention of + :py:func:`scipy.spatial.transform.Rotation.from_euler`. See + :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour). """ backend = [TransformBackends.TORCH] @@ -947,6 +952,7 @@ def __init__( align_corners: bool = False, dtype: DtypeLike | torch.dtype = torch.float32, lazy: bool = False, + rotate_order: str = "XYZ", ) -> None: LazyTransform.__init__(self, lazy=lazy) self.angle = angle @@ -955,6 +961,7 @@ def __init__( self.padding_mode: str = padding_mode self.align_corners = align_corners self.dtype = dtype + self.rotate_order = rotate_order def __call__( self, @@ -1007,6 +1014,7 @@ def __call__( _dtype, lazy=lazy_, transform_info=self.get_transform_info(), + rotate_order=self.rotate_order, ) def inverse(self, data: torch.Tensor) -> torch.Tensor: @@ -1036,6 +1044,9 @@ def inverse_transform(self, data: torch.Tensor, transform) -> torch.Tensor: out = convert_to_dst_type(out, dst=data, dtype=out.dtype)[0] if isinstance(out, MetaTensor): affine = convert_to_tensor(out.peek_pending_affine(), track_meta=False) + # Use affine matrix shape directly (not spatial_ndim) because the affine may be + # larger than the spatial dimensions (e.g., 4x4 for 2D data), and we need to match + # the actual affine matrix rank being composed mat = to_affine_nd(len(affine) - 1, transform_t) out.affine @= convert_to_dst_type(mat, affine)[0] return out @@ -1133,7 +1144,7 @@ def __call__( during initialization for this call. Defaults to None. """ img = convert_to_tensor(img, track_meta=get_track_meta()) - _zoom = ensure_tuple_rep(self.zoom, img.ndim - 1) # match the spatial image dim + _zoom = ensure_tuple_rep(self.zoom, get_spatial_ndim(img)) _mode = self.mode if mode is None else mode _padding_mode = padding_mode or self.padding_mode _align_corners = self.align_corners if align_corners is None else align_corners @@ -1521,7 +1532,7 @@ def randomize(self, data: NdarrayOrTensor) -> None: super().randomize(None) if not self._do_transform: return None - self._axis = self.R.randint(data.ndim - 1) + self._axis = self.R.randint(get_spatial_ndim(data)) def __call__(self, img: torch.Tensor, randomize: bool = True, lazy: bool | None = None) -> torch.Tensor: """ @@ -1631,13 +1642,14 @@ def randomize(self, img: NdarrayOrTensor) -> None: super().randomize(None) if not self._do_transform: return None + _sp = get_spatial_ndim(img) self._zoom = [self.R.uniform(l, h) for l, h in zip(self.min_zoom, self.max_zoom)] if len(self._zoom) == 1: # to keep the spatial shape ratio, use same random zoom factor for all dims - self._zoom = ensure_tuple_rep(self._zoom[0], img.ndim - 1) - elif len(self._zoom) == 2 and img.ndim > 3: + self._zoom = ensure_tuple_rep(self._zoom[0], _sp) + elif len(self._zoom) == 2 and _sp > 2: # if 2 zoom factors provided for 3D data, use the first factor for H and W dims, second factor for D dim - self._zoom = ensure_tuple_rep(self._zoom[0], img.ndim - 2) + ensure_tuple(self._zoom[-1]) + self._zoom = ensure_tuple_rep(self._zoom[0], _sp - 1) + ensure_tuple(self._zoom[-1]) def __call__( self, @@ -1735,6 +1747,10 @@ class AffineGrid(LazyTransform): dimensions + 1. lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False + rotate_order: for 3D inputs, the order in which the axes are rotated about when building the + rotation from ``rotate_params``, following the convention of + :py:func:`scipy.spatial.transform.Rotation.from_euler`. See + :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour). """ backend = [TransformBackends.TORCH] @@ -1750,6 +1766,7 @@ def __init__( align_corners: bool = False, affine: NdarrayOrTensor | None = None, lazy: bool = False, + rotate_order: str = "XYZ", ) -> None: LazyTransform.__init__(self, lazy=lazy) self.rotate_params = rotate_params @@ -1761,6 +1778,7 @@ def __init__( self.dtype = _dtype if _dtype in (torch.float16, torch.float64, None) else torch.float32 self.align_corners = align_corners self.affine = affine + self.rotate_order = rotate_order def __call__( self, spatial_size: Sequence[int] | None = None, grid: torch.Tensor | None = None, lazy: bool | None = None @@ -1802,7 +1820,7 @@ def __call__( if self.affine is None: affine = torch.eye(spatial_dims + 1, device=_device) if self.rotate_params: - affine @= create_rotate(spatial_dims, self.rotate_params, device=_device, backend=_b) # type: ignore[assignment] + affine @= create_rotate(spatial_dims, self.rotate_params, device=_device, backend=_b, rotate_order=self.rotate_order) # type: ignore[assignment] if self.shear_params: affine @= create_shear(spatial_dims, self.shear_params, device=_device, backend=_b) # type: ignore[assignment] if self.translate_params: @@ -2183,6 +2201,13 @@ class Affine(InvertibleTransform, LazyTransform): This transform is capable of lazy execution. See the :ref:`Lazy Resampling topic` for more information. + + Note: + This transform assumes that the origin of the coordinate system is at the spatial center + of the image. When applying transformations (rotation, scaling, etc.), they are performed + relative to this center point. If you need transformations around a different origin, + you may need to compose this transform with translation operations or adjust your affine + matrix accordingly. """ backend = list(set(AffineGrid.backend) & set(Resample.backend)) @@ -2203,6 +2228,7 @@ def __init__( align_corners: bool = False, image_only: bool = False, lazy: bool = False, + rotate_order: str = "XYZ", ) -> None: """ The affine transformations are applied in rotate, shear, translate, scale order. @@ -2245,10 +2271,12 @@ def __init__( When `mode` is an integer, using numpy/cupy backends, this argument accepts {'reflect', 'grid-mirror', 'constant', 'grid-constant', 'nearest', 'mirror', 'grid-wrap', 'wrap'}. See also: https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.map_coordinates.html - normalized: indicating whether the provided `affine` is defined to include a normalization - transform converting the coordinates from `[-(size-1)/2, (size-1)/2]` (defined in ``create_grid``) to - `[0, size - 1]` or `[-1, 1]` in order to be compatible with the underlying resampling API. - If `normalized=False`, additional coordinate normalization will be applied before resampling. + normalized: indicates whether the provided `affine` matrix already includes coordinate + normalization. Set to ``True`` if your affine matrix is designed to work with normalized + coordinates (e.g., from image processing libraries that use normalized coordinate systems). + Set to ``False`` (default) if your affine matrix works with pixel/voxel coordinates centered + at the image center. When ``False``, MONAI will automatically apply the necessary coordinate + transformations. Most users should use the default ``False``. See also: :py:func:`monai.networks.utils.normalize_transform`. device: device on which the tensor will be allocated. dtype: data type for resampling computation. Defaults to ``float32``. @@ -2259,6 +2287,10 @@ def __init__( image_only: if True return only the image volume, otherwise return (image, affine). lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False + rotate_order: for 3D inputs, the order in which the axes are rotated about when building the rotation + from ``rotate_params``, following the convention of + :py:func:`scipy.spatial.transform.Rotation.from_euler`. See + :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour). """ LazyTransform.__init__(self, lazy=lazy) self.affine_grid = AffineGrid( @@ -2271,6 +2303,7 @@ def __init__( align_corners=align_corners, device=device, lazy=lazy, + rotate_order=rotate_order, ) self.image_only = image_only self.norm_coord = not normalized @@ -2340,6 +2373,25 @@ def __call__( @classmethod def compute_w_affine(cls, spatial_rank, mat, img_size, sp_size, align_corners: bool = False): + """ + Compute the affine matrix for transforming image coordinates, accounting for + center-based coordinate system. + + This function adjusts the provided affine transformation matrix to work with images + where transformations are applied relative to the image center rather than the origin. + It composes the input matrix with translation operations that shift between + corner-based and center-based coordinate systems. + + Args: + spatial_rank: number of spatial dimensions (e.g., 2 for 2D, 3 for 3D). + mat: the base affine transformation matrix to be adjusted. + img_size: spatial dimensions of the input image. + sp_size: spatial dimensions of the output (transformed) image. + align_corners: if True, align the corners of the initial and transformed volumes. + + Returns: + The adjusted affine matrix that can be applied to image coordinates. + """ r = int(spatial_rank) mat = to_affine_nd(r, mat) shift_1 = create_translate(r, [float(d - 1) / 2 for d in img_size[:r]]) @@ -2376,6 +2428,8 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor: out = MetaTensor(out) out.meta = data.meta # type: ignore affine = convert_data_type(out.peek_pending_affine(), torch.Tensor)[0] + # Use affine matrix shape directly (not spatial_ndim) to ensure matrix composition compatibility + # when affine is larger than spatial dimensions (e.g., 4x4 for 2D data) xform, *_ = convert_to_dst_type( Affine.compute_w_affine(len(affine) - 1, inv_affine, data.shape[1:], orig_size), affine ) @@ -2645,6 +2699,8 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor: out = MetaTensor(out) out.meta = data.meta # type: ignore affine = convert_data_type(out.peek_pending_affine(), torch.Tensor)[0] + # Use affine matrix shape directly (not spatial_ndim) to ensure matrix composition compatibility + # when affine is larger than spatial dimensions (e.g., 4x4 for 2D data) xform, *_ = convert_to_dst_type( Affine.compute_w_affine(len(affine) - 1, inv_affine, data.shape[1:], orig_size), affine ) @@ -3059,10 +3115,11 @@ def __call__( raise ValueError("the spatial size of `img` does not match with the length of `distort_steps`") all_ranges = [] - num_cells = ensure_tuple_rep(self.num_cells, len(img.shape) - 1) + _sp = get_spatial_ndim(img) + num_cells = ensure_tuple_rep(self.num_cells, _sp) if isinstance(img, MetaTensor) and img.pending_operations: warnings.warn("MetaTensor img has pending operations, transform may return incorrect results.") - for dim_idx, dim_size in enumerate(img.shape[1:]): + for dim_idx, dim_size in enumerate(img.shape[1 : 1 + _sp]): dim_distort_steps = distort_steps[dim_idx] ranges = torch.zeros(dim_size, dtype=torch.float32) cell_size = dim_size // num_cells[dim_idx] diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py index 51ad0435fc0..197b80ef807 100644 --- a/monai/transforms/spatial/dictionary.py +++ b/monai/transforms/spatial/dictionary.py @@ -917,6 +917,7 @@ def __init__( align_corners: bool = False, allow_missing_keys: bool = False, lazy: bool = False, + rotate_order: str = "XYZ", ) -> None: """ Args: @@ -969,6 +970,10 @@ def __init__( allow_missing_keys: don't raise exception if key is missing. lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False + rotate_order: for 3D inputs, the order in which the axes are rotated about when building the rotation + from ``rotate_params``, following the convention of + :py:func:`scipy.spatial.transform.Rotation.from_euler`. See + :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour). See also: - :py:class:`monai.transforms.compose.MapTransform` @@ -988,6 +993,7 @@ def __init__( dtype=dtype, # type: ignore align_corners=align_corners, lazy=lazy, + rotate_order=rotate_order, ) self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys)) @@ -1752,6 +1758,9 @@ class Rotated(MapTransform, InvertibleTransform, LazyTransform): allow_missing_keys: don't raise exception if key is missing. lazy: a flag to indicate whether this transform should execute lazily or not. Defaults to False + rotate_order: for 3D inputs, the order in which the axes are rotated about, following the convention of + :py:func:`scipy.spatial.transform.Rotation.from_euler`. See + :py:func:`monai.transforms.utils.create_rotate`. Defaults to ``"XYZ"`` (the legacy behaviour). """ backend = Rotate.backend @@ -1767,10 +1776,11 @@ def __init__( dtype: Sequence[DtypeLike | torch.dtype] | DtypeLike | torch.dtype = np.float32, allow_missing_keys: bool = False, lazy: bool = False, + rotate_order: str = "XYZ", ) -> None: MapTransform.__init__(self, keys, allow_missing_keys) LazyTransform.__init__(self, lazy=lazy) - self.rotator = Rotate(angle=angle, keep_size=keep_size, lazy=lazy) + self.rotator = Rotate(angle=angle, keep_size=keep_size, lazy=lazy, rotate_order=rotate_order) self.mode = ensure_tuple_rep(mode, len(self.keys)) self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys)) diff --git a/monai/transforms/spatial/functional.py b/monai/transforms/spatial/functional.py index d976e27916c..c44d979927e 100644 --- a/monai/transforms/spatial/functional.py +++ b/monai/transforms/spatial/functional.py @@ -26,7 +26,7 @@ from monai.config.type_definitions import NdarrayOrTensor from monai.data.box_utils import get_boxmode from monai.data.meta_obj import get_track_meta -from monai.data.meta_tensor import MetaTensor +from monai.data.meta_tensor import MetaTensor, get_spatial_ndim from monai.data.utils import AFFINE_TOL, compute_shape_offset, to_affine_nd from monai.networks.layers import AffineTransform from monai.transforms.croppad.array import ResizeWithPadOrCrop @@ -140,9 +140,10 @@ def spatial_resample( src_affine: torch.Tensor = img.peek_pending_affine() if isinstance(img, MetaTensor) else torch.eye(4) img = convert_to_tensor(data=img, track_meta=get_track_meta()) # ensure spatial rank is <= 3 - spatial_rank = min(len(img.shape) - 1, src_affine.shape[0] - 1, 3) + max_rank = max(int(img.ndim) - 1, 1) + spatial_rank = min(get_spatial_ndim(img), max_rank, 3) if (not isinstance(spatial_size, int) or spatial_size != -1) and spatial_size is not None: - spatial_rank = min(len(ensure_tuple(spatial_size)), 3) # infer spatial rank based on spatial_size + spatial_rank = min(len(ensure_tuple(spatial_size)), max_rank, 3) # infer spatial rank based on spatial_size src_affine = to_affine_nd(spatial_rank, src_affine).to(torch.float64) dst_affine = to_affine_nd(spatial_rank, dst_affine) if dst_affine is not None else src_affine dst_affine = convert_to_dst_type(dst_affine, src_affine)[0] @@ -382,7 +383,9 @@ def resize( return out.copy_meta_from(meta_info) if isinstance(out, MetaTensor) else out -def rotate(img, angle, output_shape, mode, padding_mode, align_corners, dtype, lazy, transform_info): +def rotate( + img, angle, output_shape, mode, padding_mode, align_corners, dtype, lazy, transform_info, rotate_order="XYZ" +): """ Functional implementation of rotate. This function operates eagerly or lazily according to @@ -404,6 +407,9 @@ def rotate(img, angle, output_shape, mode, padding_mode, align_corners, dtype, l the output data type is always ``float32``. lazy: a flag that indicates whether the operation should be performed lazily or not transform_info: a dictionary with the relevant information pertaining to an applied transform. + rotate_order: the order in which the axes are rotated about for 3D inputs, following the convention of + :py:func:`scipy.spatial.transform.Rotation.from_euler`. See :py:func:`monai.transforms.utils.create_rotate`. + Defaults to ``"XYZ"`` (the legacy behaviour). Ignored for 2D inputs. """ @@ -412,7 +418,7 @@ def rotate(img, angle, output_shape, mode, padding_mode, align_corners, dtype, l if input_ndim not in (2, 3): raise ValueError(f"Unsupported image dimension: {input_ndim}, available options are [2, 3].") _angle = ensure_tuple_rep(angle, 1 if input_ndim == 2 else 3) - transform = create_rotate(input_ndim, _angle) + transform = create_rotate(input_ndim, _angle, rotate_order=rotate_order) if output_shape is None: corners = np.asarray(np.meshgrid(*[(0, dim) for dim in im_shape], indexing="ij")).reshape((len(im_shape), -1)) corners = transform[:-1, :-1] @ corners # type: ignore diff --git a/monai/transforms/transform.py b/monai/transforms/transform.py index 40f95d47d61..624d1f48e08 100644 --- a/monai/transforms/transform.py +++ b/monai/transforms/transform.py @@ -93,8 +93,10 @@ def _apply_transform( data = apply_pending_transforms_in_order(transform, data, lazy, overrides, logger_name) if isinstance(data, tuple) and unpack_parameters: + # pyrefly: ignore [not-callable] return transform(*data, lazy=lazy) if isinstance(transform, LazyTrait) else transform(*data) + # pyrefly: ignore [not-callable] return transform(data, lazy=lazy) if isinstance(transform, LazyTrait) else transform(data) diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index ed4b149e6b8..297b243ff91 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -30,7 +30,7 @@ from monai.config import DtypeLike from monai.config.type_definitions import NdarrayOrTensor from monai.data.meta_obj import get_track_meta -from monai.data.meta_tensor import MetaTensor +from monai.data.meta_tensor import MetaTensor, _normalize_spatial_ndim, get_spatial_ndim from monai.data.utils import is_no_channel, no_collation, orientation_ras_lps from monai.networks.layers.simplelayers import ( ApplyFilter, @@ -314,23 +314,28 @@ def __call__(self, img: torch.Tensor) -> list[torch.Tensor]: """ Apply the transform to `img`. """ - n_out = img.shape[self.dim] + dim = self.dim if self.dim >= 0 else self.dim + img.ndim + n_out = img.shape[dim] if isinstance(img, torch.Tensor): - outputs = list(torch.split(img, 1, self.dim)) + outputs = list(torch.split(img, 1, dim)) else: - outputs = np.split(img, n_out, self.dim) + outputs = np.split(img, n_out, dim) for idx, item in enumerate(outputs): if not self.keepdim: - outputs[idx] = item.squeeze(self.dim) + outputs[idx] = item.squeeze(dim) if self.update_meta and isinstance(img, MetaTensor): - if not isinstance(item, MetaTensor): - item = MetaTensor(item, meta=img.meta) - if self.dim == 0: # don't update affine if channel dim + out = outputs[idx] + if not isinstance(out, MetaTensor): + out = MetaTensor(out, meta=img.meta) + outputs[idx] = out + if dim == 0: # don't update affine if channel dim + if not self.keepdim: + out.spatial_ndim = _normalize_spatial_ndim(out.spatial_ndim, out.ndim) continue - ndim = len(item.affine) - shift = torch.eye(ndim, device=item.affine.device, dtype=item.affine.dtype) - shift[self.dim - 1, -1] = idx - item.affine = item.affine @ shift + ndim = len(out.affine) + shift = torch.eye(ndim, device=out.affine.device, dtype=out.affine.dtype) + shift[dim - 1, -1] = idx + out.affine = out.affine @ shift return outputs @@ -571,6 +576,13 @@ def __call__(self, img): class Transpose(Transform): """ Transposes the input image based on the given `indices` dimension ordering. + + .. note:: + This transform does not update the affine matrix in the metadata. As a result, + affine-dependent transforms applied after (e.g. :py:class:`monai.transforms.Spacing`) + may produce unexpected results, because the affine no longer corresponds to the + transposed data. To reorient medical images in an affine-aware way, use + :py:class:`monai.transforms.Orientation` instead. """ backend = [TransformBackends.TORCH] @@ -1521,8 +1533,9 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: Args: img: data to be transformed, assuming `img` is channel first. """ - if max(self.spatial_dims) > img.ndim - 2 or min(self.spatial_dims) < 0: - raise ValueError(f"`spatial_dims` values must be within [0, {img.ndim - 2}]") + _sp = get_spatial_ndim(img) + if max(self.spatial_dims) > _sp - 1 or min(self.spatial_dims) < 0: + raise ValueError(f"`spatial_dims` values must be within [0, {_sp - 1}]") spatial_size = img.shape[1:] coord_channels = np.array(np.meshgrid(*tuple(np.linspace(-0.5, 0.5, s) for s in spatial_size), indexing="ij")) @@ -1690,7 +1703,7 @@ def __call__( applied_operations = img.applied_operations img_, prev_type, device = convert_data_type(img, torch.Tensor) - ndim = img_.ndim - 1 # assumes channel first format + ndim = get_spatial_ndim(img) if isinstance(self.filter, str): self.filter = self._get_filter_from_string(self.filter, self.filter_size, ndim) # type: ignore diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py index 7dd24a38805..4edf2ce45f4 100644 --- a/monai/transforms/utility/dictionary.py +++ b/monai/transforms/utility/dictionary.py @@ -636,6 +636,13 @@ def __call__(self, data: Mapping[Hashable, Any]) -> dict[Hashable, Any]: class Transposed(MapTransform, InvertibleTransform): """ Dictionary-based wrapper of :py:class:`monai.transforms.Transpose`. + + .. note:: + This transform does not update the affine matrix in the metadata. As a result, + affine-dependent transforms applied after (e.g. :py:class:`monai.transforms.Spacingd`) + may produce unexpected results, because the affine no longer corresponds to the + transposed data. To reorient medical images in an affine-aware way, use + :py:class:`monai.transforms.Orientationd` instead. """ backend = Transpose.backend @@ -752,6 +759,7 @@ def __call__(self, data): sub_keys = d[key].keys() if self.sub_keys is None else self.sub_keys # move all the sub-keys to the top level + # pyrefly: ignore [not-iterable] for sk in sub_keys: # set the top-level key for the sub-key sk_top = f"{self.prefix}_{sk}" if self.prefix else sk diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 86f9d1c3e46..0b8a65b0fb3 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -864,6 +864,7 @@ def create_rotate( radians: Sequence[float] | float, device: torch.device | None = None, backend: str = TransformBackends.NUMPY, + rotate_order: str = "XYZ", ) -> NdarrayOrTensor: """ create a 2D or 3D rotation matrix @@ -872,19 +873,33 @@ def create_rotate( spatial_dims: {``2``, ``3``} spatial rank radians: rotation radians when spatial_dims == 3, the `radians` sequence corresponds to - rotation in the 1st, 2nd, and 3rd dim respectively. + rotation about the axes named by ``rotate_order``, in the order they are listed. device: device to compute and store the output (when the backend is "torch"). backend: APIs to use, ``numpy`` or ``torch``. + rotate_order: the order in which the axes are rotated about when ``spatial_dims == 3``, + following the convention of :py:func:`scipy.spatial.transform.Rotation.from_euler`. + A string of up to three characters from ``{'x', 'y', 'z'}`` (or ``{'X', 'Y', 'Z'}``), + where ``radians[i]`` is the angle applied about ``rotate_order[i]``. Lower case letters + select extrinsic rotations (about the original fixed axes); upper case letters select + intrinsic rotations (about the moving, body-fixed axes). The default ``"XYZ"`` + reproduces the legacy behaviour (intrinsic x, then y, then z). Ignored when + ``spatial_dims == 2``. Raises: ValueError: When ``radians`` is empty. ValueError: When ``spatial_dims`` is not one of [2, 3]. + ValueError: When ``rotate_order`` is not a valid Euler axis sequence. """ _backend = look_up_option(backend, TransformBackends) if _backend == TransformBackends.NUMPY: return _create_rotate( - spatial_dims=spatial_dims, radians=radians, sin_func=np.sin, cos_func=np.cos, eye_func=np.eye + spatial_dims=spatial_dims, + radians=radians, + sin_func=np.sin, + cos_func=np.cos, + eye_func=np.eye, + order=rotate_order, ) if _backend == TransformBackends.TORCH: return _create_rotate( @@ -893,16 +908,46 @@ def create_rotate( sin_func=lambda th: torch.sin(torch.as_tensor(th, dtype=torch.float32, device=device)), cos_func=lambda th: torch.cos(torch.as_tensor(th, dtype=torch.float32, device=device)), eye_func=lambda rank: torch.eye(rank, device=device), + order=rotate_order, ) raise ValueError(f"backend {backend} is not supported") +def _validate_euler_order(order: str, num_radians: int) -> None: + """ + Validate a scipy-style Euler axis sequence. + + Args: + order: the user-facing ``rotate_order`` value, a 1-3 character axis sequence. + num_radians: number of rotation angles the sequence must accommodate. + + Raises: + ValueError: when ``order`` is not a valid Euler axis sequence for ``num_radians`` angles. + """ + if not isinstance(order, str): + raise ValueError(f"`rotate_order` must be a string, got {type(order).__name__}.") + if not 1 <= len(order) <= 3: + raise ValueError(f"`rotate_order` must contain between 1 and 3 axes, got '{order}'.") + if not (order.islower() or order.isupper()): + raise ValueError( + f"`rotate_order` must be all lower case (extrinsic) or all upper case (intrinsic), got '{order}'." + ) + lowered = order.lower() + if any(axis not in "xyz" for axis in lowered): + raise ValueError(f"`rotate_order` axes must be from 'x', 'y', 'z' (any case), got '{order}'.") + if any(lowered[i] == lowered[i + 1] for i in range(len(lowered) - 1)): + raise ValueError(f"`rotate_order` must not repeat the same axis consecutively, got '{order}'.") + if len(order) < num_radians: + raise ValueError(f"`rotate_order` '{order}' is too short for {num_radians} rotation angle(s).") + + def _create_rotate( spatial_dims: int, radians: Sequence[float] | float, sin_func: Callable = np.sin, cos_func: Callable = np.cos, eye_func: Callable = np.eye, + order: str = "XYZ", ) -> NdarrayOrTensor: radians = ensure_tuple(radians) if spatial_dims == 2: @@ -915,30 +960,25 @@ def _create_rotate( raise ValueError("radians must be non empty.") if spatial_dims == 3: - affine = None - if len(radians) >= 1: - sin_, cos_ = sin_func(radians[0]), cos_func(radians[0]) - affine = eye_func(4) - affine[1, 1], affine[1, 2] = cos_, -sin_ - affine[2, 1], affine[2, 2] = sin_, cos_ - if len(radians) >= 2: - sin_, cos_ = sin_func(radians[1]), cos_func(radians[1]) - if affine is None: - raise ValueError("Affine should be a matrix.") - _affine = eye_func(4) - _affine[0, 0], _affine[0, 2] = cos_, sin_ - _affine[2, 0], _affine[2, 2] = -sin_, cos_ - affine = affine @ _affine - if len(radians) >= 3: - sin_, cos_ = sin_func(radians[2]), cos_func(radians[2]) - if affine is None: - raise ValueError("Affine should be a matrix.") - _affine = eye_func(4) - _affine[0, 0], _affine[0, 1] = cos_, -sin_ - _affine[1, 0], _affine[1, 1] = sin_, cos_ - affine = affine @ _affine - if affine is None: + if len(radians) < 1: raise ValueError("radians must be non empty.") + _validate_euler_order(order, len(radians)) + intrinsic = order.isupper() + affine = eye_func(4) + for axis, radian in zip(order.lower(), radians): + sin_, cos_ = sin_func(radian), cos_func(radian) + _affine = eye_func(4) + if axis == "x": + _affine[1, 1], _affine[1, 2] = cos_, -sin_ + _affine[2, 1], _affine[2, 2] = sin_, cos_ + elif axis == "y": + _affine[0, 0], _affine[0, 2] = cos_, sin_ + _affine[2, 0], _affine[2, 2] = -sin_, cos_ + else: # axis == "z" + _affine[0, 0], _affine[0, 1] = cos_, -sin_ + _affine[1, 0], _affine[1, 1] = sin_, cos_ + # intrinsic rotations post-multiply (body-fixed axes); extrinsic pre-multiply (world axes) + affine = affine @ _affine if intrinsic else _affine @ affine return affine # type: ignore raise ValueError(f"Unsupported spatial_dims: {spatial_dims}, available options are [2, 3].") @@ -1184,15 +1224,15 @@ def get_largest_connected_component_mask( if num_features <= num_components: out = img_.astype(bool) else: - # ignore background - nonzeros = features[lib.nonzero(features)] - # get number voxels per feature (bincount). argsort[::-1] to get indices - # of largest components. - features_to_keep = lib.argsort(lib.bincount(nonzeros))[::-1] - # only keep the first n non-background indices - features_to_keep = features_to_keep[:num_components] - # generate labelfield. True if in list of features to keep - out = lib.isin(features, features_to_keep) + # bincount counts every label; index 0 is background, so drop it before ranking + counts = lib.bincount(features.reshape(-1)) + counts[0] = 0 + # argsort[::-1] gives labels of the largest components; keep the first n + features_to_keep = lib.argsort(counts)[::-1][:num_components] + # boolean lookup-table gather over the label field, cheaper than isin + keep = lib.zeros(counts.shape[0], dtype=bool) + keep[features_to_keep] = True + out = keep[features] return convert_to_dst_type(out, dst=img, dtype=out.dtype)[0] @@ -1244,10 +1284,14 @@ def keep_merge_components_with_points( features_neg, _ = label(img_neg_, connectivity=3, return_num=True) outs = np.zeros_like(img_pos_) + # pyrefly: ignore [missing-attribute] for bs in range(point_coords.shape[0]): + # pyrefly: ignore [bad-index] for i, p in enumerate(point_coords[bs]): + # pyrefly: ignore [bad-index] if point_labels[bs, i] in pos_val: features = features_pos + # pyrefly: ignore [bad-index] elif point_labels[bs, i] in neg_val: features = features_neg else: @@ -1455,8 +1499,10 @@ def remove_small_objects( raise RuntimeError("Skimage required.") if by_measure: + # pyrefly: ignore [missing-attribute] sr = len(img.shape[1:]) if isinstance(img, monai.data.MetaTensor): + # pyrefly: ignore [missing-attribute] _pixdim = img.pixdim elif pixdim is not None: _pixdim = ensure_tuple_rep(pixdim, sr) @@ -1659,7 +1705,6 @@ def extreme_points_to_image( rescale_max: maximum value of output data. """ # points to image - # points_image = torch.zeros(label.shape[1:], dtype=torch.float) points_image = torch.zeros_like(torch.as_tensor(label[0]), dtype=torch.float) for p in points: points_image[p] = 1.0 diff --git a/monai/transforms/utils_pytorch_numpy_unification.py b/monai/transforms/utils_pytorch_numpy_unification.py index 1bc9c206d84..db6ebc26e61 100644 --- a/monai/transforms/utils_pytorch_numpy_unification.py +++ b/monai/transforms/utils_pytorch_numpy_unification.py @@ -478,6 +478,7 @@ def max(x: NdarrayTensor, dim: int | tuple | None = None, **kwargs) -> NdarrayTe else: ret = torch.max(x, int(dim), **kwargs) # type: ignore + # pyrefly: ignore [bad-index] return ret[0] if isinstance(ret, tuple) else ret @@ -544,6 +545,7 @@ def min(x: NdarrayTensor, dim: int | tuple | None = None, **kwargs) -> NdarrayTe else: ret = torch.min(x, int(dim), **kwargs) # type: ignore + # pyrefly: ignore [bad-index] return ret[0] if isinstance(ret, tuple) else ret diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py index 3efc9b5e7fd..944f34aab77 100644 --- a/monai/utils/__init__.py +++ b/monai/utils/__init__.py @@ -89,6 +89,7 @@ is_sqrt, issequenceiterable, list_to_dict, + path_to_sqlite_uri, path_to_uri, pprint_edges, progress_bar, @@ -137,6 +138,7 @@ torch_profiler_time_cpu_gpu, torch_profiler_time_end_to_end, ) +from .safeeval import SAFE_TYPES, safe_eval from .state_cacher import StateCacher from .tf32 import detect_default_tf32, has_ampere_or_later from .type_conversion import ( diff --git a/monai/utils/dist.py b/monai/utils/dist.py index 47da2bee6ed..9c321e464e2 100644 --- a/monai/utils/dist.py +++ b/monai/utils/dist.py @@ -197,5 +197,6 @@ def __init__(self, rank: int | None = None, filter_fn: Callable = lambda rank: r ) self.rank = 0 + # pyrefly: ignore [bad-override] def filter(self, *_args): return self.filter_fn(self.rank) diff --git a/monai/utils/enums.py b/monai/utils/enums.py index be00b27d73e..7be796e6b87 100644 --- a/monai/utils/enums.py +++ b/monai/utils/enums.py @@ -390,6 +390,7 @@ def orig_meta(key: str | None = None) -> str: @staticmethod def transforms(key: str | None = None) -> str: + # pyrefly: ignore [unsupported-operation] return PostFix._get_str(key, TraceKeys.KEY_SUFFIX[1:]) diff --git a/monai/utils/misc.py b/monai/utils/misc.py index ed48d4b37d7..f11e070f9df 100644 --- a/monai/utils/misc.py +++ b/monai/utils/misc.py @@ -27,6 +27,7 @@ from math import log10 from pathlib import Path from typing import TYPE_CHECKING, Any, TypeVar, cast, overload +from urllib.parse import quote import numpy as np import torch @@ -69,6 +70,7 @@ "save_obj", "label_union", "path_to_uri", + "path_to_sqlite_uri", "pprint_edges", "check_key_duplicates", "CheckKeyDuplicatesYamlLoader", @@ -727,6 +729,23 @@ def path_to_uri(path: PathLike) -> str: return Path(path).absolute().as_uri() +def path_to_sqlite_uri(path: PathLike) -> str: + """ + Convert a database file path to a SQLite connection URI, e.g. for use as an MLflow + ``tracking_uri``. If not an absolute path, it is converted to an absolute path first. + + A forward-slash (POSIX) path is used so the URI is valid on Windows as well as POSIX: + on Windows this yields ``sqlite:///C:/path/db.sqlite`` and on POSIX ``sqlite:////path/db.sqlite``. + URI-special characters in the path (e.g. ``?``, ``#``) are percent-encoded so they are not + misparsed as query/fragment components by SQLAlchemy. + + Args: + path: input database file path, can be a string or `Path` object. + + """ + return f"sqlite:///{quote(Path(path).absolute().as_posix(), safe='/:')}" + + def pprint_edges(val: Any, n_lines: int = 20) -> str: """ Pretty print the head and tail ``n_lines`` of ``val``, and omit the middle part if the part has more than 3 lines. @@ -919,11 +938,13 @@ def is_sqrt(num: Sequence[int] | int) -> bool: def unsqueeze_right(arr: NT, ndim: int) -> NT: """Append 1-sized dimensions to `arr` to create a result with `ndim` dimensions.""" + # pyrefly: ignore [bad-index, missing-attribute] return arr[(...,) + (None,) * (ndim - arr.ndim)] def unsqueeze_left(arr: NT, ndim: int) -> NT: """Prepend 1-sized dimensions to `arr` to create a result with `ndim` dimensions.""" + # pyrefly: ignore [bad-index, missing-attribute] return arr[(None,) * (ndim - arr.ndim)] diff --git a/monai/utils/ordering.py b/monai/utils/ordering.py index 1be61f98abe..6daf5d45828 100644 --- a/monai/utils/ordering.py +++ b/monai/utils/ordering.py @@ -148,7 +148,7 @@ def _order_template(self, template: np.ndarray) -> np.ndarray: else: rows, columns, depths = (template.shape[0], template.shape[1], template.shape[2]) - sequence = eval(f"self.{self.ordering_type}_idx")(rows, columns, depths) + sequence = getattr(self, f"{self.ordering_type}_idx")(rows, columns, depths) ordering = np.array([template[tuple(e)] for e in sequence]) diff --git a/monai/utils/profiling.py b/monai/utils/profiling.py index 5eda00459ee..e3259532989 100644 --- a/monai/utils/profiling.py +++ b/monai/utils/profiling.py @@ -57,7 +57,7 @@ def torch_profiler_full(func): @wraps(func) def wrapper(*args, **kwargs): - with torch.autograd.profiler.profile(use_cuda=True) as prof: + with torch.autograd.profiler.profile() as prof: result = func(*args, **kwargs) print(prof, flush=True) @@ -76,7 +76,7 @@ def torch_profiler_time_cpu_gpu(func): @wraps(func) def wrapper(*args, **kwargs): - with torch.autograd.profiler.profile(use_cuda=True) as prof: + with torch.autograd.profiler.profile() as prof: result = func(*args, **kwargs) cpu_time = prof.self_cpu_time_total @@ -377,8 +377,6 @@ def get_times_summary(self, times_in_s=True): def get_times_summary_pd(self, times_in_s=True): """Returns the same information as `get_times_summary` but in a Pandas DataFrame.""" - import pandas as pd - summ = self.get_times_summary(times_in_s) suffix = "s" if times_in_s else "ns" columns = ["Count", f"Total Time ({suffix})", "Avg", "Std", "Min", "Max"] diff --git a/monai/utils/safeeval.py b/monai/utils/safeeval.py new file mode 100644 index 00000000000..e8f0bcb1aee --- /dev/null +++ b/monai/utils/safeeval.py @@ -0,0 +1,108 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import ast +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np + +__all__ = ["SAFE_TYPES", "safe_eval"] + +# default set of safe AST node types +SAFE_TYPES: Sequence[type] = ( + ast.Expression, + ast.Name, + ast.Load, + ast.Constant, + ast.BinOp, + ast.UnaryOp, + ast.Add, + ast.Sub, + ast.Mult, + ast.Div, + ast.FloorDiv, + ast.Pow, + ast.Mod, + ast.USub, + ast.UAdd, +) + + +class _RewriteConstNp(ast.NodeTransformer): + """Replaces int and float constants in the tree with those wrapped in Numpy types.""" + + def __init__(self, int_type_str: str, float_type_str: str): + self.int_type_str = int_type_str + self.float_type_str = float_type_str + + def visit_Constant(self, node): + if isinstance(node.value, bool) or not isinstance(node.value, (int, float)): + return node + type_str = self.int_type_str if isinstance(node.value, int) else self.float_type_str + func_node = ast.parse(type_str, mode="eval").body + call_node = ast.Call(func=func_node, args=[ast.Constant(value=node.value)], keywords=[]) + return ast.copy_location(call_node, node) + + +def safe_eval( + expr: str, + globals_vars: Mapping[str, Any] | None = None, + locals_vars: Mapping[str, object] | None = None, + allowed_types: Sequence[type] = SAFE_TYPES, + rewrite_np: bool = False, + int_type_str: str = "np.int32", + float_type_str: str = "np.float32", +) -> Any: + """ + Evaluate the Python expression `expr` using `eval`, but only if it is a safe expression in that its parsed AST + contains nodes whose types are given in `allowed_types`. This ensures unsafe node types are excluded, if these + are present in the AST a ValueError is raised. The default set of such types in `SAFE_TYPES` ensures only + expressions with constants and names can be evaluated, so excludes attribute access, indexing, and calls. Code + injection is infeasible through such expressions, so this is a safe and secure way of evaluating simple expressions. + + If `rewrite_np` is True, int and float constants in the given expression will be wrapped with Numpy types as given + by `int_type_str` and `float_type_str`. These are expected to be constructor names prefixed with `np.` as Numpy + will be present in the expression global variables under that name. The values can be changed to other types if + needed, such as "int64". One advantage of doing this is to avoid denial-of-service attacks by attempting to evaluate + an expression which is incredibly slow under native Python but fast (though potentially erroneous) under Numpy. + + Args: + expr: expression to evaluate, this will be stripped before parsing to avoid indentation complaints + globals_vars: global variable mapping, this will be treated as read-only for this function, unlike `eval` + locals_vars: local variable mapping + allowed_types: sequence of allowed AST types which can be found in `expr` when parsed + rewrite_np: if True, wrap int or float literals in Numpy types + int_type_str: int Numpy wrapping type string + float_type_str: float Numpy wrapping type string + + Raises: + ValueError: raised when any node in the AST parsed from `expr` has a type not in `allowed_types` + + Returns: + The evaluated expression value, using `eval` with `globals_vars` and `locals_vars` + """ + parsed = ast.parse(expr.strip(), mode="eval") + + # collect nodes in the AST which aren't permitted and unparse them for inclusion in the exception message + disallowed = [ast.unparse(n) for n in ast.walk(parsed) if not isinstance(n, tuple(allowed_types))] + + if disallowed: + raise ValueError(f"Unsafe expression `{expr}` not evaluated, contains disallowed components: {disallowed}") + + if rewrite_np: + parsed = _RewriteConstNp(int_type_str, float_type_str).visit(parsed) + ast.fix_missing_locations(parsed) + locals_vars = {**(locals_vars or {}), "np": np} + + return eval(compile(parsed, "", "eval"), dict(globals_vars) if globals_vars else None, locals_vars) diff --git a/monai/utils/type_conversion.py b/monai/utils/type_conversion.py index b5dfb580c5a..2d7374c9d03 100644 --- a/monai/utils/type_conversion.py +++ b/monai/utils/type_conversion.py @@ -333,10 +333,12 @@ def convert_data_type( orig_device = data.device if isinstance(data, torch.Tensor) else None + # pyrefly: ignore [bad-assignment] output_type = output_type or orig_type dtype_ = get_equivalent_dtype(dtype, output_type) data_: NdarrayTensor + # pyrefly: ignore [bad-argument-type] if issubclass(output_type, torch.Tensor): track_meta = issubclass(output_type, monai.data.MetaTensor) data_ = convert_to_tensor( diff --git a/monai/visualize/img2tensorboard.py b/monai/visualize/img2tensorboard.py index 30fd4560433..a46b725113d 100644 --- a/monai/visualize/img2tensorboard.py +++ b/monai/visualize/img2tensorboard.py @@ -172,6 +172,7 @@ def plot_2d_or_3d_image( max_frames: if plot 3D RGB image as video in TensorBoardX, set the FPS to `max_frames`. tag: tag of the plotted image on TensorBoard. """ + # pyrefly: ignore [bad-index] data_index = data[index] # as the `d` data has no batch dim, reduce the spatial dim index if positive frame_dim = frame_dim - 1 if frame_dim > 0 else frame_dim diff --git a/monai/visualize/utils.py b/monai/visualize/utils.py index e79fbba8478..e13208b0cec 100644 --- a/monai/visualize/utils.py +++ b/monai/visualize/utils.py @@ -211,7 +211,7 @@ def get_label_rgb(cmap: str, label: NdarrayOrTensor) -> NdarrayOrTensor: _cmap = plt.colormaps.get_cmap(cmap) label_np, *_ = convert_data_type(label, np.ndarray) label_rgb_np = _cmap(label_np[0]) - label_rgb_np = np.moveaxis(label_rgb_np, -1, 0)[:3] + label_rgb_np = np.moveaxis(label_rgb_np, -1, 0)[:3] # pyrefly: ignore [bad-specialization] label_rgb, *_ = convert_to_dst_type(label_rgb_np, label) return label_rgb diff --git a/monai/visualize/visualizer.py b/monai/visualize/visualizer.py index 023e4444062..1f7c7e3eda0 100644 --- a/monai/visualize/visualizer.py +++ b/monai/visualize/visualizer.py @@ -11,7 +11,7 @@ from __future__ import annotations -from collections.abc import Callable, Sized +from collections.abc import Callable, Sequence import torch import torch.nn.functional as F @@ -21,7 +21,9 @@ __all__ = ["default_upsampler"] -def default_upsampler(spatial_size: Sized, align_corners: bool = False) -> Callable[[torch.Tensor], torch.Tensor]: +def default_upsampler( + spatial_size: Sequence[int], align_corners: bool = False +) -> Callable[[torch.Tensor], torch.Tensor]: """ A linear interpolation method for upsampling the feature map. The output of this function is a callable `func`, @@ -32,6 +34,6 @@ def up(x): linear_mode = [InterpolateMode.LINEAR, InterpolateMode.BILINEAR, InterpolateMode.TRILINEAR] interp_mode = linear_mode[len(spatial_size) - 1] smode = str(interp_mode.value) - return F.interpolate(x, size=spatial_size, mode=smode, align_corners=align_corners) # type: ignore + return F.interpolate(x, size=spatial_size, mode=smode, align_corners=align_corners) return up diff --git a/pyproject.toml b/pyproject.toml index 325622b66a7..684d9050025 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,265 @@ + [build-system] requires = [ + "setuptools>=78.1.1", "wheel", - "setuptools", + "versioneer[toml]", "more-itertools>=8.0", + "ninja", + "packaging", + "torch>=2.8.0", + "numpy>=1.24,<3.0", + "backports.tarfile" # see https://github.com/Project-MONAI/MONAI/issues/8791 +] +build-backend = "setuptools.build_meta" + +[project] +name = "monai" +description = "AI Toolkit for Healthcare Imaging" +readme = { file = "README.md", content-type = "text/markdown" } +requires-python = ">=3.10" +license = { text = "Apache License 2.0" } +authors = [{ name = "MONAI Consortium", email = "monai.contact@gmail.com" }] +classifiers = [ + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "Intended Audience :: Healthcare Industry", + "Programming Language :: C++", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Scientific/Engineering :: Medical Science Apps.", + "Topic :: Scientific/Engineering :: Information Analysis", + "Topic :: Software Development", + "Topic :: Software Development :: Libraries", + "Typing :: Typed", +] +dependencies = [ "torch>=2.8.0", + "numpy>=1.24,<3.0" +] +dynamic = ["version"] + +[project.urls] +Homepage = "https://project-monai.github.io/" +Documentation = "https://monai.readthedocs.io/" +"Bug Tracker" = "https://github.com/Project-MONAI/MONAI/issues" +"Source Code" = "https://github.com/Project-MONAI/MONAI" + +[project.optional-dependencies] +# All dependencies are included here except some omitted for compatibility. Testing dependencies are typically not +# needed and so present only in "testing". Ensure requirement changes in other lists are reflected here as well. +all = [ + "clearml>=1.10.0rc0", + "cucim-cu12; platform_system == 'Linux' and python_version <= '3.10'", + "cucim-cu13; platform_system == 'Linux' and python_version >= '3.11'", + "einops", + "filelock<3.12.0", + "fire", + "gdown>=4.7.3", + "h5py", + "huggingface_hub", + "imagecodecs; platform_system == 'Linux' or platform_system == 'Darwin'", + "itk>=5.2", + "jsonschema", + "lmdb", + "lpips==0.1.4", + "matplotlib>=3.6.3", + "MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support", + "mlflow>=3.15.2", + "nibabel", "ninja", - "packaging" + "nni; platform_system == 'Linux' and 'arm' not in platform_machine and 'aarch' not in platform_machine", + "nvidia-ml-py", + "onnx_graphsurgeon", + "onnx>=1.13.0", + "onnxruntime; python_version <= '3.10'", + "onnxscript", + "openslide-bin", + "openslide-python", + "optuna", + "pandas", + "pillow!=8.3.0", + "polygraphy", + "psutil", + "pyamg>=5.0.0,<5.3.0", + "pybind11", + "pydicom", + "pynrrd", + "pytorch-ignite", + "pyyaml", + "requests", + "segment_anything @ git+https://github.com/facebookresearch/segment-anything.git@6fdee8f2727f4506cfbbe553e23b895e27956588", + "scikit-image>=0.19.0", + "scipy>=1.12.0", + "tensorboard>=2.12.0", + "tensorboardX", + "tifffile; platform_system == 'Linux' or platform_system == 'Darwin'", + "torchio", + "torchvision", + "tqdm>=4.47.0", + "transformers>=5.5.0", + "zarr" +] +clearml = ["clearml>=1.10.0rc0"] +cucim = [ + "cucim-cu12; platform_system == 'Linux' and python_version <= '3.10'", + "cucim-cu13; platform_system == 'Linux' and python_version >= '3.11'" +] +cupy = ["cupy-cuda13x!=14.1.0"] # not in all, the choice between cuda12x and cuda13x that can't be resolved here +einops = ["einops"] +fire = ["fire"] +gdown = ["gdown>=4.7.3"] +h5py = ["h5py"] +huggingface_hub = ["huggingface_hub"] +hyena = ["nvsubquadratic>=0.1.1", "omegaconf", "einops"] # omitted from all for compatibility +ignite = ["pytorch-ignite"] +imagecodecs = ["imagecodecs; platform_system == 'Linux' or platform_system == 'Darwin'"] +itk = ["itk>=5.2"] +jsonschema = ["jsonschema"] +lmdb = ["lmdb"] +lpips = ["lpips==0.1.4"] +matplotlib = ["matplotlib>=3.6.3"] +metrics_reloaded = ["MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support"] +mlflow = ["mlflow>=3.15.2"] +nibabel = ["nibabel"] +nni = [ + "nni; platform_system == 'Linux' and 'arm' not in platform_machine and 'aarch' not in platform_machine", + "filelock<3.12.0", # https://github.com/microsoft/nni/issues/5523 + "typeguard<3" # https://github.com/microsoft/nni/issues/5457 +] +onnx = ["onnx>=1.13.0", "onnxruntime; python_version <= '3.10'", "onnx_graphsurgeon", "onnxscript"] +openslide = ["openslide-python", "openslide-bin"] +optuna = ["optuna"] +pandas = ["pandas"] +pillow = ["pillow!=8.3.0"] # https://github.com/python-pillow/Pillow/issues/5571 +polygraphy = ["polygraphy"] +psutil = ["psutil"] +pyamg = ["pyamg>=5.0.0,<5.3.0"] +pybind11 = ["pybind11"] +pydicom = ["pydicom"] +pynrrd = ["pynrrd"] +pynvml = ["nvidia-ml-py"] +pyyaml = ["pyyaml"] +requests = ["requests"] +segment_anything = ["segment_anything @ git+https://github.com/facebookresearch/segment-anything.git@6fdee8f2727f4506cfbbe553e23b895e27956588"] +scipy = ["scipy>=1.12.0"] +skimage = ["scikit-image>=0.19.0"] +tensorboard = ["tensorboard>=2.12.0"] # https://github.com/Project-MONAI/MONAI/issues/7434 +tensorboardX = ["tensorboardX"] +tifffile = ["tifffile; platform_system == 'Linux' or platform_system == 'Darwin'"] +torchio = ["torchio"] +torchvision = ["torchvision"] +tqdm = ["tqdm>=4.47.0"] +transformers = ["transformers>=5.5.0"] # 5.x needs the transchex BertLayer/BertConfig updates; re-verify the NGC image float8 concern +zarr = ["zarr"] +# these dependencies are for testing/building only, they aren't needed for regular use so don't appear in "all" +testing = [ + "black>=26.5.1", + "coverage>=5.5", + "isort>9.0.0", + "mccabe", + "packaging", + "parameterized", + "pep8-naming", + "pre-commit", + "pycodestyle", + "pyflakes", + "pyrefly>=1.0.0", + "ruff>=0.16.5", + "tomli", # used in print_dependencies.py for Python<3.11 + "types-PyYAML", + "types-setuptools" +] + +[tool.setuptools] +license-files = ["LICENSE"] + +[tool.setuptools.dynamic] +version = {attr = "monai.__version__"} + +[tool.versioneer] +VCS = "git" +style = "pep440" +versionfile_source = "monai/_version.py" +versionfile_build = "monai/_version.py" +tag_prefix = "" +parentdir_prefix = "" + +[tool.isort] +known_first_party = ["monai"] +profile = "black" +line_length = 120 +skip = [".git", ".eggs", "venv", ".venv", "versioneer.py", "_version.py", "conf.py", "monai/__init__.py"] +skip_glob = ["*.pyi"] +add_imports = ["from __future__ import annotations"] +append_only = true + +[tool.mypy] +ignore_missing_imports = true +no_implicit_optional = true +warn_redundant_casts = true +warn_unused_ignores = false +warn_return_any = true +strict_equality = true +show_column_numbers = true +show_error_codes = true +pretty = false +warn_unused_configs = true +extra_checks = true +exclude = ["venv/"] + +[[tool.mypy.overrides]] +module = ["versioneer", "monai._version", "monai.eggs"] +ignore_errors = true + +[[tool.mypy.overrides]] +module = ["monai.*"] +check_untyped_defs = true +disallow_untyped_decorators = true + +[[tool.mypy.overrides]] +module = [ + "monai._extensions.*", + "monai.apps.*", + "monai.auto3dseg.*", + "monai.bundle.*", + "monai.config.*", + "monai.engines.*", + "monai.fl.*", + "monai.handlers.*", + "monai.inferers.*", + "monai.losses.*", + "monai.metrics.*", + "monai.optimizers.*", + "monai.utils.*", + "monai.visualize.*" ] +disallow_incomplete_defs = true + +[tool.coverage.run] +concurrency = ["multiprocessing"] +source = ["."] +data_file = ".coverage/.coverage" +omit = ["setup.py"] + +[tool.coverage.report] +exclude_lines = [ + "pragma: no cover", + "if TYPE_CHECKING:", + "raise NotImplementedError", + "if __name__ == .__main__.:", +] +show_missing = true +skip_covered = true + +[tool.coverage.xml] +output = "coverage.xml" [tool.black] line-length = 120 @@ -20,11 +273,11 @@ exclude = ''' \.eggs | \.git | \.hg - | \.mypy_cache | \.tox | \.venv - | venv + | \.mypy_cache | \.pytype + | venv | _build | buck-out | build @@ -36,39 +289,35 @@ exclude = ''' ) ''' -[tool.pycln] -all = true -exclude = "monai/bundle/__main__.py" - [tool.ruff] line-length = 120 target-version = "py310" [tool.ruff.lint] select = [ - "B", # flake8-bugbear - https://docs.astral.sh/ruff/rules/#flake8-bugbear-b - "C90", # mccabe (complexity) - https://docs.astral.sh/ruff/rules/#mccabe-c90 - "E", # pycodestyle errors - https://docs.astral.sh/ruff/rules/#error-e - "F", # pyflakes - https://docs.astral.sh/ruff/rules/#pyflakes-f - "N", # pep8-naming - https://docs.astral.sh/ruff/rules/#pep8-naming-n - "PIE", # flake8-pie - https://docs.astral.sh/ruff/rules/#flake8-pie-pie - "TID", # flake8-tidy-imports - https://docs.astral.sh/ruff/rules/#flake8-tidy-imports-tid - "W", # pycodestyle warnings - https://docs.astral.sh/ruff/rules/#warning-w - "NPY", # NumPy specific rules - "UP", # pyupgrade - "RUF100", # aka yesqa + "B", # flake8-bugbear - https://docs.astral.sh/ruff/rules/#flake8-bugbear-b + "C90", # mccabe (complexity) - https://docs.astral.sh/ruff/rules/#mccabe-c90 + "E", # pycodestyle errors - https://docs.astral.sh/ruff/rules/#error-e + "F", # pyflakes - https://docs.astral.sh/ruff/rules/#pyflakes-f + "N", # pep8-naming - https://docs.astral.sh/ruff/rules/#pep8-naming-n + "PIE", # flake8-pie - https://docs.astral.sh/ruff/rules/#flake8-pie-pie + "TID", # flake8-tidy-imports - https://docs.astral.sh/ruff/rules/#flake8-tidy-imports-tid + "W", # pycodestyle warnings - https://docs.astral.sh/ruff/rules/#warning-w + "NPY", # NumPy specific rules - https://docs.astral.sh/ruff/rules/#numpy-specific-rules-npy + "UP", # pyupgrade - https://docs.astral.sh/ruff/rules/#pyupgrade-up + "RUF100", # aka yesqa - https://docs.astral.sh/ruff/rules/unused-noqa/ + "F401", # unused imports - https://docs.astral.sh/ruff/rules/unused-import/ ] extend-ignore = [ - "E741", # ambiguous variable name - "F401", # unused import + "E741", # ambiguous variable name "NPY002", # numpy-legacy-random - "E203", # whitespace before ':' (pycodestyle) - "E501", # line too long (pycodestyle) - "C408", # unnecessary collection call (flake8-comprehensions) - "N812", # lowercase imported as non lowercase (pep8-naming) - "B023", # function uses loop variable (flake8-bugbear) - "B905", # zip() without an explicit strict= parameter (flake8-bugbear) - "B028", # no explicit stacklevel keyword argument found (flake8-bugbear) + "E203", # whitespace before ':' (pycodestyle) + "E501", # line too long (pycodestyle) + "C408", # unnecessary collection call (flake8-comprehensions) + "N812", # lowercase imported as non lowercase (pep8-naming) + "B023", # function uses loop variable (flake8-bugbear) + "B905", # zip() without an explicit strict= parameter (flake8-bugbear) + "B028", # no explicit stacklevel keyword argument found (flake8-bugbear) ] [tool.ruff.lint.per-file-ignores] @@ -81,6 +330,8 @@ extend-ignore = [ "monai/apps/detection/utils/ATSS_matcher.py" = [ "N999" ] +"__init__.py" = ["F401"] # TODO: change importation in __init__.py files to suit F401 +"monai/bundle/__main__.py" = ["F401"] [tool.ruff.lint.mccabe] max-complexity = 50 # todo lower this treshold when yesqa id replaced with Ruff's RUF100 @@ -116,3 +367,57 @@ precise_return = true protocols = true # Experimental: Only load submodules that are explicitly imported. strict_import = false + +[tool.pyrefly] +# Check only the monai package +project-includes = ["monai/"] + +# Exclude auto-generated and vendored files +project-excludes = [ + "**/venv/**", + "**/.venv/**", + "versioneer.py", + "monai/_version.py", +] + +# Match CI environment +python-version = "3.10" +python-platform = "linux" + +# "legacy" preset provides a smooth migration from previous type checkers +preset = "legacy" + +# Check unannotated defs (previously enforced in mypy config) +check-unannotated-defs = true + +[tool.pyrefly.errors] +# Ignore missing imports +missing-import = "ignore" + +# Suppress unused-ignore warnings +unused-ignore = "ignore" + +# Suppress implicit-import globally (MONAI style uses lazy imports) +implicit-import = "ignore" + +# Downgrade errors in unannotated/dynamic code to warnings +# (pre-existing issues, not new — will fix incrementally) +bad-assignment = "warn" +bad-return = "warn" +bad-argument-type = "warn" +invalid-annotation = "ignore" +not-iterable = "warn" +not-callable = "warn" +bad-index = "warn" + +# Pre-existing errors not flagged by previous type checkers +# Suppress for a smooth migration; fix incrementally +missing-attribute = "ignore" +bad-override = "ignore" +no-matching-overload = "ignore" +unsupported-operation = "ignore" +unnecessary-type-conversion = "ignore" +missing-module-attribute = "ignore" +not-a-type = "ignore" +invalid-yield = "ignore" +deprecated = "ignore" diff --git a/requirements-dev.txt b/requirements-dev.txt deleted file mode 100644 index 08fcdc2b0e7..00000000000 --- a/requirements-dev.txt +++ /dev/null @@ -1,65 +0,0 @@ -# Full requirements for developments --r requirements-min.txt -pytorch-ignite -gdown>=4.7.3 -scipy>=1.12.0 -itk>=5.2 -nibabel -pillow!=8.3.0 # https://github.com/python-pillow/Pillow/issues/5571 -tensorboard>=2.12.0 # https://github.com/Project-MONAI/MONAI/issues/7434 -scikit-image>=0.19.0 -tqdm>=4.47.0 -lmdb -mccabe -pep8-naming -pycodestyle -pyflakes -black>=26.3.1 -isort>=5.1, <6, !=6.0.0 -ruff>=0.14.11,<0.15 -pybind11 -types-setuptools -mypy>=1.5.0, <1.12.0 -ninja -torchio -torchvision -psutil -cucim-cu12; platform_system == "Linux" and python_version <= "3.10" -cucim-cu13; platform_system == "Linux" and python_version >= '3.11' -openslide-python -openslide-bin -imagecodecs; platform_system == "Linux" or platform_system == "Darwin" -tifffile; platform_system == "Linux" or platform_system == "Darwin" -pandas -requests -einops -transformers>=4.53.0 -mlflow>=2.12.2,<3.13 -clearml>=1.10.0rc0 -matplotlib>=3.6.3 -tensorboardX -types-PyYAML -pyyaml -fire -jsonschema -pynrrd -pre-commit -pydicom -h5py -nni==2.10.1; platform_system == "Linux" and "arm" not in platform_machine and "aarch" not in platform_machine -optuna -git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded -onnx>=1.13.0 -onnxscript -onnxruntime -typeguard<3 # https://github.com/microsoft/nni/issues/5457 -filelock<3.12.0 # https://github.com/microsoft/nni/issues/5523 -zarr -lpips==0.1.4 -nvidia-ml-py -huggingface_hub -pyamg>=5.0.0, <5.3.0 -git+https://github.com/facebookresearch/segment-anything.git@6fdee8f2727f4506cfbbe553e23b895e27956588 -onnx_graphsurgeon -polygraphy -pytest # FIXME: added to get around cupy 14.1.0 creating the requirement through polygraphy and trt_compiler somehow diff --git a/requirements-min.txt b/requirements-min.txt deleted file mode 100644 index ddda9064a6b..00000000000 --- a/requirements-min.txt +++ /dev/null @@ -1,8 +0,0 @@ -# Requirements for minimal tests --r requirements.txt -setuptools>=50.3.0,<66.0.0,!=60.6.0 ; python_version < "3.12" -setuptools>=70.2.0,<=79.0.1; python_version >= "3.12" -coverage>=5.5 -parameterized -packaging -backports.tarfile diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 7d283182a40..00000000000 --- a/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -torch>=2.8.0 -numpy>=1.24,<3.0 diff --git a/runtests.sh b/runtests.sh index 431e9298c35..0fc18b36ecd 100755 --- a/runtests.sh +++ b/runtests.sh @@ -49,7 +49,7 @@ doRuffFix=false doClangFormat=false doCopyRight=false doPytypeFormat=false -doMypyFormat=false +doPyreflyFormat=false doCleanup=false doDistTests=false doPrecommit=false @@ -61,7 +61,7 @@ PY_EXE=${MONAI_PY_EXE:-$(which python)} function print_usage { echo "runtests.sh [--codeformat] [--autofix] [--black] [--isort] [--pylint] [--ruff]" - echo " [--clangformat] [--precommit] [--pytype] [-j number] [--mypy]" + echo " [--clangformat] [--precommit] [--pytype] [-j number] [--pyrefly]" echo " [--unittests] [--disttests] [--coverage] [--quick] [--min] [--net] [--build] [--list_tests]" echo " [--dryrun] [--copyright] [--clean] [--help] [--version] [--path] [--formatfix]" echo "" @@ -87,9 +87,9 @@ function print_usage { echo " --precommit : perform source code format check and fix using \"pre-commit\"" echo "" echo "Python type check options:" - echo " --pytype : perform \"pytype\" static type checks" - echo " -j, --jobs : number of parallel jobs to run \"pytype\" (default $NUM_PARALLEL)" - echo " --mypy : perform \"mypy\" static type checks" + echo " --pytype : perform \"pytype\" static type checks (deprecated, may be removed in future)" + echo " -j, --jobs : number of parallel jobs to run \"pytype\" (default $NUM_PARALLEL) (deprecated)" + echo " --pyrefly : perform \"pyrefly\" static type checks" echo "" echo "MONAI unit testing options:" echo " -u, --unittests : perform unit testing" @@ -137,8 +137,14 @@ function print_version { } function install_deps { - echo "Pip installing MONAI development dependencies and compile MONAI cpp extensions..." - ${cmdPrefix}"${PY_EXE}" -m pip install --no-build-isolation -r requirements-dev.txt + echo "Pip installing MONAI development dependencies..." + # needed for Python<3.11 + ${cmdPrefix}"${PY_EXE}" -m pip install -U tomli + # create a temporary requirements file and install using it + REQ=$(mktemp --tmpdir XXX.txt) + trap 'rm -f -- "$REQ"' EXIT + ${cmdPrefix}"${PY_EXE}" monai/config/print_dependencies.py all testing > "$REQ" + ${cmdPrefix}"${PY_EXE}" -m pip install -r "$REQ" } function compile_cpp { @@ -196,8 +202,8 @@ function clean_py { find ${TO_CLEAN} -depth -maxdepth 1 -type d -name "monai.egg-info" -exec rm -r "{}" + find ${TO_CLEAN} -depth -maxdepth 1 -type d -name "build" -exec rm -r "{}" + find ${TO_CLEAN} -depth -maxdepth 1 -type d -name "dist" -exec rm -r "{}" + - find ${TO_CLEAN} -depth -maxdepth 1 -type d -name ".mypy_cache" -exec rm -r "{}" + find ${TO_CLEAN} -depth -maxdepth 1 -type d -name ".pytype" -exec rm -r "{}" + + find ${TO_CLEAN} -depth -maxdepth 1 -type d -name ".pyrefly_cache" -exec rm -r "{}" + find ${TO_CLEAN} -depth -maxdepth 1 -type d -name ".coverage" -exec rm -r "{}" + find ${TO_CLEAN} -depth -maxdepth 1 -type d -name "__pycache__" -exec rm -r "{}" + } @@ -215,7 +221,7 @@ function print_style_fail_msg() { echo "${red}Check failed!${noColor}" if [ "$homedir" = "$currentdir" ] then - echo "Please run auto style fixes: ${green}./runtests.sh --autofix${noColor}" + echo "Please run auto style fixes if necessary: ${green}./runtests.sh --autofix${noColor}" else : fi } @@ -271,6 +277,7 @@ do doIsortFormat=true # doPylintFormat=true # https://github.com/Project-MONAI/MONAI/issues/7094 doRuffFormat=true + doPyreflyFormat=true doCopyRight=true ;; --disttests) @@ -314,10 +321,11 @@ do doPrecommit=true ;; --pytype) + echo "${yellow}WARNING: --pytype is deprecated and may be removed in a future release.${noColor}" doPytypeFormat=true ;; - --mypy) - doMypyFormat=true + --pyrefly) + doPyreflyFormat=true ;; -j|--jobs) NUM_PARALLEL=$2 @@ -611,7 +619,9 @@ fi if [ $doPytypeFormat = true ] then set +e # disable exit on failure so that diagnostics can be given on failure + echo "${yellow}WARNING: pytype is deprecated and may be removed in a future release.${noColor}" echo "${separator}${blue}pytype${noColor}" + # ensure that the necessary packages for code format testing are installed if ! is_pip_installed pytype then @@ -639,26 +649,27 @@ then fi -if [ $doMypyFormat = true ] +if [ $doPyreflyFormat = true ] then set +e # disable exit on failure so that diagnostics can be given on failure - echo "${separator}${blue}mypy${noColor}" + echo "${separator}${blue}pyrefly${noColor}" # ensure that the necessary packages for code format testing are installed - if ! is_pip_installed mypy + if ! is_pip_installed pyrefly then install_deps fi - ${cmdPrefix}"${PY_EXE}" -m mypy --version - ${cmdPrefix}"${PY_EXE}" -m mypy "$homedir" + ${cmdPrefix}"${PY_EXE}" -m pyrefly --version + # Run without file arguments to respect project-includes/excludes from pyproject.toml + ${cmdPrefix}"${PY_EXE}" -m pyrefly check - mypy_status=$? - if [ ${mypy_status} -ne 0 ] + pyrefly_status=$? + if [ ${pyrefly_status} -ne 0 ] then - : # mypy output already follows format - exit ${mypy_status} + echo "${red}failed!${noColor}" + exit ${pyrefly_status} else - : # mypy output already follows format + echo "${green}passed!${noColor}" fi set -e # enable exit on failure fi diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index d987141d0b8..00000000000 --- a/setup.cfg +++ /dev/null @@ -1,269 +0,0 @@ -[metadata] -name = monai -author = MONAI Consortium -author_email = monai.contact@gmail.com -url = https://project-monai.github.io/ -description = AI Toolkit for Healthcare Imaging -long_description = file:README.md -long_description_content_type = text/markdown; charset=UTF-8 -platforms = OS Independent -license = Apache License 2.0 -license_files = - LICENSE -project_urls = - Documentation=https://monai.readthedocs.io/ - Bug Tracker=https://github.com/Project-MONAI/MONAI/issues - Source Code=https://github.com/Project-MONAI/MONAI -classifiers = - Intended Audience :: Developers - Intended Audience :: Education - Intended Audience :: Science/Research - Intended Audience :: Healthcare Industry - Programming Language :: C++ - Programming Language :: Python :: 3 - Programming Language :: Python :: 3.10 - Programming Language :: Python :: 3.11 - Programming Language :: Python :: 3.12 - Programming Language :: Python :: 3.13 - Topic :: Scientific/Engineering - Topic :: Scientific/Engineering :: Artificial Intelligence - Topic :: Scientific/Engineering :: Medical Science Apps. - Topic :: Scientific/Engineering :: Information Analysis - Topic :: Software Development - Topic :: Software Development :: Libraries - Typing :: Typed - -[options] -python_requires = >= 3.10 -# for compiling and develop setup only -# no need to specify the versions so that we could -# compile for multiple targeted versions. -setup_requires = - torch - ninja - packaging -install_requires = - torch>=2.8.0 - numpy>=1.24,<3.0 - -[options.extras_require] -all = - nibabel - ninja - scikit-image>=0.14.2 - scipy>=1.12.0 - pillow - tensorboard - gdown>=4.7.3 - pytorch-ignite==0.4.11 - torchio - torchvision - itk>=5.2 - tqdm>=4.47.0 - lmdb - psutil - cucim-cu12; platform_system == "Linux" and python_version <= '3.10' - cucim-cu13; platform_system == "Linux" and python_version >= '3.11' - openslide-python - openslide-bin - tifffile; platform_system == "Linux" or platform_system == "Darwin" - imagecodecs; platform_system == "Linux" or platform_system == "Darwin" - pandas - einops - transformers>=4.53.0 - mlflow>=2.12.2,<3.13 - clearml>=1.10.0rc0 - matplotlib>=3.6.3 - tensorboardX - pyyaml - fire - jsonschema - pynrrd - pydicom - h5py - nni; platform_system == "Linux" and "arm" not in platform_machine and "aarch" not in platform_machine - optuna - onnx>=1.13.0 - onnxruntime - zarr - lpips==0.1.4 - nvidia-ml-py - huggingface_hub - pyamg>=5.0.0, <5.3.0 -nibabel = - nibabel -ninja = - ninja -skimage = - scikit-image>=0.14.2 -scipy = - scipy>=1.12.0 -pillow = - pillow!=8.3.0 -tensorboard = - tensorboard -gdown = - gdown>=4.7.3 -ignite = - pytorch-ignite==0.4.11 -torchio = - torchio -torchvision = - torchvision -itk = - itk>=5.2 -tqdm = - tqdm>=4.47.0 -lmdb = - lmdb -psutil = - psutil -cucim = - cucim-cu12; platform_system == "Linux" and python_version <= '3.10' - cucim-cu13; platform_system == "Linux" and python_version >= '3.11' -openslide = - openslide-python - openslide-bin -tifffile = - tifffile; platform_system == "Linux" or platform_system == "Darwin" -imagecodecs = - imagecodecs; platform_system == "Linux" or platform_system == "Darwin" -pandas = - pandas -einops = - einops -transformers = - transformers>=4.36.0, <4.41.0; python_version <= '3.10' -mlflow = - mlflow>=2.12.2,<3.13 -matplotlib = - matplotlib>=3.6.3 -clearml = - clearml -tensorboardX = - tensorboardX -pyyaml = - pyyaml -fire = - fire -packaging = - packaging -jsonschema = - jsonschema -pynrrd = - pynrrd -pydicom = - pydicom -h5py = - h5py -nni = - nni; platform_system == "Linux" and "arm" not in platform_machine and "aarch" not in platform_machine -optuna = - optuna -onnx = - onnx>=1.13.0 - onnxruntime; python_version <= '3.10' -zarr = - zarr -lpips = - lpips==0.1.4 -pynvml = - nvidia-ml-py -polygraphy = - polygraphy - -# # workaround https://github.com/Project-MONAI/MONAI/issues/5882 -# MetricsReloaded = - # MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded -huggingface_hub = - huggingface_hub -pyamg = - pyamg>=5.0.0, <5.3.0 -# segment-anything = -# segment-anything @ git+https://github.com/facebookresearch/segment-anything@6fdee8f2727f4506cfbbe553e23b895e27956588#egg=segment-anything - -[isort] -known_first_party = monai -profile = black -line_length = 120 -skip = .git, .eggs, venv, .venv, versioneer.py, _version.py, conf.py, monai/__init__.py -skip_glob = *.pyi -add_imports = from __future__ import annotations -append_only = true - -[versioneer] -VCS = git -style = pep440 -versionfile_source = monai/_version.py -versionfile_build = monai/_version.py -tag_prefix = -parentdir_prefix = - -[mypy] -# Suppresses error messages about imports that cannot be resolved. -ignore_missing_imports = True -# Changes the treatment of arguments with a default value of None by not implicitly making their type Optional. -no_implicit_optional = True -# Warns about casting an expression to its inferred type. -warn_redundant_casts = True -# No error on unneeded # type: ignore comments. -warn_unused_ignores = False -# Shows a warning when returning a value with type Any from a function declared with a non-Any return type. -warn_return_any = True -# Prohibit equality checks, identity checks, and container checks between non-overlapping types. -strict_equality = True -# Shows column numbers in error messages. -show_column_numbers = True -# Shows error codes in error messages. -show_error_codes = True -# Use visually nicer output in error messages: use soft word wrap, show source code snippets, and show error location markers. -pretty = False -# Warns about per-module sections in the config file that do not match any files processed when invoking mypy. -warn_unused_configs = True -# Make arguments prepended via Concatenate be truly positional-only. -extra_checks = True -# Allows variables to be redefined with an arbitrary type, -# as long as the redefinition is in the same block and nesting level as the original definition. -# allow_redefinition = True - -exclude = venv/ - -[mypy-versioneer] -# Ignores all non-fatal errors. -ignore_errors = True - -[mypy-monai._version] -# Ignores all non-fatal errors. -ignore_errors = True - -[mypy-monai.eggs] -# Ignores all non-fatal errors. -ignore_errors = True - -[mypy-monai.*] -# Also check the body of functions with no types in their type signature. -check_untyped_defs = True -# Warns about usage of untyped decorators. -disallow_untyped_decorators = True - -[mypy-monai.visualize.*,monai.utils.*,monai.optimizers.*,monai.losses.*,monai.inferers.*,monai.config.*,monai._extensions.*,monai.fl.*,monai.engines.*,monai.handlers.*,monai.auto3dseg.*,monai.bundle.*,monai.metrics.*,monai.apps.*] -disallow_incomplete_defs = True - -[coverage:run] -concurrency = multiprocessing -source = . -data_file = .coverage/.coverage -omit = setup.py - -[coverage:report] -exclude_lines = - pragma: no cover - if TYPE_CHECKING: - # Don't complain if tests don't hit code: - raise NotImplementedError - if __name__ == .__main__.: -show_missing = True -skip_covered = True - -[coverage:xml] -output = coverage.xml diff --git a/setup.py b/setup.py index 4d9badca415..2ebc7a9ba63 100644 --- a/setup.py +++ b/setup.py @@ -30,6 +30,7 @@ BUILD_CPP = BUILD_CUDA = False TORCH_VERSION = 0 + try: import torch @@ -126,7 +127,7 @@ def get_extensions(): ext_modules = [ extension( name="monai._C", - sources=sources, + sources=list(map(os.path.relpath, sources)), include_dirs=include_dirs, define_macros=define_macros, extra_compile_args=extra_compile_args, diff --git a/tests/apps/nnunet/__init__.py b/tests/apps/nnunet/__init__.py new file mode 100644 index 00000000000..1e97f894078 --- /dev/null +++ b/tests/apps/nnunet/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/apps/nnunet/test_nnunetv2_runner_command.py b/tests/apps/nnunet/test_nnunetv2_runner_command.py new file mode 100644 index 00000000000..506c30fad0b --- /dev/null +++ b/tests/apps/nnunet/test_nnunetv2_runner_command.py @@ -0,0 +1,78 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest +from unittest import mock + +from monai.apps.nnunet.nnunetv2_runner import nnUNetV2Runner + + +def _make_runner(export_validation_probabilities=False): + runner = nnUNetV2Runner.__new__(nnUNetV2Runner) + runner.dataset_name_or_id = "001" + runner.trainer_class_name = "nnUNetTrainer" + runner.export_validation_probabilities = export_validation_probabilities + return runner + + +class TestTrainSingleModelCommand(unittest.TestCase): + def test_store_true_flags_emit_bare(self): + runner = _make_runner() + cmd, _ = runner.train_single_model_command( + "3d_fullres", 0, 0, {"c": True, "val": True, "use_compressed": True, "disable_checkpointing": True} + ) + for flag in ("--c", "--val", "--use_compressed", "--disable_checkpointing"): + self.assertIn(flag, cmd) + self.assertNotIn("True", cmd) + + def test_store_true_flags_false_omitted(self): + runner = _make_runner() + cmd, _ = runner.train_single_model_command( + "3d_fullres", 0, 0, {"c": False, "val": False, "use_compressed": False, "disable_checkpointing": False} + ) + for flag in ("--c", "--val", "--use_compressed", "--disable_checkpointing"): + self.assertNotIn(flag, cmd) + self.assertNotIn("False", cmd) + + def test_pretrained_weights_truthy_included(self): + runner = _make_runner() + cmd, _ = runner.train_single_model_command("3d_fullres", 0, 0, {"pretrained_weights": "/path/to/weights.pth"}) + self.assertIn("-pretrained_weights", cmd) + self.assertIn("/path/to/weights.pth", cmd) + + def test_pretrained_weights_falsy_omitted(self): + runner = _make_runner() + cmd, _ = runner.train_single_model_command("3d_fullres", 0, 0, {"pretrained_weights": False}) + self.assertNotIn("-pretrained_weights", cmd) + self.assertNotIn("False", cmd) + + def test_value_kwargs_unaffected(self): + runner = _make_runner() + cmd, _ = runner.train_single_model_command("3d_fullres", 0, 0, {"npz": "something"}) + self.assertIn("--npz", cmd) + self.assertIn("something", cmd) + + +class TestValidateSingleModelCommand(unittest.TestCase): + def test_validate_emits_bare_val_flag(self): + runner = _make_runner() + with mock.patch("monai.apps.nnunet.nnunetv2_runner.run_cmd") as run_cmd: + runner.validate_single_model("3d_fullres", 0) + cmd = run_cmd.call_args.args[0] + self.assertIn("--val", cmd) + self.assertNotIn("--only_run_validation", cmd) + self.assertNotIn("True", cmd) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/apps/test_auto3dseg.py b/tests/apps/test_auto3dseg.py index 57e05d1ee6f..c310afc76ab 100644 --- a/tests/apps/test_auto3dseg.py +++ b/tests/apps/test_auto3dseg.py @@ -11,9 +11,11 @@ from __future__ import annotations +import json import os import tempfile import unittest +import warnings from copy import deepcopy from numbers import Number @@ -36,6 +38,7 @@ SampleOperations, SegSummarizer, SummaryOperations, + algo_from_json, datafold_read, verify_report_format, ) @@ -177,6 +180,20 @@ def __call__(self, data): return d +class _DummyAlgo: + """Minimal stand-in for an Auto3DSeg Algo object used in warning tests.""" + + def __init__(self) -> None: + self.template_path: str | None = None + self.output_path = os.getcwd() + + def load_state_dict(self, state: dict) -> None: + pass + + def get_output_path(self) -> str: + return self.output_path + + class TestDataAnalyzer(unittest.TestCase): def setUp(self): self.test_dir = tempfile.TemporaryDirectory() @@ -619,5 +636,23 @@ def tearDown(self) -> None: self.test_dir.cleanup() +class TestAlgoFromJsonSecurityWarning(unittest.TestCase): + def test_warns_about_untrusted_target(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + algo_file = os.path.join(tmpdir, "algo_object.json") + with open(algo_file, "w", encoding="utf-8") as f: + json.dump({"_target_": f"{__name__}._DummyAlgo"}, f) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + algo_from_json(algo_file) + + messages = [str(w.message) for w in caught] + self.assertTrue( + any("algo_object.json" in msg and "trust" in msg for msg in messages), + f"Keywords 'algo_object.json' and 'trust' not found in warning messages: {messages}", + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/apps/test_check_hash.py b/tests/apps/test_check_hash.py index 263c18703cc..75d768b0e59 100644 --- a/tests/apps/test_check_hash.py +++ b/tests/apps/test_check_hash.py @@ -11,6 +11,7 @@ from __future__ import annotations +import hashlib import os import tempfile import unittest @@ -48,6 +49,23 @@ def test_hash_type_error(self): with tempfile.TemporaryDirectory() as tempdir: check_hash(tempdir, "test_hash", "test_type") + def test_warns_when_val_is_none(self): + test_image = np.ones((5, 5, 3)) + with tempfile.TemporaryDirectory() as tempdir: + filename = os.path.join(tempdir, "test_file.png") + test_image.tofile(filename) + with self.assertWarns(UserWarning): + result = check_hash(filename, None) + self.assertTrue(result) + + def test_default_hash_type_is_sha256(self): + test_image = np.ones((5, 5, 3)) + with tempfile.TemporaryDirectory() as tempdir: + filename = os.path.join(tempdir, "test_file.png") + test_image.tofile(filename) + sha256 = hashlib.sha256(test_image.tobytes()).hexdigest() + self.assertTrue(check_hash(filename, sha256)) + if __name__ == "__main__": unittest.main() diff --git a/tests/apps/test_download_and_extract.py b/tests/apps/test_download_and_extract.py index 6d16a727351..71a810983f4 100644 --- a/tests/apps/test_download_and_extract.py +++ b/tests/apps/test_download_and_extract.py @@ -16,49 +16,68 @@ import unittest import zipfile from pathlib import Path -from urllib.error import ContentTooShortError, HTTPError from parameterized import parameterized from monai.apps import download_and_extract, download_url, extractall +from monai.apps.utils import HashCheckError from tests.test_utils import SkipIfNoModule, skip_if_downloading_fails, skip_if_quick, testing_data_config @SkipIfNoModule("requests") class TestDownloadAndExtract(unittest.TestCase): + def setUp(self): + self.testing_dir = Path(__file__).parents[1] / "testing_data" + self.config = testing_data_config("images", "mednist") + self.url = self.config["url"] + self.hash_val = self.config["hash_val"] + self.hash_type = self.config["hash_type"] + @skip_if_quick - def test_actions(self): - testing_dir = Path(__file__).parents[1] / "testing_data" - config_dict = testing_data_config("images", "mednist") - url = config_dict["url"] - filepath = Path(testing_dir) / "MedNIST.tar.gz" - output_dir = Path(testing_dir) - hash_val, hash_type = config_dict["hash_val"], config_dict["hash_type"] + def test_download_and_extract_success(self): + """End-to-end: download and extract should succeed with correct hash.""" + filepath = self.testing_dir / "MedNIST.tar.gz" + output_dir = self.testing_dir + with skip_if_downloading_fails(): - download_and_extract(url, filepath, output_dir, hash_val=hash_val, hash_type=hash_type) - download_and_extract(url, filepath, output_dir, hash_val=hash_val, hash_type=hash_type) + download_and_extract(self.url, filepath, output_dir, hash_val=self.hash_val, hash_type=self.hash_type) - wrong_md5 = "0" - with self.assertLogs(logger="monai.apps", level="ERROR"): - try: - download_url(url, filepath, wrong_md5) - except (ContentTooShortError, HTTPError, RuntimeError) as e: - if isinstance(e, RuntimeError): - # FIXME: skip MD5 check as current downloading method may fail - self.assertTrue(str(e).startswith("md5 check")) - return # skipping this test due the network connection errors - - try: - extractall(filepath, output_dir, wrong_md5) - except RuntimeError as e: - self.assertTrue(str(e).startswith("md5 check")) + self.assertTrue(filepath.exists(), "Downloaded file does not exist") + self.assertTrue(any(output_dir.iterdir()), "Extraction output is empty") + + @skip_if_quick + def test_download_url_hash_mismatch(self): + """download_url should raise HashCheckError on hash mismatch.""" + filepath = self.testing_dir / "MedNIST.tar.gz" + + with skip_if_downloading_fails(): + # First ensure file is downloaded correctly + download_url(self.url, filepath, hash_val=self.hash_val, hash_type=self.hash_type) + + # Now test incorrect hash + with self.assertRaises(HashCheckError): + download_url(self.url, filepath, hash_val="0" * len(self.hash_val), hash_type=self.hash_type) @skip_if_quick - @parameterized.expand((("icon", "tar"), ("favicon", "zip"))) - def test_default(self, key, file_type): + def test_extractall_hash_mismatch(self): + """extractall should raise HashCheckError when hash is incorrect.""" + filepath = self.testing_dir / "MedNIST.tar.gz" + output_dir = self.testing_dir + + with skip_if_downloading_fails(): + download_url(self.url, filepath, hash_val=self.hash_val, hash_type=self.hash_type) + + with self.assertRaises(HashCheckError): + extractall(filepath, output_dir, hash_val="0" * len(self.hash_val), hash_type=self.hash_type) + + @skip_if_quick + @parameterized.expand([("icon", "tar"), ("favicon", "zip")]) + def test_download_and_extract_various_formats(self, key, file_type): + """Verify different archive formats download and extract correctly.""" with tempfile.TemporaryDirectory() as tmp_dir: + img_spec = testing_data_config("images", key) + with skip_if_downloading_fails(): - img_spec = testing_data_config("images", key) download_and_extract( img_spec["url"], output_dir=tmp_dir, @@ -67,6 +86,8 @@ def test_default(self, key, file_type): file_type=file_type, ) + self.assertTrue(any(Path(tmp_dir).iterdir()), f"Extraction failed for format: {file_type}") + class TestPathTraversalProtection(unittest.TestCase): """Test cases for path traversal attack protection in extractall function.""" diff --git a/tests/apps/test_download_url_yandex.py b/tests/apps/test_download_url_yandex.py index 54d39b06ff7..b29119bf074 100644 --- a/tests/apps/test_download_url_yandex.py +++ b/tests/apps/test_download_url_yandex.py @@ -18,10 +18,6 @@ from monai.apps.utils import download_url -YANDEX_MODEL_URL = ( - "https://cloud-api.yandex.net/v1/disk/public/resources/download?" - "public_key=https%3A%2F%2Fdisk.yandex.ru%2Fd%2Fxs0gzlj2_irgWA" -) YANDEX_MODEL_FLAWED_URL = ( "https://cloud-api.yandex.net/v1/disk/public/resources/download?" "public_key=https%3A%2F%2Fdisk.yandex.ru%2Fd%2Fxs0gzlj2_irgWA-url-with-error" @@ -30,11 +26,6 @@ class TestDownloadUrlYandex(unittest.TestCase): - @unittest.skip("data source unstable") - def test_verify(self): - with tempfile.TemporaryDirectory() as tempdir: - download_url(url=YANDEX_MODEL_URL, filepath=os.path.join(tempdir, "model.pt")) - def test_verify_error(self): with tempfile.TemporaryDirectory() as tempdir: with self.assertRaises(HTTPError): diff --git a/tests/bundle/test_bundle_download.py b/tests/bundle/test_bundle_download.py index bb213cebd99..a8259ab1632 100644 --- a/tests/bundle/test_bundle_download.py +++ b/tests/bundle/test_bundle_download.py @@ -15,6 +15,7 @@ import os import tempfile import unittest +import warnings from unittest.case import skipIf, skipUnless from unittest.mock import patch @@ -24,8 +25,8 @@ import monai.networks.nets as nets from monai.apps import check_hash -from monai.bundle import ConfigParser, create_workflow, load -from monai.bundle.scripts import _examine_monai_version, _list_latest_versions, download +from monai.bundle import ConfigParser, create_workflow, load, run +from monai.bundle.scripts import _examine_monai_version, _list_latest_versions, download, download_large_files from monai.utils import optional_import from tests.test_utils import ( assert_allclose, @@ -95,6 +96,15 @@ {"model.pt": "27952767e2e154e3b0ee65defc5aed38", "model.ts": "97746870fe591f69ac09827175b00675"}, ] + +# (source, repo) pairs covering every `source` accepted by `load()`/`download()`. `repo` only +# matters for sources that read it ("github", "huggingface_hub", "ngc_private"); it's unused +# otherwise but keeps the call shape realistic for each source. +TEST_CASE_SOURCE_GITHUB = ["github", "attacker/repo"] +TEST_CASE_SOURCE_MONAIHOSTING = ["monaihosting", None] +TEST_CASE_SOURCE_NGC = ["ngc", None] +TEST_CASE_SOURCE_HUGGINGFACE_HUB = ["huggingface_hub", "attacker/repo"] + TEST_CASE_NGC_1 = [ "spleen_ct_segmentation", "0.3.7", @@ -156,7 +166,7 @@ def test_github_download_bundle(self, bundle_name, version): file_path = os.path.join(tempdir, "test_bundle", file) self.assertTrue(os.path.exists(file_path)) if file == "network.json": - self.assertTrue(check_hash(filepath=file_path, val=hash_val)) + self.assertTrue(check_hash(filepath=file_path, val=hash_val, hash_type="md5")) @parameterized.expand([TEST_CASE_3]) @skip_if_quick @@ -174,8 +184,8 @@ def test_url_download_bundle(self, bundle_files, bundle_name, url, hash_val): for file in bundle_files: file_path = os.path.join(tempdir, bundle_name, file) self.assertTrue(os.path.exists(file_path)) - if file == "network.json": - self.assertTrue(check_hash(filepath=file_path, val=hash_val)) + if file == "network.json": + self.assertTrue(check_hash(filepath=file_path, val=hash_val, hash_type="md5")) @parameterized.expand([TEST_CASE_4]) @skip_if_quick @@ -435,6 +445,15 @@ def test_load_ts_module(self, bundle_files, bundle_name, version, repo, device, class TestDownloadLargefiles(unittest.TestCase): + + def test_large_files_rejects_path_traversal(self): + with tempfile.TemporaryDirectory() as tempdir: + large_files_path = os.path.join(tempdir, "large_files.yaml") + with open(large_files_path, "w") as f: + f.write("large_files:\n" " - path: ../evil.pt\n" " url: https://example.com/evil.pt\n") + with self.assertRaises(ValueError): + download_large_files(bundle_path=tempdir) + @parameterized.expand([TEST_CASE_10]) @skip_if_quick def test_url_download_large_files(self, bundle_files, bundle_name, url, hash_val): @@ -459,7 +478,7 @@ def test_url_download_large_files(self, bundle_files, bundle_name, url, hash_val command_line_tests(cmd) for file in ["model.pt", "model.ts"]: file_path = os.path.join(tempdir, bundle_name, f"models/{file}") - self.assertTrue(check_hash(filepath=file_path, val=hash_val[file])) + self.assertTrue(check_hash(filepath=file_path, val=hash_val[file], hash_type="md5")) @skip_if_windows @@ -474,7 +493,7 @@ def test_ngc_download_bundle(self, bundle_name, version, remove_prefix, download ) full_file_path = os.path.join(tempdir, download_name, file_path) self.assertTrue(os.path.exists(full_file_path)) - self.assertTrue(check_hash(filepath=full_file_path, val=hash_val)) + self.assertTrue(check_hash(filepath=full_file_path, val=hash_val, hash_type="md5")) model = load( name=bundle_name, source="ngc", version=version, bundle_dir=tempdir, remove_prefix=remove_prefix @@ -488,5 +507,71 @@ def test_ngc_download_bundle(self, bundle_name, version, remove_prefix, download ) +class TestLoadWarnsOnConfigExecution(unittest.TestCase): + """Regression tests for GHSA-873f-pvrv-4x83: `load()`/`create_workflow()` parse and execute a + bundle's own config (arbitrary `_target_`/`$`-expression content) whenever `model` is `None`. + There is no opt-in flag -- MONAI has no way to establish whether a bundle is actually + trustworthy, so a flag would only teach callers to always pass it and ignore the risk. Instead, + a `UserWarning` is raised every time this happens, in both `load()` (via `create_workflow()`) + and `run()` (also via `create_workflow()`).""" + + def _stage_malicious_bundle(self, tempdir: str, marker: str) -> str: + name = "evil_bundle" + bundle_root = os.path.join(tempdir, name) + os.makedirs(os.path.join(bundle_root, "configs")) + os.makedirs(os.path.join(bundle_root, "models")) + torch.save({"state_dict": {}}, os.path.join(bundle_root, "models", "model.pt")) + # writes the marker directly via `pathlib` instead of shelling out through `os.system` -- + # `!r` yields a Python-source-safe literal (handling spaces and Windows backslashes alike) + # with no shell involved to reintroduce quoting/splitting issues. + payload = f"$__import__('pathlib').Path({marker!r}).write_text('pwned')" + # included under both keys so the payload runs whether the config is consumed via + # `network_def` (the `load()` tests) or via `initialize` (the `run()` test). + malicious_config = {"network_def": payload, "initialize": [payload]} + with open(os.path.join(bundle_root, "configs", "train.json"), "w") as f: + json.dump(malicious_config, f) + return name + + @parameterized.expand( + [TEST_CASE_SOURCE_GITHUB, TEST_CASE_SOURCE_MONAIHOSTING, TEST_CASE_SOURCE_NGC, TEST_CASE_SOURCE_HUGGINGFACE_HUB] + ) + def test_default_warns_and_executes_config(self, source, repo): + # `source`/`repo` only steer where `download()` would fetch from -- irrelevant here since + # the bundle is already staged on disk, so `load()` never calls `download()`. Parameterized + # anyway to confirm the warning fires the same way regardless of `source`. + with tempfile.TemporaryDirectory() as tempdir: + marker = os.path.join(tempdir, "PWNED") + name = self._stage_malicious_bundle(tempdir, marker) + with self.assertWarnsRegex(UserWarning, r"GHSA-873f-pvrv-4x83"): + with self.assertRaises(AttributeError): + # the malicious config is missing metadata.json and returns a plain `int` for + # `network_def`, so the workflow construction fails after the payload has already + # run -- this mirrors the advisory's own PoC, where the failure happens *after* RCE. + load(name=name, bundle_dir=tempdir, source=source, repo=repo) + self.assertTrue(os.path.exists(marker)) + + def test_explicit_model_skips_config_parsing(self): + with tempfile.TemporaryDirectory() as tempdir: + marker = os.path.join(tempdir, "PWNED") + name = self._stage_malicious_bundle(tempdir, marker) + model = nets.UNet(spatial_dims=2, in_channels=1, out_channels=1, channels=(4, 8), strides=(2,)) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + load(name=name, model=model, bundle_dir=tempdir, source="github", repo="attacker/repo") + self.assertFalse(os.path.exists(marker)) + + def test_run_warns_on_config_execution(self): + with tempfile.TemporaryDirectory() as tempdir: + marker = os.path.join(tempdir, "PWNED") + name = self._stage_malicious_bundle(tempdir, marker) + config_file = os.path.join(tempdir, name, "configs", "train.json") + with self.assertWarnsRegex(UserWarning, r"GHSA-873f-pvrv-4x83"): + with self.assertRaises(ValueError): + # no "run" ID is defined, so `workflow.run()` fails after `initialize()` has + # already evaluated the payload above. + run(config_file=config_file) + self.assertTrue(os.path.exists(marker)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/bundle/test_bundle_workflow.py b/tests/bundle/test_bundle_workflow.py index ceb034ecff8..2fa9b8eac48 100644 --- a/tests/bundle/test_bundle_workflow.py +++ b/tests/bundle/test_bundle_workflow.py @@ -11,11 +11,14 @@ from __future__ import annotations +import json +import logging import os import shutil import sys import tempfile import unittest +import warnings from copy import deepcopy from pathlib import Path @@ -268,5 +271,84 @@ def test_create_pythonic_workflow(self): workflow.finalize() +class TestConfigWorkflowWarnsOnLoggingConf(unittest.TestCase): + """Regression test for GHSA-wvpx-5qmp-46g3: `ConfigWorkflow` defaults `logging_file` to the + bundle's own "configs/logging.conf" and hands it to `logging.config.fileConfig`, which `eval()`s + the INI's `class=`/`args=` fields. It fires in `__init__`, before `initialize()` or `run()`, and + lives in a plain INI rather than the MONAI `$`-DSL, so it is easy to miss when reviewing a + bundle. Applying it is still not blocked -- as for GHSA-873f-pvrv-4x83, MONAI has no way to + establish whether a bundle is trustworthy -- but applying it now raises a `UserWarning`.""" + + def setUp(self): + # `fileConfig` reconfigures logging process-wide. Snapshot the root logger and restore it + # afterwards so these tests cannot leak a handler into the rest of the suite. + root = logging.getLogger() + level, handlers, filters = root.level, root.handlers[:], root.filters[:] + disabled = logging.root.manager.disable + + def _restore(): + # Detach whatever is on the root logger now, closing anything `fileConfig` installed so + # it does not linger in logging's handler registry, then put the snapshot back. Under + # `tests/runner.py` the root logger starts with no handlers, so there is nothing for + # `fileConfig` to have closed on the way in. + for handler in root.handlers[:]: + root.removeHandler(handler) + if handler not in handlers: + handler.close() + root.setLevel(level) + root.filters[:] = filters + for handler in handlers: + root.addHandler(handler) + logging.disable(disabled) + + self.addCleanup(_restore) + + def test_default_logging_conf_warns_and_executes(self): + with tempfile.TemporaryDirectory() as tempdir: + configs = os.path.join(tempdir, "configs") + os.makedirs(configs) + marker = os.path.join(tempdir, "PWNED") + with open(os.path.join(configs, "train.json"), "w") as f: + json.dump({"initialize": []}, f) + # `fileConfig` eval()s the `class=` field, so the tuple subscript runs the payload and + # still yields a usable handler class. + with open(os.path.join(configs, "logging.conf"), "w") as f: + f.write( + "[loggers]\nkeys=root\n[handlers]\nkeys=h\n[formatters]\nkeys=f\n" + "[logger_root]\nlevel=NOTSET\nhandlers=h\n" + "[handler_h]\n" + f"class=(__import__('pathlib').Path({marker!r}).write_text('pwned'), " + "__import__('logging').StreamHandler)[1]\nargs=()\nformatter=f\n" + "[formatter_f]\nformat=%(message)s\n" + ) + with self.assertWarnsRegex(UserWarning, r"GHSA-wvpx-5qmp-46g3"): + ConfigWorkflow(config_file=os.path.join(configs, "train.json"), workflow_type="train") + self.assertTrue(os.path.exists(marker)) + + def test_no_warning_when_logging_disabled(self): + """No warning when `fileConfig` is never reached -- the file exists but is opted out of.""" + with tempfile.TemporaryDirectory() as tempdir: + configs = os.path.join(tempdir, "configs") + os.makedirs(configs) + marker = os.path.join(tempdir, "PWNED") + with open(os.path.join(configs, "train.json"), "w") as f: + json.dump({"initialize": []}, f) + with open(os.path.join(configs, "logging.conf"), "w") as f: + f.write( + "[loggers]\nkeys=root\n[handlers]\nkeys=h\n[formatters]\nkeys=f\n" + "[logger_root]\nlevel=NOTSET\nhandlers=h\n" + "[handler_h]\n" + f"class=(__import__('pathlib').Path({marker!r}).write_text('pwned'), " + "__import__('logging').StreamHandler)[1]\nargs=()\nformatter=f\n" + "[formatter_f]\nformat=%(message)s\n" + ) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + ConfigWorkflow( + config_file=os.path.join(configs, "train.json"), workflow_type="train", logging_file=False + ) + self.assertFalse(os.path.exists(marker)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/bundle/test_config_parser.py b/tests/bundle/test_config_parser.py index 546957ba7e4..924287a2998 100644 --- a/tests/bundle/test_config_parser.py +++ b/tests/bundle/test_config_parser.py @@ -487,6 +487,24 @@ def test_chained_ref_backed_proxy_write_through(self): del parser.alias["y"] self.assertNotIn("y", parser.get_parsed_content("target")) + def test_ref_backed_proxy_attribute_read(self): + # Dot-notation must agree with bracket-notation on a proxy reached via $@ref: + # "alias::x" has no id in the resolver, but "x" is a key of the aliased node, so + # both notations must resolve it (parser.alias.x raised AttributeError before this + # fix, while parser.alias["x"] returned the value). + parser = ConfigParser(config={"target": {"x": 1, "y": 2}, "alias": "$@target"}, globals={"monai": "monai"}) + self.assertEqual(parser.alias.x, parser.alias["x"]) + self.assertEqual(parser.alias.x, 1) + # a key absent from the container still falls back to the container's own methods + self.assertEqual(sorted(parser.alias.keys()), ["x", "y"]) + + def test_chained_ref_backed_proxy_attribute_read(self): + # dot-notation must follow the full ref chain, as _backing_id() does for writes. + parser = ConfigParser( + config={"target": {"x": 1}, "mid": "$@target", "alias": "$@mid"}, globals={"monai": "monai"} + ) + self.assertEqual(parser.alias.x, 1) + def test_raw_is_read_only(self): with self.assertRaises(AttributeError): self.parser.A._raw = {"something": "else"} diff --git a/tests/config/test_print_dependencies.py b/tests/config/test_print_dependencies.py new file mode 100644 index 00000000000..bbf8c4c7cda --- /dev/null +++ b/tests/config/test_print_dependencies.py @@ -0,0 +1,80 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import os +import unittest +from contextlib import redirect_stdout +from io import StringIO +from tempfile import NamedTemporaryFile +from unittest.mock import patch + +from parameterized import parameterized + +from monai.config.print_dependencies import parse_dependencies, print_dependencies_argv + +TEST_TOML = """ +[build-system] +requires = ["setuptools", "wheel"] + +[project] +name = "test" +dependencies = ["torch", "numpy"] + +[project.optional-dependencies] +all = ["something", "another"] +testing = ["coverage", "black"] +""" + +PARSE_CASES = [ + ([], ["numpy", "torch"]), + (["testing"], ["black", "coverage", "numpy", "torch"]), + (["build-system"], ["numpy", "setuptools", "torch", "wheel"]), + (["*"], ["another", "black", "coverage", "numpy", "something", "torch"]), +] + + +class TestPrintDependencies(unittest.TestCase): + def setUp(self): + self.toml = NamedTemporaryFile("w", delete=False) + self.toml.write(TEST_TOML) + self.toml.close() + + def tearDown(self): + os.unlink(self.toml.name) + + @parameterized.expand(PARSE_CASES) + def test_parse_dependencies(self, sections, outputs): + deps = parse_dependencies(self.toml.name, sections) + self.assertEqual(outputs, deps) + + def test_missing_section(self): + with self.assertRaises(KeyError): + parse_dependencies(self.toml.name, ["nonexistent_section"]) + + def test_print_dependencies(self): + out = StringIO() + with redirect_stdout(out), patch("monai.config.print_dependencies.TOML_FILE", self.toml.name): + + with self.subTest("Test correct print"): + with patch("sys.argv", ["", "all", "build-system", "*"]): + print_dependencies_argv() + + self.assertGreater(out.tell(), 0) + + with self.subTest("Test missing section"): + with patch("sys.argv", ["", "nonexistent_section"]), self.assertRaises(KeyError): + print_dependencies_argv() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/data/meta_tensor/test_meta_tensor.py b/tests/data/meta_tensor/test_meta_tensor.py index c0e53fd24c9..5b63d5d7731 100644 --- a/tests/data/meta_tensor/test_meta_tensor.py +++ b/tests/data/meta_tensor/test_meta_tensor.py @@ -68,6 +68,7 @@ def check_ids(self, a, b, should_match): def check_meta(self, a: MetaTensor, b: MetaTensor) -> None: self.assertEqual(a.is_batch, b.is_batch) + self.assertEqual(a.spatial_ndim, b.spatial_ndim) meta_a, meta_b = a.meta, b.meta # need to split affine from rest of metadata aff_a = meta_a.get("affine", None) @@ -434,8 +435,12 @@ def test_astype(self): for np_types in ("float32", "np.float32", "numpy.float32", np.float32, float, "int", np.uint16): self.assertIsInstance(t.astype(np_types), np.ndarray) for pt_types in ("torch.float", torch.float, "torch.float64"): - self.assertIsInstance(t.astype(pt_types), torch.Tensor) - self.assertIsInstance(t.astype("torch.float", device="cpu"), torch.Tensor) + result = t.astype(pt_types) + self.assertIsInstance(result, MetaTensor) + self.assertEqual(result.meta.get("fname"), "filename") + result = t.astype("torch.float", device="cpu") + self.assertIsInstance(result, MetaTensor) + self.assertEqual(result.meta.get("fname"), "filename") def test_transforms(self): key = "im" diff --git a/tests/data/meta_tensor/test_spatial_ndim.py b/tests/data/meta_tensor/test_spatial_ndim.py new file mode 100644 index 00000000000..9e36603109a --- /dev/null +++ b/tests/data/meta_tensor/test_spatial_ndim.py @@ -0,0 +1,201 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest +from copy import deepcopy +from unittest import skipUnless + +import numpy as np +import torch +from parameterized import parameterized + +from monai.data import MetaTensor, get_spatial_ndim +from monai.data.utils import collate_meta_tensor_fn, decollate_batch +from monai.transforms import Affine, LabelToContour, RandAffine, RandZoom, Resize, Rotate, SqueezeDim +from monai.transforms.utility.array import SplitDim +from monai.utils import optional_import + +einops, has_einops = optional_import("einops") + +# (shape, affine, expected_spatial_ndim) +CONSTRUCTION_CASES = [ + ((1, 10, 10, 10), None, 3), # default eye(4) + ((1, 10, 10), torch.eye(3), 2), # eye(3) + ((1, 10), torch.eye(2), 1), # eye(2) +] + +# (description, op, expected_spatial_ndim) -- op takes a 2D MetaTensor and returns a new one +PRESERVATION_CASES = [ + ("reshape", lambda t: t.reshape(1, 100), 2), + ("unsqueeze", lambda t: t.unsqueeze(0), 2), + ("squeeze", lambda t: t.unsqueeze(1).squeeze(1), 2), + ("clone", lambda t: t.clone(), 2), + ("deepcopy", lambda t: deepcopy(t), 2), +] + + +class TestSpatialNdim(unittest.TestCase): + @parameterized.expand(CONSTRUCTION_CASES) + def test_construction(self, shape, affine, expected): + kwargs = {"affine": affine} if affine is not None else {} + t = MetaTensor(torch.randn(*shape), **kwargs) + self.assertEqual(t.spatial_ndim, expected) + + @parameterized.expand(PRESERVATION_CASES) + def test_preserved_through_op(self, _desc, op, expected): + t = MetaTensor(torch.randn(1, 10, 10), affine=torch.eye(3)) + t2 = op(t) + self.assertEqual(t2.spatial_ndim, expected) + + def test_setter_and_validation(self): + t = MetaTensor(torch.randn(1, 10, 10, 10)) + t.spatial_ndim = 2 + self.assertEqual(t.spatial_ndim, 2) + for bad in (0, -1): + with self.assertRaises(ValueError): + t.spatial_ndim = bad + + def test_affine_setter_syncs(self): + t = MetaTensor(torch.randn(1, 10, 10, 10)) + t.affine = torch.eye(3) + self.assertEqual(t.spatial_ndim, 2) + + def test_copy_from_meta_tensor(self): + t1 = MetaTensor(torch.randn(1, 10, 10), affine=torch.eye(3)) + self.assertEqual(MetaTensor(t1).spatial_ndim, 2) + + def test_collate_and_decollate(self): + t1 = MetaTensor(torch.randn(1, 10, 10), affine=torch.eye(3)) + t2 = MetaTensor(torch.randn(1, 10, 10), affine=torch.eye(3)) + batch = collate_meta_tensor_fn([t1, t2]) + self.assertEqual(batch.spatial_ndim, 2) + for item in decollate_batch(batch): + self.assertIsInstance(item, MetaTensor) + self.assertEqual(item.spatial_ndim, 2) + + def test_derived_properties(self): + """peek_pending_rank, peek_pending_shape, and pixdim all respect spatial_ndim.""" + aff = torch.diag(torch.tensor([2.0, 3.0, 1.0], dtype=torch.float64)) + t = MetaTensor(torch.randn(1, 10, 10), affine=aff) + self.assertEqual(t.peek_pending_rank(), 2) + self.assertEqual(t.peek_pending_shape(), (10, 10)) + self.assertEqual(len(t.pixdim), 2) + + def test_squeeze_dim_transform(self): + t = MetaTensor(torch.randn(1, 10, 1, 10)) + result = SqueezeDim(dim=2)(t) + self.assertEqual(result.spatial_ndim, result.affine.shape[-1] - 1) + + def test_splitdim_channel_dim_no_decrement(self): + t = MetaTensor(torch.randn(3, 8, 7)) + for item in SplitDim(dim=0, keepdim=False)(t): + if isinstance(item, MetaTensor): + self.assertEqual(item.spatial_ndim, 1) + + def test_lazy_apply_pending_2d(self): + """apply_pending uses spatial_ndim for 2D data instead of hardcoded 3.""" + from monai.transforms.lazy.functional import apply_pending + from monai.utils.enums import LazyAttr + + t = MetaTensor(torch.randn(1, 10, 10), affine=torch.eye(3)) + self.assertEqual(t.spatial_ndim, 2) + # Push a pending 2D affine operation + pending_op = { + LazyAttr.AFFINE: torch.eye(3, dtype=torch.float64), + LazyAttr.SHAPE: (10, 10), + LazyAttr.INTERP_MODE: "bilinear", + LazyAttr.PADDING_MODE: "zeros", + } + t.push_pending_operation(pending_op) + result, applied = apply_pending(t, overrides={"mode": "bilinear"}) + self.assertIsInstance(result, MetaTensor) + self.assertEqual(len(applied), 1) + + def test_batch_slice_clamps_spatial_ndim(self): + t = MetaTensor(torch.randn(10, 6, 5, 7), affine=torch.eye(4)) + t.is_batch = True + t.meta["affine"] = torch.eye(4)[None].repeat(10, 1, 1) + self.assertEqual(t.spatial_ndim, 3) + sliced = t[0] + self.assertEqual(sliced.shape, (6, 5, 7)) + self.assertEqual(sliced.spatial_ndim, 2) + self.assertEqual(get_spatial_ndim(sliced), 2) + + def test_label_to_contour_batch_slice_2d(self): + t = MetaTensor(torch.randint(0, 2, (10, 6, 5, 7)).float(), affine=torch.eye(4)) + t.is_batch = True + t.meta["affine"] = torch.eye(4)[None].repeat(10, 1, 1) + sliced = t[0] + out = LabelToContour()(sliced) + self.assertEqual(out.shape, sliced.shape) + + def test_rand_zoom_batch_slice_2d(self): + t = MetaTensor(torch.randn(10, 1, 64, 64), affine=torch.eye(4)) + t.is_batch = True + t.meta["affine"] = torch.eye(4)[None].repeat(10, 1, 1) + sliced = t[0] + zoom = RandZoom(prob=1.0, min_zoom=0.6, max_zoom=1.2) + zoom.set_random_state(seed=0) + zoom.randomize(sliced) + self.assertEqual(len(zoom._zoom), 2) + out = zoom(sliced) + self.assertEqual(out.ndim, sliced.ndim) + + @skipUnless(has_einops, "Requires einops") + def test_einops_rearrange_then_resize(self): + """Reproduce the exact #6397 bug: einops.rearrange -> Resize.""" + from einops import rearrange + + x = MetaTensor(torch.randn(1, 1, 64, 64, 3)) + x.is_batch = True + x.meta["affine"] = torch.eye(4)[None] + x_ = rearrange(x, "b c h w d -> (b c) h w d") + self.assertIsInstance(x_, MetaTensor) + self.assertEqual(x_.spatial_ndim, 3) + out = Resize(spatial_size=(32, 32, 3), mode="trilinear", align_corners=True)(x_) + self.assertEqual(out.shape[-3:], (32, 32, 3)) + + def test_affine_inverse_2d_metatensor(self): + """Affine.inverse on 2D data: 4x4 affine with spatial_ndim=2.""" + img = MetaTensor(torch.randn(1, 32, 32), affine=torch.eye(4)) + self.assertEqual(img.spatial_ndim, 2) + xform = Affine(rotate_params=(np.pi / 6,), padding_mode="zeros", image_only=True) + result = xform(img) + inv = xform.inverse(result) + self.assertEqual(inv.shape, img.shape) + self.assertEqual(len(inv.applied_operations), 0) + + def test_rotate_inverse_2d_metatensor(self): + """Rotate.inverse on 2D data: 4x4 affine with spatial_ndim=2.""" + img = MetaTensor(torch.randn(1, 32, 32), affine=torch.eye(4)) + self.assertEqual(img.spatial_ndim, 2) + xform = Rotate(angle=(np.pi / 4,), padding_mode="zeros") + result = xform(img) + inv = xform.inverse(result) + self.assertEqual(inv.shape, img.shape) + self.assertEqual(len(inv.applied_operations), 0) + + def test_rand_affine_inverse_2d_metatensor(self): + """RandAffine.inverse on 2D data: 4x4 affine with spatial_ndim=2.""" + img = MetaTensor(torch.randn(1, 32, 32), affine=torch.eye(4)) + self.assertEqual(img.spatial_ndim, 2) + xform = RandAffine(prob=1.0, rotate_range=(np.pi / 6,), padding_mode="zeros") + xform.set_random_state(seed=42) + result = xform(img) + inv = xform.inverse(result) + self.assertEqual(inv.shape, img.shape) + self.assertEqual(len(inv.applied_operations), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/data/test_box_utils.py b/tests/data/test_box_utils.py index 05778f691bc..71e32701876 100644 --- a/tests/data/test_box_utils.py +++ b/tests/data/test_box_utils.py @@ -23,6 +23,7 @@ CornerCornerModeTypeB, CornerCornerModeTypeC, CornerSizeMode, + batched_nms, box_area, box_centers, box_giou, @@ -34,6 +35,7 @@ convert_box_mode, convert_box_to_standard_mode, non_max_suppression, + spatial_crop_boxes, ) from monai.utils.type_conversion import convert_data_type from tests.test_utils import TEST_NDARRAYS, assert_allclose @@ -268,6 +270,30 @@ def test_integer_truncation_bug(self): self.assertTrue(np.issubdtype(iou.dtype, np.floating)) self.assertGreater(iou[0, 0], 0.0, "IoU should not be truncated to 0") + def test_large_coordinates_are_not_dropped(self): + """Verify large-coordinate boxes are preserved by cropping and clipping.""" + boxes = torch.tensor([[41000.0, 5000.0, 45000.0, 15000.0]], dtype=torch.float32) + + cropped_boxes, keep = spatial_crop_boxes( + boxes=boxes, roi_start=[40000, 0], roi_end=[50000, 20000], remove_empty=True + ) + assert_allclose(keep, torch.tensor([True])) + assert_allclose(cropped_boxes, torch.tensor([[1000.0, 5000.0, 5000.0, 15000.0]])) + + clipped_boxes, keep = clip_boxes_to_image(boxes=boxes, spatial_size=[50000, 50000], remove_empty=True) + assert_allclose(keep, torch.tensor([True])) + assert_allclose(clipped_boxes, boxes) + + +class TestBatchedNms(unittest.TestCase): + @parameterized.expand(TEST_NDARRAYS) + def test_batched_nms_backend(self, p): + boxes = p(np.array([[0, 0, 10, 10], [1, 1, 11, 11], [100, 100, 110, 110]], dtype=np.float32)) + scores = p(np.array([0.9, 0.8, 0.7], dtype=np.float32)) + labels = p(np.array([0, 0, 1])) + keep = batched_nms(boxes, scores, labels, nms_thresh=0.5) + assert_allclose(keep, [0, 2], type_test=False) + if __name__ == "__main__": unittest.main() diff --git a/tests/data/test_init_reader.py b/tests/data/test_init_reader.py index 169fd20a5fc..35a0b9f9139 100644 --- a/tests/data/test_init_reader.py +++ b/tests/data/test_init_reader.py @@ -19,6 +19,7 @@ from monai.data import ITKReader, NibabelReader, NrrdReader, NumpyReader, PILReader, PydicomReader from monai.transforms import LoadImage, LoadImaged +from monai.utils import MetaKeys, OptionalImportError, optional_import from tests.test_utils import SkipIfNoModule @@ -29,9 +30,27 @@ def test_load_image(self): self.assertIsInstance(instance1, LoadImage) self.assertIsInstance(instance2, LoadImage) - for r in ["NibabelReader", "PILReader", "ITKReader", "NumpyReader", "NrrdReader", "PydicomReader", None]: - inst = LoadImaged("image", reader=r) - self.assertIsInstance(inst, LoadImaged) + optional_readers = { + "NibabelReader": "nibabel", + "PILReader": "PIL", + "ITKReader": "itk", + "NrrdReader": "nrrd", + "PydicomReader": "pydicom", + } + for r, module in optional_readers.items(): + with self.subTest(reader=r): + _, has_module = optional_import(module, allow_namespace_pkg=module in ("itk", "nrrd")) + if has_module: + inst = LoadImaged("image", reader=r) + self.assertIsInstance(inst, LoadImaged) + else: + with self.assertRaises(OptionalImportError): + LoadImaged("image", reader=r) + + inst = LoadImaged("image", reader="NumpyReader") + self.assertIsInstance(inst, LoadImaged) + inst = LoadImaged("image", reader=None) + self.assertIsInstance(inst, LoadImaged) @SkipIfNoModule("nibabel") @SkipIfNoModule("cupy") @@ -48,7 +67,7 @@ def test_load_image_to_gpu(self): @SkipIfNoModule("nibabel") @SkipIfNoModule("PIL") @SkipIfNoModule("nrrd") - @SkipIfNoModule("Pydicom") + @SkipIfNoModule("pydicom") def test_readers(self): inst = ITKReader() self.assertIsInstance(inst, ITKReader) @@ -100,6 +119,40 @@ def test_nibabel_reader_avoids_eager_c_order_copy(self): # (F-order) layout from nibabel should be preserved here. self.assertFalse(data.flags.c_contiguous) + @SkipIfNoModule("pydicom") + def test_pydicom_reader_get_affine_single_slice_with_last_position(self): + reader = PydicomReader() + metadata = { + "00200037": {"Value": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]}, + "00200032": {"Value": [10.0, 20.0, 30.0]}, + "00280030": {"Value": [0.5, 0.25]}, + "lastImagePositionPatient": np.array([10.0, 20.0, 30.0]), + MetaKeys.SPATIAL_SHAPE: np.array([64, 64, 1]), + } + + affine = reader._get_affine(metadata, lps_to_ras=False) + + np.testing.assert_allclose(affine[0, 2], 0.0) + np.testing.assert_allclose(affine[1, 2], 0.0) + np.testing.assert_allclose(affine[2, 2], 1.0) + + @SkipIfNoModule("pydicom") + def test_pydicom_reader_get_affine_multi_slice_uses_last_position(self): + reader = PydicomReader() + metadata = { + "00200037": {"Value": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]}, + "00200032": {"Value": [0.0, 0.0, 0.0]}, + "00280030": {"Value": [1.0, 1.0]}, + "lastImagePositionPatient": np.array([0.0, 0.0, 8.0]), + MetaKeys.SPATIAL_SHAPE: np.array([8, 8, 5]), + } + + affine = reader._get_affine(metadata, lps_to_ras=False) + + np.testing.assert_allclose(affine[0, 2], 0.0) + np.testing.assert_allclose(affine[1, 2], 0.0) + np.testing.assert_allclose(affine[2, 2], 2.0) + if __name__ == "__main__": unittest.main() diff --git a/tests/data/test_persistentdataset.py b/tests/data/test_persistentdataset.py index ca62cdb1840..c70519d98e5 100644 --- a/tests/data/test_persistentdataset.py +++ b/tests/data/test_persistentdataset.py @@ -13,8 +13,11 @@ import contextlib import os +import pickle import tempfile import unittest +from pathlib import Path +from unittest.mock import patch import nibabel as nib import numpy as np @@ -46,7 +49,7 @@ TEST_CASE_4 = [True, False, False, MetaTensor] -TEST_CASE_5 = [True, True, True, None] +TEST_CASE_5 = [True, True, False, MetaTensor] TEST_CASE_6 = [False, False, False, torch.Tensor] @@ -200,6 +203,133 @@ def test_track_meta_and_weights_only(self, track_meta, weights_only, expected_er im = test_dataset[0]["image"] self.assertIsInstance(im, expected_type) + def test_metatensor_loading(self): + """ + Thorough test of metadata loading correctly with MetaTensor. This will store a MetaTensor with safe object types + in its metadata dictionary, test the cache file exists and can be safely loaded with weights only, and that the + loaded object is another MetaTensor with the correct information + """ + meta = {"test_meta": 123, "foo": "bar", "test_tuple": (1, 2, 3)} + imt = MetaTensor(torch.rand(1, 128, 128, 128), meta=dict(meta), affine=torch.rand(4, 4)) + + with tempfile.TemporaryDirectory() as tempdir: + cache_dir = Path(tempdir, "cache", "data") + + test_data = [{"image": imt}] + + test_dataset = PersistentDataset( + data=test_data, + transform=Compose([Identity()]), + cache_dir=str(cache_dir), + track_meta=True, + weights_only=True, + ) + + im = test_dataset[0]["image"] + self.assertIsInstance(im, MetaTensor, "MetaTensor not stored in dataset.") + + for k, v in meta.items(): + self.assertIn(k, im.meta, f"Metadata key {k} missing from loaded object.") + self.assertEqual(im.meta[k], v, f"Metadata key {k} not equal ({im.meta[k]}!={v}).") + + torch.testing.assert_close(imt.affine, im.affine) + + cache_files = list(cache_dir.glob("*")) + self.assertEqual(len(cache_files), 1, "Cached file not present.") + + cache_im = torch.load(cache_files[0], weights_only=True)["image"] + + self.assertIsInstance(cache_im, MetaTensor, "MetaTensor not stored in dataset.") + + for k, v in meta.items(): + self.assertIn(k, cache_im.meta, f"Metadata key {k} missing from loaded object.") + self.assertEqual(cache_im.meta[k], v, f"Metadata key {k} not equal ({cache_im.meta[k]}!={v}).") + + # create a new dataset to be sure + test_dataset2 = PersistentDataset( + data=test_data, + transform=Compose([Identity()]), + cache_dir=str(cache_dir), + track_meta=True, + weights_only=True, + ) + + # Replace torch.load with a function returning the same thing wrapped in a tuple, this is used to indicate + # the dataset loaded the cached data rather than recomputed. + old_load = torch.load + + def _mock_load(f, weights_only): + self.assertTrue(weights_only, f"torch.load called with {weights_only=}.") + return (old_load(f, weights_only=weights_only),) + + # check the returned object is a tuple containing the expected dict, if not then _mock_load wasn't called + with patch("torch.load", _mock_load): + im2_t = test_dataset2[0] + self.assertIsInstance(im2_t, tuple, "Special tuple not returned, so mock not used.") + self.assertIsInstance(im2_t[0]["image"], MetaTensor, "MetaTensor not stored in dataset.") + + def test_metatensor_badcache(self): + """ + Test attempting to save then load a MetaTensor with an unsafe metadata item raises an exception. This creates + a MetaTensor with an object in its metadata using unsafe code in __reduce__ which gets stored in the pickle. + When attempting to load this through torch.load, pickle.UnpicklingError should be raised to force a recompute + of the cached data rather than attempting to load something unsafe. + """ + with tempfile.TemporaryDirectory() as tempdir: + cache_dir = Path(tempdir) / "cache" / "data" + + class _BadType: + def __reduce__(self): + # something more insecure than this could be done with os.system + return (os.system, (f'echo "Code injected!" > {Path(tempdir)/"out.txt"!s}',)) + + meta = {"test_meta": 123, "foo": "bar", "bad_item": _BadType()} + imt = MetaTensor(torch.rand(1, 128, 128, 128), meta=dict(meta), affine=torch.rand(4, 4)) + test_data = [{"image": imt}] + + test_dataset = PersistentDataset( + data=test_data, + transform=Compose([Identity()]), + cache_dir=str(cache_dir), + track_meta=True, + weights_only=True, + ) + + # This will trigger the _BadType class code injection because deepcopy will use __reduce__, but will still + # write the cache file as needed for the test. The alternative was to write the cache file directly with a + # computed hash value, but computing that hash without using pickle_hashing isn't trivial. + im = test_dataset[0]["image"] + + self.assertIsInstance(im, MetaTensor, "MetaTensor not stored in dataset.") + + cache_files = list(cache_dir.glob("*")) + self.assertEqual(len(cache_files), 1, "Cached file not present.") + + # loading the cache file directly will raise the pickle exception as expected + with self.assertRaises(pickle.UnpicklingError): + torch.load(cache_files[0], weights_only=True) + + # create a new dataset object just to be sure. When loading, a cache hit will occur but this will raise + # the pickle exception again and force a recompute of the cached data as well as a warning, this indicates + # the unsafe data was correctly rejected. + test_dataset2 = PersistentDataset( + data=test_data, + transform=Compose([Identity()]), + cache_dir=str(cache_dir), + track_meta=True, + weights_only=True, + ) + + # warning raised about recomputing the corrupted cache file which raised UnpicklingError + with self.assertWarns(UserWarning): + im = test_dataset2[0]["image"] + + self.assertIsInstance(im, MetaTensor, "MetaTensor not stored in dataset.") + + cache_files2 = list(cache_dir.glob("*")) + + self.assertEqual(cache_files[0], cache_files2[0], "Hashes for cached data differ.") + if __name__ == "__main__": unittest.main() diff --git a/tests/data/test_pydicom_reader.py b/tests/data/test_pydicom_reader.py new file mode 100644 index 00000000000..42fcd89185e --- /dev/null +++ b/tests/data/test_pydicom_reader.py @@ -0,0 +1,112 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest + +import numpy as np + +from monai.data import PydicomReader +from monai.utils import MetaKeys +from tests.test_utils import SkipIfNoModule + + +@SkipIfNoModule("pydicom") +class TestPydicomReaderAffine(unittest.TestCase): + def test_missing_orientation_tags_warns_and_returns_identity(self): + # Without ImageOrientationPatient (0020,0037) and ImagePositionPatient + # (0020,0032) the affine cannot be derived. The reader falls back to the + # identity matrix; regression test for #8468 ensures this is no longer + # silent so users know orientation/spacing may be wrong. + reader = PydicomReader() + with self.assertWarns(UserWarning): + affine = reader._get_affine({}) + np.testing.assert_array_equal(affine, np.eye(4)) + + def test_partial_orientation_tags_warns(self): + # Only one of the two required tags present is still insufficient. + reader = PydicomReader() + metadata = {"00200037": {"Value": [1, 0, 0, 0, 1, 0]}} # orientation only + with self.assertWarns(UserWarning): + affine = reader._get_affine(metadata) + np.testing.assert_array_equal(affine, np.eye(4)) + + def test_non_finite_pixel_spacing_raises(self): + reader = PydicomReader() + metadata = { + "00200037": {"Value": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]}, + "00200032": {"Value": [0.0, 0.0, 0.0]}, + "00280030": {"Value": [np.nan, 1.0]}, + } + with self.assertRaisesRegex(ValueError, "PixelSpacing"): + reader._get_affine(metadata, lps_to_ras=False) + + def test_non_finite_image_position_raises(self): + reader = PydicomReader() + metadata = { + "00200037": {"Value": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]}, + "00200032": {"Value": [np.inf, 0.0, 0.0]}, + "00280030": {"Value": [1.0, 1.0]}, + } + with self.assertRaisesRegex(ValueError, "ImagePositionPatient"): + reader._get_affine(metadata, lps_to_ras=False) + + def test_finite_values_return_affine(self): + reader = PydicomReader() + metadata = { + "00200037": {"Value": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]}, + "00200032": {"Value": [10.0, 20.0, 30.0]}, + "00280030": {"Value": [0.5, 0.25]}, + } + affine = reader._get_affine(metadata, lps_to_ras=False) + self.assertEqual(affine.shape, (4, 4)) + self.assertTrue(np.all(np.isfinite(affine))) + np.testing.assert_allclose(affine[0, 3], 10.0) + np.testing.assert_allclose(affine[1, 3], 20.0) + np.testing.assert_allclose(affine[2, 3], 30.0) + + def test_non_finite_orientation_raises(self): + reader = PydicomReader() + metadata = { + "00200037": {"Value": [np.nan, 0.0, 0.0, 0.0, 1.0, 0.0]}, + "00200032": {"Value": [0.0, 0.0, 0.0]}, + "00280030": {"Value": [1.0, 1.0]}, + } + with self.assertRaisesRegex(ValueError, "ImageOrientationPatient"): + reader._get_affine(metadata, lps_to_ras=False) + + def test_non_finite_last_image_position_raises(self): + reader = PydicomReader() + metadata = { + "00200037": {"Value": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]}, + "00200032": {"Value": [0.0, 0.0, 0.0]}, + "00280030": {"Value": [1.0, 1.0]}, + "lastImagePositionPatient": [0.0, 0.0, np.inf], + MetaKeys.SPATIAL_SHAPE: [1, 1, 2], + } + with self.assertRaisesRegex(ValueError, "lastImagePositionPatient"): + reader._get_affine(metadata, lps_to_ras=False) + + def test_overflow_from_finite_inputs_raises(self): + # Finite inputs whose product overflows produce a non-finite affine. + reader = PydicomReader() + metadata = { + "00200037": {"Value": [1e308, 0.0, 0.0, 1e308, 0.0, 0.0]}, + "00200032": {"Value": [0.0, 0.0, 0.0]}, + "00280030": {"Value": [1e308, 1e308]}, + } + with self.assertRaisesRegex(ValueError, "not finite"): + reader._get_affine(metadata, lps_to_ras=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/fl/monai_algo/test_fl_monai_algo.py b/tests/fl/monai_algo/test_fl_monai_algo.py index 2c1a8488cc0..06374e4af65 100644 --- a/tests/fl/monai_algo/test_fl_monai_algo.py +++ b/tests/fl/monai_algo/test_fl_monai_algo.py @@ -12,9 +12,13 @@ from __future__ import annotations import glob +import json +import logging import os import shutil +import tempfile import unittest +import warnings from copy import deepcopy from os.path import join as pathjoin from pathlib import Path @@ -23,10 +27,10 @@ from monai.bundle import ConfigParser, ConfigWorkflow from monai.bundle.utils import DEFAULT_HANDLERS_ID -from monai.fl.client.monai_algo import MonaiAlgo +from monai.fl.client.monai_algo import MonaiAlgo, MonaiAlgoStats from monai.fl.utils.constants import ExtraItems from monai.fl.utils.exchange_object import ExchangeObject -from monai.utils import path_to_uri +from monai.utils import path_to_sqlite_uri from tests.test_utils import SkipIfNoModule _root_dir = Path(__file__).resolve().parents[2] @@ -79,7 +83,7 @@ "save_execute_config": f"{_data_dir}/config_executed.json", "trainer": { "_target_": "MLFlowHandler", - "tracking_uri": path_to_uri(_data_dir) + "/mlflow_override", + "tracking_uri": path_to_sqlite_uri(os.path.join(_data_dir, "mlflow_override.db")), "output_transform": "$monai.handlers.from_engine(['loss'], first=True)", "close_on_complete": True, }, @@ -103,7 +107,7 @@ workflow_type="train", logging_file=_logging_file, tracking="mlflow", - tracking_uri=path_to_uri(_data_dir) + "/mlflow_1", + tracking_uri=path_to_sqlite_uri(os.path.join(_data_dir, "mlflow_1.db")), experiment_name="monai_eval1", ), "config_filters_filename": os.path.join(_data_dir, "config_fl_filters.json"), @@ -119,7 +123,7 @@ ], "eval_kwargs": { "tracking": "mlflow", - "tracking_uri": path_to_uri(_data_dir) + "/mlflow_2", + "tracking_uri": path_to_sqlite_uri(os.path.join(_data_dir, "mlflow_2.db")), "experiment_name": "monai_eval2", }, "eval_workflow_name": "training", @@ -179,6 +183,38 @@ ] +def _dispose_sqlite_engines(): + """Dispose MLflow's open SQLAlchemy SQLite engines so the test ``.db`` files can be removed. + + MLflow keeps a SQLite connection open for the lifetime of its client; on Windows that + locks the database file and breaks cleanup. ``MLFlowHandler.close()`` releases it, but a + workflow may finish without closing every handler, so dispose defensively here before + deleting the files. Scoped to the test's ``mlflow*.db`` backends so unrelated (e.g. + in-memory) sqlite engines elsewhere in the process are left untouched. + """ + import gc + + try: + from sqlalchemy.engine import Engine + except ImportError: + return + gc.collect() + for obj in gc.get_objects(): + # gc.get_objects() can include dead weakref proxies, whose isinstance() raises + # ReferenceError, so guard the whole inspection (ReferenceError is an Exception). + try: + if not isinstance(obj, Engine): + continue + url = obj.url + db = url.database if url.get_backend_name() == "sqlite" else None + # the test backends are all files named ``mlflow*.db``; match those only so + # unrelated (e.g. in-memory) sqlite engines in the process are left untouched. + if db and os.path.basename(db).startswith("mlflow"): + obj.dispose() + except Exception: + pass + + @SkipIfNoModule("ignite") @SkipIfNoModule("mlflow") class TestFLMonaiAlgo(unittest.TestCase): @@ -202,8 +238,11 @@ def test_train(self, input_params): # test experiment management if "save_execute_config" in algo.train_workflow.parser: - self.assertTrue(os.path.exists(f"{_data_dir}/mlflow_override")) - shutil.rmtree(f"{_data_dir}/mlflow_override") + _dispose_sqlite_engines() # release SQLite handles so the db file can be removed on Windows + self.assertTrue(os.path.exists(f"{_data_dir}/mlflow_override.db")) + os.remove(f"{_data_dir}/mlflow_override.db") + if os.path.isdir(f"{_data_dir}/mlruns"): + shutil.rmtree(f"{_data_dir}/mlruns") self.assertTrue(os.path.exists(f"{_data_dir}/config_executed.json")) os.remove(f"{_data_dir}/config_executed.json") @@ -225,9 +264,12 @@ def test_evaluate(self, input_params): # test experiment management if "save_execute_config" in algo.eval_workflow.parser: + _dispose_sqlite_engines() # release SQLite handles so the db files can be removed on Windows self.assertGreater(len(list(glob.glob(f"{_data_dir}/mlflow_*"))), 0) for f in list(glob.glob(f"{_data_dir}/mlflow_*")): - shutil.rmtree(f) + shutil.rmtree(f) if os.path.isdir(f) else os.remove(f) + if os.path.isdir(f"{_data_dir}/mlruns"): + shutil.rmtree(f"{_data_dir}/mlruns") self.assertGreater(len(list(glob.glob(f"{_data_dir}/eval/config_*"))), 0) for f in list(glob.glob(f"{_data_dir}/eval/config_*")): os.remove(f) @@ -247,5 +289,142 @@ def test_get_weights(self, input_params): self.assertIsInstance(weights, ExchangeObject) +@SkipIfNoModule("ignite") +class TestFLMonaiAlgoWarnsOnProvisionedConfig(unittest.TestCase): + """Regression tests for GHSA-x6pr-233j-x5cw: `MonaiAlgo`/`MonaiAlgoStats` execute a bundle whose + whole app directory -- configs included -- is provisioned by the FL system, and the aggregation + server dispatches the tasks that run it with no per-round human interaction. `MonaiAlgo` builds + its `ConfigWorkflow` directly rather than through `create_workflow()`, so the warning added for + GHSA-873f-pvrv-4x83 never fired on this path. + + Executing the config is still not blocked -- MONAI has no way to establish whether a bundle is + trustworthy, so a flag would only teach operators to set it once and forget it -- but a + `UserWarning` is now raised, and the one sink with no functional role in FL, the bundle's own + "configs/logging.conf", is no longer applied unless the FL system asks for it via + `ExtraItems.LOGGING_FILE` (GHSA-wvpx-5qmp-46g3).""" + + def setUp(self): + # `fileConfig` reconfigures logging process-wide. Snapshot the root logger and restore it + # afterwards so these tests cannot leak a handler into the rest of the suite. + root = logging.getLogger() + level, handlers, filters = root.level, root.handlers[:], root.filters[:] + disabled = logging.root.manager.disable + + def _restore(): + # Detach whatever is on the root logger now, closing anything `fileConfig` installed so + # it does not linger in logging's handler registry, then put the snapshot back. Under + # `tests/runner.py` the root logger starts with no handlers, so there is nothing for + # `fileConfig` to have closed on the way in. + for handler in root.handlers[:]: + root.removeHandler(handler) + if handler not in handlers: + handler.close() + root.setLevel(level) + root.filters[:] = filters + for handler in handlers: + root.addHandler(handler) + logging.disable(disabled) + + self.addCleanup(_restore) + + def _stage_malicious_app(self, tempdir: str) -> tuple[str, str, str]: + """Write an FL app whose config and logging.conf each drop a distinct marker file.""" + app_root = os.path.join(tempdir, "app") + os.makedirs(os.path.join(app_root, "configs")) + config_marker = os.path.join(tempdir, "CONFIG_PWNED") + logging_marker = os.path.join(tempdir, "LOGGING_PWNED") + # write the markers via `pathlib` instead of shelling out through `os.system` -- `!r` yields + # a Python-source-safe literal (handling spaces and Windows backslashes alike) with no shell + # involved to reintroduce quoting/splitting issues. + payload = f"$__import__('pathlib').Path({config_marker!r}).write_text('pwned')" + with open(os.path.join(app_root, "configs", "train.json"), "w") as f: + json.dump({"initialize": [payload]}, f) + # `fileConfig` eval()s the `class=` field, so the tuple subscript runs the payload and still + # yields a usable handler class. + with open(os.path.join(app_root, "configs", "logging.conf"), "w") as f: + f.write( + "[loggers]\nkeys=root\n[handlers]\nkeys=h\n[formatters]\nkeys=f\n" + "[logger_root]\nlevel=NOTSET\nhandlers=h\n" + "[handler_h]\n" + f"class=(__import__('pathlib').Path({logging_marker!r}).write_text('pwned'), " + "__import__('logging').StreamHandler)[1]\nargs=()\nformatter=f\n" + "[formatter_f]\nformat=%(message)s\n" + ) + return app_root, config_marker, logging_marker + + @staticmethod + def _algo(algo_class): + # `bundle_root=""` so the whole path comes from the server-supplied `APP_ROOT`, exactly as in + # the advisory's PoC. `MonaiAlgo` additionally defaults to building an evaluate workflow. + kwargs = {"bundle_root": "", "config_train_filename": "configs/train.json"} + if algo_class is MonaiAlgo: + kwargs["config_evaluate_filename"] = None + return algo_class(**kwargs) + + @parameterized.expand([[MonaiAlgoStats], [MonaiAlgo]]) + def test_warns_and_executes_provisioned_config(self, algo_class): + with tempfile.TemporaryDirectory() as tempdir: + app_root, config_marker, logging_marker = self._stage_malicious_app(tempdir) + algo = self._algo(algo_class) + with self.assertWarnsRegex(UserWarning, r"GHSA-x6pr-233j-x5cw"): + # the staged config defines only `initialize`, so resolving the `bundle_root` + # property fails *after* the payload has already run -- as in the advisory's own + # PoC, where the failure happens after code execution. + with self.assertRaises(KeyError): + algo.initialize(extra={ExtraItems.CLIENT_NAME: "test_fl", ExtraItems.APP_ROOT: app_root}) + # executing the config is deliberately still not blocked + self.assertTrue(os.path.exists(config_marker)) + # ... but the server's logging.conf is no longer handed to `fileConfig` + self.assertFalse(os.path.exists(logging_marker)) + + @parameterized.expand([[MonaiAlgoStats], [MonaiAlgo]]) + def test_explicit_none_logging_file_does_not_apply_provisioned_conf(self, algo_class): + """`None` was the pre-fix default, so an FL system may well pass the key explicitly with that + value. `ConfigWorkflow` reads `None` as "fall back to the bundle's own configs/logging.conf", + which would hand the server's INI straight to `fileConfig`; it has to mean disabled here.""" + with tempfile.TemporaryDirectory() as tempdir: + app_root, _, logging_marker = self._stage_malicious_app(tempdir) + algo = self._algo(algo_class) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with self.assertRaises(KeyError): + algo.initialize( + extra={ + ExtraItems.CLIENT_NAME: "test_fl", + ExtraItems.APP_ROOT: app_root, + ExtraItems.LOGGING_FILE: None, + } + ) + self.assertFalse(any("GHSA-wvpx-5qmp-46g3" in str(w.message) for w in caught)) + self.assertFalse(os.path.exists(logging_marker)) + + @parameterized.expand([[MonaiAlgoStats], [MonaiAlgo]]) + def test_logging_file_opt_in_applies_provisioned_conf(self, algo_class): + with tempfile.TemporaryDirectory() as tempdir: + app_root, _, logging_marker = self._stage_malicious_app(tempdir) + algo = self._algo(algo_class) + with self.assertWarnsRegex(UserWarning, r"GHSA-wvpx-5qmp-46g3"): + with self.assertRaises(KeyError): + algo.initialize( + extra={ + ExtraItems.CLIENT_NAME: "test_fl", + ExtraItems.APP_ROOT: app_root, + ExtraItems.LOGGING_FILE: os.path.join(app_root, "configs", "logging.conf"), + } + ) + self.assertTrue(os.path.exists(logging_marker)) + + def test_no_logging_warning_when_logging_disabled(self): + """The `fileConfig` warning must not fire when nothing is actually executed.""" + with tempfile.TemporaryDirectory() as tempdir: + app_root, _, _ = self._stage_malicious_app(tempdir) + algo = self._algo(MonaiAlgoStats) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with self.assertRaises(KeyError): + algo.initialize(extra={ExtraItems.CLIENT_NAME: "test_fl", ExtraItems.APP_ROOT: app_root}) + self.assertFalse(any("GHSA-wvpx-5qmp-46g3" in str(w.message) for w in caught)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/handlers/test_handler_mlflow.py b/tests/handlers/test_handler_mlflow.py index 80630e6f5a2..a396227eb9a 100644 --- a/tests/handlers/test_handler_mlflow.py +++ b/tests/handlers/test_handler_mlflow.py @@ -13,11 +13,10 @@ import glob import os -import shutil import tempfile import unittest from concurrent.futures import ThreadPoolExecutor -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np from ignite.engine import Engine, Events @@ -26,11 +25,9 @@ from monai.apps import download_and_extract from monai.bundle import ConfigWorkflow, download from monai.handlers import MLFlowHandler -from monai.utils import optional_import, path_to_uri +from monai.utils import path_to_sqlite_uri, path_to_uri from tests.test_utils import skip_if_downloading_fails, skip_if_quick -_, has_dataset_tracking = optional_import("mlflow", "2.4.0") - def get_event_filter(e): def event_filter(_, event): @@ -41,9 +38,7 @@ def event_filter(_, event): return event_filter -def dummy_train(tracking_folder): - tempdir = tempfile.mkdtemp() - +def dummy_train(tracking_folder, tempdir): # set up engine def _train_func(engine, batch): return [batch + 1.0] @@ -55,7 +50,7 @@ def _train_func(engine, batch): handler = MLFlowHandler( iteration_log=False, epoch_log=True, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], close_on_complete=True, ) @@ -65,14 +60,6 @@ def _train_func(engine, batch): class TestHandlerMLFlow(unittest.TestCase): - def setUp(self): - self.tmpdir_list = [] - - def tearDown(self): - for tmpdir in self.tmpdir_list: - if tmpdir and os.path.exists(tmpdir): - shutil.rmtree(tmpdir) - def test_multi_run(self): with tempfile.TemporaryDirectory() as tempdir: # set up the train function for engine @@ -95,7 +82,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=False, epoch_log=True, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], close_on_complete=True, ) @@ -106,6 +93,124 @@ def _update_metric(engine): # the run count should equal to the times of creating engine self.assertEqual(create_engine_times, run_cnt) + def test_default_tracking_uri_is_sqlite(self): + """Verify the handler defaults to a local SQLite backend, not the file store, without a tracking URI.""" + with tempfile.TemporaryDirectory() as tempdir: + cwd = os.getcwd() + os.chdir(tempdir) + handler = None + try: + handler = MLFlowHandler(iteration_log=False, epoch_log=False) + self.assertTrue(handler.client.tracking_uri.startswith("sqlite:///")) + self.assertTrue(handler.client.tracking_uri.endswith("mlruns.db")) + # artifacts should still default to a `./mlruns`-style directory + self.assertIsNotNone(handler.artifact_location) + self.assertTrue(handler.artifact_location.endswith("mlruns")) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + os.chdir(cwd) + + def test_remote_tracking_uri_leaves_artifact_location_unset(self): + """Verify a remote tracking URI gets no local artifact location injected.""" + handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri="http://localhost:5000") + self.assertEqual(handler.client.tracking_uri, "http://localhost:5000") + self.assertIsNone(handler.artifact_location) + + def test_file_store_tracking_uri_is_rejected(self): + """Verify local paths and file:// URIs are rejected with an actionable error.""" + for uri in ("/tmp/mlruns", path_to_uri(os.path.join("some", "dir"))): + with self.assertRaises(ValueError): + MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=uri) + + def test_explicit_sqlite_tracking_uri_colocates_artifacts(self): + """Verify an explicit SQLite tracking URI co-locates artifacts next to the database.""" + with tempfile.TemporaryDirectory() as tempdir: + uri = path_to_sqlite_uri(os.path.join(tempdir, "sub", "mlruns.db")) + handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=uri) + try: + self.assertEqual(handler.client.tracking_uri, uri) + self.assertIsNotNone(handler.artifact_location) + self.assertTrue(handler.artifact_location.endswith("mlruns")) + finally: + handler.close() # release the SQLite handle so Windows can delete the db + + def test_env_var_sqlite_tracking_uri_colocates_artifacts(self): + """Verify a SQLite ``MLFLOW_TRACKING_URI`` env var co-locates artifacts next to the db.""" + with tempfile.TemporaryDirectory() as tempdir: + uri = path_to_sqlite_uri(os.path.join(tempdir, "sub", "mlruns.db")) + handler = None + with patch.dict(os.environ, {"MLFLOW_TRACKING_URI": uri}): + try: + handler = MLFlowHandler(iteration_log=False, epoch_log=False) + self.assertTrue(handler.client.tracking_uri.endswith("mlruns.db")) + self.assertIsNotNone(handler.artifact_location) + self.assertTrue(handler.artifact_location.endswith("mlruns")) + # co-located with the db file (the `sub` dir), not a cwd-relative `./mlruns` + self.assertIn("sub", handler.artifact_location) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + + def test_env_var_tracking_uri_takes_priority_over_argument(self): + """Verify ``MLFLOW_TRACKING_URI`` overrides an explicit ``tracking_uri`` argument.""" + with tempfile.TemporaryDirectory() as tempdir: + env_uri = path_to_sqlite_uri(os.path.join(tempdir, "env.db")) + arg_uri = path_to_sqlite_uri(os.path.join(tempdir, "arg.db")) + handler = None + with patch.dict(os.environ, {"MLFLOW_TRACKING_URI": env_uri}): + try: + handler = MLFlowHandler(iteration_log=False, epoch_log=False, tracking_uri=arg_uri) + self.assertTrue(handler.client.tracking_uri.endswith("env.db")) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + + def test_explicit_artifact_location_is_used(self): + """Verify an explicit artifact location is preserved with the default SQLite backend.""" + with tempfile.TemporaryDirectory() as tempdir: + cwd = os.getcwd() + os.chdir(tempdir) + handler = None + try: + art = path_to_uri(os.path.join(tempdir, "artifacts")) + handler = MLFlowHandler(iteration_log=False, epoch_log=False, artifact_location=art) + self.assertEqual(handler.artifact_location, art) + finally: + if handler is not None: + handler.close() # release the SQLite handle so Windows can delete the db + os.chdir(cwd) + + def test_default_sqlite_run_flow(self): + """Verify a basic run flow works end-to-end with the default SQLite backend.""" + with tempfile.TemporaryDirectory() as tempdir: + cwd = os.getcwd() + os.chdir(tempdir) + try: + + def _train_func(engine, batch): + return [batch + 1.0] + + engine = Engine(_train_func) + + @engine.on(Events.EPOCH_COMPLETED) + def _update_metric(engine): + current_metric = engine.state.metrics.get("acc", 0.1) + engine.state.metrics["acc"] = current_metric + 0.1 + + # close_on_complete=False so cur_run stays available after the run for the metric + # check below; the run is closed explicitly afterwards. + handler = MLFlowHandler(iteration_log=False, epoch_log=True, close_on_complete=False) + handler.attach(engine) + engine.run(range(3), max_epochs=2) + cur_run = handler.client.get_run(handler.cur_run.info.run_id) + self.assertTrue("acc" in cur_run.data.metrics.keys()) + handler.close() + # the default backend should have created a SQLite database file in the cwd + self.assertTrue(os.path.exists(os.path.join(tempdir, "mlruns.db"))) + finally: + os.chdir(cwd) + def test_metrics_track(self): experiment_param = {"backbone": "efficientnet_b0"} with tempfile.TemporaryDirectory() as tempdir: @@ -137,7 +242,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=False, epoch_log=True, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], experiment_param=experiment_param, artifacts=[artifact_path], @@ -173,7 +278,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=False, epoch_log=epoch_log, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], experiment_param=experiment_param, close_on_complete=True, @@ -212,7 +317,7 @@ def _update_metric(engine): handler = MLFlowHandler( iteration_log=iteration_log, epoch_log=False, - tracking_uri=path_to_uri(test_path), + tracking_uri=path_to_sqlite_uri(test_path), state_attributes=["test"], experiment_param=experiment_param, close_on_complete=True, @@ -232,18 +337,17 @@ def _update_metric(engine): def test_multi_thread(self): test_uri_list = ["monai_mlflow_test1", "monai_mlflow_test2"] - with ThreadPoolExecutor(2, "Training") as executor: - futures = {} - for t in test_uri_list: - futures[t] = executor.submit(dummy_train, t) + with tempfile.TemporaryDirectory() as tempdir: + with ThreadPoolExecutor(2, "Training") as executor: + futures = {} + for t in test_uri_list: + futures[t] = executor.submit(dummy_train, t, tempdir) - for _, future in futures.items(): - res = future.result() - self.tmpdir_list.append(res) - self.assertTrue(len(glob.glob(res)) > 0) + for _, future in futures.items(): + res = future.result() + self.assertTrue(len(glob.glob(res)) > 0) @skip_if_quick - @unittest.skipUnless(has_dataset_tracking, reason="Requires mlflow version >= 2.4.0.") def test_dataset_tracking(self): test_bundle_name = "endoscopic_tool_segmentation" with tempfile.TemporaryDirectory() as tempdir: @@ -271,7 +375,7 @@ def test_dataset_tracking(self): final_id="finalize", ) - tracking_path = os.path.join(bundle_root, "eval") + tracking_path = os.path.join(tempdir, "mlflow_dataset.db") workflow.bundle_root = bundle_root workflow.dataset_dir = data_dir workflow.initialize() @@ -280,7 +384,7 @@ def test_dataset_tracking(self): iteration_log=False, epoch_log=False, dataset_dict={"test": infer_dataset}, - tracking_uri=path_to_uri(tracking_path), + tracking_uri=path_to_sqlite_uri(tracking_path), ) mlflow_handler.attach(workflow.evaluator) workflow.run() diff --git a/tests/handlers/test_write_metrics_reports.py b/tests/handlers/test_write_metrics_reports.py index 1013f15d85d..07cf46c122b 100644 --- a/tests/handlers/test_write_metrics_reports.py +++ b/tests/handlers/test_write_metrics_reports.py @@ -63,6 +63,28 @@ def test_content(self): self.assertTrue(os.path.exists(os.path.join(tempdir, "metric4_raw.csv"))) self.assertTrue(os.path.exists(os.path.join(tempdir, "metric4_summary.csv"))) + def test_multi_metric_details_headers(self): + with tempfile.TemporaryDirectory() as tempdir: + write_metrics_reports( + save_dir=Path(tempdir), + images=["img1", "img2"], + metrics=None, + metric_details={ + "m1": torch.tensor([[1, 2, 3], [4, 5, 6]]), + "m2": torch.tensor([[7, 8], [9, 10]]), + "m3": torch.tensor([[11, 12, 13, 14], [15, 16, 17, 18]]), + }, + summary_ops=None, + deli=",", + output_type="csv", + ) + for name, nclass in [("m1", 3), ("m2", 2), ("m3", 4)]: + path = os.path.join(tempdir, f"{name}_raw.csv") + self.assertTrue(os.path.exists(path)) + with open(path) as f: + header = f.readline().strip().split(",") + self.assertEqual(header, ["filename"] + [f"class{i}" for i in range(nclass)] + ["mean"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/inferers/test_sliding_window_inference.py b/tests/inferers/test_sliding_window_inference.py index 5a624c787f9..70ebb61639f 100644 --- a/tests/inferers/test_sliding_window_inference.py +++ b/tests/inferers/test_sliding_window_inference.py @@ -32,7 +32,6 @@ [(1, 3, 16, 15, 7), (4, -1, 7), 3, 0.25, "constant", torch.device("cpu:0")], # 3D small roi [(2, 3, 16, 15, 7), (4, -1, 7), 3, 0.25, "constant", torch.device("cpu:0")], # 3D small roi [(3, 3, 16, 15, 7), (4, -1, 7), 3, 0.25, "constant", torch.device("cpu:0")], # 3D small roi - [(2, 3, 16, 15, 7), (4, -1, 7), 3, 0.25, "constant", torch.device("cpu:0")], # 3D small roi [(1, 3, 16, 15, 7), (4, 10, 7), 3, 0.25, "constant", torch.device("cpu:0")], # 3D small roi [(1, 3, 16, 15, 7), (20, 22, 23), 10, 0.25, "constant", torch.device("cpu:0")], # 3D large roi [(2, 3, 15, 7), (2, 6), 1000, 0.25, "constant", torch.device("cpu:0")], # 2D small roi, large batch diff --git a/tests/integration/test_integration_bundle_run.py b/tests/integration/test_integration_bundle_run.py index 7f366d47458..67f44562593 100644 --- a/tests/integration/test_integration_bundle_run.py +++ b/tests/integration/test_integration_bundle_run.py @@ -29,7 +29,7 @@ from monai.bundle import ConfigParser from monai.bundle.utils import DEFAULT_HANDLERS_ID from monai.transforms import LoadImage -from monai.utils import path_to_uri +from monai.utils import path_to_sqlite_uri from tests.test_utils import command_line_tests TESTS_PATH = Path(__file__).parents[1] @@ -175,7 +175,7 @@ def test_shape(self, config_file, expected_shape): "no_epoch": True, # test override config in the settings file "evaluator": { "_target_": "MLFlowHandler", - "tracking_uri": "$monai.utils.path_to_uri(@output_dir) + '/mlflow_override1'", + "tracking_uri": "$monai.utils.path_to_sqlite_uri(@output_dir + '/mlflow_override1.db')", "iteration_log": "@no_epoch", }, }, @@ -208,16 +208,17 @@ def test_shape(self, config_file, expected_shape): command_line_tests(la + ["--args_file", def_args_file] + ["--tracking", settings_file]) loader = LoadImage(image_only=True) self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_seg.nii.gz")).shape, expected_shape) - self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override1")) + self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override1.db")) - tracking_uri = path_to_uri(tempdir) + "/mlflow_override2" # test override experiment management configs + # test override experiment management configs + tracking_uri = path_to_sqlite_uri(os.path.join(tempdir, "mlflow_override2.db")) # here test the script with `google fire` tool as CLI cmd = "-m fire monai.bundle.scripts run --tracking mlflow --evaluator#amp False" cmd += f" --tracking_uri {tracking_uri} {override} --output_dir {tempdir} --device {device}" la = ["coverage", "run"] + cmd.split(" ") + ["--meta_file", meta_file] + ["--config_file", config_file] command_line_tests(la) self.assertTupleEqual(loader(os.path.join(tempdir, "image", "image_trans.nii.gz")).shape, expected_shape) - self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override2")) + self.assertTrue(os.path.exists(f"{tempdir}/mlflow_override2.db")) # test the saved execution configs self.assertTrue(len(glob(f"{tempdir}/config_*.json")), 2) diff --git a/tests/integration/test_integration_nnunetv2_runner.py b/tests/integration/test_integration_nnunetv2_runner.py index 1da131c8908..f4cf0b4fb11 100644 --- a/tests/integration/test_integration_nnunetv2_runner.py +++ b/tests/integration/test_integration_nnunetv2_runner.py @@ -107,41 +107,41 @@ def setUp(self) -> None: self.good_yml2 = os.path.join(test_path, "good2.yml") self.inject_yml = os.path.join(test_path, "test.yml") - good_yml_content1 = """ + good_yml_content1 = f""" dataset_name_or_id: Dataset123 - dataroot: ./data - datalist: ./lists/task4.json - work_dir: ./work - nnunet_raw: ./nnUNet_raw - nnunet_preprocessed: ./nnUNet_preprocessed - nnunet_results: ./nnUNet_results + dataroot: {test_path}/data + datalist: {test_path}/lists/task4.json + work_dir: {test_path}/work + nnunet_raw: {test_path}/nnUNet_raw + nnunet_preprocessed: {test_path}/nnUNet_preprocessed + nnunet_results: {test_path}/nnUNet_results """ with open(self.good_yml1, "w") as o: o.write(dedent(good_yml_content1)) - good_yml_content2 = """ + good_yml_content2 = f""" dataset_name_or_id: 123 - dataroot: ./data - datalist: ./lists/task4.json - work_dir: ./work - nnunet_raw: ./nnUNet_raw - nnunet_preprocessed: ./nnUNet_preprocessed - nnunet_results: ./nnUNet_results + dataroot: {test_path}/data + datalist: {test_path}/lists/task4.json + work_dir: {test_path}/work + nnunet_raw: {test_path}/nnUNet_raw + nnunet_preprocessed: {test_path}/nnUNet_preprocessed + nnunet_results: {test_path}/nnUNet_results """ with open(self.good_yml2, "w") as o: o.write(dedent(good_yml_content2)) # define a config file with code-injecting dataset name - injecting_yml_content = """ - dataset_name_or_id: '4 & echo "This is exploited" > "./test.txt" & rem' - dataroot: ./data - datalist: ./lists/task4.json - work_dir: ./work - nnunet_raw: ./nnUNet_raw - nnunet_preprocessed: ./nnUNet_preprocessed - nnunet_results: ./nnUNet_results + injecting_yml_content = f""" + dataset_name_or_id: '4 & echo "This is exploited" > "{test_path}/test.txt" & rem' + dataroot: {test_path}/data + datalist: {test_path}/lists/task4.json + work_dir: {test_path}/work + nnunet_raw: {test_path}/nnUNet_raw + nnunet_preprocessed: {test_path}/nnUNet_preprocessed + nnunet_results: {test_path}/nnUNet_results """ with open(self.inject_yml, "w") as o: diff --git a/tests/losses/deform/test_bending_energy.py b/tests/losses/deform/test_bending_energy.py index 2e8ab32dbd8..5e713b3e476 100644 --- a/tests/losses/deform/test_bending_energy.py +++ b/tests/losses/deform/test_bending_energy.py @@ -23,6 +23,7 @@ TEST_CASES = [ [{}, {"pred": torch.ones((1, 3, 5, 5, 5), device=device)}, 0.0], + [{}, {"pred": torch.ones((1, 3, 3, 3, 3), device=device)}, 0.0], [{}, {"pred": torch.arange(0, 5, device=device)[None, None, None, None, :].expand(1, 3, 5, 5, 5)}, 0.0], [ {"normalize": False}, @@ -64,11 +65,11 @@ def test_ill_shape(self): with self.assertRaisesRegex(ValueError, "Expecting 3-d, 4-d or 5-d"): loss.forward(torch.ones((1, 4, 5, 5, 5, 5), device=device)) with self.assertRaisesRegex(ValueError, "All spatial dimensions"): - loss.forward(torch.ones((1, 3, 4, 5, 5), device=device)) + loss.forward(torch.ones((1, 3, 2, 5, 5), device=device)) with self.assertRaisesRegex(ValueError, "All spatial dimensions"): - loss.forward(torch.ones((1, 3, 5, 4, 5))) + loss.forward(torch.ones((1, 3, 5, 2, 5))) with self.assertRaisesRegex(ValueError, "All spatial dimensions"): - loss.forward(torch.ones((1, 3, 5, 5, 4))) + loss.forward(torch.ones((1, 3, 5, 5, 2))) # number of vector components unequal to number of spatial dims with self.assertRaisesRegex(ValueError, "Number of vector components"): diff --git a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py index a16499ac113..19a60f72193 100644 --- a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py +++ b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py @@ -164,5 +164,51 @@ def test_ill_opts(self, num_bins, reduction, expected_exception, expected_messag GlobalMutualInformationLoss(num_bins=num_bins, reduction=reduction)(pred, target) +class TestGlobalMutualInformationLossBuffers(unittest.TestCase): + def test_gaussian_kernel_registers_buffers(self): + """Verify gaussian kernel registers preterm and bin_centers as non-trainable, non-persistent buffers.""" + loss = GlobalMutualInformationLoss(kernel_type="gaussian") + self.assertIn("preterm", loss._buffers) + self.assertIn("bin_centers", loss._buffers) + self.assertFalse(loss.preterm.requires_grad) + self.assertFalse(loss.bin_centers.requires_grad) + self.assertEqual(loss.bin_centers.ndim, 3) + state = loss.state_dict() + self.assertNotIn("preterm", state) + self.assertNotIn("bin_centers", state) + + def test_bspline_kernel_has_no_gaussian_buffers(self): + """Verify b-spline kernel does not populate gaussian-specific buffers.""" + loss = GlobalMutualInformationLoss(kernel_type="b-spline") + self.assertIsNone(loss.preterm) + self.assertIsNone(loss.bin_centers) + state = loss.state_dict() + self.assertNotIn("preterm", state) + self.assertNotIn("bin_centers", state) + + def test_gaussian_kernel_forward_correct(self): + """Verify gaussian kernel forward pass returns a scalar loss tensor.""" + pred = torch.rand(2, 1, 8, 8, dtype=torch.float32) + target = torch.rand(2, 1, 8, 8, dtype=torch.float32) + loss = GlobalMutualInformationLoss(kernel_type="gaussian") + result = loss(pred, target) + self.assertEqual(result.shape, torch.Size([])) + + def test_gaussian_buffers_move_with_module(self): + """Verify preterm and bin_centers buffers move to the target device with the module.""" + loss = GlobalMutualInformationLoss(kernel_type="gaussian") + self.assertEqual(loss.preterm.device.type, "cpu") + self.assertEqual(loss.bin_centers.device.type, "cpu") + if not torch.cuda.is_available(): + self.skipTest("CUDA not available") + loss = loss.cuda() + self.assertEqual(loss.preterm.device.type, "cuda") + self.assertEqual(loss.bin_centers.device.type, "cuda") + pred = torch.rand(2, 1, 8, 8, device="cuda") + target = torch.rand(2, 1, 8, 8, device="cuda") + result = loss(pred, target) + self.assertEqual(result.device.type, "cuda") + + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_boundary_loss.py b/tests/losses/test_boundary_loss.py new file mode 100644 index 00000000000..156d4767c01 --- /dev/null +++ b/tests/losses/test_boundary_loss.py @@ -0,0 +1,334 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest +from unittest.case import skipUnless + +import torch +from parameterized import parameterized + +from monai.losses import BoundaryLoss +from monai.utils import optional_import + +_, has_scipy = optional_import("scipy") + +# Reusable test tensors +ONES_2D = {"input": torch.ones((2, 2, 8, 8)), "target": torch.ones((2, 2, 8, 8))} +ONES_3D = {"input": torch.ones((2, 2, 8, 8, 8)), "target": torch.ones((2, 2, 8, 8, 8))} + +# Perfect match: target is a 2x2 square, input matches exactly +PERFECT_MATCH = { + "input": torch.tensor( + [[[[1.0, 1.0, 0.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]] + ), + "target": torch.tensor( + [[[[1.0, 1.0, 0.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]] + ), +} + +# Partial overlap: two 2x2 squares shifted by 1 pixel +PARTIAL_OVERLAP = { + "input": torch.tensor( + [[[[1.0, 1.0, 0.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]] + ), + "target": torch.tensor( + [[[[0.0, 1.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]] + ), +} + +# Empty foreground class: target has no foreground in class 1 +EMPTY_FOREGROUND = { + "input": torch.tensor( + [[[[0.9, 0.9, 0.9], [0.9, 0.9, 0.9], [0.9, 0.9, 0.9]], [[0.1, 0.1, 0.1], [0.1, 0.1, 0.1], [0.1, 0.1, 0.1]]]] + ), + "target": torch.tensor( + [[[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]] + ), +} + +TEST_CASES = [] +for device in ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]: + # Basic 2D test with sigmoid + TEST_CASES.append( + [ + {"include_background": True, "sigmoid": True}, + { + "input": torch.tensor([[[[2.0, -2.0], [-2.0, 2.0]]]], device=device), + "target": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]], device=device), + }, + None, # Just check it runs, value depends on distance map + ] + ) + # Basic 3D test with sigmoid + TEST_CASES.append( + [ + {"include_background": True, "sigmoid": True}, + { + "input": torch.tensor([[[[[2.0, -2.0], [-2.0, 2.0]], [[2.0, -2.0], [-2.0, 2.0]]]]], device=device), + "target": torch.tensor([[[[[1.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 1.0]]]]], device=device), + }, + None, + ] + ) + # Multi-class 2D with softmax + TEST_CASES.append( + [ + {"include_background": True, "softmax": True}, + { + "input": torch.tensor([[[[2.0, 0.0], [0.0, 2.0]], [[-2.0, 0.0], [0.0, -2.0]]]], device=device), + "target": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]]], device=device), + }, + None, + ] + ) + # With to_onehot_y + TEST_CASES.append( + [ + {"include_background": True, "to_onehot_y": True, "softmax": True}, + { + "input": torch.tensor([[[[2.0, 0.0], [0.0, 2.0]], [[-2.0, 0.0], [0.0, -2.0]]]], device=device), + "target": torch.tensor([[[[0, 0], [0, 1]]]], device=device), + }, + None, + ] + ) + # With reduction="none" + TEST_CASES.append( + [ + {"include_background": True, "sigmoid": True, "reduction": "none"}, + { + "input": torch.tensor([[[[2.0, -2.0], [-2.0, 2.0]]]], device=device), + "target": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]], device=device), + }, + None, + ] + ) + # With reduction="sum" + TEST_CASES.append( + [ + {"include_background": True, "sigmoid": True, "reduction": "sum"}, + { + "input": torch.tensor([[[[2.0, -2.0], [-2.0, 2.0]]]], device=device), + "target": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]], device=device), + }, + None, + ] + ) + # Exclude background + TEST_CASES.append( + [ + {"include_background": False, "sigmoid": True}, + { + "input": torch.tensor([[[[2.0, -2.0], [-2.0, 2.0]], [[-2.0, 2.0], [2.0, -2.0]]]], device=device), + "target": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]]], device=device), + }, + None, + ] + ) + # With other_act + TEST_CASES.append( + [ + {"include_background": True, "other_act": torch.tanh}, + { + "input": torch.tensor([[[[1.0, -1.0], [-1.0, 1.0]]]], device=device), + "target": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]], device=device), + }, + None, + ] + ) + + +def _describe_test_case(test_func, test_number, params): + input_param, input_data, _ = params.args + return f"params:{input_param}, shape:{input_data['input'].shape}, device:{input_data['input'].device}" + + +@skipUnless(has_scipy, "Scipy required") +class TestBoundaryLoss(unittest.TestCase): + + @parameterized.expand(TEST_CASES, doc_func=_describe_test_case) + def test_runs(self, input_param, input_data, _): + """Test that the loss runs without errors for various configurations.""" + loss = BoundaryLoss(**input_param) + result = loss(**input_data) + # Just verify it's a scalar tensor and finite + self.assertTrue(torch.isfinite(result).all()) + + def test_perfect_match(self): + """Test that perfect predictions yield lower loss than imperfect ones.""" + loss_fn = BoundaryLoss() + perfect_loss = loss_fn(PERFECT_MATCH["input"], PERFECT_MATCH["target"]) + partial_loss = loss_fn(PARTIAL_OVERLAP["input"], PARTIAL_OVERLAP["target"]) + # Perfect match should have lower loss than partial overlap + self.assertLess(perfect_loss.item(), partial_loss.item()) + + def test_reduction_shapes(self): + """Test that different reductions produce expected shapes.""" + input_tensor = torch.ones((4, 2, 8, 8)) + target = torch.ones((4, 2, 8, 8)) + + self.assertEqual(BoundaryLoss(reduction="mean")(input_tensor, target).shape, torch.Size([])) + self.assertEqual(BoundaryLoss(reduction="sum")(input_tensor, target).shape, torch.Size([])) + # With include_background=True and 2 classes, shape should be (4, 2) + self.assertEqual(BoundaryLoss(reduction="none")(input_tensor, target).shape, torch.Size([4, 2])) + + def test_reduction_shapes_exclude_background(self): + """Test shapes when background is excluded.""" + input_tensor = torch.ones((4, 3, 8, 8)) + target = torch.ones((4, 3, 8, 8)) + + # With include_background=False, shape should be (4, 2) for 3 classes + self.assertEqual( + BoundaryLoss(reduction="none", include_background=False)(input_tensor, target).shape, torch.Size([4, 2]) + ) + + def test_single_channel_options_warn_and_are_ignored(self): + """Test that single-channel-only options follow other MONAI loss behavior.""" + input_tensor = torch.randn((1, 1, 4, 4), requires_grad=True) + target = torch.zeros((1, 1, 4, 4)) + target[..., 1:3, 1:3] = 1 + + with self.assertWarns(Warning): + loss = BoundaryLoss(softmax=True)(input_tensor, target) + loss.backward() + self.assertGreater(input_tensor.grad.abs().sum().item(), 0.0) + + with self.assertWarns(Warning): + result = BoundaryLoss(include_background=False)(input_tensor.detach(), target) + self.assertTrue(torch.isfinite(result)) + + with self.assertWarns(Warning): + result = BoundaryLoss(to_onehot_y=True)(input_tensor.detach(), target) + self.assertTrue(torch.isfinite(result)) + + def test_to_onehot_y_accepts_channel_free_target(self): + """Test target labels can omit the singleton channel dimension.""" + input_tensor = torch.randn((2, 3, 4, 4)) + target = torch.randint(0, 3, size=(2, 4, 4)) + result = BoundaryLoss(to_onehot_y=True, softmax=True)(input_tensor, target) + self.assertTrue(torch.isfinite(result)) + + def test_degenerate_target_distance_map_is_zero(self): + """Test that empty and full classes don't create edge-biased distance maps.""" + loss_fn = BoundaryLoss() + empty_target = torch.zeros((1, 1, 4, 4)) + full_target = torch.ones((1, 1, 4, 4)) + + self.assertTrue(torch.equal(loss_fn.compute_distance_map(empty_target), torch.zeros_like(empty_target))) + self.assertTrue(torch.equal(loss_fn.compute_distance_map(full_target), torch.zeros_like(full_target))) + + def test_batch_reduction_changes_none_shape_and_values(self): + """Test that batch=True reduces the batch dimension before final reduction.""" + input_tensor = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]], [[[0.0, 1.0], [1.0, 0.0]]]]) + target = torch.tensor([[[[1.0, 1.0], [1.0, 1.0]]], [[[1.0, 0.0], [1.0, 0.0]]]]) + + batch_false = BoundaryLoss(reduction="none", batch=False)(input_tensor, target) + batch_true = BoundaryLoss(reduction="none", batch=True)(input_tensor, target) + + self.assertEqual(batch_false.shape, torch.Size([2, 1])) + self.assertEqual(batch_true.shape, torch.Size([1])) + self.assertTrue(torch.allclose(batch_true, batch_false.mean(dim=0))) + + def test_ill_shape(self): + """Test that mismatched shapes raise an error.""" + loss = BoundaryLoss() + with self.assertRaisesRegex(AssertionError, "shapes do not match"): + loss(torch.ones((1, 1, 2, 3)), torch.ones((1, 4, 5, 6))) + + def test_ill_opts(self): + """Test that invalid options raise errors.""" + with self.assertRaisesRegex(ValueError, ""): + BoundaryLoss(sigmoid=True, softmax=True) + with self.assertRaisesRegex(ValueError, ""): + BoundaryLoss(sigmoid=True, other_act=torch.tanh) + with self.assertRaisesRegex(ValueError, ""): + BoundaryLoss(softmax=True, other_act=torch.tanh) + with self.assertRaisesRegex(ValueError, ""): + BoundaryLoss(sigmoid=True, softmax=True, other_act=torch.tanh) + + chn_input = torch.ones((1, 1, 3, 3)) + chn_target = torch.ones((1, 1, 3, 3)) + with self.assertRaisesRegex(ValueError, ""): + BoundaryLoss(reduction="unknown")(chn_input, chn_target) + + def test_invalid_other_act_type(self): + """Test that non-callable other_act raises TypeError.""" + with self.assertRaises(TypeError): + BoundaryLoss(other_act="invalid") + + def test_empty_foreground(self): + """Test that empty foreground classes don't crash the loss.""" + loss_fn = BoundaryLoss(sigmoid=False) + result = loss_fn(EMPTY_FOREGROUND["input"], EMPTY_FOREGROUND["target"]) + self.assertTrue(torch.isfinite(result)) + + def test_dimension_validation(self): + """Test that unsupported dimensions raise errors.""" + loss = BoundaryLoss() + with self.assertRaises(ValueError): + # 1D input should fail + loss(torch.ones((1, 1, 10)), torch.ones((1, 1, 10))) + with self.assertRaises(ValueError): + # 4D input (5D with batch+channel) should fail + loss(torch.ones((1, 1, 2, 2, 2, 2)), torch.ones((1, 1, 2, 2, 2, 2))) + + def test_distance_map_computation(self): + """Test that distance maps are computed correctly for a simple case.""" + # Simple 3x3 case: foreground in center pixel + target = torch.zeros((1, 1, 3, 3)) + target[0, 0, 1, 1] = 1.0 # Center pixel is foreground + + loss_fn = BoundaryLoss() + distance_map = loss_fn.compute_distance_map(target) + + # Center pixel is on the boundary (single-pixel object), so distance should be 0 or near 0 + self.assertAlmostEqual(distance_map[0, 0, 1, 1].item(), 0.0, places=5) + + # Corners should be positive (outside foreground) + self.assertGreater(distance_map[0, 0, 0, 0].item(), 0) + self.assertGreater(distance_map[0, 0, 0, 2].item(), 0) + self.assertGreater(distance_map[0, 0, 2, 0].item(), 0) + self.assertGreater(distance_map[0, 0, 2, 2].item(), 0) + + # Neighbors of center should also be positive (outside foreground) + self.assertGreater(distance_map[0, 0, 0, 1].item(), 0) + self.assertGreater(distance_map[0, 0, 1, 0].item(), 0) + + def test_loss_gradient_flow(self): + """Test that gradients flow through the loss.""" + input_tensor = torch.randn((2, 2, 8, 8), requires_grad=True) + target = torch.ones((2, 2, 8, 8)) + + loss_fn = BoundaryLoss(sigmoid=True) + loss = loss_fn(input_tensor, target) + loss.backward() + + self.assertIsNotNone(input_tensor.grad) + self.assertTrue(torch.isfinite(input_tensor.grad).all()) + + def test_consistency_with_hausdorff_loss(self): + """Test that BoundaryLoss behaves differently from HausdorffDTLoss on the same input.""" + from monai.losses import HausdorffDTLoss + + input_tensor = torch.tensor([[[[2.0, -2.0], [-2.0, 2.0]]]]) + target = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) + + bl_loss = BoundaryLoss(sigmoid=True)(input_tensor, target) + hd_loss = HausdorffDTLoss(sigmoid=True)(input_tensor, target) + + # They should produce different values (different formulations) + self.assertNotAlmostEqual(bl_loss.item(), hd_loss.item(), places=3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/losses/test_cldice_loss.py b/tests/losses/test_cldice_loss.py index cb17cb81ad0..23c22dd395e 100644 --- a/tests/losses/test_cldice_loss.py +++ b/tests/losses/test_cldice_loss.py @@ -114,6 +114,29 @@ def test_invalid_iter_value(self): with self.assertRaises(ValueError): SoftclDiceLoss(iter_=-1) + def test_zero_input_is_finite(self): + loss = SoftclDiceLoss(smooth=1e-7, smooth_dr=1e-5) + result = loss(torch.zeros((1, 2, 4, 4)), torch.zeros((1, 2, 4, 4))) + self.assertTrue(torch.isfinite(result).all()) + + def test_non_default_smooth_dr_changes_result(self): + input_tensor = torch.zeros((1, 2, 4, 4)) + target = torch.zeros((1, 2, 4, 4)) + loss_a = SoftclDiceLoss(smooth=1e-7, smooth_dr=1e-3) + loss_b = SoftclDiceLoss(smooth=1e-7, smooth_dr=1e-5) + result_a = loss_a(input_tensor, target) + result_b = loss_b(input_tensor, target) + self.assertTrue(torch.isfinite(result_a).all()) + self.assertTrue(torch.isfinite(result_b).all()) + self.assertNotAlmostEqual(result_a.item(), result_b.item(), places=5) + + def test_non_overlapping_input_is_finite(self): + loss = SoftclDiceLoss(smooth=1e-7, smooth_dr=1e-5) + input_tensor = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + target = torch.tensor([[[[0.0, 0.0], [0.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + result = loss(input_tensor, target) + self.assertTrue(torch.isfinite(result).all()) + class TestSoftDiceclDiceLoss(unittest.TestCase): @parameterized.expand(COMBINED_CASES) @@ -146,6 +169,29 @@ def test_invalid_alpha_negative(self): with self.assertRaises(ValueError): SoftDiceclDiceLoss(alpha=-0.5) + def test_zero_input_is_finite(self): + loss = SoftDiceclDiceLoss(smooth=1e-7, smooth_dr=1e-5) + result = loss(torch.zeros((1, 2, 4, 4)), torch.zeros((1, 2, 4, 4))) + self.assertTrue(torch.isfinite(result).all()) + + def test_non_default_smooth_dr_changes_result(self): + input_tensor = torch.zeros((1, 2, 4, 4)) + target = torch.zeros((1, 2, 4, 4)) + loss_a = SoftDiceclDiceLoss(smooth=1e-7, smooth_dr=1e-3) + loss_b = SoftDiceclDiceLoss(smooth=1e-7, smooth_dr=1e-5) + result_a = loss_a(input_tensor, target) + result_b = loss_b(input_tensor, target) + self.assertTrue(torch.isfinite(result_a).all()) + self.assertTrue(torch.isfinite(result_b).all()) + self.assertNotAlmostEqual(result_a.item(), result_b.item(), places=5) + + def test_non_overlapping_input_is_finite(self): + loss = SoftDiceclDiceLoss(smooth=1e-7, smooth_dr=1e-5) + input_tensor = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + target = torch.tensor([[[[0.0, 0.0], [0.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + result = loss(input_tensor, target) + self.assertTrue(torch.isfinite(result).all()) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_dice_loss.py b/tests/losses/test_dice_loss.py index 66c038783a6..d8fb5e11957 100644 --- a/tests/losses/test_dice_loss.py +++ b/tests/losses/test_dice_loss.py @@ -104,11 +104,6 @@ }, 1.534853, ], - [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) - {"include_background": True, "sigmoid": True, "smooth_nr": 1e-6, "smooth_dr": 1e-6}, - {"input": torch.tensor([[[[1.0, -1.0], [-1.0, 1.0]]]]), "target": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]])}, - 0.307576, - ], [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) {"include_background": True, "sigmoid": True, "squared_pred": True}, {"input": torch.tensor([[[[1.0, -1.0], [-1.0, 1.0]]]]), "target": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]])}, diff --git a/tests/losses/test_generalized_dice_loss.py b/tests/losses/test_generalized_dice_loss.py index 8549e874822..7d60b99932e 100644 --- a/tests/losses/test_generalized_dice_loss.py +++ b/tests/losses/test_generalized_dice_loss.py @@ -112,11 +112,6 @@ }, 0.0, ], - [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) - {"include_background": True, "sigmoid": True, "smooth_nr": 1e-6, "smooth_dr": 1e-6}, - {"input": torch.tensor([[[[1.0, -1.0], [-1.0, 1.0]]]]), "target": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]])}, - 0.307576, - ], [ # shape: (1, 2, 4), (1, 1, 4) { "include_background": True, diff --git a/tests/losses/test_perceptual_loss.py b/tests/losses/test_perceptual_loss.py index 8d94fdc1aef..051a79fef13 100644 --- a/tests/losses/test_perceptual_loss.py +++ b/tests/losses/test_perceptual_loss.py @@ -17,6 +17,7 @@ from parameterized import parameterized from monai.losses import PerceptualLoss +from monai.losses.perceptual import normalize_tensor from monai.utils import optional_import from tests.test_utils import assert_allclose, skip_if_downloading_fails, skip_if_quick @@ -126,6 +127,16 @@ def test_non_medicalnet_3d_without_fake_3d(self, network_type): with self.assertRaises(ValueError): PerceptualLoss(spatial_dims=3, network_type=network_type, is_fake_3d=False) + def test_normalize_tensor_zero_norm_finite_gradient(self): + # regression test for #8412: a zero-norm feature vector (e.g. from identical + # input/target features) must not produce NaN gradients via SqrtBackward. + x = torch.zeros(2, 4, 8, 8, requires_grad=True) + out = normalize_tensor(x) + out.sum().backward() + self.assertFalse(torch.isnan(out).any()) + self.assertIsNotNone(x.grad) + self.assertFalse(torch.isnan(x.grad).any()) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_unified_focal_loss.py b/tests/losses/test_unified_focal_loss.py index 3b868a560e8..845d359cc79 100644 --- a/tests/losses/test_unified_focal_loss.py +++ b/tests/losses/test_unified_focal_loss.py @@ -26,14 +26,7 @@ "y_true": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]), }, 0.0, - ], - [ # shape: (2, 1, 2, 2), (2, 1, 2, 2) - { - "y_pred": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]), - "y_true": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]), - }, - 0.0, - ], + ] ] diff --git a/tests/metrics/test_absolute_volume_difference.py b/tests/metrics/test_absolute_volume_difference.py new file mode 100644 index 00000000000..37ee61bcfba --- /dev/null +++ b/tests/metrics/test_absolute_volume_difference.py @@ -0,0 +1,161 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest + +import torch + +from monai.metrics import AbsoluteVolumeDifferenceMetric, compute_absolute_volume_difference + + +class TestComputeAbsoluteVolumeDifference(unittest.TestCase): + """Tests for the standalone compute_absolute_volume_difference function.""" + + def test_perfect_prediction_returns_zero(self): + """Identical prediction and ground truth should yield AVD of zero for all classes.""" + # identical masks → AVD = 0 for every class + y = torch.zeros(2, 3, 4, 4) + y[:, 1, :2, :2] = 1.0 + y[:, 2, 2:, 2:] = 1.0 + result = compute_absolute_volume_difference(y_pred=y, y=y, ignore_empty=False) + self.assertEqual(result.shape, torch.Size([2, 3])) + self.assertTrue(torch.all(result == 0.0)) + + def test_known_volume_difference(self): + """AVD should equal the absolute difference in foreground voxel counts between prediction and GT.""" + # batch=1, 2 classes (background + foreground), 1D spatial of length 10 + y_pred = torch.zeros(1, 2, 10) + y_true = torch.zeros(1, 2, 10) + y_pred[0, 1, :7] = 1.0 # 7 foreground voxels predicted + y_true[0, 1, :4] = 1.0 # 4 foreground voxels in GT + result = compute_absolute_volume_difference(y_pred=y_pred, y=y_true, ignore_empty=False) + # channel 0: both all-zeros → AVD = 0 + # channel 1: |7 - 4| = 3 + self.assertAlmostEqual(result[0, 0].item(), 0.0) + self.assertAlmostEqual(result[0, 1].item(), 3.0) + + def test_ignore_background(self): + """Setting include_background=False should strip the first channel and reduce output shape accordingly.""" + y_pred = torch.zeros(2, 3, 8, 8) + y_true = torch.zeros(2, 3, 8, 8) + y_pred[:, 1, :3, :3] = 1.0 + y_true[:, 1, :4, :4] = 1.0 + result = compute_absolute_volume_difference(y_pred=y_pred, y=y_true, include_background=False) + # background channel stripped → shape [2, 2] + self.assertEqual(result.shape, torch.Size([2, 2])) + + def test_ignore_empty_sets_nan(self): + """Channels with no ground-truth foreground voxels should be NaN when ignore_empty=True.""" + # channel 1 has no GT voxels → should be NaN when ignore_empty=True + y_pred = torch.zeros(1, 2, 6) + y_true = torch.zeros(1, 2, 6) + y_pred[0, 0, :3] = 1.0 + result = compute_absolute_volume_difference(y_pred=y_pred, y=y_true, ignore_empty=True) + # channel 0: GT is empty → NaN + self.assertTrue(torch.isnan(result[0, 0])) + # channel 1: GT is empty → NaN + self.assertTrue(torch.isnan(result[0, 1])) + + def test_ignore_empty_false_returns_pred_volume(self): + """With ignore_empty=False and empty GT, AVD should equal the predicted volume.""" + # when GT is all zero and ignore_empty=False, AVD = |V_pred - 0| = V_pred + y_pred = torch.zeros(1, 2, 6) + y_true = torch.zeros(1, 2, 6) + y_pred[0, 1, :5] = 1.0 + result = compute_absolute_volume_difference(y_pred=y_pred, y=y_true, ignore_empty=False) + self.assertAlmostEqual(result[0, 1].item(), 5.0) + + def test_shape_mismatch_raises(self): + """Mismatched y_pred and y shapes should raise a ValueError.""" + with self.assertRaises(ValueError): + compute_absolute_volume_difference(y_pred=torch.zeros(2, 3, 8, 8), y=torch.zeros(2, 3, 4, 4)) + + def test_too_few_dims_raises(self): + """Input tensors with fewer than 3 dimensions should raise a ValueError.""" + with self.assertRaises(ValueError): + compute_absolute_volume_difference(y_pred=torch.zeros(2, 3), y=torch.zeros(2, 3)) + + def test_3d_volumes(self): + """AVD should correctly count voxel differences in 3-D spatial inputs.""" + # 3-D spatial (D, H, W) + y_pred = torch.zeros(1, 2, 8, 8, 8) + y_true = torch.zeros(1, 2, 8, 8, 8) + y_pred[0, 1, :4, :4, :4] = 1.0 # 64 voxels + y_true[0, 1, :3, :3, :3] = 1.0 # 27 voxels + result = compute_absolute_volume_difference(y_pred=y_pred, y=y_true, ignore_empty=False) + self.assertAlmostEqual(result[0, 1].item(), 37.0) + + def test_output_shape_multi_class(self): + """Output shape should be [batch_size, num_classes] for multi-class inputs.""" + y = torch.randint(0, 2, (4, 5, 16, 16)).float() + result = compute_absolute_volume_difference(y_pred=y, y=y, ignore_empty=False) + self.assertEqual(result.shape, torch.Size([4, 5])) + + +class TestAbsoluteVolumeDifferenceMetric(unittest.TestCase): + """Tests for the AbsoluteVolumeDifferenceMetric class (cumulative interface).""" + + def test_aggregate_mean(self): + """Mean reduction over accumulated batches should return the correct per-class AVD.""" + y_pred = torch.zeros(2, 2, 8, 8) + y_true = torch.zeros(2, 2, 8, 8) + y_pred[:, 1, :6, :6] = 1.0 # 36 voxels per batch item + y_true[:, 1, :4, :4] = 1.0 # 16 voxels per batch item + metric = AbsoluteVolumeDifferenceMetric(include_background=False, reduction="mean", ignore_empty=False) + metric(y_pred, y_true) + agg = metric.aggregate() + # single foreground channel, AVD = 20 for both batch items → mean = 20 + self.assertAlmostEqual(agg.item(), 20.0) + metric.reset() + + def test_aggregate_returns_not_nans_when_requested(self): + """When get_not_nans=True, aggregate should return a (metric, not_nans) tuple.""" + y_pred = torch.zeros(2, 2, 4, 4) + y_true = torch.zeros(2, 2, 4, 4) + y_pred[:, 1, :2, :2] = 1.0 + y_true[:, 1, :2, :2] = 1.0 + metric = AbsoluteVolumeDifferenceMetric(include_background=False, get_not_nans=True) + metric(y_pred, y_true) + result, not_nans = metric.aggregate() + self.assertIsInstance(result, torch.Tensor) + self.assertIsInstance(not_nans, torch.Tensor) + metric.reset() + + def test_cumulative_accumulation(self): + """Multiple forward calls before aggregate should use all accumulated data correctly.""" + # calling the metric twice and aggregating should use all accumulated data + metric = AbsoluteVolumeDifferenceMetric(include_background=False, reduction="mean", ignore_empty=False) + for _ in range(3): + y_pred = torch.zeros(1, 2, 8) + y_true = torch.zeros(1, 2, 8) + y_pred[0, 1, :6] = 1.0 + y_true[0, 1, :4] = 1.0 + metric(y_pred, y_true) + agg = metric.aggregate() + self.assertAlmostEqual(agg.item(), 2.0) + metric.reset() + + def test_reset_clears_buffer(self): + """Calling reset() should clear the buffer so a subsequent aggregate() raises.""" + metric = AbsoluteVolumeDifferenceMetric(ignore_empty=False) + y = torch.zeros(1, 2, 4) + y[0, 1, :2] = 1.0 + metric(y, y) + metric.reset() + # after reset the buffer should be empty; calling aggregate raises + with self.assertRaises(ValueError): + metric.aggregate() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/metrics/test_compute_froc.py b/tests/metrics/test_compute_froc.py index 4dc05073665..aa889ddb073 100644 --- a/tests/metrics/test_compute_froc.py +++ b/tests/metrics/test_compute_froc.py @@ -60,6 +60,34 @@ 3, ] +TEST_CASE_EXCLUDE_ABSENT = [ + { + "probs": torch.tensor([1, 0.6, 0.8]), + "y_coord": torch.tensor([0, 2, 3]), + "x_coord": torch.tensor([3, 0, 1]), + "evaluation_mask": np.array([[0, 0, 1, 1], [2, 2, 0, 0], [0, 3, 3, 0], [0, 3, 3, 3]]), + "labels_to_exclude": [5], + "resolution_level": 0, + }, + np.array([0.6]), + np.array([1, 0, 0.8]), + 3, +] + +TEST_CASE_EXCLUDE_DUPLICATE = [ + { + "probs": torch.tensor([1, 0.6, 0.8]), + "y_coord": torch.tensor([0, 2, 3]), + "x_coord": torch.tensor([3, 0, 1]), + "evaluation_mask": np.array([[0, 0, 1, 1], [2, 2, 0, 0], [0, 3, 3, 0], [0, 3, 3, 3]]), + "labels_to_exclude": [2, 2], + "resolution_level": 0, + }, + np.array([0.6]), + np.array([1, 0, 0.8]), + 2, +] + TEST_CASE_4 = [ { "fp_probs": np.array([0.8, 0.6]), @@ -112,7 +140,9 @@ class TestComputeFpTp(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) + @parameterized.expand( + [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_EXCLUDE_ABSENT, TEST_CASE_EXCLUDE_DUPLICATE] + ) def test_value(self, input_data, expected_fp, expected_tp, expected_num): fp_probs, tp_probs, num_tumors = compute_fp_tp_probs(**input_data) np.testing.assert_allclose(fp_probs, expected_fp, rtol=1e-5) diff --git a/tests/metrics/test_surface_distance.py b/tests/metrics/test_surface_distance.py index 85db389f80a..3461e44a5bb 100644 --- a/tests/metrics/test_surface_distance.py +++ b/tests/metrics/test_surface_distance.py @@ -18,6 +18,10 @@ from parameterized import parameterized from monai.metrics import SurfaceDistanceMetric +from monai.metrics.utils import get_mask_edges, get_surface_distance +from monai.utils import optional_import + +distance_transform_edt, has_scipy = optional_import("scipy.ndimage", name="distance_transform_edt") _device = "cuda:0" if torch.cuda.is_available() else "cpu" @@ -182,5 +186,44 @@ def test_nans(self, input_data): np.testing.assert_allclose(0, not_nans, rtol=1e-5) +KDTREE_SPACINGS = [["isotropic_default", None], ["isotropic", (1.0, 1.0, 1.0)], ["anisotropic", (1.0, 2.5, 0.5)]] + + +def _edge_masks(seed=0): + # two offset spheres plus a few scattered false positives in the prediction, so the + # surfaces are non-trivially apart and an outlier expands the cropped bounding box. + gt = create_spherical_seg_3d(radius=20, centre=(30, 30, 30)) + pred = create_spherical_seg_3d(radius=20, centre=(32, 31, 30)) + rng = np.random.RandomState(seed) + for _ in range(5): + pred[tuple(rng.randint(0, s) for s in pred.shape)] = 1 + edges_pred, edges_gt = get_mask_edges(pred, gt) + return np.asarray(edges_pred, dtype=bool), np.asarray(edges_gt, dtype=bool) + + +@unittest.skipUnless(has_scipy, "Requires scipy.") +class TestSurfaceDistanceKDTreeMatchesEDT(unittest.TestCase): + @parameterized.expand(KDTREE_SPACINGS) + def test_cpu_kdtree_euclidean_distances_match_dense_edt(self, _name, spacing): + edges_pred, edges_gt = _edge_masks() + result = np.asarray(get_surface_distance(edges_pred, edges_gt, distance_metric="euclidean", spacing=spacing)) + reference = distance_transform_edt(~edges_gt, sampling=spacing)[edges_pred] + # same multiset of distances (downstream metrics only use max/percentile/mean) + np.testing.assert_allclose(np.sort(result), np.sort(reference), rtol=1e-5, atol=1e-5) + self.assertEqual(result.dtype, np.float32) + self.assertEqual(result.shape, reference.shape) + + def test_torch_input_preserves_type_device_and_matches_dense_edt(self): + edges_pred, edges_gt = _edge_masks() + spacing = (1.0, 2.5, 0.5) + seg_pred, seg_gt = torch.as_tensor(edges_pred), torch.as_tensor(edges_gt) + result = get_surface_distance(seg_pred, seg_gt, distance_metric="euclidean", spacing=spacing) + self.assertIsInstance(result, torch.Tensor) + self.assertEqual(result.dtype, torch.float32) + self.assertEqual(result.device, seg_pred.device) + reference = distance_transform_edt(~edges_gt, sampling=spacing)[edges_pred] + np.testing.assert_allclose(np.sort(result.cpu().numpy()), np.sort(reference), rtol=1e-5, atol=1e-5) + + if __name__ == "__main__": unittest.main() diff --git a/tests/min_tests.py b/tests/min_tests.py index 2d68f099a7f..25a42fe4b5d 100644 --- a/tests/min_tests.py +++ b/tests/min_tests.py @@ -20,10 +20,10 @@ def run_testsuit(): """ Load test cases by excluding those need external dependencies. - The loaded cases should work with "requirements-min.txt":: + The loaded cases should work with testing requirements:: # in the monai repo folder: - pip install -r requirements-min.txt + pip install -e .[testing] QUICKTEST=true python -m tests.min_tests :return: a test suite @@ -112,6 +112,8 @@ def run_testsuit(): "test_hausdorff_distance", "test_header_correct", "test_hilbert_transform", + "test_hyena_block", + "test_hyena_nd_unetr", "test_hovernet_loss", "test_image_dataset", "test_image_rw", @@ -145,6 +147,7 @@ def run_testsuit(): "test_mlp", "test_nifti_header_revise", "test_nifti_rw", + "test_navit", "test_nuclick_transforms", "test_nrrd_reader", "test_occlusion_sensitivity", diff --git a/tests/networks/blocks/test_hyena_block.py b/tests/networks/blocks/test_hyena_block.py new file mode 100644 index 00000000000..5f3835e2dfa --- /dev/null +++ b/tests/networks/blocks/test_hyena_block.py @@ -0,0 +1,336 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +import unittest +from unittest import skipUnless + +import torch +import torch.nn as nn +from parameterized import parameterized + +from monai.networks.blocks.hyena import ( + DepthwiseFFTConv2d, + DepthwiseFFTConv3d, + HyenaMixer, + HyenaTransformerBlock, + is_nvsubquadratic_available, +) + +HAS_NVSUBQ = is_nvsubquadratic_available() +HAS_CUDA = torch.cuda.is_available() + + +# --------------------------------------------------------------------------- +# DepthwiseFFTConv{2,3}d — no nvsubquadratic dependency +# --------------------------------------------------------------------------- + + +class TestDepthwiseFFTConvShape(unittest.TestCase): + """The FFT conv must preserve spatial dimensions for any depthwise config.""" + + @parameterized.expand( + [ + ("3d_k3_d16", (2, 8, 16, 16, 16), 3, 3), + ("3d_k1_d10", (1, 16, 10, 10, 10), 1, 3), + ("3d_k5_d12", (1, 8, 12, 12, 12), 5, 3), + ("2d_k3_d32", (2, 8, 32, 32), 3, 2), + ] + ) + def test_output_shape(self, _name, input_shape, kernel_size, spatial_dims): + channels = input_shape[1] + cls = DepthwiseFFTConv3d if spatial_dims == 3 else DepthwiseFFTConv2d + conv = cls(channels, channels, kernel_size=kernel_size, groups=channels, padding=kernel_size // 2) + x = torch.randn(*input_shape) + self.assertEqual(conv(x).shape, x.shape) + + +class TestDepthwiseFFTConvNumerics(unittest.TestCase): + """FFT conv must match the equivalent ``nn.Conv{2,3}d`` numerically.""" + + @parameterized.expand([("d8_s12", 8, 12), ("d16_s8", 16, 8), ("d32_s6", 32, 6)]) + def test_matches_conv3d(self, _name, channels, spatial): + ref = nn.Conv3d(channels, channels, kernel_size=3, groups=channels, padding=1, bias=False) + fft = DepthwiseFFTConv3d(channels, channels, kernel_size=3, groups=channels, padding=1) + with torch.no_grad(): + fft.weight.copy_(ref.weight) + x = torch.randn(2, channels, spatial, spatial, spatial) + with torch.no_grad(): + torch.testing.assert_close(fft(x), ref(x), atol=1e-4, rtol=1e-4) + + def test_matches_conv2d(self): + channels, spatial = 8, 16 + ref = nn.Conv2d(channels, channels, kernel_size=3, groups=channels, padding=1, bias=False) + fft = DepthwiseFFTConv2d(channels, channels, kernel_size=3, groups=channels, padding=1) + with torch.no_grad(): + fft.weight.copy_(ref.weight) + x = torch.randn(2, channels, spatial, spatial) + with torch.no_grad(): + torch.testing.assert_close(fft(x), ref(x), atol=1e-4, rtol=1e-4) + + +class TestDepthwiseFFTConvDtype(unittest.TestCase): + """Output dtype must match input dtype (AMP transparency).""" + + @parameterized.expand([("fp16", torch.float16), ("bf16", torch.bfloat16)]) + def test_amp_dtype_preserved(self, _name, dtype): + conv = DepthwiseFFTConv3d(8, 8, kernel_size=3, groups=8, padding=1) + x = torch.randn(1, 8, 8, 8, 8, dtype=dtype) + out = conv(x) + self.assertEqual(out.dtype, dtype) + self.assertEqual(out.shape, x.shape) + + def test_float32_preserved(self): + conv = DepthwiseFFTConv3d(8, 8, kernel_size=3, groups=8, padding=1) + x = torch.randn(1, 8, 8, 8, 8) + self.assertEqual(conv(x).dtype, torch.float32) + + +class TestDepthwiseFFTConvGradients(unittest.TestCase): + """Backward pass must produce gradients on both input and weight.""" + + def test_gradients_flow_3d(self): + conv = DepthwiseFFTConv3d(8, 8, kernel_size=3, groups=8, padding=1) + x = torch.randn(1, 8, 10, 10, 10, requires_grad=True) + conv(x).sum().backward() + self.assertIsNotNone(x.grad) + self.assertIsNotNone(conv.weight.grad) + + def test_gradients_flow_2d(self): + conv = DepthwiseFFTConv2d(8, 8, kernel_size=3, groups=8, padding=1) + x = torch.randn(1, 8, 16, 16, requires_grad=True) + conv(x).sum().backward() + self.assertIsNotNone(x.grad) + self.assertIsNotNone(conv.weight.grad) + + +class TestDepthwiseFFTConvConstruction(unittest.TestCase): + """Reject configurations the FFT path cannot represent.""" + + def test_rejects_non_depthwise(self): + with self.assertRaises(ValueError): + DepthwiseFFTConv3d(8, 8, kernel_size=3, groups=1, padding=1) + + def test_rejects_bias(self): + with self.assertRaises(ValueError): + DepthwiseFFTConv3d(8, 8, kernel_size=3, groups=8, padding=1, bias=True) + + def test_rejects_non_same_padding(self): + # forward() crops to the input size assuming padding == kernel_size // 2. + with self.assertRaisesRegex(ValueError, "same"): + DepthwiseFFTConv3d(8, 8, kernel_size=3, groups=8, padding=0) + + def test_rejects_even_kernel(self): + with self.assertRaisesRegex(ValueError, "same"): + DepthwiseFFTConv3d(8, 8, kernel_size=4, groups=8, padding=2) + + def test_weight_shape(self): + conv = DepthwiseFFTConv3d(16, 16, kernel_size=3, groups=16, padding=1) + self.assertEqual(conv.weight.shape, (16, 1, 3, 3, 3)) + + def test_weight_initialised(self): + conv = DepthwiseFFTConv3d(64, 64, kernel_size=3, groups=64, padding=1) + # kaiming_uniform with fan_in = 1 * 3^3 = 27 → bound ≈ 1/sqrt(27) ≈ 0.19 + self.assertGreater(conv.weight.abs().max().item(), 0.0) + self.assertLess(conv.weight.abs().max().item(), 5.0 / math.sqrt(27)) + + +# --------------------------------------------------------------------------- +# HyenaMixer configuration validation — no CUDA required (construction only, +# but nvsubquadratic must be present to reach the validation branch) +# --------------------------------------------------------------------------- + + +@skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") +class TestHyenaMixerConfigValidation(unittest.TestCase): + def test_rejects_circular_with_double_grid(self): + with self.assertRaisesRegex(ValueError, "circular.*single"): + HyenaMixer(dim=12, spatial_dims=3, fft_padding="circular", grid_type="double") + + def test_rejects_chunked_with_circular(self): + with self.assertRaisesRegex(ValueError, "chunked.*zero|zero.*chunked"): + HyenaMixer(dim=12, spatial_dims=3, fft_padding="circular", use_chunked_fftconv=True) + + def test_rejects_bad_fft_padding(self): + with self.assertRaisesRegex(ValueError, "fft_padding"): + HyenaMixer(dim=12, spatial_dims=3, fft_padding="reflective") + + def test_rejects_bad_grid_type(self): + with self.assertRaisesRegex(ValueError, "grid_type"): + HyenaMixer(dim=12, spatial_dims=3, grid_type="triple") + + def test_rejects_bad_spatial_dims(self): + with self.assertRaisesRegex(ValueError, "spatial_dims"): + HyenaMixer(dim=12, spatial_dims=4) + + def test_zero_double_chunked_constructs(self): + m = HyenaMixer(dim=12, spatial_dims=3, fft_padding="zero", grid_type="double", use_chunked_fftconv=True) + self.assertEqual(m.dim, 12) + + +class TestHyenaMixerOptionalDep(unittest.TestCase): + """When ``nvsubquadratic`` is missing, ``HyenaMixer`` must raise a clear ImportError.""" + + @skipUnless(not HAS_NVSUBQ, "Only runs when nvsubquadratic is absent") + def test_raises_import_error(self): + with self.assertRaisesRegex(ImportError, "nvsubquadratic"): + HyenaMixer(dim=12, spatial_dims=3) + + +# --------------------------------------------------------------------------- +# Forward shape — channels-last [B, *spatial, C] preserved +# --------------------------------------------------------------------------- + + +@skipUnless(HAS_NVSUBQ and HAS_CUDA, "Requires nvsubquadratic and CUDA") +class TestHyenaMixerForward(unittest.TestCase): + device = "cuda" + + def test_3d_forward_shape(self): + m = HyenaMixer(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 8, 8, 8, 12, device=self.device) + self.assertEqual(m(x).shape, x.shape) + + def test_2d_forward_shape(self): + m = HyenaMixer(dim=8, spatial_dims=2).to(self.device) + x = torch.randn(2, 16, 16, 8, device=self.device) + self.assertEqual(m(x).shape, x.shape) + + def test_zero_padding_forward(self): + m = HyenaMixer(dim=12, spatial_dims=3, fft_padding="zero", grid_type="single").to(self.device) + x = torch.randn(2, 8, 8, 8, 12, device=self.device) + self.assertEqual(m(x).shape, x.shape) + + def test_zero_double_chunked_forward(self): + m = HyenaMixer(dim=12, spatial_dims=3, fft_padding="zero", grid_type="double", use_chunked_fftconv=True).to( + self.device + ) + x = torch.randn(2, 8, 8, 8, 12, device=self.device) + self.assertEqual(m(x).shape, x.shape) + + +@skipUnless(HAS_NVSUBQ and HAS_CUDA, "Requires nvsubquadratic and CUDA") +class TestHyenaMixerGradients(unittest.TestCase): + device = "cuda" + + def test_qkv_and_out_proj_get_grads(self): + m = HyenaMixer(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 6, 6, 6, 12, device=self.device, requires_grad=True) + m(x).sum().backward() + self.assertIsNotNone(m.qkv_proj.weight.grad) + self.assertIsNotNone(m.out_proj.weight.grad) + self.assertIsNotNone(x.grad) + + def test_mixer_internal_params_get_grads(self): + m = HyenaMixer(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 6, 6, 6, 12, device=self.device) + m(x).sum().backward() + with_grad = [ + name for name, p in m.mixer.named_parameters() if p.grad is not None and p.grad.abs().sum().item() > 0 + ] + self.assertGreater(len(with_grad), 0, "no mixer-internal params received a gradient") + + +@skipUnless(HAS_NVSUBQ and HAS_CUDA, "Requires nvsubquadratic and CUDA") +class TestHyenaMixerAMP(unittest.TestCase): + """Under ``torch.autocast`` the output dtype must match the autocast dtype.""" + + device = "cuda" + + @parameterized.expand([("fp16", torch.float16), ("bf16", torch.bfloat16)]) + def test_autocast_output_dtype(self, _name, dtype): + m = HyenaMixer(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 6, 6, 6, 12, device=self.device) + with torch.autocast("cuda", dtype=dtype): + out = m(x) + self.assertEqual(out.dtype, dtype) + self.assertEqual(out.shape, x.shape) + + def test_float32_preserved(self): + m = HyenaMixer(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 6, 6, 6, 12, device=self.device) + self.assertEqual(m(x).dtype, torch.float32) + + +@skipUnless(HAS_NVSUBQ and HAS_CUDA, "Requires nvsubquadratic and CUDA") +class TestHyenaMixerDeterminism(unittest.TestCase): + device = "cuda" + + def test_same_seed_same_output(self): + torch.manual_seed(0) + m1 = HyenaMixer(dim=12, spatial_dims=3).to(self.device) + torch.manual_seed(0) + m2 = HyenaMixer(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(1, 6, 6, 6, 12, device=self.device) + with torch.no_grad(): + y1, y2 = m1(x), m2(x) + torch.testing.assert_close(y1, y2, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# HyenaTransformerBlock — full residual forward path +# --------------------------------------------------------------------------- + + +@skipUnless(HAS_NVSUBQ and HAS_CUDA, "Requires nvsubquadratic and CUDA") +class TestHyenaTransformerBlock(unittest.TestCase): + device = "cuda" + + def test_3d_forward_shape(self): + blk = HyenaTransformerBlock(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 6, 6, 6, 12, device=self.device) + self.assertEqual(blk(x).shape, x.shape) + + def test_2d_forward_shape(self): + blk = HyenaTransformerBlock(dim=8, spatial_dims=2).to(self.device) + x = torch.randn(2, 16, 16, 8, device=self.device) + self.assertEqual(blk(x).shape, x.shape) + + def test_grad_flow_through_block(self): + blk = HyenaTransformerBlock(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 6, 6, 6, 12, device=self.device) + blk(x).sum().backward() + self.assertIsNotNone(blk.mixer.qkv_proj.weight.grad) + self.assertIsNotNone(blk.mixer.out_proj.weight.grad) + mlp_params_with_grad = [p for p in blk.mlp.parameters() if p.grad is not None] + self.assertGreater(len(mlp_params_with_grad), 0) + + def test_mask_matrix_accepted_and_ignored(self): + """``mask_matrix`` is accepted (signature parity with Swin) but ignored.""" + blk = HyenaTransformerBlock(dim=12, spatial_dims=3).to(self.device) + x = torch.randn(2, 6, 6, 6, 12, device=self.device) + with torch.no_grad(): + y1 = blk(x) + y2 = blk(x, mask_matrix=torch.ones(1, device=self.device)) + torch.testing.assert_close(y1, y2) + + +@skipUnless(HAS_NVSUBQ and HAS_CUDA, "Requires nvsubquadratic and CUDA") +class TestHyenaMixerFFTShortConv(unittest.TestCase): + """The use_fft_short_conv=True path swaps Conv3d for DepthwiseFFTConv3d.""" + + device = "cuda" + + def test_3d_constructs_and_runs(self): + m = HyenaMixer(dim=12, spatial_dims=3, use_fft_short_conv=True).to(self.device) + x = torch.randn(2, 8, 8, 8, 12, device=self.device) + self.assertEqual(m(x).shape, x.shape) + + def test_3d_with_short_conv_chunks(self): + m = HyenaMixer(dim=12, spatial_dims=3, use_fft_short_conv=True, short_conv_fft_chunk_size=4).to(self.device) + x = torch.randn(2, 8, 8, 8, 12, device=self.device) + self.assertEqual(m(x).shape, x.shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/networks/blocks/warp/test_warp.py b/tests/networks/blocks/warp/test_warp.py index 93af5597908..1f23664234b 100644 --- a/tests/networks/blocks/warp/test_warp.py +++ b/tests/networks/blocks/warp/test_warp.py @@ -12,6 +12,7 @@ import unittest from pathlib import Path +from unittest import mock import numpy as np import torch @@ -138,6 +139,43 @@ def test_ill_shape(self): with self.assertRaisesRegex(ValueError, ""): warp_layer(image=torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), ddf=torch.zeros(1, 2, 3, 3)) + def test_jitter(self): + ddf = torch.zeros(1, 2, 4, 5) + grid = Warp(jitter=True).get_reference_grid(ddf, jitter=True, seed=0) + self.assertTrue(grid.is_floating_point()) + self.assertFalse(torch.equal(grid, grid.round())) + + grid = Warp().get_reference_grid(ddf, jitter=False) + self.assertTrue(torch.equal(grid, grid.round())) + + same = Warp().get_reference_grid(ddf, jitter=True, seed=7) + repeat = Warp().get_reference_grid(ddf, jitter=True, seed=7) + other = Warp().get_reference_grid(ddf, jitter=True, seed=8) + self.assertTrue(torch.equal(same, repeat)) + self.assertFalse(torch.equal(same, other)) + + @mock.patch("monai.networks.blocks.warp.USE_COMPILED", False) + def test_singleton_spatial_dim(self): + """ + Regression test for a singleton spatial dimension (a single-slice volume or a + single-row/column image), where the grid normalization ``* 2 / (dim - 1)`` previously + divided by zero. + + The native ``grid_sample`` path is forced via ``USE_COMPILED=False`` because only that + branch normalizes the grid; the csrc ``grid_pull`` path is unaffected. ``padding_mode`` + is ``"zeros"`` so an out-of-range (pre-fix ``nan``) coordinate maps to 0 and exposes the + bug, whereas ``"border"``/``"reflection"`` would clamp onto the lone voxel and mask it. + For a zero displacement field the warped output must contain no ``nan`` and must equal + the input image. + """ + for shape, ndim in [((1, 1, 1, 4, 4), 3), ((1, 1, 1, 5), 2), ((1, 1, 5, 1), 2)]: + image = torch.rand(*shape) + ddf = torch.zeros(shape[0], ndim, *shape[2:]) + warp_layer = Warp(mode="bilinear", padding_mode="zeros") + result = warp_layer(image, ddf) + self.assertFalse(torch.isnan(result).any(), f"NaN in warp output for shape {shape}") + np.testing.assert_allclose(result.cpu().numpy(), image.cpu().numpy(), rtol=1e-4, atol=1e-4) + def test_grad(self): for b in GridSampleMode: for p in GridSamplePadMode: diff --git a/tests/networks/nets/test_autoencoderkl.py b/tests/networks/nets/test_autoencoderkl.py index af0c55d6ece..33972c7ece3 100644 --- a/tests/networks/nets/test_autoencoderkl.py +++ b/tests/networks/nets/test_autoencoderkl.py @@ -99,21 +99,6 @@ (1, 1, 16, 16), (1, 4, 4, 4), ], - [ - { - "spatial_dims": 2, - "in_channels": 1, - "out_channels": 1, - "channels": (4, 4, 4), - "latent_channels": 4, - "attention_levels": (False, False, False), - "num_res_blocks": 1, - "norm_num_groups": 4, - }, - (1, 1, 16, 16), - (1, 1, 16, 16), - (1, 4, 4, 4), - ], [ { "spatial_dims": 2, diff --git a/tests/networks/nets/test_hyena_nd_unetr.py b/tests/networks/nets/test_hyena_nd_unetr.py new file mode 100644 index 00000000000..4fdb7356f10 --- /dev/null +++ b/tests/networks/nets/test_hyena_nd_unetr.py @@ -0,0 +1,137 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest +from unittest import skipUnless + +import torch +from parameterized import parameterized + +from monai.networks.blocks.hyena import HyenaTransformerBlock, is_nvsubquadratic_available +from monai.networks.nets.hyena_nd_unetr import PAPER_VARIANTS, HyenaNDUNETR +from monai.networks.nets.swin_unetr import SwinTransformerBlock, SwinUNETR +from tests.test_utils import skip_if_no_cuda + +HAS_NVSUBQ = is_nvsubquadratic_available() + + +PAPER_VARIANT_CASES = [ + ("HHHH", (True, True, True, True)), + ("HAHA", (True, False, True, False)), + ("HHAA", (True, True, False, False)), +] + + +def _block_type_at_stage(model, stage_idx): + layer_attr = ["layers1", "layers2", "layers3", "layers4"][stage_idx] + return type(getattr(model.swinViT, layer_attr)[0].blocks[0]) + + +class TestHyenaNDUNETRConstructorContract(unittest.TestCase): + """``HyenaNDUNETR.__init__`` enforces an explicit, non-empty ``hyena_stages``.""" + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_explicit_stages_required(self): + with self.assertRaisesRegex(ValueError, "requires `hyena_stages`"): + HyenaNDUNETR(in_channels=1, out_channels=14, feature_size=12, hyena_stages=None) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_wrong_length_stages_rejected(self): + with self.assertRaisesRegex(ValueError, "length 4"): + HyenaNDUNETR(in_channels=1, out_channels=14, feature_size=12, hyena_stages=(True, True)) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_all_false_stages_rejected(self): + with self.assertRaisesRegex(ValueError, "at least one stage"): + HyenaNDUNETR(in_channels=1, out_channels=14, feature_size=12, hyena_stages=(False, False, False, False)) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_use_hyena_kwarg_rejected(self): + """The subclass forces use_hyena=True; caller may not override via kwargs.""" + with self.assertRaisesRegex(TypeError, "use_hyena"): + HyenaNDUNETR( + in_channels=1, out_channels=14, feature_size=12, hyena_stages=(True, True, False, False), use_hyena=True + ) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_subclass_of_swin_unetr(self): + m = HyenaNDUNETR(in_channels=1, out_channels=14, feature_size=12, hyena_stages=(True, True, False, False)) + self.assertIsInstance(m, SwinUNETR) + # The forced kwargs land on the instance via SwinUNETR.__init__. + self.assertTrue(m.use_hyena) + self.assertEqual(m.hyena_stages, (True, True, False, False)) + + +class TestHyenaNDUNETRFromPaperVariant(unittest.TestCase): + """``get_variant`` maps {HHHH, HAHA, HHAA} to the correct stage pattern.""" + + @parameterized.expand(PAPER_VARIANT_CASES) + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_returns_expected_stages(self, name, expected_stages): + m = HyenaNDUNETR.get_variant(name, in_channels=1, out_channels=14, feature_size=12) + self.assertEqual(m.hyena_stages, expected_stages) + for stage_idx, want_hyena in enumerate(expected_stages): + block_type = _block_type_at_stage(m, stage_idx) + if want_hyena: + self.assertIs(block_type, HyenaTransformerBlock) + else: + self.assertIs(block_type, SwinTransformerBlock) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_case_insensitive(self): + m_upper = HyenaNDUNETR.get_variant("HHAA", in_channels=1, out_channels=14, feature_size=12) + m_lower = HyenaNDUNETR.get_variant("hhaa", in_channels=1, out_channels=14, feature_size=12) + self.assertEqual(m_upper.hyena_stages, m_lower.hyena_stages) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_aaaa_rejected(self): + """AAAA is plain SwinUNETR and intentionally not exposed via this constructor.""" + with self.assertRaisesRegex(ValueError, "Unknown paper variant"): + HyenaNDUNETR.get_variant("AAAA", in_channels=1, out_channels=14, feature_size=12) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_unknown_variant_rejected(self): + with self.assertRaisesRegex(ValueError, "Unknown paper variant"): + HyenaNDUNETR.get_variant("HAAA", in_channels=1, out_channels=14, feature_size=12) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_redundant_hyena_stages_kwarg_rejected(self): + with self.assertRaisesRegex(ValueError, "do not also pass hyena_stages"): + HyenaNDUNETR.get_variant( + "HHAA", in_channels=1, out_channels=14, feature_size=12, hyena_stages=(True, False, True, False) + ) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_paper_variants_table_matches_constants(self): + """Guard against the table in PAPER_VARIANTS drifting.""" + self.assertEqual(PAPER_VARIANTS["HHHH"], (True, True, True, True)) + self.assertEqual(PAPER_VARIANTS["HAHA"], (True, False, True, False)) + self.assertEqual(PAPER_VARIANTS["HHAA"], (True, True, False, False)) + + +class TestHyenaNDUNETRForward(unittest.TestCase): + """End-to-end forward over the three paper variants. CUDA required.""" + + @parameterized.expand(PAPER_VARIANT_CASES) + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + @skip_if_no_cuda + def test_forward_shape(self, name, _stages): + m = HyenaNDUNETR.get_variant(name, in_channels=1, out_channels=14, feature_size=12).cuda().eval() + x = torch.randn(1, 1, 64, 64, 64, device="cuda") + with torch.no_grad(): + out = m(x) + self.assertEqual(out.shape, (1, 14, 64, 64, 64)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/networks/nets/test_navit.py b/tests/networks/nets/test_navit.py new file mode 100644 index 00000000000..d6c33d6baa0 --- /dev/null +++ b/tests/networks/nets/test_navit.py @@ -0,0 +1,233 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest + +import torch +from parameterized import parameterized + +from monai.networks import eval_mode +from monai.networks.nets.navit import NaViT +from tests.test_utils import skip_if_quick + +# Shared default kwargs to reduce duplication across test cases. +DEFAULT_2D_KWARGS = { + "image_size": 64, + "patch_size": 16, + "num_classes": 10, + "hidden_size": 128, + "mlp_dim": 256, + "num_layers": 2, + "num_heads": 4, + "in_channels": 3, + "spatial_dims": 2, +} + +DEFAULT_3D_KWARGS = { + "image_size": 64, + "patch_size": 16, + "num_classes": 2, + "hidden_size": 256, + "mlp_dim": 512, + "num_layers": 2, + "num_heads": 8, + "in_channels": 1, + "spatial_dims": 3, +} + +# Each entry: (init_kwargs, batched_images_spec, expected_output_shape) +# batched_images_spec is a list of groups; each group is a list of image shape tuples. +TEST_CASES_SHAPE = [ + # 2D single image + (DEFAULT_2D_KWARGS, [[(3, 64, 64)]], (1, 10)), + # 2D multiple images in one group + ({**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1}, [[(1, 64, 64), (1, 32, 32), (1, 64, 32)]], (3, 5)), + # 2D multiple groups + ( + {**DEFAULT_2D_KWARGS, "image_size": 96, "num_classes": 8, "hidden_size": 192, "mlp_dim": 384, "num_heads": 6}, + [[(3, 96, 96), (3, 64, 64)], [(3, 80, 80)]], + (3, 8), + ), + # 3D single volume + (DEFAULT_3D_KWARGS, [[(1, 64, 64, 64)]], (1, 2)), + # 3D multiple volumes, multiple groups + ( + {**DEFAULT_3D_KWARGS, "image_size": 96, "num_classes": 3}, + [[(1, 96, 96, 96), (1, 64, 64, 64)], [(1, 80, 96, 80)]], + (3, 3), + ), + # token dropout (float) + ({**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1, "token_dropout_prob": 0.2}, [[(1, 64, 64)]], (1, 5)), + # custom dim_head + ({**DEFAULT_2D_KWARGS, "num_classes": 4, "in_channels": 1, "dim_head": 64}, [[(1, 64, 64)]], (1, 4)), + # qkv_bias enabled + ({**DEFAULT_2D_KWARGS, "num_classes": 3, "in_channels": 1, "qkv_bias": True}, [[(1, 64, 64)]], (1, 3)), + # anisotropic image_size 2D + ({**DEFAULT_2D_KWARGS, "image_size": (64, 128), "num_classes": 4, "in_channels": 1}, [[(1, 64, 128)]], (1, 4)), + # anisotropic image_size 3D + ( + {**DEFAULT_3D_KWARGS, "image_size": (64, 64, 96), "hidden_size": 192, "mlp_dim": 384, "num_heads": 6}, + [[(1, 64, 64, 96)]], + (1, 2), + ), +] + +# Invalid constructor arguments that should raise ValueError +TEST_CASES_ILL_ARG = [ + # spatial_dims not 2 or 3 + {**DEFAULT_2D_KWARGS, "in_channels": 1, "spatial_dims": 4}, + # hidden_size not divisible by num_heads + {**DEFAULT_2D_KWARGS, "in_channels": 1, "hidden_size": 100, "num_heads": 7}, + # dropout_rate out of [0, 1] + {**DEFAULT_2D_KWARGS, "in_channels": 1, "dropout_rate": 1.5}, + # emb_dropout_rate out of [0, 1] + {**DEFAULT_2D_KWARGS, "in_channels": 1, "emb_dropout_rate": -0.1}, + # token_dropout_prob out of (0, 1) as float + {**DEFAULT_2D_KWARGS, "in_channels": 1, "token_dropout_prob": 1.5}, + # num_heads zero + {**DEFAULT_2D_KWARGS, "in_channels": 1, "num_heads": 0}, + # image_size not divisible by patch_size + {**DEFAULT_2D_KWARGS, "in_channels": 1, "image_size": 50}, +] + +# Forward-validation cases: (description, image_tensor_shape) +# All use the same base net with in_channels=1, spatial_dims=2 +TEST_CASES_FORWARD_VALIDATION = [ + # wrong number of input channels (3 instead of 1) + ("wrong_channels", (3, 64, 64)), + # wrong number of spatial dimensions (3D image for 2D net) + ("wrong_spatial_dims", (1, 64, 64, 64)), + # spatial size not divisible by patch_size + ("patch_size_not_divisible", (1, 50, 64)), +] + + +@skip_if_quick +class TestNaViT(unittest.TestCase): + + @parameterized.expand(TEST_CASES_SHAPE) + def test_shape(self, input_param, batched_images_spec, expected_shape): + """Test output shape for various configurations.""" + net = NaViT(**input_param) + with eval_mode(net): + batched_images = [[torch.randn(*img_shape) for img_shape in group] for group in batched_images_spec] + result = net(batched_images) + self.assertEqual(result.shape, expected_shape) + + @parameterized.expand([(kwargs,) for kwargs in TEST_CASES_ILL_ARG]) + def test_ill_arg(self, input_param): + """Test that invalid constructor arguments raise ValueError.""" + with self.assertRaises(ValueError): + NaViT(**input_param) + + @parameterized.expand(TEST_CASES_FORWARD_VALIDATION) + def test_forward_validation(self, _, image_shape): + """Forward pass should raise ValueError for invalid input tensors.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "in_channels": 1, "num_classes": 2}) + net.eval() + with self.assertRaises(ValueError): + net([[torch.randn(*image_shape)]]) + + def test_auto_grouping(self): + """Auto-packing with group_images=True should produce correct total output size.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1}) + net.eval() + flat_images = [torch.randn(1, 64, 64) for _ in range(4)] + result = net(flat_images, group_images=True, group_max_seq_len=32) + self.assertEqual(result.shape, (4, 5)) + + def test_token_dropout_callable_invoked_during_training(self): + """Token dropout callable is invoked during training and produces correct shape.""" + call_log: list[tuple] = [] + + def recording_dropout(h, w): + call_log.append((h, w)) + return 0.25 + + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1}, token_dropout_prob=recording_dropout) + net.train() + result = net([[torch.randn(1, 64, 64)]]) + self.assertEqual(result.shape, (1, 5)) + self.assertGreater(len(call_log), 0, "Token dropout callable was not invoked during training.") + + def test_token_dropout_callable_not_invoked_during_eval(self): + """Token dropout callable is NOT invoked during eval.""" + call_log: list[tuple] = [] + + def recording_dropout(h, w): + call_log.append((h, w)) + return 0.25 + + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1}, token_dropout_prob=recording_dropout) + net.eval() + net([[torch.randn(1, 64, 64)]]) + self.assertEqual(len(call_log), 0, "Token dropout callable was invoked during eval mode.") + + def test_token_dropout_produces_different_outputs_in_training(self): + """With token dropout, different RNG seeds produce different training outputs.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1}, token_dropout_prob=0.5) + net.train() + input_data = [[torch.randn(1, 64, 64)]] + torch.manual_seed(0) + out1 = net(input_data) + torch.manual_seed(42) + out2 = net(input_data) + self.assertFalse( + torch.allclose(out1, out2), + "Token dropout should produce different outputs with different RNG seeds during training.", + ) + + def test_token_dropout_disabled_in_eval(self): + """Token dropout should not be applied during eval, producing deterministic output.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1}, token_dropout_prob=0.5) + net.eval() + input_data = [[torch.randn(1, 64, 64)]] + out1 = net(input_data) + out2 = net(input_data) + self.assertTrue(torch.allclose(out1, out2)) + + def test_eval_mode_deterministic(self): + """In eval mode, outputs should be identical across calls.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 5, "in_channels": 1}) + net.eval() + input_data = [[torch.randn(1, 64, 64)]] + out1 = net(input_data) + out2 = net(input_data) + self.assertTrue(torch.allclose(out1, out2)) + + def test_gradient_flow(self): + """All trainable parameters should receive gradients after a backward pass.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 3, "in_channels": 1}) + net.train() + output = net([[torch.randn(1, 64, 64)]]) + output.sum().backward() + for name, param in net.named_parameters(): + if param.requires_grad: + self.assertIsNotNone(param.grad, f"No gradient for parameter: {name}") + + def test_all_parameters_trainable(self): + """All parameters should be trainable by default.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 3, "in_channels": 1}) + frozen = [n for n, p in net.named_parameters() if not p.requires_grad] + self.assertEqual(frozen, [], f"Found frozen parameters: {frozen}") + + def test_variable_resolution_beyond_reference(self): + """Images larger than reference image_size should work via positional encoding clamping.""" + net = NaViT(**{**DEFAULT_2D_KWARGS, "num_classes": 2, "in_channels": 1}) + net.eval() + result = net([[torch.randn(1, 96, 96)]]) + self.assertEqual(result.shape, (1, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/networks/nets/test_spade_autoencoderkl.py b/tests/networks/nets/test_spade_autoencoderkl.py index 9353ceedc2e..c9a17be55fd 100644 --- a/tests/networks/nets/test_spade_autoencoderkl.py +++ b/tests/networks/nets/test_spade_autoencoderkl.py @@ -99,23 +99,6 @@ (1, 1, 16, 16), (1, 4, 4, 4), ], - [ - { - "spatial_dims": 2, - "label_nc": 3, - "in_channels": 1, - "out_channels": 1, - "channels": (4, 4, 4), - "latent_channels": 4, - "attention_levels": (False, False, False), - "num_res_blocks": 1, - "norm_num_groups": 4, - }, - (1, 1, 16, 16), - (1, 3, 16, 16), - (1, 1, 16, 16), - (1, 4, 4, 4), - ], [ { "spatial_dims": 2, diff --git a/tests/networks/nets/test_swin_unetr.py b/tests/networks/nets/test_swin_unetr.py index ba94aab4f90..b570270abb0 100644 --- a/tests/networks/nets/test_swin_unetr.py +++ b/tests/networks/nets/test_swin_unetr.py @@ -21,7 +21,15 @@ from monai.apps import download_url from monai.networks import eval_mode -from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR, filter_swinunetr +from monai.networks.blocks.hyena import HyenaTransformerBlock, is_nvsubquadratic_available +from monai.networks.nets.swin_unetr import ( + PatchMerging, + PatchMergingV2, + SwinTransformer, + SwinTransformerBlock, + SwinUNETR, + filter_swinunetr, +) from monai.networks.utils import copy_model_state from monai.utils import optional_import from tests.test_utils import ( @@ -34,6 +42,8 @@ ) einops, has_einops = optional_import("einops") +HAS_NVSUBQ = is_nvsubquadratic_available() +HAS_CUDA = torch.cuda.is_available() test_merging_mode = ["mergingv2", "merging", PatchMerging, PatchMergingV2] checkpoint_vals = [True, False] @@ -101,6 +111,19 @@ def test_invalid_input_shape(self): with self.assertRaises(ValueError): net_2d(torch.randn(1, 1, 48, 33)) # 33 is not divisible by 32 + @skipUnless(has_einops, "Requires einops") + def test_flash_attention(self): + input_param = {"in_channels": 1, "out_channels": 2, "feature_size": 12, "spatial_dims": 3} + net_ref = SwinUNETR(use_flash_attention=False, **input_param).double() + net_flash = SwinUNETR(use_flash_attention=True, **input_param).double() + net_flash.load_state_dict(net_ref.state_dict()) + x = torch.randn(1, 1, 64, 64, 64, dtype=torch.float64) + with eval_mode(net_ref, net_flash): + ref = net_ref.swinViT(x, net_ref.normalize) + out = net_flash.swinViT(x, net_flash.normalize) + for a, b in zip(ref, out, strict=True): + assert_allclose(a, b, atol=1e-6, rtol=1e-6, type_test=False) + def test_patch_merging(self): dim = 10 t = PatchMerging(dim)(torch.zeros((1, 21, 20, 20, dim))) @@ -126,5 +149,202 @@ def test_filter_swinunetr(self, input_param, key, value): self.assertTrue(len(loaded) == 157 and len(not_loaded) == 2) +# Backward-compat reference for SwinUNETR(use_hyena=False), feature_size=12, img_size=64^3, +# seeds (model=0, input=1), CPU. Captured before the HyenaND port; the default code path must +# keep reproducing this within tolerance. Tolerance-based (not a byte hash) so it tolerates +# benign cross-platform float drift while still catching a real change to the non-Hyena path. +HYENA_BACKCOMPAT_REF = torch.tensor( + [ + -0.069162, + -0.209673, + 0.543457, + -0.111868, + 0.474825, + 0.031108, + 0.191482, + -0.167401, + 0.091668, + 0.272223, + -0.084950, + -0.042126, + ] +) + + +def _build_hyena_unetr(use_hyena=False, hyena_stages=None, feature_size=12, out_channels=14): + return SwinUNETR( + in_channels=1, + out_channels=out_channels, + feature_size=feature_size, + use_hyena=use_hyena, + hyena_stages=hyena_stages, + ) + + +def _block_type_at_stage(model, stage_idx): + layer_attr = ["layers1", "layers2", "layers3", "layers4"][stage_idx] + return type(getattr(model.swinViT, layer_attr)[0].blocks[0]) + + +HYENA_VARIANT_CASES = [ + ("AAAA", False, None), + ("HHHH", True, None), + ("HAHA", True, (True, False, True, False)), + ("HHAA", True, (True, True, False, False)), +] + + +class TestSwinUNETRHyenaBackCompat(unittest.TestCase): + """The non-Hyena code path must keep reproducing its pre-port output (within tolerance).""" + + @skipUnless(has_einops, "Requires einops") + def test_default_path_unchanged(self): + """SwinUNETR with no hyena kwargs reproduces the pre-port reference output. + + Runs on CPU so it executes in environments without a GPU and is stable across + platforms; ``assert_close`` tolerates benign float drift while still flagging a real + change to the default (non-Hyena) code path. + """ + torch.manual_seed(0) + net = SwinUNETR(in_channels=1, out_channels=14, feature_size=12).eval() + torch.manual_seed(1) + x = torch.randn(1, 1, 64, 64, 64) + with torch.no_grad(): + out = net(x) + self.assertEqual(out.shape, (1, 14, 64, 64, 64)) + assert_allclose( + out.flatten()[: HYENA_BACKCOMPAT_REF.numel()], HYENA_BACKCOMPAT_REF, atol=1e-4, rtol=1e-4, type_test=False + ) + + +class TestSwinUNETRHyenaStages(unittest.TestCase): + """``hyena_stages`` must place :class:`HyenaTransformerBlock` at flagged stages and + :class:`SwinTransformerBlock` everywhere else. Construction-only; no CUDA required.""" + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_haha_pattern(self): + m = _build_hyena_unetr(use_hyena=True, hyena_stages=(True, False, True, False)) + self.assertIs(_block_type_at_stage(m, 0), HyenaTransformerBlock) + self.assertIs(_block_type_at_stage(m, 1), SwinTransformerBlock) + self.assertIs(_block_type_at_stage(m, 2), HyenaTransformerBlock) + self.assertIs(_block_type_at_stage(m, 3), SwinTransformerBlock) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_hhaa_pattern(self): + m = _build_hyena_unetr(use_hyena=True, hyena_stages=(True, True, False, False)) + self.assertIs(_block_type_at_stage(m, 0), HyenaTransformerBlock) + self.assertIs(_block_type_at_stage(m, 1), HyenaTransformerBlock) + self.assertIs(_block_type_at_stage(m, 2), SwinTransformerBlock) + self.assertIs(_block_type_at_stage(m, 3), SwinTransformerBlock) + + def test_aaaa_pattern_default(self): + m = _build_hyena_unetr(use_hyena=False) + for i in range(4): + self.assertIs(_block_type_at_stage(m, i), SwinTransformerBlock) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_hhhh_pattern_default(self): + m = _build_hyena_unetr(use_hyena=True) + for i in range(4): + self.assertIs(_block_type_at_stage(m, i), HyenaTransformerBlock) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_wrong_length_hyena_stages_raises(self): + with self.assertRaisesRegex(ValueError, "hyena_stages must have length"): + _build_hyena_unetr(use_hyena=True, hyena_stages=(True, True)) + + +class TestSwinUNETRHyenaForward(unittest.TestCase): + """Forward shape across the four paper variants. CUDA required.""" + + @parameterized.expand(HYENA_VARIANT_CASES) + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + @skip_if_no_cuda + def test_forward_shape(self, _name, use_hyena, hyena_stages): + m = _build_hyena_unetr(use_hyena=use_hyena, hyena_stages=hyena_stages).cuda() + x = torch.randn(1, 1, 64, 64, 64, device="cuda") + with torch.no_grad(): + out = m(x) + self.assertEqual(out.shape, (1, 14, 64, 64, 64)) + + +class TestSwinUNETRHyenaGradient(unittest.TestCase): + """Backward through the HHAA variant must produce grads on at least 90 percent of params.""" + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + @skip_if_no_cuda + def test_hhaa_backward(self): + m = _build_hyena_unetr(use_hyena=True, hyena_stages=(True, True, False, False)).cuda() + x = torch.randn(1, 1, 64, 64, 64, device="cuda") + m(x).sum().backward() + total = list(m.parameters()) + with_grad = [p for p in total if p.grad is not None] + coverage = len(with_grad) / len(total) + self.assertGreater(coverage, 0.9, f"only {coverage:.1%} of params received gradients") + + +class TestSwinTransformerRoPEDivisibility(unittest.TestCase): + """3D Hyena requires embed_dim * 2^layer % 6 == 0; 2D requires % 4.""" + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_3d_rejects_non_divisible_embed_dim(self): + with self.assertRaisesRegex(ValueError, "divisible by 6"): + SwinTransformer( + in_chans=1, + embed_dim=14, + window_size=(2, 2, 2), + patch_size=(2, 2, 2), + depths=(2, 2, 2, 2), + num_heads=(3, 6, 12, 24), + spatial_dims=3, + use_hyena=True, + ) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_2d_rejects_non_divisible_embed_dim(self): + with self.assertRaisesRegex(ValueError, "divisible by 4"): + SwinTransformer( + in_chans=1, + embed_dim=14, + window_size=(2, 2), + patch_size=(2, 2), + depths=(2, 2, 2, 2), + num_heads=(3, 6, 12, 24), + spatial_dims=2, + use_hyena=True, + ) + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + def test_per_stage_skips_check_for_attention_stages(self): + """Per-stage False suppresses the check for that stage; remaining Hyena stages still fire.""" + with self.assertRaisesRegex(ValueError, "divisible by 6"): + SwinTransformer( + in_chans=1, + embed_dim=14, + window_size=(2, 2, 2), + patch_size=(2, 2, 2), + depths=(2, 2, 2, 2), + num_heads=(3, 6, 12, 24), + spatial_dims=3, + use_hyena=True, + hyena_stages=(False, True, False, False), + ) + + +class TestSwinUNETRHyenaSlidingWindow(unittest.TestCase): + """The production inference path: sliding-window inference over HHAA must succeed.""" + + @skipUnless(HAS_NVSUBQ, "Requires nvsubquadratic") + @skip_if_no_cuda + def test_swi_hhaa(self): + from monai.inferers import sliding_window_inference + + m = _build_hyena_unetr(use_hyena=True, hyena_stages=(True, True, False, False)).cuda().eval() + x = torch.randn(1, 1, 96, 96, 96, device="cuda") + with torch.no_grad(): + out = sliding_window_inference(inputs=x, roi_size=(64, 64, 64), sw_batch_size=2, predictor=m, overlap=0.25) + self.assertEqual(out.shape, (1, 14, 96, 96, 96)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/test_lazy_onnx_import.py b/tests/networks/test_lazy_onnx_import.py new file mode 100644 index 00000000000..d66c7389f64 --- /dev/null +++ b/tests/networks/test_lazy_onnx_import.py @@ -0,0 +1,51 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest + + +class TestLazyOnnxImport(unittest.TestCase): + """Regression test for #8455. + + ``onnx``/``onnx.reference``/``onnxruntime`` used to be imported at module + scope in ``monai/networks/utils.py`` and ``monai/bundle/scripts.py`` via + ``optional_import``, which imports eagerly. Because ``import monai`` + auto-loads ``monai.networks``, a broken or hanging onnx install (e.g. + onnx 1.18 on Windows) would take down ``import monai`` with no error. + + The fix moves those imports inside the functions that use them, so neither + module binds onnx at module scope any more. Assert that directly: it is + deterministic and independent of which other optional packages happen to be + installed (some of them import onnx transitively, so checking + ``sys.modules`` after ``import monai`` is not a reliable signal). + """ + + def test_utils_does_not_bind_onnx_at_module_scope(self): + import monai.networks.utils as utils + + for attr in ("onnx", "onnxreference", "onnxruntime"): + self.assertFalse( + hasattr(utils, attr), + f"monai.networks.utils must not import {attr} at module scope (regression for #8455)", + ) + + def test_scripts_does_not_bind_onnx_at_module_scope(self): + import monai.bundle.scripts as scripts + + self.assertFalse( + hasattr(scripts, "onnx"), "monai.bundle.scripts must not import onnx at module scope (regression for #8455)" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_timedcall_dist.py b/tests/test_timedcall_dist.py index 28b4ab93067..863c0990db6 100644 --- a/tests/test_timedcall_dist.py +++ b/tests/test_timedcall_dist.py @@ -19,7 +19,7 @@ from tests.test_utils import TimedCall -@TimedCall(seconds=20 if sys.platform == "linux" else 60, force_quit=False) +@TimedCall(seconds=20 if sys.platform == "linux" else 60, force_quit=True) def case_1_seconds(arg=None): time.sleep(1) return "good" if not arg else arg diff --git a/tests/test_utils.py b/tests/test_utils.py index 05f7cb88d9e..320a23d0cd0 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -64,6 +64,8 @@ quick_test_var = "QUICKTEST" _tf32_enabled = None _test_data_config: dict = {} +# Fix dynamic warningregistry logs noise in python unit/pytest configurations +warnings.filterwarnings("ignore", message="Accessing.*__warningregistry__") MODULE_PATH = Path(__file__).resolve().parents[1] @@ -79,7 +81,7 @@ "unexpected EOF", # incomplete download "network issue", "gdown dependency", # gdown not installed - "md5 check", + "hash check", # check hash value of downloaded file "limit", # HTTP Error 503: Egress is over the account limit "authenticate", "timed out", # urlopen error [Errno 110] Connection timed out @@ -184,37 +186,6 @@ def skip_if_downloading_fails(): raise rt_e -SAMPLE_TIFF = "https://huggingface.co/datasets/MONAI/testing_data/resolve/main/CMU-1.tiff" -SAMPLE_TIFF_HASH = "73a7e89bc15576587c3d68e55d9bf92f09690280166240b48ff4b48230b13bcd" -SAMPLE_TIFF_HASH_TYPE = "sha256" - - -class TestDownloadUrl(unittest.TestCase): - """Exercise ``download_url`` success and hash-mismatch paths.""" - - def test_download_url(self): - """Download a sample TIFF and validate hash handling. - - Raises: - RuntimeError: When the downloaded file's hash does not match. - """ - with tempfile.TemporaryDirectory() as tempdir: - with skip_if_downloading_fails(): - download_url( - url=SAMPLE_TIFF, - filepath=os.path.join(tempdir, "model.tiff"), - hash_val=SAMPLE_TIFF_HASH, - hash_type=SAMPLE_TIFF_HASH_TYPE, - ) - with self.assertRaises(RuntimeError): - download_url( - url=SAMPLE_TIFF, - filepath=os.path.join(tempdir, "model_bad.tiff"), - hash_val="0" * 64, - hash_type=SAMPLE_TIFF_HASH_TYPE, - ) - - def test_pretrained_networks(network, input_param, device): with skip_if_downloading_fails(): return network(**input_param).to(device) diff --git a/tests/transforms/test_adaptors.py b/tests/transforms/test_adaptors.py index 2495fdc72e7..36f81f60cfb 100644 --- a/tests/transforms/test_adaptors.py +++ b/tests/transforms/test_adaptors.py @@ -125,6 +125,15 @@ def foo(a): dres = adaptor(foo, {"a": "b"}, {"b": "a"})(d) self.assertEqual(dres["b"], 4) + def test_kwargs_with_dict_inputs(self): + + def foo(**kwargs): + return {k: v * 2 for k, v in kwargs.items()} + + d = {"x": 3} + dres = adaptor(foo, {"out": "out"}, {"x": "out"})(d) + self.assertEqual(dres["out"], 6) + class TestApplyAlias(unittest.TestCase): diff --git a/tests/transforms/test_affine.py b/tests/transforms/test_affine.py index 5384db0f507..8a29e31de87 100644 --- a/tests/transforms/test_affine.py +++ b/tests/transforms/test_affine.py @@ -199,6 +199,48 @@ def test_affine(self, input_param, input_data, expected_val): ) +class TestComputeWAffine(unittest.TestCase): + def test_identity_2d(self): + """Identity matrix with same input/output size should produce pure translation to/from center.""" + mat = np.eye(3) + img_size = (4, 4) + sp_size = (4, 4) + result = Affine.compute_w_affine(2, mat, img_size, sp_size) + # For identity transform with same sizes, result should be identity + assert_allclose(result, np.eye(3), atol=1e-6) + + def test_identity_3d(self): + """Identity matrix in 3D with same input/output size.""" + mat = np.eye(4) + img_size = (6, 6, 6) + sp_size = (6, 6, 6) + result = Affine.compute_w_affine(3, mat, img_size, sp_size) + assert_allclose(result, np.eye(4), atol=1e-6) + + def test_different_sizes(self): + """When img_size != sp_size, result should include net translation.""" + mat = np.eye(3) + img_size = (4, 4) + sp_size = (8, 8) + result = Affine.compute_w_affine(2, mat, img_size, sp_size) + # Translation should account for the shift: (4-1)/2 - (8-1)/2 = 1.5 - 3.5 = -2.0 + expected_translation = np.array([(d1 - 1) / 2 - (d2 - 1) / 2 for d1, d2 in zip(img_size, sp_size)]) + assert_allclose(result[:2, 2], expected_translation, atol=1e-6) + + def test_output_shape(self): + """Output should be (r+1) x (r+1) matrix.""" + for r in [2, 3]: + mat = np.eye(r + 1) + result = Affine.compute_w_affine(r, mat, (4,) * r, (4,) * r) + self.assertEqual(result.shape, (r + 1, r + 1)) + + def test_torch_input(self): + """Method should accept torch tensor input.""" + mat = torch.eye(3) + result = Affine.compute_w_affine(2, mat, (4, 4), (4, 4)) + assert_allclose(result, np.eye(3), atol=1e-6) + + @unittest.skipUnless(optional_import("scipy")[1], "Requires scipy library.") class TestAffineConsistency(unittest.TestCase): @parameterized.expand([[7], [8], [9]]) diff --git a/tests/transforms/test_border_pad.py b/tests/transforms/test_border_pad.py index d0ea112d3af..adb01629f56 100644 --- a/tests/transforms/test_border_pad.py +++ b/tests/transforms/test_border_pad.py @@ -23,7 +23,6 @@ [{"spatial_border": 2}, (3, 8, 8, 4), (3, 12, 12, 8)], [{"spatial_border": [1, 2, 3]}, (3, 8, 8, 4), (3, 10, 12, 10)], [{"spatial_border": [1, 2, 3, 4, 5, 6]}, (3, 8, 8, 4), (3, 11, 15, 15)], - [{"spatial_border": [1, 2, 3, 4, 5, 6]}, (3, 8, 8, 4), (3, 11, 15, 15)], ] diff --git a/tests/transforms/test_border_padd.py b/tests/transforms/test_border_padd.py index c7eb3da762c..0b3ef058ab8 100644 --- a/tests/transforms/test_border_padd.py +++ b/tests/transforms/test_border_padd.py @@ -23,8 +23,6 @@ [{"keys": "img", "spatial_border": 2}, (3, 8, 8, 4), (3, 12, 12, 8)], [{"keys": "img", "spatial_border": [1, 2, 3]}, (3, 8, 8, 4), (3, 10, 12, 10)], [{"keys": "img", "spatial_border": [1, 2, 3, 4, 5, 6]}, (3, 8, 8, 4), (3, 11, 15, 15)], - [{"keys": "img", "spatial_border": 2}, (3, 8, 8, 4), (3, 12, 12, 8)], - [{"keys": "img", "spatial_border": 2}, (3, 8, 8, 4), (3, 12, 12, 8)], ] diff --git a/tests/transforms/test_center_spatial_crop.py b/tests/transforms/test_center_spatial_crop.py index c0da043ecbb..b2cb756cc53 100644 --- a/tests/transforms/test_center_spatial_crop.py +++ b/tests/transforms/test_center_spatial_crop.py @@ -14,15 +14,17 @@ import unittest import numpy as np +import torch from parameterized import parameterized +from monai.data.meta_obj import get_track_meta, set_track_meta from monai.transforms import CenterSpatialCrop +from monai.transforms.croppad.array import Crop from tests.croppers import CropTest TEST_SHAPES = [ [{"roi_size": [2, 2, -1]}, (3, 3, 3, 3), (3, 2, 2, 3), True], [{"roi_size": [2, 2, 2]}, (3, 3, 3, 3), (3, 2, 2, 2), True], - [{"roi_size": [2, 2, 2]}, (3, 3, 3, 3), (3, 2, 2, 2), True], [{"roi_size": [2, 1, 2]}, (3, 3, 3, 3), (3, 2, 1, 2), False], [{"roi_size": [2, 1, 3]}, (3, 3, 1, 3), (3, 2, 1, 3), True], ] @@ -51,6 +53,28 @@ def test_value(self, input_param, input_arr, expected_arr): def test_pending_ops(self, input_param, input_shape, _, align_corners): self.crop_test_pending_ops(input_param, input_shape, align_corners) + def test_compute_slices_broadcast(self): + self.assertEqual(Crop.compute_slices(roi_center=2, roi_size=(4, 6, 8)), (slice(0, 4), slice(0, 6), slice(0, 8))) + self.assertEqual(Crop.compute_slices(roi_start=1, roi_end=(3, 5, 7)), (slice(1, 3), slice(1, 5), slice(1, 7))) + with self.assertRaises(ValueError): + Crop.compute_slices(roi_center=(2, 3), roi_size=(4, 5, 6)) + with self.assertRaises(ValueError): + Crop.compute_slices(roi_start=(1, 2), roi_end=(3, 5, 7)) + with self.assertRaises(TypeError): + Crop.compute_slices(roi_center="10", roi_size=(4, 6)) + + def test_torch_compile(self): + prev_track_meta = get_track_meta() + set_track_meta(False) + try: + # eager backend traces the transform without needing the Inductor C++ compiler + cropper = torch.compile(CenterSpatialCrop(roi_size=(1, 16, 16)), backend="eager") + img = torch.rand(1, 1, 32, 32, dtype=torch.float32) + self.assertEqual(tuple(cropper(img).shape), (1, 1, 16, 16)) + finally: + set_track_meta(prev_track_meta) + torch._dynamo.reset() + if __name__ == "__main__": unittest.main() diff --git a/tests/transforms/test_clip_intensity_percentiles.py b/tests/transforms/test_clip_intensity_percentiles.py index 18ed47dbaa0..12d93da47de 100644 --- a/tests/transforms/test_clip_intensity_percentiles.py +++ b/tests/transforms/test_clip_intensity_percentiles.py @@ -192,5 +192,29 @@ def test_channel_wise(self, p): assert_allclose(result[i], p(expected), type_test="tensor", rtol=1e-4, atol=0) +class TestClipIntensityPercentilesClippingValues(unittest.TestCase): + def test_clipping_values_repeated_channel_wise_calls(self): + clipper = ClipIntensityPercentiles(lower=0, upper=100, channel_wise=True, return_clipping_values=True) + first = clipper(torch.tensor([[[0.0, 1.0]], [[10.0, 20.0]]])) + first_clipping_values = list(first.meta["clipping_values"]) + + second = clipper(torch.tensor([[[100.0, 200.0]], [[1000.0, 2000.0]]])) + + self.assertEqual(first_clipping_values, [(0.0, 1.0), (10.0, 20.0)]) + self.assertEqual(first.meta["clipping_values"], first_clipping_values) + self.assertEqual(second.meta["clipping_values"], [(100.0, 200.0), (1000.0, 2000.0)]) + + def test_clipping_values_repeated_non_channel_wise_calls(self): + clipper = ClipIntensityPercentiles(lower=0, upper=100, return_clipping_values=True) + first = clipper(torch.tensor([[[0.0, 1.0]]])) + first_clipping_values = list(first.meta["clipping_values"]) + + second = clipper(torch.tensor([[[100.0, 200.0]]])) + + self.assertEqual(first_clipping_values, [(0.0, 1.0)]) + self.assertEqual(first.meta["clipping_values"], first_clipping_values) + self.assertEqual(second.meta["clipping_values"], [(100.0, 200.0)]) + + if __name__ == "__main__": unittest.main() diff --git a/tests/transforms/test_create_rotate_order.py b/tests/transforms/test_create_rotate_order.py new file mode 100644 index 00000000000..1052c487518 --- /dev/null +++ b/tests/transforms/test_create_rotate_order.py @@ -0,0 +1,99 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest + +import numpy as np +import torch +from parameterized import parameterized + +from monai.transforms import Affine, Rotate +from monai.transforms.utils import create_rotate +from monai.utils import optional_import + +Rotation, has_scipy = optional_import("scipy.spatial.transform", name="Rotation") + +RADIANS = (0.3, -0.7, 1.1) +SEQUENCES = ["xyz", "zyx", "zxy", "yxz", "yzx", "xzy", "XYZ", "ZYX", "ZXY", "xzx", "ZXZ"] +BAD_ORDERS = ["abc", "Xy", "xx", "wxyz", "", "xyzz"] + + +def _legacy_rotate_3d(radians): + affine = np.eye(4) + a = np.eye(4) + a[1, 1], a[1, 2], a[2, 1], a[2, 2] = np.cos(radians[0]), -np.sin(radians[0]), np.sin(radians[0]), np.cos(radians[0]) + affine = affine @ a + a = np.eye(4) + a[0, 0], a[0, 2], a[2, 0], a[2, 2] = np.cos(radians[1]), np.sin(radians[1]), -np.sin(radians[1]), np.cos(radians[1]) + affine = affine @ a + a = np.eye(4) + a[0, 0], a[0, 1], a[1, 0], a[1, 1] = np.cos(radians[2]), -np.sin(radians[2]), np.sin(radians[2]), np.cos(radians[2]) + return affine @ a + + +class TestCreateRotateOrder(unittest.TestCase): + def test_default_matches_legacy(self): + legacy = _legacy_rotate_3d(RADIANS) + np.testing.assert_allclose(np.asarray(create_rotate(3, RADIANS)), legacy, atol=1e-6) + np.testing.assert_allclose(np.asarray(create_rotate(3, RADIANS, rotate_order="XYZ")), legacy, atol=1e-6) + + @parameterized.expand([(s,) for s in SEQUENCES]) + @unittest.skipUnless(has_scipy, "requires scipy") + def test_matches_scipy(self, order): + radians = RADIANS[: len(order)] + expected = Rotation.from_euler(order, radians).as_matrix() + np_mat = np.asarray(create_rotate(3, radians, rotate_order=order))[:3, :3] + torch_mat = create_rotate(3, radians, rotate_order=order, backend="torch").cpu().numpy()[:3, :3] + np.testing.assert_allclose(np_mat, expected, atol=1e-6) + np.testing.assert_allclose(torch_mat, expected, atol=1e-5) + + @parameterized.expand([(b,) for b in BAD_ORDERS]) + def test_invalid_order_raises(self, order): + with self.assertRaises(ValueError): + create_rotate(3, RADIANS, rotate_order=order) + + def test_order_too_short_for_radians(self): + with self.assertRaises(ValueError): + create_rotate(3, RADIANS, rotate_order="xy") + + def test_2d_ignores_order(self): + np.testing.assert_allclose( + np.asarray(create_rotate(2, [0.5])), np.asarray(create_rotate(2, [0.5], rotate_order="x")), atol=1e-6 + ) + + def test_transform_order_changes_output(self): + img = torch.arange(8 * 9 * 10, dtype=torch.float32).reshape(1, 8, 9, 10) + default = Rotate(angle=RADIANS, rotate_order="XYZ")(img) + reordered = Rotate(angle=RADIANS, rotate_order="zyx")(img) + self.assertFalse(torch.allclose(default, reordered)) + + def test_transform_invertible_with_order(self): + img = torch.arange(10 * 10 * 10, dtype=torch.float32).reshape(1, 10, 10, 10) + rotate = Rotate(angle=RADIANS, rotate_order="zyx", keep_size=True) + out = rotate(img) + inv = rotate.inverse(out) + self.assertEqual(tuple(inv.shape), tuple(img.shape)) + # rotation is lossy, so check the inverse undoes most of it rather than an exact round-trip + err_inv = np.abs(np.asarray(inv.cpu()) - np.asarray(img)).mean() + err_rot = np.abs(np.asarray(out.cpu()) - np.asarray(img)).mean() + self.assertLess(err_inv, err_rot) + + def test_affine_propagates_order(self): + img = torch.arange(6 * 7 * 8, dtype=torch.float32).reshape(1, 6, 7, 8) + affine_default = Affine(rotate_params=RADIANS, image_only=True)(img) + affine_reordered = Affine(rotate_params=RADIANS, rotate_order="zyx", image_only=True)(img) + self.assertFalse(torch.allclose(affine_default, affine_reordered)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/transforms/test_load_image.py b/tests/transforms/test_load_image.py index 4a470a624c9..e7ebec0f971 100644 --- a/tests/transforms/test_load_image.py +++ b/tests/transforms/test_load_image.py @@ -15,7 +15,9 @@ import shutil import tempfile import unittest +import warnings from pathlib import Path +from unittest.mock import patch import nibabel as nib import numpy as np @@ -24,11 +26,11 @@ from PIL import Image from monai.apps import download_and_extract -from monai.data import NibabelReader, PydicomReader +from monai.data import ImageReader, NibabelReader, PydicomReader from monai.data.meta_obj import get_track_meta, set_track_meta from monai.data.meta_tensor import MetaTensor from monai.transforms import LoadImage -from monai.utils import optional_import +from monai.utils import OptionalImportError, optional_import from tests.test_utils import SkipIfNoModule, assert_allclose, skip_if_downloading_fails, testing_data_config itk, has_itk = optional_import("itk", allow_namespace_pkg=True) @@ -52,6 +54,38 @@ def get_data(self, _obj): return np.zeros((1, 1, 1)), {"name": "my test"} +class _MissingDependencyReader(ImageReader): + """a test reader that simulates a missing optional dependency""" + + def __init__(self): + raise OptionalImportError("mock missing dependency") + + def verify_suffix(self, _filename): + return True + + def read(self, _data, **_kwargs): + return None + + def get_data(self, _img): + return np.zeros((1, 1)), {} + + +class _FallbackReader(ImageReader): + """a test reader that should not be used after an explicit reader import failure""" + + read_called = False + + def verify_suffix(self, _filename): + return True + + def read(self, data, **_kwargs): + type(self).read_called = True + return data + + def get_data(self, _img): + return np.zeros((1, 1)), {"name": "fallback"} + + TEST_CASE_1 = [{}, ["test_image.nii.gz"], (128, 128, 128)] TEST_CASE_2 = [{}, ["test_image.nii.gz"], (128, 128, 128)] @@ -184,6 +218,25 @@ def get_data(self, _obj): TESTS_META.append([{"reader": "ITKReader", "fallback_only": False}, (128, 128, 128), track_meta]) +class TestLoadImageReaderSelection(unittest.TestCase): + def test_explicit_string_reader_missing_dependency_raises(self): + """test explicitly requested string readers don't fall back when their dependency is missing""" + _FallbackReader.read_called = False + readers = {"missingreader": _MissingDependencyReader, "fallbackreader": _FallbackReader} + with patch("monai.transforms.io.array.SUPPORTED_READERS", readers): + loader = LoadImage() + self.assertEqual(len(loader.readers), 1) + self.assertIsInstance(loader.readers[0], _FallbackReader) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with self.assertRaises(OptionalImportError): + LoadImage(reader="missingreader") + + self.assertEqual(len(caught), 0) + self.assertFalse(_FallbackReader.read_called) + + @unittest.skipUnless(has_itk, "itk not installed") class TestLoadImage(unittest.TestCase): @classmethod diff --git a/tests/transforms/test_spatial_padd.py b/tests/transforms/test_spatial_padd.py index 10bf9587381..1b05f2e3afa 100644 --- a/tests/transforms/test_spatial_padd.py +++ b/tests/transforms/test_spatial_padd.py @@ -21,7 +21,6 @@ TESTS = [ [{"keys": ["img"], "spatial_size": [15, 8, 8], "method": "symmetric"}, (3, 8, 8, 5), (3, 15, 8, 8)], [{"keys": ["img"], "spatial_size": [15, 8, 8], "method": "end"}, (3, 8, 8, 5), (3, 15, 8, 8)], - [{"keys": ["img"], "spatial_size": [15, 8, 8], "method": "end"}, (3, 8, 8, 5), (3, 15, 8, 8)], [{"keys": ["img"], "spatial_size": [15, 8, -1], "method": "end"}, (3, 8, 5, 4), (3, 15, 8, 4)], ] diff --git a/tests/transforms/test_squeezedim.py b/tests/transforms/test_squeezedim.py index 5fd333d8217..8e838629f46 100644 --- a/tests/transforms/test_squeezedim.py +++ b/tests/transforms/test_squeezedim.py @@ -38,6 +38,7 @@ def test_shape(self, input_param, test_data, expected_shape): self.assertTupleEqual(result.shape, expected_shape) if "dim" in input_param and input_param["dim"] == 2 and isinstance(result, MetaTensor): assert_allclose(result.affine.shape, [3, 3]) + self.assertEqual(result.spatial_ndim, result.affine.shape[-1] - 1) @parameterized.expand(TESTS_FAIL) def test_invalid_inputs(self, exception, input_param, test_data): diff --git a/tests/transforms/utility/test_apply_transform_to_pointsd.py b/tests/transforms/utility/test_apply_transform_to_pointsd.py index 978113931cf..91aab663f79 100644 --- a/tests/transforms/utility/test_apply_transform_to_pointsd.py +++ b/tests/transforms/utility/test_apply_transform_to_pointsd.py @@ -57,7 +57,6 @@ POINT_3D_WORLD, ], [MetaTensor(DATA_3D, affine=AFFINE_2), POINT_3D_WORLD, None, True, True, POINT_3D_IMAGE_RAS], - [MetaTensor(DATA_3D, affine=AFFINE_2), POINT_3D_WORLD, None, True, True, POINT_3D_IMAGE_RAS], ] TEST_CASES_SEQUENCE = [ [ diff --git a/tests/transforms/utility/test_splitdim.py b/tests/transforms/utility/test_splitdim.py index 31d9983a2b8..090d55a6a54 100644 --- a/tests/transforms/utility/test_splitdim.py +++ b/tests/transforms/utility/test_splitdim.py @@ -16,6 +16,7 @@ import numpy as np from parameterized import parameterized +from monai.data import MetaTensor from monai.transforms.utility.array import SplitDim from tests.test_utils import TEST_NDARRAYS @@ -47,6 +48,44 @@ def test_singleton(self): out = SplitDim(dim=1)(arr) self.assertEqual(out[0].shape, shape) + def test_spatial_ndim_decremented(self): + """spatial_ndim decremented for keepdim=False on spatial dim.""" + import torch + + arr = MetaTensor(torch.randn(2, 3, 8, 7)) + self.assertEqual(arr.spatial_ndim, 3) + out = SplitDim(dim=1, keepdim=False)(arr) + for item in out: + self.assertIsInstance(item, MetaTensor) + self.assertEqual(item.spatial_ndim, 2) + + def test_spatial_ndim_negative_dim(self): + """spatial_ndim decremented for keepdim=False with negative dim.""" + import torch + + arr = MetaTensor(torch.randn(2, 3, 8, 7)) + self.assertEqual(arr.spatial_ndim, 3) + out = SplitDim(dim=-1, keepdim=False)(arr) + for item in out: + self.assertIsInstance(item, MetaTensor) + self.assertEqual(item.spatial_ndim, 2) + + def test_spatial_ndim_channel_dim_no_decrement(self): + """spatial_ndim clamped to the new tensor rank for keepdim=False on channel dim (dim=0).""" + import torch + + arr = MetaTensor(torch.randn(3, 8, 7)) + self.assertEqual(arr.spatial_ndim, 2) + out = SplitDim(dim=0, keepdim=False)(arr) + for item in out: + self.assertIsInstance(item, MetaTensor) + self.assertEqual(item.spatial_ndim, 1) + + out_keep = SplitDim(dim=0, keepdim=True)(arr) + for item in out_keep: + self.assertIsInstance(item, MetaTensor) + self.assertEqual(item.spatial_ndim, 2) + if __name__ == "__main__": unittest.main() diff --git a/tests/utils/enums/test_wsireader.py b/tests/utils/enums/test_wsireader.py index 2c6498234e7..ee6462b5791 100644 --- a/tests/utils/enums/test_wsireader.py +++ b/tests/utils/enums/test_wsireader.py @@ -11,6 +11,7 @@ from __future__ import annotations +import gc import os import unittest from pathlib import Path @@ -474,6 +475,21 @@ class WSIReaderTests: class Tests(unittest.TestCase): backend = None + def tearDown(self): + """Force deterministic cleanup of any backend WSI handles. + + ``LoadImage`` calls ``reader.read`` and then discards the returned + object after ``get_data`` (see ``monai/transforms/io/array.py``); + for WSI readers that object is a ``TiffFile`` / ``OpenSlide`` / + ``CuImage`` instance that owns an open file descriptor. Without + forcing a collection here, the temp TIFFs used by these tests + stay open long enough for the interpreter to emit + ``ResourceWarning: unclosed file ...``. Running ``gc.collect`` + invokes the corresponding ``__del__`` finalizers, which all close + the underlying handle. + """ + gc.collect() + @parameterized.expand([TEST_CASE_WHOLE_0]) def test_read_whole_image(self, file_path, level, expected_shape): reader = WSIReader(self.backend, level=level) diff --git a/tests/utils/misc/test_monai_utils_misc.py b/tests/utils/misc/test_monai_utils_misc.py index f4eb5d3956d..2c9e2f89a71 100644 --- a/tests/utils/misc/test_monai_utils_misc.py +++ b/tests/utils/misc/test_monai_utils_misc.py @@ -16,7 +16,13 @@ from parameterized import parameterized -from monai.utils.misc import MONAIEnvVars, check_kwargs_exist_in_class_init, run_cmd, to_tuple_of_dictionaries +from monai.utils.misc import ( + MONAIEnvVars, + check_kwargs_exist_in_class_init, + path_to_sqlite_uri, + run_cmd, + to_tuple_of_dictionaries, +) TO_TUPLE_OF_DICTIONARIES_TEST_CASES = [ ({}, tuple(), tuple()), @@ -99,5 +105,14 @@ def test_run_cmd(self): self.assertNotIn("\\t", str(cm.exception)) +class TestPathToSqliteUri(unittest.TestCase): + def test_path_to_sqlite_uri(self): + """Verify a sqlite:/// URI is built from the absolute path with special chars escaped.""" + self.assertTrue(path_to_sqlite_uri("/tmp/mlruns.db").startswith("sqlite:///")) + self.assertTrue(path_to_sqlite_uri("/tmp/mlruns.db").endswith("mlruns.db")) + self.assertIn("%3F", path_to_sqlite_uri("a?b.db")) + self.assertIn("%23", path_to_sqlite_uri("a#b.db")) + + if __name__ == "__main__": unittest.main() diff --git a/tests/utils/test_alias.py b/tests/utils/test_alias.py index e7abff3d890..8ec1f8ae007 100644 --- a/tests/utils/test_alias.py +++ b/tests/utils/test_alias.py @@ -23,7 +23,10 @@ class TestModuleAlias(unittest.TestCase): - """check that 'import monai.xx.file_name' returns a module""" + """ + Check that 'import monai.xx.file_name' returns a module. Note that this test will fail if a module has the same name + as a member of that module (or any other) which is imported in a `__init__.py` file. + """ def test_files(self): src_dir = os.path.dirname(TESTS_PATH) diff --git a/tests/utils/test_safe_eval.py b/tests/utils/test_safe_eval.py new file mode 100644 index 00000000000..d578ced9acc --- /dev/null +++ b/tests/utils/test_safe_eval.py @@ -0,0 +1,97 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import ast +import unittest + +import numpy as np +from parameterized import parameterized + +from monai.utils import safe_eval + +GOOD_EXPRS = [ + ("1+2", None, None, 3), + (" 1 + 2 ", None, None, 3), + ("1+2+x", {"x": 4}, None, 7), + ("1+2+x", None, {"x": 4}, 7), + ("1*2+x", {"x": 4}, None, 6), + ("(1+2)*3", None, None, 9), + ("foo+bar", {"foo": 1030}, {"bar": 204}, 1234), +] + +BAD_EXPRS = [("foo()",), ("foo.bar",), ("foo[123]",), ("(1,2)",), ("[3,4]",), ("int.__class__.__init__.__globals__",)] + + +class TestSafeEval(unittest.TestCase): + @parameterized.expand(GOOD_EXPRS) + def test_good_exprs(self, expr, globals_vars, locals_vars, expected): + """Test valid expressions with globals/locals evaluate to correct values.""" + result = safe_eval(expr, globals_vars, locals_vars) + self.assertEqual(result, expected) + + @parameterized.expand(GOOD_EXPRS) + def test_good_exprs_np(self, expr, globals_vars, locals_vars, expected): + """Test valid expressions with globals/locals evaluate to correct values with Numpy wrapping.""" + result = safe_eval(expr, globals_vars, locals_vars, rewrite_np=True) + self.assertEqual(result, expected) + + @parameterized.expand(BAD_EXPRS) + def test_bad_exprs(self, expr): + """Test bad expressions correctly raise ValueError.""" + with self.assertRaises(ValueError): + safe_eval(expr) + + with self.assertRaises(ValueError): + safe_eval(expr, rewrite_np=True) + + def test_allowed_types(self): + """Test restricting the allowed list of types.""" + allowed = [ast.Expression, ast.Constant, ast.BinOp, ast.Add] + result = safe_eval("1+2", allowed_types=allowed) + self.assertEqual(result, 3) + + with self.assertRaises(ValueError): + safe_eval("1*2", allowed_types=allowed) + + def test_rewrite_np_produces_numpy_types(self): + """Test that rewrite_np wraps literals in numpy types.""" + result = safe_eval("2 + 3", rewrite_np=True) + self.assertIsInstance(result, np.integer) + + result = safe_eval("2.5 + 1.5", rewrite_np=True) + self.assertIsInstance(result, np.floating) + + def test_rewrite_np_large_exponent(self): + """Test that rewrite_np prevents slow native-Python exponentiation.""" + # Under native Python, 9**9**9 produces a ~369-million-digit integer; + # under np.int32 it overflows and completes almost instantly. + result = safe_eval("9**9**9", rewrite_np=True) + self.assertIsInstance(result, np.integer) + + def test_rewrite_np_preserves_bool(self): + """Test that rewrite_np does not wrap bool constants.""" + result = safe_eval("True", rewrite_np=True) + self.assertIs(result, True) + + result = safe_eval("False", rewrite_np=True) + self.assertIs(result, False) + + def test_rewrite_np_inf_constant(self): + """Test that rewrite_np handles overflowing infinity literals.""" + result = safe_eval("1e309", rewrite_np=True) + self.assertIsInstance(result, np.floating) + self.assertTrue(np.isinf(result)) + + +if __name__ == "__main__": + unittest.main() diff --git a/versioneer.py b/versioneer.py index 5d0a606c91d..1e3753e63fb 100644 --- a/versioneer.py +++ b/versioneer.py @@ -1,4 +1,5 @@ -# Version: 0.23 + +# Version: 0.29 """The Versioneer - like a rocketeer, but for versions. @@ -8,12 +9,12 @@ * like a rocketeer, but for versions! * https://github.com/python-versioneer/python-versioneer * Brian Warner -* License: Public Domain (CC0-1.0) -* Compatible with: Python 3.7, 3.8, 3.9, 3.10 and pypy3 +* License: Public Domain (Unlicense) +* Compatible with: Python 3.7, 3.8, 3.9, 3.10, 3.11 and pypy3 * [![Latest Version][pypi-image]][pypi-url] * [![Build Status][travis-image]][travis-url] -This is a tool for managing a recorded version number in distutils/setuptools-based +This is a tool for managing a recorded version number in setuptools-based python projects. The goal is to remove the tedious and error-prone "update the embedded version string" step from your release process. Making a new release should be as easy as recording a new tag in your version-control @@ -22,10 +23,38 @@ ## Quick Install +Versioneer provides two installation modes. The "classic" vendored mode installs +a copy of versioneer into your repository. The experimental build-time dependency mode +is intended to allow you to skip this step and simplify the process of upgrading. + +### Vendored mode + +* `pip install versioneer` to somewhere in your $PATH + * A [conda-forge recipe](https://github.com/conda-forge/versioneer-feedstock) is + available, so you can also use `conda install -c conda-forge versioneer` +* add a `[tool.versioneer]` section to your `pyproject.toml` or a + `[versioneer]` section to your `setup.cfg` (see [Install](INSTALL.md)) + * Note that you will need to add `tomli; python_version < "3.11"` to your + build-time dependencies if you use `pyproject.toml` +* run `versioneer install --vendor` in your source tree, commit the results +* verify version information with `python setup.py version` + +### Build-time dependency mode + * `pip install versioneer` to somewhere in your $PATH -* add a `[versioneer]` section to your setup.cfg (see [Install](INSTALL.md)) -* run `versioneer install` in your source tree, commit the results -* Verify version information with `python setup.py version` + * A [conda-forge recipe](https://github.com/conda-forge/versioneer-feedstock) is + available, so you can also use `conda install -c conda-forge versioneer` +* add a `[tool.versioneer]` section to your `pyproject.toml` or a + `[versioneer]` section to your `setup.cfg` (see [Install](INSTALL.md)) +* add `versioneer` (with `[toml]` extra, if configuring in `pyproject.toml`) + to the `requires` key of the `build-system` table in `pyproject.toml`: + ```toml + [build-system] + requires = ["setuptools", "versioneer[toml]"] + build-backend = "setuptools.build_meta" + ``` +* run `versioneer install --no-vendor` in your source tree, commit the results +* verify version information with `python setup.py version` ## Version Identifiers @@ -230,9 +259,10 @@ To upgrade your project to a new release of Versioneer, do the following: * install the new Versioneer (`pip install -U versioneer` or equivalent) -* edit `setup.cfg`, if necessary, to include any new configuration settings - indicated by the release notes. See [UPGRADING](./UPGRADING.md) for details. -* re-run `versioneer install` in your source tree, to replace +* edit `setup.cfg` and `pyproject.toml`, if necessary, + to include any new configuration settings indicated by the release notes. + See [UPGRADING](./UPGRADING.md) for details. +* re-run `versioneer install --[no-]vendor` in your source tree, to replace `SRC/_version.py` * commit any changed files @@ -262,9 +292,8 @@ To make Versioneer easier to embed, all its code is dedicated to the public domain. The `_version.py` that it creates is also in the public domain. -Specifically, both are released under the Creative Commons "Public Domain -Dedication" license (CC0-1.0), as described in -https://creativecommons.org/publicdomain/zero/1.0/ . +Specifically, both are released under the "Unlicense", as described in +https://unlicense.org/. [pypi-image]: https://img.shields.io/pypi/v/versioneer.svg [pypi-url]: https://pypi.python.org/pypi/versioneer/ @@ -273,7 +302,6 @@ [travis-url]: https://travis-ci.com/github/python-versioneer/python-versioneer """ - # pylint:disable=invalid-name,import-outside-toplevel,missing-function-docstring # pylint:disable=missing-class-docstring,too-many-branches,too-many-statements # pylint:disable=raise-missing-from,too-many-lines,too-many-locals,import-error @@ -287,15 +315,34 @@ import re import subprocess import sys -from typing import Callable, Dict +from pathlib import Path +from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union +from typing import NoReturn import functools +have_tomllib = True +if sys.version_info >= (3, 11): + import tomllib +else: + try: + import tomli as tomllib + except ImportError: + have_tomllib = False + class VersioneerConfig: """Container for Versioneer configuration parameters.""" + VCS: str + style: str + tag_prefix: str + versionfile_source: str + versionfile_build: Optional[str] + parentdir_prefix: Optional[str] + verbose: Optional[bool] + -def get_root(): +def get_root() -> str: """Get the project root directory. We require that all commands are run from the project root, i.e. the @@ -303,20 +350,28 @@ def get_root(): """ root = os.path.realpath(os.path.abspath(os.getcwd())) setup_py = os.path.join(root, "setup.py") + pyproject_toml = os.path.join(root, "pyproject.toml") versioneer_py = os.path.join(root, "versioneer.py") - if not (os.path.exists(setup_py) or os.path.exists(versioneer_py)): + if not ( + os.path.exists(setup_py) + or os.path.exists(pyproject_toml) + or os.path.exists(versioneer_py) + ): # allow 'python path/to/setup.py COMMAND' root = os.path.dirname(os.path.realpath(os.path.abspath(sys.argv[0]))) setup_py = os.path.join(root, "setup.py") + pyproject_toml = os.path.join(root, "pyproject.toml") versioneer_py = os.path.join(root, "versioneer.py") - if not (os.path.exists(setup_py) or os.path.exists(versioneer_py)): - err = ( - "Versioneer was unable to run the project root directory. " - "Versioneer requires setup.py to be executed from " - "its immediate directory (like 'python setup.py COMMAND'), " - "or in a way that lets it use sys.argv[0] to find the root " - "(like 'python path/to/setup.py COMMAND')." - ) + if not ( + os.path.exists(setup_py) + or os.path.exists(pyproject_toml) + or os.path.exists(versioneer_py) + ): + err = ("Versioneer was unable to run the project root directory. " + "Versioneer requires setup.py to be executed from " + "its immediate directory (like 'python setup.py COMMAND'), " + "or in a way that lets it use sys.argv[0] to find the root " + "(like 'python path/to/setup.py COMMAND').") raise VersioneerBadRootError(err) try: # Certain runtime workflows (setup.py install/develop in a setuptools @@ -328,38 +383,59 @@ def get_root(): my_path = os.path.realpath(os.path.abspath(__file__)) me_dir = os.path.normcase(os.path.splitext(my_path)[0]) vsr_dir = os.path.normcase(os.path.splitext(versioneer_py)[0]) - if me_dir != vsr_dir: - print("Warning: build in %s is using versioneer.py from %s" % (os.path.dirname(my_path), versioneer_py)) + if me_dir != vsr_dir and "VERSIONEER_PEP518" not in globals(): + print("Warning: build in %s is using versioneer.py from %s" + % (os.path.dirname(my_path), versioneer_py)) except NameError: pass return root -def get_config_from_root(root): +def get_config_from_root(root: str) -> VersioneerConfig: """Read the project setup.cfg file to determine Versioneer config.""" # This might raise OSError (if setup.cfg is missing), or # configparser.NoSectionError (if it lacks a [versioneer] section), or # configparser.NoOptionError (if it lacks "VCS="). See the docstring at # the top of versioneer.py for instructions on writing your setup.cfg . - setup_cfg = os.path.join(root, "setup.cfg") - parser = configparser.ConfigParser() - with open(setup_cfg, "r") as cfg_file: - parser.read_file(cfg_file) - VCS = parser.get("versioneer", "VCS") # mandatory - - # Dict-like interface for non-mandatory entries - section = parser["versioneer"] + root_pth = Path(root) + pyproject_toml = root_pth / "pyproject.toml" + setup_cfg = root_pth / "setup.cfg" + section: Union[Dict[str, Any], configparser.SectionProxy, None] = None + if pyproject_toml.exists() and have_tomllib: + try: + with open(pyproject_toml, 'rb') as fobj: + pp = tomllib.load(fobj) + section = pp['tool']['versioneer'] + except (tomllib.TOMLDecodeError, KeyError) as e: + print(f"Failed to load config from {pyproject_toml}: {e}") + print("Try to load it from setup.cfg") + if not section: + parser = configparser.ConfigParser() + with open(setup_cfg) as cfg_file: + parser.read_file(cfg_file) + parser.get("versioneer", "VCS") # raise error if missing + + section = parser["versioneer"] + + # `cast`` really shouldn't be used, but its simplest for the + # common VersioneerConfig users at the moment. We verify against + # `None` values elsewhere where it matters cfg = VersioneerConfig() - cfg.VCS = VCS + cfg.VCS = section['VCS'] cfg.style = section.get("style", "") - cfg.versionfile_source = section.get("versionfile_source") + cfg.versionfile_source = cast(str, section.get("versionfile_source")) cfg.versionfile_build = section.get("versionfile_build") - cfg.tag_prefix = section.get("tag_prefix") + cfg.tag_prefix = cast(str, section.get("tag_prefix")) if cfg.tag_prefix in ("''", '""', None): cfg.tag_prefix = "" cfg.parentdir_prefix = section.get("parentdir_prefix") - cfg.verbose = section.get("verbose") + if isinstance(section, configparser.SectionProxy): + # Make sure configparser translates to bool + cfg.verbose = section.getboolean("verbose") + else: + cfg.verbose = section.get("verbose") + return cfg @@ -372,23 +448,28 @@ class NotThisMethod(Exception): HANDLERS: Dict[str, Dict[str, Callable]] = {} -def register_vcs_handler(vcs, method): # decorator +def register_vcs_handler(vcs: str, method: str) -> Callable: # decorator """Create decorator to mark a method as the handler of a VCS.""" - - def decorate(f): + def decorate(f: Callable) -> Callable: """Store f in HANDLERS[vcs][method].""" HANDLERS.setdefault(vcs, {})[method] = f return f - return decorate -def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=None): +def run_command( + commands: List[str], + args: List[str], + cwd: Optional[str] = None, + verbose: bool = False, + hide_stderr: bool = False, + env: Optional[Dict[str, str]] = None, +) -> Tuple[Optional[str], Optional[int]]: """Call the given command(s).""" assert isinstance(commands, list) process = None - popen_kwargs = {} + popen_kwargs: Dict[str, Any] = {} if sys.platform == "win32": # This hides the console window if pythonw.exe is used startupinfo = subprocess.STARTUPINFO() @@ -399,17 +480,12 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env= try: dispcmd = str([command] + args) # remember shell=False, so use git.cmd on windows, not just git - process = subprocess.Popen( - [command] + args, - cwd=cwd, - env=env, - stdout=subprocess.PIPE, - stderr=(subprocess.PIPE if hide_stderr else None), - **popen_kwargs, - ) + process = subprocess.Popen([command] + args, cwd=cwd, env=env, + stdout=subprocess.PIPE, + stderr=(subprocess.PIPE if hide_stderr + else None), **popen_kwargs) break - except OSError: - e = sys.exc_info()[1] + except OSError as e: if e.errno == errno.ENOENT: continue if verbose: @@ -429,17 +505,16 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env= return stdout, process.returncode -LONG_VERSION_PY[ - "git" -] = r''' +LONG_VERSION_PY['git'] = r''' # This file helps to compute a version number in source trees obtained from # git-archive tarball (such as those provided by githubs download-from-tag # feature). Distribution tarballs (built by setup.py sdist) and build # directories (produced by setup.py build) will contain a much shorter file # that just contains the computed version number. -# This file is released into the public domain. Generated by -# versioneer-0.23 (https://github.com/python-versioneer/python-versioneer) +# This file is released into the public domain. +# Generated by versioneer-0.29 +# https://github.com/python-versioneer/python-versioneer """Git implementation of _version.py.""" @@ -448,11 +523,11 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env= import re import subprocess import sys -from typing import Callable, Dict +from typing import Any, Callable, Dict, List, Optional, Tuple import functools -def get_keywords(): +def get_keywords() -> Dict[str, str]: """Get the keywords needed to look up the version information.""" # these strings will be replaced by git during git-archive. # setup.py/versioneer.py will grep for the variable names, so they must @@ -468,8 +543,15 @@ def get_keywords(): class VersioneerConfig: """Container for Versioneer configuration parameters.""" + VCS: str + style: str + tag_prefix: str + parentdir_prefix: str + versionfile_source: str + verbose: bool + -def get_config(): +def get_config() -> VersioneerConfig: """Create, populate and return the VersioneerConfig() object.""" # these strings are filled in when 'setup.py versioneer' creates # _version.py @@ -491,9 +573,9 @@ class NotThisMethod(Exception): HANDLERS: Dict[str, Dict[str, Callable]] = {} -def register_vcs_handler(vcs, method): # decorator +def register_vcs_handler(vcs: str, method: str) -> Callable: # decorator """Create decorator to mark a method as the handler of a VCS.""" - def decorate(f): + def decorate(f: Callable) -> Callable: """Store f in HANDLERS[vcs][method].""" if vcs not in HANDLERS: HANDLERS[vcs] = {} @@ -502,13 +584,19 @@ def decorate(f): return decorate -def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, - env=None): +def run_command( + commands: List[str], + args: List[str], + cwd: Optional[str] = None, + verbose: bool = False, + hide_stderr: bool = False, + env: Optional[Dict[str, str]] = None, +) -> Tuple[Optional[str], Optional[int]]: """Call the given command(s).""" assert isinstance(commands, list) process = None - popen_kwargs = {} + popen_kwargs: Dict[str, Any] = {} if sys.platform == "win32": # This hides the console window if pythonw.exe is used startupinfo = subprocess.STARTUPINFO() @@ -524,8 +612,7 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, stderr=(subprocess.PIPE if hide_stderr else None), **popen_kwargs) break - except OSError: - e = sys.exc_info()[1] + except OSError as e: if e.errno == errno.ENOENT: continue if verbose: @@ -545,7 +632,11 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, return stdout, process.returncode -def versions_from_parentdir(parentdir_prefix, root, verbose): +def versions_from_parentdir( + parentdir_prefix: str, + root: str, + verbose: bool, +) -> Dict[str, Any]: """Try to determine the version from the parent directory name. Source tarballs conventionally unpack into a directory that includes both @@ -570,13 +661,13 @@ def versions_from_parentdir(parentdir_prefix, root, verbose): @register_vcs_handler("git", "get_keywords") -def git_get_keywords(versionfile_abs): +def git_get_keywords(versionfile_abs: str) -> Dict[str, str]: """Extract version information from the given file.""" # the code embedded in _version.py can just fetch the value of these # keywords. When used from setup.py, we don't want to import _version.py, # so we do it with a regexp instead. This function is not used from # _version.py. - keywords = {} + keywords: Dict[str, str] = {} try: with open(versionfile_abs, "r") as fobj: for line in fobj: @@ -598,7 +689,11 @@ def git_get_keywords(versionfile_abs): @register_vcs_handler("git", "keywords") -def git_versions_from_keywords(keywords, tag_prefix, verbose): +def git_versions_from_keywords( + keywords: Dict[str, str], + tag_prefix: str, + verbose: bool, +) -> Dict[str, Any]: """Get version information from git keywords.""" if "refnames" not in keywords: raise NotThisMethod("Short version file found") @@ -662,7 +757,12 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose): @register_vcs_handler("git", "pieces_from_vcs") -def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): +def git_pieces_from_vcs( + tag_prefix: str, + root: str, + verbose: bool, + runner: Callable = run_command +) -> Dict[str, Any]: """Get version from 'git describe' in the root of the source tree. This only gets called if the git-archive 'subst' keywords were *not* @@ -681,7 +781,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): runner = functools.partial(runner, env=env) _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, - hide_stderr=True) + hide_stderr=not verbose) if rc != 0: if verbose: print("Directory %%s not under git control" %% root) @@ -692,7 +792,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): describe_out, rc = runner(GITS, [ "describe", "--tags", "--dirty", "--always", "--long", "--match", f"{tag_prefix}[[:digit:]]*" - ], cwd=root) + ], cwd=root) # --long was added in git-1.5.5 if describe_out is None: raise NotThisMethod("'git describe' failed") @@ -702,7 +802,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): raise NotThisMethod("'git rev-parse' failed") full_out = full_out.strip() - pieces = {} + pieces: Dict[str, Any] = {} pieces["long"] = full_out pieces["short"] = full_out[:7] # maybe improved later pieces["error"] = None @@ -794,14 +894,14 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): return pieces -def plus_or_dot(pieces): +def plus_or_dot(pieces: Dict[str, Any]) -> str: """Return a + if we don't already have one, else return a .""" if "+" in pieces.get("closest-tag", ""): return "." return "+" -def render_pep440(pieces): +def render_pep440(pieces: Dict[str, Any]) -> str: """Build up version string, with post-release "local version identifier". Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you @@ -826,7 +926,7 @@ def render_pep440(pieces): return rendered -def render_pep440_branch(pieces): +def render_pep440_branch(pieces: Dict[str, Any]) -> str: """TAG[[.dev0]+DISTANCE.gHEX[.dirty]] . The ".dev0" means not master branch. Note that .dev0 sorts backwards @@ -856,7 +956,7 @@ def render_pep440_branch(pieces): return rendered -def pep440_split_post(ver): +def pep440_split_post(ver: str) -> Tuple[str, Optional[int]]: """Split pep440 version string at the post-release segment. Returns the release segments before the post-release and the @@ -866,7 +966,7 @@ def pep440_split_post(ver): return vc[0], int(vc[1] or 0) if len(vc) == 2 else None -def render_pep440_pre(pieces): +def render_pep440_pre(pieces: Dict[str, Any]) -> str: """TAG[.postN.devDISTANCE] -- No -dirty. Exceptions: @@ -890,7 +990,7 @@ def render_pep440_pre(pieces): return rendered -def render_pep440_post(pieces): +def render_pep440_post(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]+gHEX] . The ".dev0" means dirty. Note that .dev0 sorts backwards @@ -917,7 +1017,7 @@ def render_pep440_post(pieces): return rendered -def render_pep440_post_branch(pieces): +def render_pep440_post_branch(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] . The ".dev0" means not master branch. @@ -946,7 +1046,7 @@ def render_pep440_post_branch(pieces): return rendered -def render_pep440_old(pieces): +def render_pep440_old(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]] . The ".dev0" means dirty. @@ -968,7 +1068,7 @@ def render_pep440_old(pieces): return rendered -def render_git_describe(pieces): +def render_git_describe(pieces: Dict[str, Any]) -> str: """TAG[-DISTANCE-gHEX][-dirty]. Like 'git describe --tags --dirty --always'. @@ -988,7 +1088,7 @@ def render_git_describe(pieces): return rendered -def render_git_describe_long(pieces): +def render_git_describe_long(pieces: Dict[str, Any]) -> str: """TAG-DISTANCE-gHEX[-dirty]. Like 'git describe --tags --dirty --always -long'. @@ -1008,7 +1108,7 @@ def render_git_describe_long(pieces): return rendered -def render(pieces, style): +def render(pieces: Dict[str, Any], style: str) -> Dict[str, Any]: """Render the given version pieces into the requested style.""" if pieces["error"]: return {"version": "unknown", @@ -1044,7 +1144,7 @@ def render(pieces, style): "date": pieces.get("date")} -def get_versions(): +def get_versions() -> Dict[str, Any]: """Get version information or return default if unable to do so.""" # I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have # __file__, we can work backwards from there to the root. Some @@ -1092,13 +1192,13 @@ def get_versions(): @register_vcs_handler("git", "get_keywords") -def git_get_keywords(versionfile_abs): +def git_get_keywords(versionfile_abs: str) -> Dict[str, str]: """Extract version information from the given file.""" # the code embedded in _version.py can just fetch the value of these # keywords. When used from setup.py, we don't want to import _version.py, # so we do it with a regexp instead. This function is not used from # _version.py. - keywords = {} + keywords: Dict[str, str] = {} try: with open(versionfile_abs, "r") as fobj: for line in fobj: @@ -1120,7 +1220,11 @@ def git_get_keywords(versionfile_abs): @register_vcs_handler("git", "keywords") -def git_versions_from_keywords(keywords, tag_prefix, verbose): +def git_versions_from_keywords( + keywords: Dict[str, str], + tag_prefix: str, + verbose: bool, +) -> Dict[str, Any]: """Get version information from git keywords.""" if "refnames" not in keywords: raise NotThisMethod("Short version file found") @@ -1146,7 +1250,7 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose): # starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of # just "foo-1.0". If we see a "tag: " prefix, prefer those. TAG = "tag: " - tags = {r[len(TAG) :] for r in refs if r.startswith(TAG)} + tags = {r[len(TAG):] for r in refs if r.startswith(TAG)} if not tags: # Either we're using git < 1.8.3, or there really are no tags. We use # a heuristic: assume all version tags have a digit. The old git %d @@ -1155,7 +1259,7 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose): # between branches and tags. By ignoring refnames without digits, we # filter out many common branch names like "release" and # "stabilization", as well as "HEAD" and "master". - tags = {r for r in refs if re.search(r"\d", r)} + tags = {r for r in refs if re.search(r'\d', r)} if verbose: print("discarding '%s', no digits" % ",".join(refs - tags)) if verbose: @@ -1163,35 +1267,33 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose): for ref in sorted(tags): # sorting will prefer e.g. "2.0" over "2.0rc1" if ref.startswith(tag_prefix): - r = ref[len(tag_prefix) :] + r = ref[len(tag_prefix):] # Filter out refs that exactly match prefix or that don't start # with a number once the prefix is stripped (mostly a concern # when prefix is '') - if not re.match(r"\d", r): + if not re.match(r'\d', r): continue if verbose: print("picking %s" % r) - return { - "version": r, - "full-revisionid": keywords["full"].strip(), - "dirty": False, - "error": None, - "date": date, - } + return {"version": r, + "full-revisionid": keywords["full"].strip(), + "dirty": False, "error": None, + "date": date} # no suitable tags, so version is "0+unknown", but full hex is still there if verbose: print("no suitable tags, using unknown + full revision id") - return { - "version": "0+unknown", - "full-revisionid": keywords["full"].strip(), - "dirty": False, - "error": "no suitable tags", - "date": None, - } + return {"version": "0+unknown", + "full-revisionid": keywords["full"].strip(), + "dirty": False, "error": "no suitable tags", "date": None} @register_vcs_handler("git", "pieces_from_vcs") -def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): +def git_pieces_from_vcs( + tag_prefix: str, + root: str, + verbose: bool, + runner: Callable = run_command +) -> Dict[str, Any]: """Get version from 'git describe' in the root of the source tree. This only gets called if the git-archive 'subst' keywords were *not* @@ -1209,7 +1311,8 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): env.pop("GIT_DIR", None) runner = functools.partial(runner, env=env) - _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=True) + _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, + hide_stderr=not verbose) if rc != 0: if verbose: print("Directory %s not under git control" % root) @@ -1217,9 +1320,10 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): # if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty] # if there isn't one, this yields HEX[-dirty] (no NUM) - describe_out, rc = runner( - GITS, ["describe", "--tags", "--dirty", "--always", "--long", "--match", f"{tag_prefix}[[:digit:]]*"], cwd=root - ) + describe_out, rc = runner(GITS, [ + "describe", "--tags", "--dirty", "--always", "--long", + "--match", f"{tag_prefix}[[:digit:]]*" + ], cwd=root) # --long was added in git-1.5.5 if describe_out is None: raise NotThisMethod("'git describe' failed") @@ -1229,12 +1333,13 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): raise NotThisMethod("'git rev-parse' failed") full_out = full_out.strip() - pieces = {} + pieces: Dict[str, Any] = {} pieces["long"] = full_out pieces["short"] = full_out[:7] # maybe improved later pieces["error"] = None - branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"], cwd=root) + branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"], + cwd=root) # --abbrev-ref was added in git-1.6.3 if rc != 0 or branch_name is None: raise NotThisMethod("'git rev-parse --abbrev-ref' returned error") @@ -1274,16 +1379,17 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): dirty = git_describe.endswith("-dirty") pieces["dirty"] = dirty if dirty: - git_describe = git_describe[: git_describe.rindex("-dirty")] + git_describe = git_describe[:git_describe.rindex("-dirty")] # now we have TAG-NUM-gHEX or HEX if "-" in git_describe: # TAG-NUM-gHEX - mo = re.search(r"^(.+)-(\d+)-g([0-9a-f]+)$", git_describe) + mo = re.search(r'^(.+)-(\d+)-g([0-9a-f]+)$', git_describe) if not mo: # unparsable. Maybe git-describe is misbehaving? - pieces["error"] = "unable to parse git-describe output: '%s'" % describe_out + pieces["error"] = ("unable to parse git-describe output: '%s'" + % describe_out) return pieces # tag @@ -1292,9 +1398,10 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): if verbose: fmt = "tag '%s' doesn't start with prefix '%s'" print(fmt % (full_tag, tag_prefix)) - pieces["error"] = "tag '%s' doesn't start with prefix '%s'" % (full_tag, tag_prefix) + pieces["error"] = ("tag '%s' doesn't start with prefix '%s'" + % (full_tag, tag_prefix)) return pieces - pieces["closest-tag"] = full_tag[len(tag_prefix) :] + pieces["closest-tag"] = full_tag[len(tag_prefix):] # distance: number of commits since tag pieces["distance"] = int(mo.group(2)) @@ -1318,7 +1425,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command): return pieces -def do_vcs_install(versionfile_source, ipy): +def do_vcs_install(versionfile_source: str, ipy: Optional[str]) -> None: """Git-specific installation logic for Versioneer. For Git, this means creating/changing .gitattributes to mark _version.py @@ -1330,14 +1437,15 @@ def do_vcs_install(versionfile_source, ipy): files = [versionfile_source] if ipy: files.append(ipy) - try: - my_path = __file__ - if my_path.endswith(".pyc") or my_path.endswith(".pyo"): - my_path = os.path.splitext(my_path)[0] + ".py" - versioneer_file = os.path.relpath(my_path) - except NameError: - versioneer_file = "versioneer.py" - files.append(versioneer_file) + if "VERSIONEER_PEP518" not in globals(): + try: + my_path = __file__ + if my_path.endswith((".pyc", ".pyo")): + my_path = os.path.splitext(my_path)[0] + ".py" + versioneer_file = os.path.relpath(my_path) + except NameError: + versioneer_file = "versioneer.py" + files.append(versioneer_file) present = False try: with open(".gitattributes", "r") as fobj: @@ -1355,7 +1463,11 @@ def do_vcs_install(versionfile_source, ipy): run_command(GITS, ["add", "--"] + files) -def versions_from_parentdir(parentdir_prefix, root, verbose): +def versions_from_parentdir( + parentdir_prefix: str, + root: str, + verbose: bool, +) -> Dict[str, Any]: """Try to determine the version from the parent directory name. Source tarballs conventionally unpack into a directory that includes both @@ -1367,23 +1479,20 @@ def versions_from_parentdir(parentdir_prefix, root, verbose): for _ in range(3): dirname = os.path.basename(root) if dirname.startswith(parentdir_prefix): - return { - "version": dirname[len(parentdir_prefix) :], - "full-revisionid": None, - "dirty": False, - "error": None, - "date": None, - } + return {"version": dirname[len(parentdir_prefix):], + "full-revisionid": None, + "dirty": False, "error": None, "date": None} rootdirs.append(root) root = os.path.dirname(root) # up a level if verbose: - print("Tried directories %s but none started with prefix %s" % (str(rootdirs), parentdir_prefix)) + print("Tried directories %s but none started with prefix %s" % + (str(rootdirs), parentdir_prefix)) raise NotThisMethod("rootdir doesn't start with parentdir_prefix") SHORT_VERSION_PY = """ -# This file was generated by 'versioneer.py' (0.23) from +# This file was generated by 'versioneer.py' (0.29) from # revision-control system data, or from the parent directory name of an # unpacked source archive. Distribution tarballs contain a pre-generated copy # of this file. @@ -1400,39 +1509,41 @@ def get_versions(): """ -def versions_from_file(filename): +def versions_from_file(filename: str) -> Dict[str, Any]: """Try to determine the version from _version.py if present.""" try: with open(filename) as f: contents = f.read() except OSError: raise NotThisMethod("unable to read _version.py") - mo = re.search(r"version_json = '''\n(.*)''' # END VERSION_JSON", contents, re.M | re.S) + mo = re.search(r"version_json = '''\n(.*)''' # END VERSION_JSON", + contents, re.M | re.S) if not mo: - mo = re.search(r"version_json = '''\r\n(.*)''' # END VERSION_JSON", contents, re.M | re.S) + mo = re.search(r"version_json = '''\r\n(.*)''' # END VERSION_JSON", + contents, re.M | re.S) if not mo: raise NotThisMethod("no version_json in _version.py") return json.loads(mo.group(1)) -def write_to_version_file(filename, versions): +def write_to_version_file(filename: str, versions: Dict[str, Any]) -> None: """Write the given version number to the given _version.py file.""" - os.unlink(filename) - contents = json.dumps(versions, sort_keys=True, indent=1, separators=(",", ": ")) + contents = json.dumps(versions, sort_keys=True, + indent=1, separators=(",", ": ")) with open(filename, "w") as f: f.write(SHORT_VERSION_PY % contents) print("set %s to '%s'" % (filename, versions["version"])) -def plus_or_dot(pieces): +def plus_or_dot(pieces: Dict[str, Any]) -> str: """Return a + if we don't already have one, else return a .""" if "+" in pieces.get("closest-tag", ""): return "." return "+" -def render_pep440(pieces): +def render_pep440(pieces: Dict[str, Any]) -> str: """Build up version string, with post-release "local version identifier". Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you @@ -1450,13 +1561,14 @@ def render_pep440(pieces): rendered += ".dirty" else: # exception #1 - rendered = "0+untagged.%d.g%s" % (pieces["distance"], pieces["short"]) + rendered = "0+untagged.%d.g%s" % (pieces["distance"], + pieces["short"]) if pieces["dirty"]: rendered += ".dirty" return rendered -def render_pep440_branch(pieces): +def render_pep440_branch(pieces: Dict[str, Any]) -> str: """TAG[[.dev0]+DISTANCE.gHEX[.dirty]] . The ".dev0" means not master branch. Note that .dev0 sorts backwards @@ -1479,13 +1591,14 @@ def render_pep440_branch(pieces): rendered = "0" if pieces["branch"] != "master": rendered += ".dev0" - rendered += "+untagged.%d.g%s" % (pieces["distance"], pieces["short"]) + rendered += "+untagged.%d.g%s" % (pieces["distance"], + pieces["short"]) if pieces["dirty"]: rendered += ".dirty" return rendered -def pep440_split_post(ver): +def pep440_split_post(ver: str) -> Tuple[str, Optional[int]]: """Split pep440 version string at the post-release segment. Returns the release segments before the post-release and the @@ -1495,7 +1608,7 @@ def pep440_split_post(ver): return vc[0], int(vc[1] or 0) if len(vc) == 2 else None -def render_pep440_pre(pieces): +def render_pep440_pre(pieces: Dict[str, Any]) -> str: """TAG[.postN.devDISTANCE] -- No -dirty. Exceptions: @@ -1519,7 +1632,7 @@ def render_pep440_pre(pieces): return rendered -def render_pep440_post(pieces): +def render_pep440_post(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]+gHEX] . The ".dev0" means dirty. Note that .dev0 sorts backwards @@ -1546,7 +1659,7 @@ def render_pep440_post(pieces): return rendered -def render_pep440_post_branch(pieces): +def render_pep440_post_branch(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] . The ".dev0" means not master branch. @@ -1575,7 +1688,7 @@ def render_pep440_post_branch(pieces): return rendered -def render_pep440_old(pieces): +def render_pep440_old(pieces: Dict[str, Any]) -> str: """TAG[.postDISTANCE[.dev0]] . The ".dev0" means dirty. @@ -1597,7 +1710,7 @@ def render_pep440_old(pieces): return rendered -def render_git_describe(pieces): +def render_git_describe(pieces: Dict[str, Any]) -> str: """TAG[-DISTANCE-gHEX][-dirty]. Like 'git describe --tags --dirty --always'. @@ -1617,7 +1730,7 @@ def render_git_describe(pieces): return rendered -def render_git_describe_long(pieces): +def render_git_describe_long(pieces: Dict[str, Any]) -> str: """TAG-DISTANCE-gHEX[-dirty]. Like 'git describe --tags --dirty --always -long'. @@ -1637,16 +1750,14 @@ def render_git_describe_long(pieces): return rendered -def render(pieces, style): +def render(pieces: Dict[str, Any], style: str) -> Dict[str, Any]: """Render the given version pieces into the requested style.""" if pieces["error"]: - return { - "version": "unknown", - "full-revisionid": pieces.get("long"), - "dirty": None, - "error": pieces["error"], - "date": None, - } + return {"version": "unknown", + "full-revisionid": pieces.get("long"), + "dirty": None, + "error": pieces["error"], + "date": None} if not style or style == "default": style = "pep440" # the default @@ -1670,20 +1781,16 @@ def render(pieces, style): else: raise ValueError("unknown style '%s'" % style) - return { - "version": rendered, - "full-revisionid": pieces["long"], - "dirty": pieces["dirty"], - "error": None, - "date": pieces.get("date"), - } + return {"version": rendered, "full-revisionid": pieces["long"], + "dirty": pieces["dirty"], "error": None, + "date": pieces.get("date")} class VersioneerBadRootError(Exception): """The project root directory is unknown or missing key files.""" -def get_versions(verbose=False): +def get_versions(verbose: bool = False) -> Dict[str, Any]: """Get the project version from whatever source is available. Returns dict with two keys: 'version' and 'full'. @@ -1698,8 +1805,9 @@ def get_versions(verbose=False): assert cfg.VCS is not None, "please set [versioneer]VCS= in setup.cfg" handlers = HANDLERS.get(cfg.VCS) assert handlers, "unrecognized VCS '%s'" % cfg.VCS - verbose = verbose or cfg.verbose - assert cfg.versionfile_source is not None, "please set versioneer.versionfile_source" + verbose = verbose or bool(cfg.verbose) # `bool()` used to avoid `None` + assert cfg.versionfile_source is not None, \ + "please set versioneer.versionfile_source" assert cfg.tag_prefix is not None, "please set versioneer.tag_prefix" versionfile_abs = os.path.join(root, cfg.versionfile_source) @@ -1753,21 +1861,17 @@ def get_versions(verbose=False): if verbose: print("unable to compute version") - return { - "version": "0+unknown", - "full-revisionid": None, - "dirty": None, - "error": "unable to compute version", - "date": None, - } + return {"version": "0+unknown", "full-revisionid": None, + "dirty": None, "error": "unable to compute version", + "date": None} -def get_version(): +def get_version() -> str: """Get the short version string for this project.""" return get_versions()["version"] -def get_cmdclass(cmdclass=None): +def get_cmdclass(cmdclass: Optional[Dict[str, Any]] = None): """Get the custom setuptools subclasses used by Versioneer. If the package uses a different cmdclass (e.g. one from numpy), it @@ -1795,16 +1899,16 @@ def get_cmdclass(cmdclass=None): class cmd_version(Command): description = "report generated version string" - user_options = [] - boolean_options = [] + user_options: List[Tuple[str, str, str]] = [] + boolean_options: List[str] = [] - def initialize_options(self): + def initialize_options(self) -> None: pass - def finalize_options(self): + def finalize_options(self) -> None: pass - def run(self): + def run(self) -> None: vers = get_versions(verbose=True) print("Version: %s" % vers["version"]) print(" full-revisionid: %s" % vers.get("full-revisionid")) @@ -1812,7 +1916,6 @@ def run(self): print(" date: %s" % vers.get("date")) if vers["error"]: print(" error: %s" % vers["error"]) - cmds["version"] = cmd_version # we override "build_py" in setuptools @@ -1834,13 +1937,13 @@ def run(self): # but the build_py command is not expected to copy any files. # we override different "build_py" commands for both environments - if "build_py" in cmds: - _build_py = cmds["build_py"] + if 'build_py' in cmds: + _build_py: Any = cmds['build_py'] else: from setuptools.command.build_py import build_py as _build_py class cmd_build_py(_build_py): - def run(self): + def run(self) -> None: root = get_root() cfg = get_config_from_root(root) versions = get_versions() @@ -1852,19 +1955,19 @@ def run(self): # now locate _version.py in the new build/ directory and replace # it with an updated value if cfg.versionfile_build: - target_versionfile = os.path.join(self.build_lib, cfg.versionfile_build) + target_versionfile = os.path.join(self.build_lib, + cfg.versionfile_build) print("UPDATING %s" % target_versionfile) write_to_version_file(target_versionfile, versions) - cmds["build_py"] = cmd_build_py - if "build_ext" in cmds: - _build_ext = cmds["build_ext"] + if 'build_ext' in cmds: + _build_ext: Any = cmds['build_ext'] else: from setuptools.command.build_ext import build_ext as _build_ext class cmd_build_ext(_build_ext): - def run(self): + def run(self) -> None: root = get_root() cfg = get_config_from_root(root) versions = get_versions() @@ -1877,22 +1980,21 @@ def run(self): return # now locate _version.py in the new build/ directory and replace # it with an updated value - target_versionfile = os.path.join(self.build_lib, cfg.versionfile_build) + if not cfg.versionfile_build: + return + target_versionfile = os.path.join(self.build_lib, + cfg.versionfile_build) if not os.path.exists(target_versionfile): - print( - f"Warning: {target_versionfile} does not exist, skipping " - "version update. This can happen if you are running build_ext " - "without first running build_py." - ) + print(f"Warning: {target_versionfile} does not exist, skipping " + "version update. This can happen if you are running build_ext " + "without first running build_py.") return print("UPDATING %s" % target_versionfile) write_to_version_file(target_versionfile, versions) - cmds["build_ext"] = cmd_build_ext if "cx_Freeze" in sys.modules: # cx_freeze enabled? - from cx_Freeze.dist import build_exe as _build_exe - + from cx_Freeze.dist import build_exe as _build_exe # type: ignore # nczeczulin reports that py2exe won't like the pep440-style string # as FILEVERSION, but it can be used for PRODUCTVERSION, e.g. # setup(console=[{ @@ -1901,7 +2003,7 @@ def run(self): # ... class cmd_build_exe(_build_exe): - def run(self): + def run(self) -> None: root = get_root() cfg = get_config_from_root(root) versions = get_versions() @@ -1913,25 +2015,24 @@ def run(self): os.unlink(target_versionfile) with open(cfg.versionfile_source, "w") as f: LONG = LONG_VERSION_PY[cfg.VCS] - f.write( - LONG - % { - "DOLLAR": "$", - "STYLE": cfg.style, - "TAG_PREFIX": cfg.tag_prefix, - "PARENTDIR_PREFIX": cfg.parentdir_prefix, - "VERSIONFILE_SOURCE": cfg.versionfile_source, - } - ) - + f.write(LONG % + {"DOLLAR": "$", + "STYLE": cfg.style, + "TAG_PREFIX": cfg.tag_prefix, + "PARENTDIR_PREFIX": cfg.parentdir_prefix, + "VERSIONFILE_SOURCE": cfg.versionfile_source, + }) cmds["build_exe"] = cmd_build_exe del cmds["build_py"] - if "py2exe" in sys.modules: # py2exe enabled? - from py2exe.distutils_buildexe import py2exe as _py2exe + if 'py2exe' in sys.modules: # py2exe enabled? + try: + from py2exe.setuptools_buildexe import py2exe as _py2exe # type: ignore + except ImportError: + from py2exe.distutils_buildexe import py2exe as _py2exe # type: ignore class cmd_py2exe(_py2exe): - def run(self): + def run(self) -> None: root = get_root() cfg = get_config_from_root(root) versions = get_versions() @@ -1943,27 +2044,23 @@ def run(self): os.unlink(target_versionfile) with open(cfg.versionfile_source, "w") as f: LONG = LONG_VERSION_PY[cfg.VCS] - f.write( - LONG - % { - "DOLLAR": "$", - "STYLE": cfg.style, - "TAG_PREFIX": cfg.tag_prefix, - "PARENTDIR_PREFIX": cfg.parentdir_prefix, - "VERSIONFILE_SOURCE": cfg.versionfile_source, - } - ) - + f.write(LONG % + {"DOLLAR": "$", + "STYLE": cfg.style, + "TAG_PREFIX": cfg.tag_prefix, + "PARENTDIR_PREFIX": cfg.parentdir_prefix, + "VERSIONFILE_SOURCE": cfg.versionfile_source, + }) cmds["py2exe"] = cmd_py2exe # sdist farms its file list building out to egg_info - if "egg_info" in cmds: - _sdist = cmds["egg_info"] + if 'egg_info' in cmds: + _egg_info: Any = cmds['egg_info'] else: from setuptools.command.egg_info import egg_info as _egg_info class cmd_egg_info(_egg_info): - def find_sources(self): + def find_sources(self) -> None: # egg_info.find_sources builds the manifest list and writes it # in one shot super().find_sources() @@ -1971,7 +2068,7 @@ def find_sources(self): # Modify the filelist and normalize it root = get_root() cfg = get_config_from_root(root) - self.filelist.append("versioneer.py") + self.filelist.append('versioneer.py') if cfg.versionfile_source: # There are rare cases where versionfile_source might not be # included by default, so we must be explicit @@ -1984,23 +2081,23 @@ def find_sources(self): # We will instead replicate their final normalization (to unicode, # and POSIX-style paths) from setuptools import unicode_utils + normalized = [unicode_utils.filesys_decode(f).replace(os.sep, '/') + for f in self.filelist.files] - normalized = [unicode_utils.filesys_decode(f).replace(os.sep, "/") for f in self.filelist.files] + manifest_filename = os.path.join(self.egg_info, 'SOURCES.txt') + with open(manifest_filename, 'w') as fobj: + fobj.write('\n'.join(normalized)) - manifest_filename = os.path.join(self.egg_info, "SOURCES.txt") - with open(manifest_filename, "w") as fobj: - fobj.write("\n".join(normalized)) - - cmds["egg_info"] = cmd_egg_info + cmds['egg_info'] = cmd_egg_info # we override different "sdist" commands for both environments - if "sdist" in cmds: - _sdist = cmds["sdist"] + if 'sdist' in cmds: + _sdist: Any = cmds['sdist'] else: from setuptools.command.sdist import sdist as _sdist class cmd_sdist(_sdist): - def run(self): + def run(self) -> None: versions = get_versions() self._versioneer_generated_versions = versions # unless we update this, the command will keep using the old @@ -2008,7 +2105,7 @@ def run(self): self.distribution.metadata.version = versions["version"] return _sdist.run(self) - def make_release_tree(self, base_dir, files): + def make_release_tree(self, base_dir: str, files: List[str]) -> None: root = get_root() cfg = get_config_from_root(root) _sdist.make_release_tree(self, base_dir, files) @@ -2017,8 +2114,8 @@ def make_release_tree(self, base_dir, files): # updated value target_versionfile = os.path.join(base_dir, cfg.versionfile_source) print("UPDATING %s" % target_versionfile) - write_to_version_file(target_versionfile, self._versioneer_generated_versions) - + write_to_version_file(target_versionfile, + self._versioneer_generated_versions) cmds["sdist"] = cmd_sdist return cmds @@ -2073,14 +2170,16 @@ def make_release_tree(self, base_dir, files): """ -def do_setup(): +def do_setup() -> int: """Do main VCS-independent setup function for installing Versioneer.""" root = get_root() try: cfg = get_config_from_root(root) - except (OSError, configparser.NoSectionError, configparser.NoOptionError) as e: + except (OSError, configparser.NoSectionError, + configparser.NoOptionError) as e: if isinstance(e, (OSError, configparser.NoSectionError)): - print("Adding sample versioneer config to setup.cfg", file=sys.stderr) + print("Adding sample versioneer config to setup.cfg", + file=sys.stderr) with open(os.path.join(root, "setup.cfg"), "a") as f: f.write(SAMPLE_CONFIG) print(CONFIG_ERROR, file=sys.stderr) @@ -2089,18 +2188,16 @@ def do_setup(): print(" creating %s" % cfg.versionfile_source) with open(cfg.versionfile_source, "w") as f: LONG = LONG_VERSION_PY[cfg.VCS] - f.write( - LONG - % { - "DOLLAR": "$", - "STYLE": cfg.style, - "TAG_PREFIX": cfg.tag_prefix, - "PARENTDIR_PREFIX": cfg.parentdir_prefix, - "VERSIONFILE_SOURCE": cfg.versionfile_source, - } - ) - - ipy = os.path.join(os.path.dirname(cfg.versionfile_source), "__init__.py") + f.write(LONG % {"DOLLAR": "$", + "STYLE": cfg.style, + "TAG_PREFIX": cfg.tag_prefix, + "PARENTDIR_PREFIX": cfg.parentdir_prefix, + "VERSIONFILE_SOURCE": cfg.versionfile_source, + }) + + ipy = os.path.join(os.path.dirname(cfg.versionfile_source), + "__init__.py") + maybe_ipy: Optional[str] = ipy if os.path.exists(ipy): try: with open(ipy, "r") as f: @@ -2121,16 +2218,16 @@ def do_setup(): print(" %s unmodified" % ipy) else: print(" %s doesn't exist, ok" % ipy) - ipy = None + maybe_ipy = None # Make VCS-specific changes. For git, this means creating/changing # .gitattributes to mark _version.py for export-subst keyword # substitution. - do_vcs_install(cfg.versionfile_source, ipy) + do_vcs_install(cfg.versionfile_source, maybe_ipy) return 0 -def scan_setup_py(): +def scan_setup_py() -> int: """Validate the contents of setup.py against Versioneer's expectations.""" found = set() setters = False @@ -2167,10 +2264,14 @@ def scan_setup_py(): return errors +def setup_command() -> NoReturn: + """Set up Versioneer and exit with appropriate error code.""" + errors = do_setup() + errors += scan_setup_py() + sys.exit(1 if errors else 0) + + if __name__ == "__main__": cmd = sys.argv[1] if cmd == "setup": - errors = do_setup() - errors += scan_setup_py() - if errors: - sys.exit(1) + setup_command()